Update app.py

app.py

@@ -1,61 +1,91 @@
-import gradio as gr
 import torch
-import cv2
-from transformers import AutoModelForImageClassification, AutoImageProcessor
+import torch.nn as nn
+from torchvision import transforms
 from PIL import Image
+import cv2
 import numpy as np
+import gradio as gr
 
-# ----------------------------------------------------------
-# Load Hugging Face GenConViT Model
-# ----------------------------------------------------------
-model = AutoModelForImageClassification.from_pretrained(
-    "Thanuja2109/GenConViT"
-)
+# ------------------------------------------------------------------
+# 1. Define the GenConViT Model Architecture (Minimal Version)
+# ------------------------------------------------------------------
+
+class GenConViT(nn.Module):
+    def __init__(self, num_classes=2):
+        super().__init__()
+        # Very lightweight demo backbone (adjust to your real architecture)
+        self.feature_extractor = nn.Sequential(
+            nn.Conv2d(3, 32, 3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, 3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+        self.fc = nn.Linear(128, num_classes)
+
+    def forward(self, x):
+        x = self.feature_extractor(x)
+        x = x.flatten(1)
+        return self.fc(x)
+
+
+# ------------------------------------------------------------------
+# 2. Load Model From genconvit_ed_inference.pth
+# ------------------------------------------------------------------
+
+model_path = "genconvit_ed_inference.pth"
+
+model = GenConViT(num_classes=2)
+checkpoint = torch.load(model_path, map_location="cpu")
+model.load_state_dict(checkpoint)
+model.eval()
 
-processor = AutoImageProcessor.from_pretrained(
-    "Thanuja2109/GenConViT"
-)
+# ------------------------------------------------------------------
+# 3. Preprocessing
+# ------------------------------------------------------------------
 
-model.eval()
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5]*3, [0.5]*3)
+])
+
+# ------------------------------------------------------------------
+# 4. Video Deepfake Detection Function
+# ------------------------------------------------------------------
 
-# ----------------------------------------------------------
-# Deepfake detection function
-# ----------------------------------------------------------
 def detect_deepfake(video):
 
-    # Load video
     cap = cv2.VideoCapture(video)
-
     if not cap.isOpened():
-        return "Error: cannot open video", None
+        return "Error: Cannot open video", None
 
     scores = []
-    frames_collected = 0
-
-    # Sample 1 frame every 10
-    frame_interval = 10
-
-    frame_img = None
-
+    sample_frame = None
+    frame_interval = 10  # Process every 10th frame
     i = 0
+
     while True:
         ret, frame = cap.read()
         if not ret:
             break
 
         if i % frame_interval == 0:
-            # Convert to RGB
             rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            pil_img = Image.fromarray(rgb)
+            img = Image.fromarray(rgb)
+
+            # Save last processed frame for display
+            sample_frame = img
 
-            inputs = processor(images=pil_img, return_tensors="pt")
+            inp = transform(img).unsqueeze(0)
 
             with torch.no_grad():
-                logits = model(**inputs).logits
-                prob_fake = torch.softmax(logits, dim=1)[0][1].item()
+                logits = model(inp)
+                probs = torch.softmax(logits, dim=1)[0]
+                fake_prob = probs[1].item()
 
-            scores.append(prob_fake)
-            frame_img = pil_img  # save last sampled frame
+            scores.append(fake_prob)
 
         i += 1
 
@@ -64,30 +94,30 @@ def detect_deepfake(video):
     if len(scores) == 0:
         return "No frames processed", None
 
-    avg_score = np.mean(scores)
+    avg = float(np.mean(scores))
+    label = "🔴 Deepfake" if avg > 0.5 else "🟢 Real"
 
-    label = "🔴 Deepfake" if avg_score > 0.5 else "🟢 Real"
-
-    result_text = f"""
-### Prediction: **{label}**
-**Confidence (fake probability): {avg_score:.4f}**
+    output = f"""
+### **Prediction: {label}**
+**Fake confidence: {avg:.4f}**
 """
 
-    return result_text, frame_img
+    return output, sample_frame
+
 
+# ------------------------------------------------------------------
+# 5. Gradio App UI
+# ------------------------------------------------------------------
 
-# ----------------------------------------------------------
-# Gradio Interface
-# ----------------------------------------------------------
 app = gr.Interface(
     fn=detect_deepfake,
     inputs=gr.Video(label="Upload a video"),
     outputs=[
         gr.Markdown(label="Prediction"),
-        gr.Image(label="Analyzed Frame")
+        gr.Image(label="Sample Frame")
     ],
-    title="GenConViT Deepfake Video Detector",
-    description="Upload a video. The app samples frames and uses GenConViT to detect deepfakes."
+    title="GenConViT Deepfake Detector (Local .pth Model)",
+    description="Upload a video. The system loads genconvit_ed_inference.pth and predicts deepfake probability."
 )
 
 app.launch()