Update app.py

app.py

@@ -2,12 +2,17 @@ import gradio as gr
 from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # 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")
 
-def detect(image):
+def detect(image, confidence_threshold=0.7):
     """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")
@@ -16,6 +21,9 @@ def detect(image):
         # Convert Gradio image (filepath) to PIL Image
         pil_image = Image.open(image).convert("RGB")
         
+        # Ensure image meets ViT's expected size (224x224)
+        pil_image = pil_image.resize((224, 224), Image.Resampling.LANCZOS)
+        
         # Preprocess the image
         inputs = processor(images=pil_image, return_tensors="pt")
         
@@ -23,27 +31,31 @@ def detect(image):
         with torch.no_grad():
             outputs = model(**inputs)
             logits = outputs.logits
-            predicted_class = torch.argmax(logits, dim=1).item()
+            probabilities = torch.softmax(logits, dim=1)[0]
         
         # 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]
+        # Determine prediction based on threshold
+        predicted_label = "Fake" if confidence_fake / 100 >= confidence_threshold else "Real"
+        confidence_score = max(confidence_real, confidence_fake)
+        
+        # Log the prediction
+        logger.info(f"Real: {confidence_real:.1f}%, Fake: {confidence_fake:.1f}%, Predicted: {predicted_label}")
         
         # Prepare output
-        overall = f"{max(confidence_real, confidence_fake):.1f}% Confidence"
+        overall = f"{confidence_score:.1f}% Confidence ({predicted_label})"
         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:
+        logger.error(f"Error during analysis: {str(e)}")
         raise gr.Error(f"Analysis error: {str(e)}")
 
-# Custom CSS (unchanged from your original)
+# Custom CSS (unchanged)
 custom_css = """
 .container {
     max-width: 1200px;
@@ -85,7 +97,8 @@ MARKDOWN0 = """
 <div class="header">
     <h1>DeepFake Detection System</h1>
     <p>Advanced AI-powered analysis for identifying manipulated media<br>
-    Powered by dima806/deepfake_vs_real_image_detection model</p>
+    Powered by dima806/deepfake_vs_real_image_detection model<br>
+    Note: Accuracy may vary with recent deepfakes due to training data age</p>
 </div>
 """