Update app.py

app.py

@@ -1,27 +1,27 @@
 import io
-from PIL import Image, ImageChops, ImageStat, ExifTags
+from typing import List, Tuple, Dict, Any
+from PIL import Image
+import numpy as np
+import torch
 import gradio as gr
 
-# Possible generator keywords that may appear in EXIF metadata (extendable)
-GENERATOR_KEYWORDS = [
-    "stable diffusion", "stability.ai", "sdxl", "midjourney", "dall", "openai",
-    "novelai", "leonardo", "kaiber", "flux", "comfyui", "automatic1111", "invokeai"
-]
+# Face detector
+from facenet_pytorch import MTCNN
 
-def to_rgb_flat(img, bg=(255, 255, 255)):
-    """Ensure RGB mode; when RGBA/transparent, composite on a white background to avoid JPEG save errors."""
-    if img.mode == "RGB":
-        return img
-    if img.mode in ("RGBA", "LA", "P"):
-        bg_img = Image.new("RGB", img.size, bg)
-        if img.mode == "P":
-            img = img.convert("RGBA")
-        bg_img.paste(img, mask=img.split()[-1] if "A" in img.getbands() else None)
-        return bg_img
-    return img.convert("RGB")
-
-def resize_max(img, max_side=1024):
-    """Limit the longest side to 1024 px to reduce memory and avoid encoding errors."""
+# HF image classifier
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# ========= Config =========
+# You can change the model below to another public model on Hugging Face
+# Example: prithivMLmods/Deep-Fake-Detector-v2-Model (binary: Deepfake vs Realism)
+MODEL_ID = "prithivMLmods/Deep-Fake-Detector-v2-Model"
+DEVICE = "cpu"  # Use "cuda" if GPU is available
+MAX_SIDE = 640  # Resize to keep the longest side ≤ 640px for efficiency
+# =========================
+
+# ---- Utilities ----
+def resize_keep_ratio(img: Image.Image, max_side: int = MAX_SIDE) -> Image.Image:
+    """Resize the image while keeping aspect ratio and limit max side length."""
     w, h = img.size
     m = max(w, h)
     if m <= max_side:
@@ -29,133 +29,118 @@ def resize_max(img, max_side=1024):
     scale = max_side / float(m)
     return img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
 
