Update app.py

app.py

@@ -1,55 +1,17 @@
-import io, os, numpy as np, gradio as gr
-from PIL import Image, ImageChops, ImageDraw
+# app.py
+import io
+import numpy as np
+import gradio as gr
+from PIL import Image, ImageDraw
 import cv2
-from skimage import exposure
+import torch
+from transformers import AutoImageProcessor, ViTForImageClassification
 import mediapipe as mp
 
-# ====== HF model choice (pick one) ======
-HF_MODEL_ID = os.getenv("HF_MODEL_ID", "prithivMLmods/Deep-Fake-Detector-v2-Model")  # ViT 224
-HF_IMAGE_SIZE = int(os.getenv("HF_IMAGE_SIZE", "224"))  # 224 for v2 ViT, 512 for v1 SigLIP
-
-# ====== HF imports (lazy so app can start even if transformers missing) ======
-_hf_loaded = False
-_hf_processor = None
-_hf_model = None
-def _try_load_hf():
-    global _hf_loaded, _hf_processor, _hf_model
-    if _hf_loaded: 
-        return True
-    try:
-        from transformers import AutoImageProcessor, AutoModelForImageClassification
-        _hf_processor = AutoImageProcessor.from_pretrained(HF_MODEL_ID)
-        _hf_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_ID)
-        _hf_model.eval()
-        _hf_loaded = True
-        return True
-    except Exception as e:
-        print("HF load failed:", e)
-        _hf_loaded = False
-        return False
-
-def _hf_predict_proba(pil_rgb_face):
-    """Returns probability that image is deepfake, in [0,1]."""
-    import torch
-    with torch.no_grad():
-        inputs = _hf_processor(images=pil_rgb_face.resize((HF_IMAGE_SIZE, HF_IMAGE_SIZE)), return_tensors="pt")
-        outputs = _hf_model(**inputs)
-        logits = outputs.logits[0]
-        probs = torch.softmax(logits, dim=-1).cpu().numpy()
-    # Map label -> index; models commonly use ["Deepfake","Realism"] or ["fake","real"]
-    id2label = _hf_model.config.id2label
-    lab2idx = {v.lower(): k for k, v in _hf_model.config.label2id.items()}
-    # Try a few common names
-    deep_idx = lab2idx.get("deepfake", None)
-    if deep_idx is None:
-        deep_idx = lab2idx.get("fake", None)
-    if deep_idx is None:
-        # Heuristic: choose the class whose label name contains 'fake'
-        deep_idx = next((i for i, name in id2label.items() if "fake" in name.lower()), 0)
-    return float(probs[int(deep_idx)])
-
-# ====== Face detect / crop (your pipeline) ======
-_mp_face = mp.solutions.face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.4)
+# -------------------- Face crop utilities --------------------
+_mp_face = mp.solutions.face_detection.FaceDetection(
+    model_selection=0, min_detection_confidence=0.4
+)
 
 def crop_face(pil_img, pad=0.25):
     img = np.array(pil_img.convert("RGB"))
