Create interface.py

interface.py

@@ -0,0 +1,75 @@
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+from model import load_model
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 🔥 Download models
+model_paths = {
+    "m1": hf_hub_download("Simma7/deepfake_model", "video1.pth"),
+    "m2": hf_hub_download("Simma7/deepfake_model", "video2.pth"),
+    "m3": hf_hub_download("Simma7/deepfake_model", "video3.pt"),
+}
+
+# 🔥 Load models
+models = {
+    "m1": load_model(model_paths["m1"], "clip"),
+    "m2": load_model(model_paths["m2"], "vit"),
+    "m3": load_model(model_paths["m3"], "clip"),
+}
+
+# transforms
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
+])
+
+# extract frames
+def extract_frames(video_path, num_frames=16):
+    cap = cv2.VideoCapture(video_path)
+    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    indices = np.linspace(0, total - 1, num_frames, dtype=int)
+    frames = []
+
+    for i in indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        ret, frame = cap.read()
+        if ret:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frames.append(Image.fromarray(frame))
+
+    cap.release()
+    return frames
+
+# predict
+def predict(video_path):
+    frames = extract_frames(video_path)
+
+    all_probs = []
+
+    with torch.no_grad():
+        for frame in frames:
+            x = transform(frame).unsqueeze(0).to(DEVICE)
+
+            probs = []
+            for key, model in models.items():
+                out = model(x)
+
+                if out.shape[-1] == 1:
+                    prob = torch.sigmoid(out).item()
+                else:
+                    prob = torch.softmax(out, dim=1)[0, 1].item()
+
+                probs.append(prob)
+
+            frame_prob = sum(probs) / len(probs)
+            all_probs.append(frame_prob)
+
+    final_prob = sum(all_probs) / len(all_probs)
+
+    return final_prob