-def compute_ela_score(img, quality=95):
-    """
-    ELA (Error Level Analysis): recompress the image at given JPEG quality, then compute the mean/std of differences.
-    If fails (unsupported mode/codec), return (None, None) and let upper layer handle gracefully.
-    """
-    try:
-        img_rgb = to_rgb_flat(img)
-        img_rgb = resize_max(img_rgb, 1024)
-        buf = io.BytesIO()
-        img_rgb.save(buf, "JPEG", quality=quality, optimize=True)
-        buf.seek(0)
-        recompressed = Image.open(buf).convert("RGB")
-        ela = ImageChops.difference(img_rgb, recompressed)
-        stat = ImageStat.Stat(ela)
-        mean = float(sum(stat.mean) / len(stat.mean))
-        std = float(sum(stat.stddev) / len(stat.stddev))
-        return mean, std
-    except Exception:
-        # Retry once with lower JPEG quality (more conservative)
-        try:
-            img_rgb = to_rgb_flat(img)
-            img_rgb = resize_max(img_rgb, 1024)
-            buf = io.BytesIO()
-            img_rgb.save(buf, "JPEG", quality=85)
-            buf.seek(0)
-            recompressed = Image.open(buf).convert("RGB")
-            ela = ImageChops.difference(img_rgb, recompressed)
-            stat = ImageStat.Stat(ela)
-            mean = float(sum(stat.mean) / len(stat.mean))
-            std = float(sum(stat.stddev) / len(stat.stddev))
-            return mean, std
-        except Exception:
-            return None, None  # Fully give up on ELA and let upper layer degrade gracefully
-
-def extract_exif_flags(img):
-    """Read common EXIF fields and search for generator keywords. 
-    Ignore exceptions and return as much info as possible."""
-    exif = {}
-    try:
-        raw = img.getexif()
-        for k, v in raw.items():
-            tag = ExifTags.TAGS.get(k, str(k))
-            exif[tag] = v
-    except Exception:
-        pass
-
-    exif_str = " ".join([str(v).lower() for v in exif.values()]) if exif else ""
-    has_camera_fields = any(tag in exif for tag in ["Make", "Model", "LensModel", "DateTimeOriginal"])
-    has_generator_kw = any(kw in exif_str for kw in GENERATOR_KEYWORDS)
-    empty_exif = (len(exif) == 0)
-
-    preview = {}
-    for k in ["Make", "Model", "LensModel", "Software", "DateTimeOriginal"]:
-        if k in exif:
-            preview[k] = str(exif[k])
-
+def canonical_label(label: str) -> str:
+    """Map model-specific labels to canonical 'fake' or 'real' categories."""
+    l = label.lower()
+    if any(k in l for k in ["fake", "ai", "synthetic", "deepfake"]):
+        return "fake"
+    if any(k in l for k in ["real", "authentic", "genuine"]):
+        return "real"
+    # Default fallback if label doesn't match known keywords
+    return label
+
+def rank_probs(id2label: Dict[int, str], probs: List[float]) -> List[Tuple[str, float]]:
+    """Return sorted list of (label, probability) pairs."""
+    pairs = [(id2label[i], float(probs[i])) for i in range(len(probs))]
+    return sorted(pairs, key=lambda x: x[1], reverse=True)
+
+# ---- Load models (once) ----
+mtcnn = MTCNN(keep_all=True, device=DEVICE)
+processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+clf = AutoModelForImageClassification.from_pretrained(MODEL_ID).to(DEVICE)
+id2label = clf.config.id2label
+
+# ---- Core inference ----
+@torch.no_grad()
+def classify_pil(img: Image.Image) -> Dict[str, Any]:
+    """Run classification on a single PIL image and return ranked probabilities."""
+    inputs = processor(images=img, return_tensors="pt")
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    logits = clf(**inputs).logits
+    probs = torch.softmax(logits, dim=-1).squeeze().tolist()
+    ranked = rank_probs(id2label, probs)
+
+    # Extract approximate fake / real probabilities based on label keywords
+    fake_p, real_p = None, None
+    for lbl, p in ranked:
+        cat = canonical_label(lbl)
+        if cat == "fake" and fake_p is None:
+            fake_p = p
+        if cat == "real" and real_p is None:
+            real_p = p
     return {
-        "has_camera_fields": has_camera_fields,
-        "has_generator_kw": has_generator_kw,
-        "empty_exif": empty_exif,
-        "exif_preview": preview
+        "top": ranked[:3],
+        "fake_prob": fake_p,
+        "real_prob": real_p
     }
 
