final CAM layers fixed

app.py

@@ -1,156 +1,146 @@
-# app.py ────────────────────────────────────────────────────────────────
+# app.py ───────────────────────────────────────────────────────────────
 """
-Two‑stage local AI‑image detector
-1. haywoodsloan/ai-image-detector-deploy  → Real vs AI         (Swin‑V2)
-2. SuSy.pt                                → Likely generator   (ResNet‑based)
+Two‑stage AI‑image detector with visual explainability
 
-Includes Grad‑CAM overlays:
-    • always show heat‑map for binary decision
-    • if image is flagged AI, also show heat‑map for SuSy
+Stage‑1 : haywoodsloan/ai-image-detector-deploy (Swin‑V2)  → Real vs AI
+          ⟳ Grad‑CAM (torchcam) overlay
+Stage‑2 : SuSy.pt (torchscript ResNet)                      → Generator
+          ⟳ Saliency‑grad overlay (Captum), because hooks are disabled
 """
 
-import gradio as gr
-import numpy as np, torch, pandas as pd, matplotlib.pyplot as plt
+# ───────────────────── Imports ────────────────────────────────────────
+import torch, numpy as np, pandas as pd, matplotlib.pyplot as plt
 from PIL import Image
 from torchvision import transforms
-from skimage.feature import graycomatrix, graycoprops
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 from torchcam.methods import GradCAM
+from captum.attr import Saliency
+from skimage.feature import graycomatrix, graycoprops
+import gradio as gr
 
-# ────────────  Stage‑1 model (binary)  ────────────────────────────────
-BIN_ID      = "haywoodsloan/ai-image-detector-deploy"
-bin_proc    = AutoImageProcessor.from_pretrained(BIN_ID)
-bin_model   = AutoModelForImageClassification.from_pretrained(BIN_ID)
-bin_model.eval()
+# ─────────────────── Runtime / models ─────────────────────────────────
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+plt.set_loglevel("ERROR")
 
-CAM_LAYER_BIN  = "encoder.layers.3.blocks.1.layernorm_after"   # <- from dump
+# Stage‑1 (eager)
+BIN_ID   = "haywoodsloan/ai-image-detector-deploy"
+bin_proc = AutoImageProcessor.from_pretrained(BIN_ID)
+bin_mod  = AutoModelForImageClassification.from_pretrained(BIN_ID).to(device).eval()
+CAM_LAYER_BIN = "encoder.layers.3.blocks.1.layernorm_after"
 
-# ────────────  Stage‑2 model (SuSy)  ──────────────────────────────────
-susy_model    = torch.jit.load("SuSy.pt").eval()
-CAM_LAYER_SUSY = "feature_extractor.resnet_model.layer4.1.relu"  # <- from dump
+# Stage‑2 (scripted)
+susy_mod = torch.jit.load("SuSy.pt").to(device).eval()  # ScriptModule
+CAM_LAYER_SUSY = "feature_extractor.resnet_model.layer4.1.relu"
 
-GEN_CLASSES  = [
-    "Stable Diffusion 1.x", "DALL·E 3", "MJ V5/V6",
-    "Stable Diffusion XL",  "MJ V1/V2",
-]
+GEN_CLASSES = ["Stable Diffusion 1.x", "DALL·E 3",
+               "MJ V5/V6", "Stable Diffusion XL", "MJ V1/V2"]
 PATCH, TOP = 224, 5
 
