Update app.py

app.py

@@ -1,9 +1,10 @@
 import cv2
+import numpy as np
 from PIL import Image
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
 import gradio as gr
-from spaces import GPU  # ✅ import this to use the decorator
+from spaces import GPU  # Hugging Face ZeroGPU decorator
 
 # Load model and processor
 model_name = "dima806/ai_vs_real_image_detection"
@@ -15,22 +16,57 @@ model.eval()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
-@GPU  # ✅ This activates GPU on Hugging Face ZeroGPU Spaces
+# === Preprocessing Function ===
+def preprocess_frame(frame: np.ndarray) -> np.ndarray:
+    # Resize for consistency
+    frame = cv2.resize(frame, (512, 512), interpolation=cv2.INTER_AREA)
+
+    # Denoise slightly
+    frame = cv2.GaussianBlur(frame, (3, 3), 0)
+
+    # Sharpen
+    kernel_sharpen = np.array([[0, -1, 0], 
+                               [-1, 5, -1], 
+                               [0, -1, 0]])
+    frame = cv2.filter2D(frame, -1, kernel_sharpen)
+
+    # Brightness/contrast adjustment
+    frame = cv2.convertScaleAbs(frame, alpha=1.1, beta=10)
+
+    # Histogram equalization
+    yuv = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV)
+    yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])
+    frame = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
+
+    return frame
+
+# === Main Video Analysis Function ===
+@GPU
 def analyze_video(video):
     cap = cv2.VideoCapture(video)
-    frame_num = 0
-    frame_interval = 60  
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    max_frames_to_process = 10
+
+    if total_frames == 0:
+        return "Could not read video frames."
+
+    # Evenly sample 10 frames
+    frame_indices = [int(i * total_frames / max_frames_to_process) for i in range(max_frames_to_process)]
     frames_to_process = []
+    current_index = 0
+    frame_num = 0
 
     while cap.isOpened():
         ret, frame = cap.read()
-        if not ret:
+        if not ret or current_index >= len(frame_indices):
             break
 
-        if frame_num % frame_interval == 0:
-            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        if frame_num == frame_indices[current_index]:
+            processed = preprocess_frame(frame)
+            rgb = cv2.cvtColor(processed, cv2.COLOR_BGR2RGB)
             pil_image = Image.fromarray(rgb)
             frames_to_process.append((frame_num, pil_image))
+            current_index += 1
 
         frame_num += 1
 
@@ -42,24 +78,48 @@ def analyze_video(video):
     frame_numbers, pil_images = zip(*frames_to_process)
 
     inputs = processor(images=pil_images, return_tensors="pt")
-    inputs = {k: v.to(device) for k, v in inputs.items()}  # ✅ Move inputs to GPU if available
+    inputs = {k: v.to(device) for k, v in inputs.items()}
 
     with torch.no_grad():
         outputs = model(**inputs)
-        predictions = torch.argmax(outputs.logits, dim=1).tolist()
+        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
+        predictions = torch.argmax(probs, dim=1).tolist()
+
+    # Find index for "AI" class
+    ai_class_index = [k for k, v in model.config.id2label.items() if "AI" in v.upper()][0]
+    ai_scores = probs[:, ai_class_index].tolist()
 
+    # === Smart Verdict Logic ===
+    ai_conf_threshold = 0.7  # Must be at least 70% confident to count
+    min_ai_frames = 3        # Require 3 or more confident AI frames
+    avg_score = sum(ai_scores) / len(ai_scores)
+    confident_ai_frames = sum(1 for score in ai_scores if score >= ai_conf_threshold)
+
+    if confident_ai_frames >= min_ai_frames:
+        verdict = "Video is likely AI-generated."
+    elif avg_score >= 0.5:
+        verdict = "Uncertain — some signs of AI generation but not enough evidence."
+    else:
+        verdict = "Video is likely REAL."
+
+    # === Output Formatting ===
     results = []
-    for frame_idx, pred in zip(frame_numbers, predictions):
+    for frame_idx, pred, ai_score in zip(frame_numbers, predictions, ai_scores):
         label = model.config.id2label[pred]
-        results.append(f"Frame {frame_idx}: {label}")
+        confidence = ai_score * 100
+        results.append(f"Frame {frame_idx}: {label} (AI Confidence: {confidence:.1f}%)")
+
+    results.append(f"\nConfident AI-flagged frames: {confident_ai_frames}")
+    results.append(f"Average AI Confidence: {avg_score * 100:.1f}%")
+    results.append(f"Final Verdict: {verdict}")
 
     return "\n".join(results)
 
-# Gradio UI (note: NO gpu=True here)
+# === Gradio UI ===
 gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload a video"),
     outputs=gr.Textbox(label="Detection Results"),
     title="AI Frame Detector",
-    description="Detects whether frames in a video are AI-generated or real."
+    description="Analyzes 10 frames from a video with preprocessing to detect if it's AI-generated or real."
 ).launch()