-def ai_likelihood(img):
-    """
-    Main entry point: catch any exceptions and return JSON-friendly output instead of crashing the frontend.
-    """
-    try:
-        if img is None:
-            return {"label": "Error", "message": "No image uploaded."}
-
-        info = extract_exif_flags(img)
-        ela_mean, ela_std = compute_ela_score(img)
-
-        # Initial score (0.5 = uncertain)
-        score = 0.5
-        reasons = []
-
-        if info["has_generator_kw"]:
-            score += 0.4
-            reasons.append("Metadata contains generator keywords (e.g., Stable Diffusion/Midjourney).")
-        if info["has_camera_fields"]:
-            score -= 0.2
-            reasons.append("Camera EXIF fields found (Make/Model/Lens/DateTimeOriginal).")
-        if info["empty_exif"]:
-            score += 0.1
-            reasons.append("No EXIF found (common in exported AI images or screenshots).")
-
-        if ela_mean is not None and ela_std is not None:
-            if ela_mean < 2.0 and ela_std < 2.0:
-                score += 0.15
-                reasons.append("ELA mean/std are very low → uniform compression error (AI-like).")
-            elif ela_mean > 4.0 or ela_std > 4.0:
-                score -= 0.05
-                reasons.append("ELA mean/std are higher → natural camera/post-processing artifacts (Real-like).")
+def analyze(img: Image.Image) -> Dict[str, Any]:
+    """Main analysis pipeline: detect faces, classify each face or full image."""
+    if img is None:
+        return {"error": "No image uploaded."}
+
+    img = img.convert("RGB")
+    img = resize_keep_ratio(img, MAX_SIDE)
+
+    # 1) Face detection
+    boxes, _ = mtcnn.detect(img)
+    crops = []
+    if boxes is not None:
+        for (x1, y1, x2, y2) in boxes:
+            x1 = max(0, int(x1)); y1 = max(0, int(y1))
+            x2 = min(img.width, int(x2)); y2 = min(img.height, int(y2))
+            if x2 > x1 and y2 > y1:
+                crops.append(img.crop((x1, y1, x2, y2)))
+
+    results = []
+    if crops:
+        # 2) Classify each detected face
+        for idx, face in enumerate(crops, 1):
+            r = classify_pil(face)
+            results.append({"face": idx, **r})
+    else:
+        # 3) If no face is detected, classify the whole image
+        r = classify_pil(img)
+        results.append({"face": None, **r})
+
+    # Aggregate: use median of fake probabilities across all faces
+    fake_scores = []
+    for r in results:
+        if r.get("fake_prob") is not None:
+            fake_scores.append(r["fake_prob"])
         else:
-            reasons.append("ELA failed (unsupported format/codec); decision based on metadata only.")
+            # Fallback: use top-1 label keyword
+            top1 = r["top"][0][0]
+            fake_scores.append(1.0 if canonical_label(top1) == "fake" else 0.0)
 
-        score = max(0.0, min(1.0, score))
-        label = "Likely AI" if score >= 0.6 else ("Uncertain" if 0.4 <= score < 0.6 else "Likely Real")
+    if fake_scores:
+        overall_fake = float(np.median(fake_scores))
+    else:
+        overall_fake = 0.5
 
-        return {
-            "label": label,
-            "ai_probability": round(score, 3),
-            "ela_mean": None if ela_mean is None else round(ela_mean, 3),
-            "ela_std": None if ela_std is None else round(ela_std, 3),
-            "exif": info["exif_preview"],
-            "notes": reasons or ["No strong signals; result uncertain."]
-        }
+    label = "Likely AI/Deepfake" if overall_fake >= 0.6 else ("Uncertain" if overall_fake >= 0.4 else "Likely Real")
 
-    except Exception as e:
-        # Fallback: show error details in JSON instead of crashing frontend
-        return {"label": "Error", "message": str(e)}
+    return {
+        "label": label,
+        "overall_fake_probability": round(overall_fake, 3),
+        "faces_detected": len(crops),
+        "per_face_results": results
+    }
 
+# ---- Gradio UI ----
 with gr.Blocks() as demo:
-    gr.Markdown("""
-    # 🕵️ FakeSpotter (Heuristic Demo)
-    Upload an image to estimate whether it is **AI-generated** or **Real** using simple FREE heuristics:
-    - Metadata scan (generator keywords vs. camera EXIF)
-    - ELA (Error Level Analysis) statistics  
-    > ⚠️ Classroom demo, **not** a forensic tool.
-    """)
-    inp = gr.Image(type="pil", label="Upload image")
-    out = gr.JSON(label="Result")
-    btn = gr.Button("Analyze")
-    btn.click(ai_likelihood, inputs=inp, outputs=out)
+    gr.Markdown(
+        """
+        # 🕵️ FakeSpotter — Image Deepfake Detector (CPU)
+        Upload an image. If a face is detected, each face is analyzed; otherwise, the whole image is classified.  
+        **No EXIF is used.** Model can be swapped by editing `MODEL_ID` in the code.  
+        > Classroom demo — not a forensic tool.
+        """
+    )
+    with gr.Row():
+        inp = gr.Image(type="pil", label="Upload image")
+        out = gr.JSON(label="Results")
+    gr.Button("Analyze").click(analyze, inputs=inp, outputs=out)
 
 if __name__ == "__main__":
     demo.launch()
-