@@ -57,7 +19,10 @@ def crop_face(pil_img, pad=0.25):
     res = _mp_face.process(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
     if not res.detections:
         return pil_img
-    det = max(res.detections, key=lambda d: d.location_data.relative_bounding_box.width)
+    det = max(
+        res.detections,
+        key=lambda d: d.location_data.relative_bounding_box.width
+    )
     b = det.location_data.relative_bounding_box
     x, y, bw, bh = b.xmin, b.ymin, b.width, b.height
     x1 = int(max(0, (x - pad*bw) * w)); y1 = int(max(0, (y - pad*bh) * h))
@@ -65,73 +30,77 @@ def crop_face(pil_img, pad=0.25):
     face = Image.fromarray(img[y1:y2, x1:x2])
     return face if face.size[0] > 20 and face.size[1] > 20 else pil_img
 
-# ====== Heuristic fallback (unchanged core) ======
-def _enhance_for_display(pil_img, scale: float):
-    arr = np.array(pil_img).astype("float32") * scale
-    arr = np.clip(arr, 0, 255).astype("uint8")
-    return Image.fromarray(arr)
-
-def error_level_analysis(pil_img: Image.Image, quality: int = 90):
-    img = pil_img.convert("RGB")
-    with io.BytesIO() as buf:
-        img.save(buf, "JPEG", quality=quality); buf.seek(0)
-        comp = Image.open(buf).convert("RGB")
-    diff = ImageChops.difference(img, comp)
-    extrema = diff.getextrema(); max_diff = max([m for (_, m) in extrema])
-    scale = 255.0 / max(1, max_diff)
-    ela_vis = _enhance_for_display(diff, scale)
-    ela_np = np.array(ela_vis, dtype=np.float32)
-    mean_intensity = float(ela_np.mean() / 255.0)
-    return ela_vis, mean_intensity
-
-def ela_sweep_mean(pil_img, qualities=(95, 90, 85)):
-    vals = []
-    for q in qualities:
-        _, m = error_level_analysis(pil_img, quality=q); vals.append(m)
-    return float(max(vals)), float(np.mean(vals))
-
-def fft_high_freq_ratio(pil_img: Image.Image):
-    y = pil_img.convert("YCbCr").split()[0]
-    gray = np.array(y, dtype=np.float32)/255.0
-    h, w = gray.shape
-    wy, wx = np.hanning(h)[:, None], np.hanning(w)[None, :]
-    F = np.fft.fftshift(np.fft.fft2(gray * (wy * wx)))
-    mag = np.log1p(np.abs(F))
-    cy, cx = h//2, w//2
-    yy, xx = np.ogrid[:h, :w]; dist = np.sqrt((yy - cy)**2 + (xx - cx)**2)
-    r_low = min(h, w) * 0.08
-    low = float(mag[dist <= r_low].sum()); high = float(mag[dist > r_low].sum())
-    return None, float(high / (high + low + 1e-9))
-
-def noise_inconsistency(pil_img: Image.Image):
-    y = pil_img.convert("YCbCr").split()[0]
-    img = np.array(y, dtype=np.float32)
-    lap = cv2.Laplacian(img, cv2.CV_32F, ksize=3); lap_abs = np.abs(lap)
-    tile = 32; H, W = lap_abs.shape; vals = []
-    for yy in range(0, H, tile):
-        for xx in range(0, W, tile):
-            patch = lap_abs[yy:min(yy+tile, H), xx:min(xx+tile, W)]
-            if patch.size: vals.append(patch.var())
-    if not vals: return None, 0.0
-    vals = np.array(vals, dtype=np.float32)
-    score = float(vals.std() / (vals.mean() + 1e-9))
-    return None, float(np.tanh(score / 5.0))
-
-def combine_scores(ela_mean, hf_ratio, noise_incons_score):
-    w1, w2, w3 = 0.30, 0.40, 0.30
-    s_ela = np.clip(ela_mean * 3.0, 0, 1)
-    s_hf  = np.clip((hf_ratio - 0.65) / 0.25, 0, 1)
-    s_noi = np.clip(noise_incons_score, 0, 1)
-    conf = float(w1*s_ela + w2*s_hf + w3*s_noi)
-    label = "Likely Manipulated" if conf >= 0.65 else "Likely Authentic"
-    return label, conf
-
-# ====== Result card ======
-def _result_card(label: str, conf: float, note: str | None = None) -> str:
+def face_oval_mask(img_pil, shrink=0.80):
+    # (Optional) not used by the model; kept if you ever want to mask background
+    w, h = img_pil.size
+    mask = Image.new("L", (w, h), 0)
+    draw = ImageDraw.Draw(mask)
+    dx, dy = int((1 - shrink) * w / 2), int((1 - shrink) * h / 2)
+    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; if you run locally with GPU, you can .to("cuda")
+_hf_processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+_hf_model = ViTForImageClassification.from_pretrained(MODEL_ID)
+_hf_model.eval()
+torch.set_grad_enabled(False)
+
+_FAKE_KEYS = ("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'}.
+    """
+    # Prefer id2label
+    id2label = getattr(cfg, "id2label", None)
+    if id2label:
+        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
+    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)
+
+def _hf_predict_proba(pil_img: Image.Image) -> float:
+    """
+    Returns P(Deepfake) in [0,1] using the ViT classifier.
+    """
+    inputs = _hf_processor(images=pil_img.convert("RGB"), return_tensors="pt")
+    logits = _hf_model(**inputs).logits  # (1, C)
+    if logits.shape[-1] == 1:
+        # Unlikely for this model, but handle binary-sigmoid heads
+        return torch.sigmoid(logits.squeeze(0))[0].item()
+
+    probs = torch.softmax(logits.squeeze(0), dim=-1).cpu().numpy()
+    if _DEEP_IDX is not None and 0 <= _DEEP_IDX < probs.shape[0]:
+        return float(probs[_DEEP_IDX])
+
+    # Binary softmax fallback: assume index 1 = deepfake
+    if probs.shape[0] == 2:
+        return float(probs[1])
+
+    # Last resort: take the highest class prob (not ideal, but safe)
+    return float(probs.max())
+
+# -------------------- Output card --------------------
+def _result_card(label: str, conf: float) -> str:
     pct = max(0.0, min(1.0, conf)) * 100.0
     color = "#d84a4a" if label.startswith("Likely Manipulated") else "#2e7d32"
     bar_bg = "#e9ecef"
-    extra = f"<div style='color:#6b7280;font-size:12px;margin-top:10px;text-align:center;'>{note}</div>" if note else ""
     return f"""
     <div style="max-width:860px;margin:0 auto;">
       <div style="border:1px solid #e5e7eb;border-radius:14px;padding:18px 20px;background:#fff;
@@ -144,31 +113,23 @@ def _result_card(label: str, conf: float, note: str | None = None) -> str:
           <div style="height:100%;width:{pct:.4f}%;background:{color};"></div>
         </div>
       </div>
-      {extra}
     </div>
     """
 
-# ====== Inference ======
+# -------------------- Gradio handler --------------------
 def analyze(pil_img: Image.Image):
     if pil_img is None:
         return _result_card("Likely Authentic", 0.0)
-    face = crop_face(pil_img).convert("RGB")
-
-    if _try_load_hf():
-        prob_fake = _hf_predict_proba(face)
-        label = "Likely Manipulated" if prob_fake >= 0.5 else "Likely Authentic"
-        note  = f"HF model: {HF_MODEL_ID}"
-        return _result_card(label, prob_fake, note=note)
-
-    # Fallback heuristic (if HF model failed)
-    face = face.resize((512, 512))
-    _, ela_mean = error_level_analysis(face, quality=90)
-    _, hf_ratio  = fft_high_freq_ratio(face)
-    _, noi_score = noise_inconsistency(face)
-    label, conf = combine_scores(ela_mean, hf_ratio, noi_score)
-    return _result_card(label, conf, note="Heuristic fallback")
-
-# ====== UI ======
+
+    # 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)
+
+# -------------------- UI --------------------
 CUSTOM_CSS = """
 .gradio-container {max-width: 980px !important;}
 .sleek-card {
@@ -176,17 +137,26 @@ CUSTOM_CSS = """
   box-shadow: 0 2px 10px rgba(16,24,40,.04); padding: 18px;
 }
 """
-with gr.Blocks(title="Deepfake Detector (Pretrained HF Model)", css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("<h2 style='text-align:center;margin-bottom:6px;'>Deepfake Detector</h2>"
-                "<p style='text-align:center;color:#6b7280;'>Face-crop → pretrained classifier → single likelihood.</p>")
+
+with gr.Blocks(title="Deepfake Detector (ViT)", css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "<h2 style='text-align:center;margin-bottom:6px;'>Deepfake Detector (ViT)</h2>"
+        "<p style='text-align:center;color:#6b7280;'>Upload an image to get a single, clean likelihood estimate using a fine-tuned Vision Transformer.</p>"
+    )
     with gr.Row():
         with gr.Column(scale=6, elem_classes=["sleek-card"]):
-            inp = gr.Image(type="pil", label="Upload / Paste Image",
-                           sources=["upload", "webcam", "clipboard"],
-                           height=420, show_label=True, interactive=True)
+            inp = gr.Image(
+                type="pil",
+                label="Upload / Paste Image",
+                sources=["upload", "webcam", "clipboard"],
+                height=420,
+                show_label=True,
+                interactive=True,
+            )
             btn = gr.Button("Analyze", variant="primary", size="lg")
         with gr.Column(scale=6):
             out = gr.HTML()
+
     btn.click(analyze, inputs=inp, outputs=out)
     inp.change(analyze, inputs=inp, outputs=out)