Update app.py

app.py

@@ -1,19 +1,33 @@
 # app.py
-import io
 import numpy as np
 import gradio as gr
 from PIL import Image, ImageDraw
 import cv2
 import torch
-from transformers import AutoImageProcessor, ViTForImageClassification
-import mediapipe as mp
-
-# -------------------- Face crop utilities --------------------
-_mp_face = mp.solutions.face_detection.FaceDetection(
-    model_selection=0, min_detection_confidence=0.4
-)
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# ---- Config ----
+MODEL_ID = "SadraCoding/SDXL-Deepfake-Detector"
+THRESHOLD = 0.65  # >= -> "Likely Manipulated"
+IMAGE_SIZE = 224  # ViT input size
+
+# Optional: MediaPipe face detection (app still works if not installed)
+try:
+    import mediapipe as mp
+    _mp_face = mp.solutions.face_detection.FaceDetection(
+        model_selection=0, min_detection_confidence=0.4
+    )
+except Exception:
+    _mp_face = None
 
+# ---- Face crop ----
 def crop_face(pil_img, pad=0.25):
+    """
+    Crop the most prominent face using MediaPipe. If MP missing or no face found,
+    return the original image.
+    """
+    if _mp_face is None:
+        return pil_img
     img = np.array(pil_img.convert("RGB"))
     h, w = img.shape[:2]
     res = _mp_face.process(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
@@ -28,10 +42,12 @@ def crop_face(pil_img, pad=0.25):
     x1 = int(max(0, (x - pad*bw) * w)); y1 = int(max(0, (y - pad*bh) * h))
     x2 = int(min(w, (x + bw + pad*bw) * w)); y2 = int(min(h, (y + bh + pad*bh) * h))
     face = Image.fromarray(img[y1:y2, x1:x2])
-    return face if face.size[0] > 20 and face.size[1] > 20 else pil_img
+    if face.size[0] < 20 or face.size[1] < 20:
+        return pil_img
+    return face
 
+# (Not used for inference; kept if you want to mask background later)
 def face_oval_mask(img_pil, shrink=0.80):
-    
     w, h = img_pil.size
     mask = Image.new("L", (w, h), 0)
     draw = ImageDraw.Draw(mask)
@@ -39,69 +55,73 @@ def face_oval_mask(img_pil, shrink=0.80):
     draw.ellipse((dx, dy, w - dx, h - dy), fill=255)
     return np.array(mask, dtype=np.float32) / 255.0
 
-# -------------------- HF model: Deepfake vs Realism --------------------
-MODEL_ID = "prithivMLmods/Deep-Fake-Detector-v2-Model"
-
-# CPU by default
-_hf_processor = AutoImageProcessor.from_pretrained(MODEL_ID)
-_hf_model = ViTForImageClassification.from_pretrained(MODEL_ID)
-_hf_model.eval()
+# ---- HF model load ----
+processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
+model.eval()
 torch.set_grad_enabled(False)
 
-_FAKE_KEYS = ("fake", "deepfake", "manipulated", "spoof", "forged")
+# Resolve which index corresponds to "fake"
+_FAKE_KEYS = ("artificial", "fake", "deepfake", "manipulated", "spoof", "forged")
 
-def _deepfake_index_from_config(cfg) -> int | None:
-    """
-    Try to find the class index for 'Deepfake' from id2label/label2id.
-    This model typically has {0:'Realism', 1:'Deepfake'}.
-    """
+def _fake_index_from_config(cfg) -> int | None:
     # Prefer id2label
     id2label = getattr(cfg, "id2label", None)
     if id2label:
-        normalized = {int(k): str(v).lower() for k, v in id2label.items()}
+        try:
+            normalized = {int(k): str(v).lower() for k, v in id2label.items()}
+        except Exception:
+            # sometimes keys already ints
+            normalized = {int(k): str(v).lower() for k, v in id2label.items()}
         for idx, lab in normalized.items():
             if any(k in lab for k in _FAKE_KEYS):
                 return idx
-
-    # Fallback to label2id if present
+    # Fallback: invert label2id
     label2id = getattr(cfg, "label2id", None)
     if label2id:
         inv = {int(v): str(k).lower() for k, v in label2id.items()}
         for idx, lab in inv.items():
             if any(k in lab for k in _FAKE_KEYS):
                 return idx
-
     return None
 
-_DEEP_IDX = _deepfake_index_from_config(_hf_model.config)
+_FAKE_IDX = _fake_index_from_config(model.config)
 
-def _hf_predict_proba(pil_img: Image.Image) -> float:
+# ---- Inference ----
+def predict_fake_prob(pil_img: Image.Image) -> float:
     """
-    Returns P(Deepfake) in [0,1] using the ViT classifier.
+    Returns P(fake) in [0,1].
+    Model labels per card: 0 -> 'artificial' (fake), 1 -> 'human' (real).
     """
-    inputs = _hf_processor(images=pil_img.convert("RGB"), return_tensors="pt")
-    with torch.no_grad():
-        logits = _hf_model(**inputs).logits  # (1, C)
+    # Face-focus to reduce background bias
+    face = crop_face(pil_img)
+    face = face.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE))
+
+    inputs = processor(images=face, return_tensors="pt")
+    logits = model(**inputs).logits  # (1, C)
 
     if logits.shape[-1] == 1:
