Update app.py

app.py

@@ -1,21 +1,40 @@
 import keras
 import numpy as np
 import streamlit as st
+import random
 from PIL import Image
 import os
 from huggingface_hub import snapshot_download
 
-models = ['DDPM','Glide','Latent Diffusion','Palette','Echt','Stable Diffusion','Unseen Fake','VQ Diffusion']
+def random_crop(img, min_size=160, max_size=2048, ratio=5/8):
+    width, height = img.size
 
-local_model_path = snapshot_download(repo_id="ElBeh/ma_basemodel")
-model = keras.models.load_model(local_model_path)
+    crop_width = random.randint(min_size, min(max_size, width))
+    crop_height = int(crop_width * ratio)
 
+    if crop_height > height:
+        crop_height = height
+        crop_width = int(crop_height / ratio)
+    
+    left = random.randint(0, width - crop_width)
+    top = random.randint(0, height - crop_height)
+    right = left + crop_width
+    bottom = top + crop_height
+
+    return img.crop((left, top, right, bottom))
+    
 def get_prediction(img):
     x = np.array(img)
     x = np.expand_dims(x, axis=0)
     predictions = model.predict(x)
     return predictions[0,:]
 
+models = ['DDPM', 'Glide', 'Latent Diffusion', 'Palette', 'Stable Diffusion', 'VQ Diffusion', 'real', 'unseen_fake']
+
+local_model_path = snapshot_download(repo_id="ElBeh/ma_basemodel")
+model = keras.models.load_model(local_model_path)
+
+
 st.title("Fake Detection")
 
 file_name = st.file_uploader("Choose an image...", ['jpg'])
@@ -24,6 +43,10 @@ if file_name is not None:
     col1, col2 = st.columns(2)
 
     image = Image.open(file_name)
+    if image.size != (200, 200) or image.mode != 'RGB':
+        image = random_crop(image)
+        image = image.resize((200, 200))     
+        
     col1.image(image, use_column_width=True)  
     predictions = get_prediction(image)