Update app.py

app.py

@@ -1,30 +1,65 @@
-from transformers import AutoImageProcessor, SiglipForImageClassification
-from PIL import Image
 import gradio as gr
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
 import torch
 
-# New model details
-model_name = "prithivMLmods/deepfake-detector-model-v1"
-processor = AutoImageProcessor.from_pretrained(model_name)
-model = SiglipForImageClassification.from_pretrained(model_name)
+# Load multiple models
+model_names = [
+    "dima806/deepfake_vs_real_image_detection",
+    "BuzzFeedNews/Deepfake-Detection",
+    "umarlai/deepfake-detection-vit"
+]
+
+models = []
+processors = []
+
+for name in model_names:
+    processors.append(AutoImageProcessor.from_pretrained(name))
+    models.append(AutoModelForImageClassification.from_pretrained(name))
+
+def predict(image):
+    votes = {"Real": 0, "Fake": 0}
+    probs_list = []
+
+    for model, processor in zip(models, processors):
+        inputs = processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
+        
+        # Map outputs depending on model’s labels
+        labels = model.config.id2label
+        result = {labels[i]: float(probs[i]) for i in range(len(labels))}
 
-def detect_deepfake(image):
-    image = Image.fromarray(image).convert("RGB")
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
-    id2label = { "0": "Fake", "1": "Real" }
-    return { id2label[str(i)]: float(probs[i]) for i in range(2) }
+        # Normalize labels to Real/Fake
+        real_score = result.get("Real", result.get("REAL", result.get("0", 0)))
+        fake_score = result.get("Fake", result.get("FAKE", result.get("1", 0)))
 
+        probs_list.append({"Real": real_score, "Fake": fake_score})
+
+        if fake_score > real_score:
+            votes["Fake"] += 1
+        else:
+            votes["Real"] += 1
+
+    # Majority voting
+    final_label = "Fake" if votes["Fake"] > votes["Real"] else "Real"
+
+    # Average probability
+    avg_real = sum([p["Real"] for p in probs_list]) / len(probs_list)
+    avg_fake = sum([p["Fake"] for p in probs_list]) / len(probs_list)
+
+    return {final_label: max(avg_real, avg_fake), "Real": avg_real, "Fake": avg_fake}
+
+# UI
 demo = gr.Interface(
-    fn=detect_deepfake,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Label(num_top_classes=2),
-    title="Deepfake Detector",
-    description="Upload an image; this model predicts Real vs Fake."
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="label",
+    title="🕵️ Deepfake Detector (Ensemble)",
+    description="Uploads an image and checks if it's REAL or FAKE using 3 different models combined."
 )
 
 if __name__ == "__main__":
     demo.launch()
-
+