Update app.py

app.py

@@ -1,123 +1,54 @@
+import gradio as gr
 import torch
-import torch.nn as nn
-from torchvision import transforms
-from PIL import Image
 import cv2
 import numpy as np
-import gradio as gr
-
-# ------------------------------------------------------------------
-# 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)
+from model import GenConViT
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ------------------------------------------------------------------
-# 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)
+# Load model
+model = GenConViT().to(device)
+state = torch.load("genconvit_ed_inference.pth", map_location=device)
+model.load_state_dict(state)
 model.eval()
 
-# ------------------------------------------------------------------
-# 3. Preprocessing
-# ------------------------------------------------------------------
-
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5]*3, [0.5]*3)
-])
-
-# ------------------------------------------------------------------
-# 4. Video Deepfake Detection Function
-# ------------------------------------------------------------------
-
-def detect_deepfake(video):
+def preprocess(frame):
+    frame = cv2.resize(frame, (224, 224))
+    frame = frame[:, :, ::-1] / 255.0
+    frame = torch.tensor(frame, dtype=torch.float32).permute(2, 0, 1)
+    return frame.unsqueeze(0)
 
+def predict(video):
     cap = cv2.VideoCapture(video)
-    if not cap.isOpened():
-        return "Error: Cannot open video", None
-
     scores = []
-    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:
-            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            img = Image.fromarray(rgb)
-
-            # Save last processed frame for display
-            sample_frame = img
-
-            inp = transform(img).unsqueeze(0)
-
-            with torch.no_grad():
-                logits = model(inp)
-                probs = torch.softmax(logits, dim=1)[0]
-                fake_prob = probs[1].item()
-
-            scores.append(fake_prob)
-
-        i += 1
+        inp = preprocess(frame).to(device)
+        with torch.no_grad():
+            pred = model(inp)
+            prob = torch.softmax(pred, dim=1)[0, 1].item()
+            scores.append(prob)
 
     cap.release()
 
     if len(scores) == 0:
-        return "No frames processed", None
-
-    avg = float(np.mean(scores))
-    label = "🔴 Deepfake" if avg > 0.5 else "🟢 Real"
-
-    output = f"""
-### **Prediction: {label}**
-**Fake confidence: {avg:.4f}**
-"""
-
-    return output, sample_frame
+        return "No frames detected."
 
+    deepfake_prob = float(np.mean(scores))
+    label = "Deepfake" if deepfake_prob > 0.5 else "Real"
 
-# ------------------------------------------------------------------
-# 5. Gradio App UI
-# ------------------------------------------------------------------
+    return f"{label} (score: {deepfake_prob:.4f})"
 
-app = gr.Interface(
-    fn=detect_deepfake,
-    inputs=gr.Video(label="Upload a video"),
-    outputs=[
-        gr.Markdown(label="Prediction"),
-        gr.Image(label="Sample Frame")
-    ],
-    title="GenConViT Deepfake Detector (Local .pth Model)",
-    description="Upload a video. The system loads genconvit_ed_inference.pth and predicts deepfake probability."
+# UI
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Video(),
+    outputs="text",
+    title="GenConViT Deepfake Detector",
 )
 
-app.launch()
+demo.launch()