Update app.py

app.py

@@ -2,8 +2,9 @@ import keras
 import numpy as np
 import streamlit as st
 import random
-from PIL import Image
 import os
+from PIL import Image
+from io import BytesIO
 from huggingface_hub import snapshot_download
 
 def random_crop(img, min_size=160, max_size=2048, ratio=5/8):
@@ -22,6 +23,14 @@ def random_crop(img, min_size=160, max_size=2048, ratio=5/8):
     bottom = top + crop_height
 
     return img.crop((left, top, right, bottom))
+
+def jpg_compression(img):
+    quality = random.randint(65, 100)
+    jpeg_image = BytesIO()
+    img.convert("RGB").save(jpeg_image, 'JPEG', quality=quality)
+    jpeg_image.seek(0)
+    compressed_img = Image.open(jpeg_image)
+    return compressed_img
     
 def get_prediction(img):
     x = np.array(img)
@@ -37,15 +46,18 @@ model = keras.models.load_model(local_model_path)
 
 st.title("Fake Detection")
 
-file_name = st.file_uploader("Choose an image...", ['jpg'])
+file_name = st.file_uploader("Choose an image...")
 
 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))     
+
+    if image.format != "JPEG":
+        image = jpg_compression(image)
         
     col1.image(image, use_column_width=True)  
     predictions = get_prediction(image)