Create inference.py

inference.py

@@ -0,0 +1,120 @@
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from retinaface import RetinaFace
+from pathlib import Path
+from typing import Optional
+
+# --- Configuration ---
+CHECKPOINT_PATH = Path("pytorch_model.bin")  # Updated to match your new HF filename
+_IN_FEATURES = 1408
+_DROPOUT = 0.3
+_NUM_CLASSES = 2
+_INPUT_SIZE = 260
+_CONFIDENCE_THRESHOLD = 0.90
+_MIN_FACE_PX = 50
+_PADDING = 20
+
+# --- Transform ---
+_transform = transforms.Compose([
+    transforms.Resize((_INPUT_SIZE, _INPUT_SIZE), interpolation=transforms.InterpolationMode.BICUBIC),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# --- 1. Load Architecture ---
+def load_model() -> tuple[nn.Module, torch.device]:
+    # Universal Hardware Routing
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+    
+    net = models.efficientnet_b2(weights=None)
+    net.classifier = nn.Sequential(
+        nn.Dropout(_DROPOUT),
+        nn.Linear(_IN_FEATURES, _NUM_CLASSES),
+    )
+    
+    checkpoint = torch.load(CHECKPOINT_PATH, map_location=device, weights_only=False)
+    net.load_state_dict(checkpoint["model_state_dict"])
+    
+    net.to(device)
+    net.eval()
+    
+    return net, device
+
+MODEL, DEVICE = load_model()
+
+# --- 2. Extract & Preprocess ---
+def detect_and_crop_face(image_path: str) -> Optional[torch.Tensor]:
+    image_bgr = cv2.imread(image_path)
+    if image_bgr is None:
+        raise ValueError(f"Could not load image at {image_path}")
+
+    detections = RetinaFace.detect_faces(image_bgr)
+    if not isinstance(detections, dict):
+        return None
+
+    best_conf = -1.0
+    best_box = None
+
+    for face_data in detections.values():
+        conf = float(face_data.get("score", 0.0))
+        if conf < _CONFIDENCE_THRESHOLD:
+            continue
+
+        x1, y1, x2, y2 = face_data["facial_area"]
+        w, h = x2 - x1, y2 - y1
+        if w < _MIN_FACE_PX or h < _MIN_FACE_PX:
+            continue
+
+        if conf > best_conf:
+            best_conf = conf
+            best_box = (x1, y1, x2, y2)
+
+    if best_box is None:
+        return None
+
+    H, W = image_bgr.shape[:2]
+    x1, y1, x2, y2 = best_box
+    x1 = max(0, x1 - _PADDING)
+    y1 = max(0, y1 - _PADDING)
+    x2 = min(W, x2 + _PADDING)
+    y2 = min(H, y2 + _PADDING)
+
+    crop_bgr = image_bgr[y1:y2, x1:x2]
+    crop_rgb = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB)
+    pil_face = Image.fromarray(crop_rgb)
+    
+    return _transform(pil_face).unsqueeze(0)
+
+# --- 3. Execute Prediction ---
+def predict_deepfake(image_path: str) -> dict:
+    face_tensor = detect_and_crop_face(image_path)
+    
+    if face_tensor is None:
+        return {"error": "No face detected in the image meeting confidence thresholds."}
+        
+    face_tensor = face_tensor.to(DEVICE)
+    
+    with torch.no_grad():
+        logits = MODEL(face_tensor)
+        probs = torch.softmax(logits, dim=1)[0]
+        
+    fake_prob = probs[0].item() * 100
+    real_prob = probs[1].item() * 100
+    predicted_idx = int(torch.argmax(probs).item())
+    
+    prediction = "REAL" if predicted_idx == 1 else "FAKE"
+    
+    return {
+        "prediction": prediction,
+        "fake_confidence": f"{fake_prob:.2f}%",
+        "real_confidence": f"{real_prob:.2f}%"
+    }