-# ────────────  Heat‑map helper  ───────────────────────────────────────
-def grad_cam_overlay(model, inputs, target_layer, class_idx, orig_pil):
-    # prepare Grad‑CAM extractor
-    cam_ex = GradCAM(model, target_layer=target_layer,
-                     input_shape=next(iter(inputs.values()) if isinstance(inputs, dict) else [inputs]).shape)
-
-    # forward & backward
-    scores = model(**inputs).logits if isinstance(inputs, dict) else model(inputs)
-    scores[0, class_idx].backward()
-
-    # normalise cam
-    mask = cam_ex(class_idx)[0].cpu().numpy()
+# ─────────────── Universal overlay helper ─────────────────────────────
+def overlay_explanation(model, model_inputs, target_layer, class_idx, base_img):
+    """
+    • If model is eager (supports hooks) → Grad‑CAM via torchcam
+    • If model is ScriptModule         → absolute‑gradient saliency via Captum
+    Returns an RGBA PIL image blended with the heat‑map.
+    """
+    is_script = isinstance(model, torch.jit.ScriptModule)
+
+    # Prepare inputs for forward
+    forward_inputs = model_inputs if torch.is_tensor(model_inputs) else dict(model_inputs)
+
+    # ---------- Scripted: Captum Saliency ----------
+    if is_script:
+        model.zero_grad(set_to_none=True)
+        sal = Saliency(model)
+        if not torch.is_tensor(forward_inputs):
+            forward_inputs = forward_inputs["pixel_values"]
+        grads = sal.attribute(forward_inputs, target=class_idx).abs().mean(1, keepdim=True)
+        mask  = grads.squeeze().detach().cpu().numpy()
+    # ---------- Eager: torchcam Grad‑CAM ----------
+    else:
+        mods = dict(model.named_modules())
+        tgt  = mods.get(target_layer) or next(m for n, m in mods.items() if n.endswith(target_layer))
+        cam  = GradCAM(model, target_layer=tgt)
+        outputs = (model(forward_inputs) if torch.is_tensor(forward_inputs)
+                   else model(**forward_inputs))
+        logits = outputs.logits if hasattr(outputs, "logits") else outputs
+        mask   = cam(class_idx, logits)[0].detach().cpu().numpy()
+
+    # normalise & overlay
     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-6)
-    mask = Image.fromarray(np.uint8(plt.cm.jet(mask)[:, :, :3] * 255)).resize(orig_pil.size, Image.BICUBIC)
-    return Image.blend(orig_pil.convert("RGBA"), mask.convert("RGBA"), alpha=0.45)
+    heat = Image.fromarray((plt.cm.jet(mask)[:, :, :3] * 255).astype(np.uint8))\
+               .resize(base_img.size, Image.BICUBIC)
+    return Image.blend(base_img.convert("RGBA"), heat.convert("RGBA"), alpha=0.45)
 
-# ────────────  SuSy helper  ───────────────────────────────────────────
-to_tensor     = transforms.ToTensor()
-to_gray_pil   = transforms.Compose([transforms.PILToTensor(), transforms.Grayscale()])
+# ───────────── SuSy patch‑ranking helper ──────────────────────────────
+to_tensor = transforms.ToTensor()
+to_gray   = transforms.Compose([transforms.PILToTensor(), transforms.Grayscale()])
 
