Update app.py

app.py

@@ -53,6 +53,7 @@ def load_text_model():
 @st.cache_resource
 def load_image_model():
     import os
+    import tensorflow as tf
     model_path = "image_model.h5"
 
     if not os.path.exists(model_path):
@@ -61,11 +62,21 @@ def load_image_model():
         with open(model_path, "wb") as f:
             f.write(response.content)
 
-    # Use a lightweight PyTorch MobileNetV2 instead of TensorFlow
-    model = models.mobilenet_v2(weights=None)
-    model.classifier[1] = nn.Linear(model.last_channel, 1)
-    model.eval()
+    model = tf.keras.models.load_model(model_path)
     return model
+
+def predict_image(img, image_model):
+    import numpy as np
+    import tensorflow as tf
+    img_resized = img.resize((224, 224))
+    img_array = tf.keras.preprocessing.image.img_to_array(img_resized) / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+    prediction = image_model.predict(img_array)
+    confidence = float(prediction[0][0]) * 100
+    if confidence > 50:
+        return "AI", confidence
+    else:
+        return "REAL", 100 - confidence
     
 
 # # ---------------------------
@@ -98,21 +109,21 @@ def predict_news(text, tokenizer, text_model):
     # This model outputs: 0 = FAKE, 1 = REAL
     label = text_model.config.id2label[prediction]
     return label, confidence
-# ---------------------------
-# PREDICT IMAGE
-# ---------------------------
-def predict_image(img, image_model):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406],
-                             [0.229, 0.224, 0.225])
-    ])
-    tensor = transform(img).unsqueeze(0)
-    with torch.no_grad():
-        output = torch.sigmoid(image_model(tensor))
-    confidence = output.item() * 100
-    return confidence
+# # ---------------------------
+# # PREDICT IMAGE
+# # ---------------------------
+# def predict_image(img, image_model):
+#     transform = transforms.Compose([
+#         transforms.Resize((224, 224)),
+#         transforms.ToTensor(),
+#         transforms.Normalize([0.485, 0.456, 0.406],
+#                              [0.229, 0.224, 0.225])
+#     ])
+#     tensor = transform(img).unsqueeze(0)
+#     with torch.no_grad():
+#         output = torch.sigmoid(image_model(tensor))
+#     confidence = output.item() * 100
+#     return confidence
 
 # ---------------------------
 # TABS