✨ Two-stage AI detector (binary + generator classifier)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+SuSy.pt filter=lfs diff=lfs merge=lfs -text

SuSy.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa10fae300ee2742c7a373b6c3332c2595b461954b8f5616d2d382ef2751020e
+size 50810392

app.py

@@ -1,10 +1,112 @@
+# ─── app.py ──────────────────────────────────────────────────────────────
 import gradio as gr
+import torch, numpy as np
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from torchvision import transforms
+from skimage.feature import graycomatrix, graycoprops
 
-with gr.Blocks(fill_height=True) as demo:
-    with gr.Sidebar():
-        gr.Markdown("# Inference Provider")
-        gr.Markdown("This Space showcases the haywoodsloan/ai-image-detector-deploy model, served by the hf-inference API. Sign in with your Hugging Face account to use this API.")
-        button = gr.LoginButton("Sign in")
-    gr.load("models/haywoodsloan/ai-image-detector-deploy", accept_token=button, provider="hf-inference")
-    
-demo.launch()
+# ───── Binary (real-vs-AI) model ─────────────────────────────────────────
+BIN_CKPT = "haywoodsloan/ai-image-detector-deploy"
+bin_proc = AutoImageProcessor.from_pretrained(BIN_CKPT)
+bin_model = AutoModelForImageClassification.from_pretrained(BIN_CKPT)
+bin_model.eval()
+
+# ───── Generator-classifier (“SuSy”) model ───────────────────────────────
+susy_model = torch.jit.load("SuSy.pt")
+susy_model.eval()
+GEN_CLASSES = [
+    'DALL·E 3', 'Stable Diffusion 1.x', 'MJ V5/V6',
+    'MJ V1/V2', 'Stable Diffusion XL'
+]  # “Authentic” intentionally omitted
+
+# ─── Helper - classify with SuSy (patch-based) ───────────────────────────
+def susy_predict(image: Image.Image, top_k_patches: int = 5, patch_size: int = 224):
+    w, h = image.size
+    npx, npy = w // patch_size, h // patch_size
+    if npx == 0 or npy == 0:          # tiny images → resize once and exit early
+        image = image.resize((patch_size, patch_size), Image.LANCZOS)
+        npx = npy = 1
+
+    # split into patches ---------------------------------------------------
+    patches = np.zeros((npx * npy, patch_size, patch_size, 3), dtype=np.uint8)
+    for i in range(npx):
+        for j in range(npy):
+            x, y = i * patch_size, j * patch_size
+            patches[i * npy + j] = np.array(image.crop((x, y, x + patch_size, y + patch_size)))
+
+    # pick the most “textured” patches (GLCM contrast) ---------------------
+    scores = []
+    to_tensor = transforms.Compose([transforms.PILToTensor(), transforms.Grayscale()])
+    for p in patches:
+        g = to_tensor(Image.fromarray(p)).squeeze(0).numpy()
+        glcm = graycomatrix(g, [5], [0], 256, symmetric=True, normed=True)
+        scores.append(graycoprops(glcm, "contrast")[0, 0])
+
+    keep = patches[np.argsort(scores)[::-1][:top_k_patches]]
+    keep = torch.from_numpy(keep.transpose(0, 3, 1, 2)).float() / 255.0
+
+    # predict --------------------------------------------------------------
+    susy_model.eval()
+    with torch.no_grad():
+        probs = susy_model(keep).softmax(dim=-1).mean(dim=0).numpy()
+
+    # “Authentic” is index 0 – skip it
+    out = {cls: float(p) for cls, p in zip(GEN_CLASSES, probs[1:])}
+    return dict(sorted(out.items(), key=lambda x: x[1], reverse=True))
+
+# ─── Two-stage pipeline ***************************************************
+def pipeline(img):
+    if isinstance(img, np.ndarray):           # (inference images arrive as ndarrays)
+        img = Image.fromarray(img)
+
+    # Stage 1 – binary -----------------------------------------------------
+    with torch.no_grad():
+        inp = bin_proc(images=img, return_tensors="pt")
+        logits = bin_model(**inp).logits
+        p = torch.softmax(logits, dim=-1)[0]
+        pred_idx = int(p.argmax())
+        label = bin_model.config.id2label[pred_idx]
+        conf = float(p[pred_idx])
+
+    if label.lower() in {"ai", "fake", "synthetic", "generated"}:
+        # Stage 2 – generator detection -----------------------------------
+        gen_probs = susy_predict(img)
+        return {
+            "binary": f"AI-generated ({conf*100:.1f} %)",
+            "progress": "",                        # hide progress line
+            "generator": gen_probs,
+        }
+    else:
+        return {
+            "binary": f"Authentic ({conf*100:.1f} %)",
+            "progress": "",
+            "generator": None,                     # hides second label
+        }
+
+# ─── Gradio UI ************************************************************
+with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
+    gr.Markdown("## 🖼️ AI Fake Detector")
+    gr.Markdown(
+        "1. Upload an image → we say **Real** or **AI**.<br>"
+        "2. If it’s AI-generated, we guess the **likely generator model**."
+    )
+
+    img_in = gr.Image(type="numpy", label="Input image")
+    run_btn = gr.Button("��� Run detection")
+    bin_out = gr.Label(label="Step 1 — Real vs AI")
+    progress = gr.Markdown(visible=False)            # not used in final version
+    gen_out = gr.Label(label="Step 2 — Probable generator", visible=False)
+
+    def _on_click(img):
+        result = pipeline(img)
+        show_gen = result["generator"] is not None
+        return (
+            result["binary"],
+            gr.update(value=result["progress"], visible=False),
+            gr.update(value=result["generator"] or {}, visible=show_gen),
+        )
+
+    run_btn.click(_on_click, inputs=img_in, outputs=[bin_out, progress, gen_out])
+
+demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+gradio>=4.26.0
+torch            # will pull torchvision too
+transformers
+pillow
+scikit-image