-        # Binary sigmoid head
+        # Binary sigmoid head (unlikely for this model, but safe)
         return torch.sigmoid(logits.squeeze(0))[0].item()
 
-    # Softmax head
+    # Softmax multi-class (expected 2 classes)
     probs = torch.softmax(logits.squeeze(0), dim=-1).detach().cpu().numpy()
-    if _DEEP_IDX is not None and 0 <= _DEEP_IDX < probs.shape[0]:
-        return float(probs[_DEEP_IDX])
 
-    # Binary fallback
+    # Use explicit mapping if available
+    if _FAKE_IDX is not None and 0 <= _FAKE_IDX < probs.shape[0]:
+        return float(probs[_FAKE_IDX])
+
+    # Known mapping from the model card: 0=artificial (fake), 1=human
     if probs.shape[0] == 2:
-        return float(probs[1])
+        return float(probs[0])  # class-0 is fake
 
-    # Last resort: take max
+    # Last resort
     return float(probs.max())
 
-# -------------------- Output card --------------------
-def _result_card(label: str, conf: float) -> str:
-    pct = max(0.0, min(1.0, conf)) * 100.0
+# ---- UI helpers ----
+def result_card(prob_fake: float) -> str:
+    label = "Likely Manipulated" if prob_fake >= THRESHOLD else "Likely Authentic"
+    pct = prob_fake * 100.0
     color = "#d84a4a" if label.startswith("Likely Manipulated") else "#2e7d32"
     bar_bg = "#e9ecef"
     return f"""
@@ -115,24 +135,21 @@ def _result_card(label: str, conf: float) -> str:
         <div style="width:100%;height:10px;background:{bar_bg};border-radius:999px;overflow:hidden;">
           <div style="height:100%;width:{pct:.4f}%;background:{color};"></div>
         </div>
+        <div style="font-size:12px;color:#6b7280;margin-top:8px;">
+          Model: {MODEL_ID} · Threshold: {int(THRESHOLD*100)}%
+        </div>
       </div>
     </div>
     """
 
-# -------------------- Gradio handler --------------------
+# ---- Gradio handlers ----
 def analyze(pil_img: Image.Image):
     if pil_img is None:
-        return _result_card("Likely Authentic", 0.0)
-
-    # Focus on the face to reduce background false positives
-    face = crop_face(pil_img)
-    face = face.convert("RGB").resize((224, 224))  # ViT expects 224x224
-
-    p_fake = _hf_predict_proba(face)
-    label = "Likely Manipulated" if p_fake >= 0.65 else "Likely Authentic"
-    return _result_card(label, p_fake)
+        return result_card(0.0)
+    p_fake = predict_fake_prob(pil_img)
+    return result_card(p_fake)
 
-# -------------------- UI --------------------
+# ---- UI ----
 CUSTOM_CSS = """
 .gradio-container {max-width: 980px !important;}
 .sleek-card {
@@ -141,9 +158,10 @@ CUSTOM_CSS = """
 }
 """
 
-with gr.Blocks(title="Deepfake Detector (ViT)", css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Deepfake Detector (SDXL ViT)", css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
     gr.Markdown(
-        "<h2 style='text-align:center;margin-bottom:6px;'>Deepfake Detector (ViT)</h2>"
+        "<h2 style='text-align:center;margin-bottom:6px;'>Deepfake Detector (SDXL ViT)</h2>"
+        "<p style='text-align:center;color:#6b7280;'>MediaPipe face-crop + Vision Transformer fine-tuned for artificial vs human faces.</p>"
     )
     with gr.Row():
         with gr.Column(scale=6, elem_classes=["sleek-card"]):