Update app.py

app.py

@@ -1,32 +1,49 @@
-import os
 import gradio as gr
-import json
+from transformers import ViTForImageClassification, ViTImageProcessor
+from PIL import Image
+import torch
+
+# Load the model and processor from Hugging Face
+model = ViTForImageClassification.from_pretrained("dima806/deepfake_vs_real_image_detection")
+processor = ViTImageProcessor.from_pretrained("dima806/deepfake_vs_real_image_detection")
 
-# Since this is running in Hugging Face Spaces, we'll assume the detection logic
-# needs to be implemented here or use a simpler demo version
 def detect(image):
-    """Detect deepfake content in an image with comprehensive error handling"""
+    """Detect deepfake content in an image using dima806/deepfake_vs_real_image_detection"""
     if image is None:
         raise gr.Error("Please upload an image to analyze")
     
     try:
-        # Mock detection logic (replace with actual model inference if available)
-        # In a real implementation, you'd load your model here
-        import random
-        overall_score = random.uniform(60, 99)
-        aigen_score = random.uniform(0, 100)
-        deepfake_score = random.uniform(0, 100)
+        # Convert Gradio image (filepath) to PIL Image
+        pil_image = Image.open(image).convert("RGB")
+        
+        # Preprocess the image
+        inputs = processor(images=pil_image, return_tensors="pt")
+        
+        # Perform inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            predicted_class = torch.argmax(logits, dim=1).item()
+        
+        # Get confidence scores
+        probabilities = torch.softmax(logits, dim=1)[0]
+        confidence_real = probabilities[0].item() * 100
+        confidence_fake = probabilities[1].item() * 100
+        
+        # Map class index to label
+        label = model.config.id2label[predicted_class]
         
-        overall = f"{overall_score:.1f}% Confidence"
-        aigen = f"{aigen_score:.1f}% (AI-Generated Content Likelihood)"
-        deepfake = f"{deepfake_score:.1f}% (Face Manipulation Likelihood)"
+        # Prepare output
+        overall = f"{max(confidence_real, confidence_fake):.1f}% Confidence"
+        aigen = f"{confidence_fake:.1f}% (AI-Generated Content Likelihood)"
+        deepfake = f"{confidence_fake:.1f}% (Face Manipulation Likelihood)"
         
         return overall, aigen, deepfake
     
     except Exception as e:
         raise gr.Error(f"Analysis error: {str(e)}")
 
-# Custom CSS remains the same
+# Custom CSS (unchanged from your original)
 custom_css = """
 .container {
     max-width: 1200px;
@@ -67,10 +84,12 @@ custom_css = """
 MARKDOWN0 = """
 <div class="header">
     <h1>DeepFake Detection System</h1>
-    <p>Advanced AI-powered analysis for identifying manipulated media</p>
+    <p>Advanced AI-powered analysis for identifying manipulated media<br>
+    Powered by dima806/deepfake_vs_real_image_detection model</p>
 </div>
 """
 
+# Create Gradio interface
 with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
     gr.Markdown(MARKDOWN0)
     with gr.Row(elem_classes="container"):
@@ -88,7 +107,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
         outputs=[overall, aigen, deepfake]
     )
 
-# Launch configuration optimized for Hugging Face Spaces
+# Launch the application
 demo.launch(
     debug=True
-)
+)