-def susy_predict(img: Image.Image) -> dict:
+def susy_predict(img: Image.Image):
     w, h = img.size
     npx, npy = max(1, w // PATCH), max(1, h // PATCH)
-    patches = np.zeros((npx * npy, PATCH, PATCH, 3), dtype=np.uint8)
+    patches  = np.zeros((npx * npy, PATCH, PATCH, 3), dtype=np.uint8)
 
     for i in range(npx):
         for j in range(npy):
             x, y = i * PATCH, j * PATCH
-            patches[i * npy + j] = np.array(img.crop((x, y, x + PATCH, y + PATCH)).resize((PATCH, PATCH)))
+            patches[i*npy+j] = np.array(img.crop((x, y, x+PATCH, y+PATCH)).resize((PATCH, PATCH)))
 
     contrasts = []
     for p in patches:
-        g = to_gray_pil(Image.fromarray(p)).squeeze(0).numpy()
+        g = to_gray(Image.fromarray(p)).squeeze(0).numpy()
         glcm = graycomatrix(g, [5], [0], 256, symmetric=True, normed=True)
         contrasts.append(graycoprops(glcm, "contrast")[0, 0])
 
-    idx    = np.argsort(contrasts)[::-1][:TOP]
-    tensor = torch.from_numpy(patches[idx].transpose(0, 3, 1, 2)).float() / 255.0
-
+    idx  = np.argsort(contrasts)[::-1][:TOP]
+    tens = torch.from_numpy(patches[idx].transpose(0,3,1,2)).float()/255.0
     with torch.no_grad():
-        probs = susy_model(tensor).softmax(-1).mean(0).numpy()[1:]
+        probs = susy_mod(tens.to(device)).softmax(-1).mean(0).cpu().numpy()[1:]
     return dict(zip(GEN_CLASSES, probs))
 
-# ────────────  Main pipeline  ─────────────────────────────────────────
+# ───────────────────── Pipeline ───────────────────────────────────────
 def pipeline(img_arr):
     img = Image.fromarray(img_arr) if isinstance(img_arr, np.ndarray) else img_arr
     heatmaps = []
 
-    # Stage‑1: binary Real/AI
+    # Stage‑1
     with torch.no_grad():
-        inp_bin = bin_proc(images=img, return_tensors="pt")
-        logits  = bin_model(**inp_bin).logits
-        probs   = torch.softmax(logits, -1)[0].tolist()   # [artificial, real]
-
-    ai_conf, real_conf = probs[0], probs[1]
-
-    # Grad‑CAM for winning class
-    class_idx = 0 if ai_conf >= real_conf else 1
-    heatmaps.append(
-        grad_cam_overlay(
-            bin_model, inp_bin,
-            target_layer=CAM_LAYER_BIN,
-            class_idx=class_idx,
-            orig_pil=img
-        )
-    )
-
-    # defaults
-    msg, bar_df, bar_vis = f"Authentic ({real_conf*100:.1f} %)", None, False
-
-    # Stage‑2 if AI
+        inp_bin = bin_proc(images=img, return_tensors="pt").to(device)
+        logits  = bin_mod(**inp_bin).logits.softmax(-1)[0]   # [AI, Real]
+
+    ai_conf, real_conf = logits
+    winner_idx = 0 if ai_conf >= real_conf else 1
+    heatmaps.append(overlay_explanation(bin_mod, inp_bin, CAM_LAYER_BIN, winner_idx, img))
+
+    verdict = f"Authentic ({real_conf*100:.1f} %)"
+    bar_df, show_bar = None, False
+
+    # Stage‑2 (only if AI)
     if ai_conf > real_conf:
-        msg = f"AI‑generated ({ai_conf*100:.1f} %)"
+        verdict = f"AI‑generated ({ai_conf*100:.1f} %)"
         gen_probs = susy_predict(img)
-        bar_df  = pd.DataFrame({"class": gen_probs.keys(), "prob": gen_probs.values()})
-        bar_vis = True
+        bar_df = pd.DataFrame({"class": gen_probs.keys(), "prob": gen_probs.values()})
+        show_bar = True
 
-        # SuSy heat‑map: choose most‑probable generator class
         with torch.no_grad():
-            t_inp = to_tensor(img.resize((224, 224))).unsqueeze(0)
-            logits_susy = susy_model(t_inp)
-        susy_class = logits_susy[0, 1:].argmax().item() + 1  # skip 'real'
-
-        heatmaps.append(
-            grad_cam_overlay(
-                susy_model, t_inp,
-                target_layer=CAM_LAYER_SUSY,
-                class_idx=susy_class,
-                orig_pil=img
-            )
-        )
-
-    return msg, gr.update(value=bar_df, visible=bar_vis), heatmaps
-
-# ────────────  Gradio UI  ─────────────────────────────────────────────
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🖼️ Local AI Fake Detector")
+            susy_in = to_tensor(img.resize((224,224))).unsqueeze(0).to(device)
+            g_idx   = susy_mod(susy_in)[0,1:].argmax().item() + 1
+        heatmaps.append(overlay_explanation(susy_mod, susy_in, CAM_LAYER_SUSY, g_idx, img))
+
+    return verdict, gr.update(value=bar_df, visible=show_bar), heatmaps
 
+# ───────────────────────── UI ─────────────────────────────────────────
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🖼️ Two‑Stage AI Fake Detector — Explained with Heat‑maps")
     with gr.Row():
         img_in = gr.Image(type="numpy", label="Upload image")
         btn    = gr.Button("Detect")
 
-    txt_bin = gr.Textbox(label="Step 1 — Real vs AI", interactive=False)
-
-    bar_gen = gr.BarPlot(
-        x="class", y="prob",
-        title="Step 2 — Likely generator",
-        y_label="probability",
-        visible=False
-    )
-
-    gal_cam = gr.Gallery(
-        label="Model attention heat‑maps",
-        columns=2, height=300, visible=True
-    )
+    txt_out = gr.Textbox(label="Verdict", interactive=False)
+    bar_out = gr.BarPlot(x="class", y="prob", title="Likely generator",
+                         y_label="probability", visible=False)
+    gal_out = gr.Gallery(label="Heat‑maps", columns=2, height=320)
 
-    btn.click(pipeline, inputs=img_in, outputs=[txt_bin, bar_gen, gal_cam])
+    btn.click(pipeline, inputs=img_in, outputs=[txt_out, bar_out, gal_out])
 
 demo.launch()

requirements.txt

@@ -8,4 +8,7 @@ pydantic==2.10.6
 wheel
 huggingface_hub>=0.22
 pandas
-torchcam>=0.4
+torchcam>=0.4
+matplotlib>=3.8
+timm>=0.9.12
+captum>=0.7