Update deepfake_vs_real.py

deepfake_vs_real.py

@@ -1,5 +1,4 @@
 import gradio as gr
-import spaces
 from transformers import AutoImageProcessor
 from transformers import SiglipForImageClassification
 from transformers.image_utils import load_image
@@ -7,13 +6,12 @@ from PIL import Image
 import torch
 
 # Load model and processor
-model_name = "prithivMLmods/Deepfake-Detect-Siglip2"
+model_name = "prithivMLmods/Deepfake-vs-Real-8000"
 model = SiglipForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name)
 
-@spaces.GPU
-def deepfake_detection(image):
-    """Classifies an image as Fake or Real."""
+def deepfake_classification(image):
+    """Predicts whether an image is a Deepfake or Real."""
     image = Image.fromarray(image).convert("RGB")
     inputs = processor(images=image, return_tensors="pt")
     
@@ -22,18 +20,20 @@ def deepfake_detection(image):
         logits = outputs.logits
         probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
     
-    labels = model.config.id2label
-    predictions = {labels[i]: round(probs[i], 3) for i in range(len(probs))}
+    labels = {
+        "0": "Deepfake", "1": "Real one"
+    }
+    predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
     
     return predictions
 
 # Create Gradio interface
 iface = gr.Interface(
-    fn=deepfake_detection,
+    fn=deepfake_classification,
     inputs=gr.Image(type="numpy"),
-    outputs=gr.Label(label="Detection Result"),
-    title="Deepfake Detection Model",
-    description="Upload an image to determine if it is Fake or Real."
+    outputs=gr.Label(label="Prediction Scores"),
+    title="Deepfake vs. Real Image Classification",
+    description="Upload an image to determine if it's a Deepfake or a Real one."
 )
 
 # Launch the app