"""AI Image Detector — Gradio demo. Upload any image; get REAL vs AI-GENERATED with a calibrated confidence and a heatmap of what the model looked at. Three things here exist specifically so the demo behaves on arbitrary uploads, and all three are easy to get wrong (CLAUDE.md R9/R10): 1. **Identical preprocessing to training.** The exact same canonical pipeline is imported from the code repo — not reimplemented. A demo that resizes differently from training silently loses several points of accuracy. 2. **5-crop averaging.** A single centre crop of a 4000px photo throws away most of the evidence. Averaging logits over four corners plus the centre is markedly steadier. 3. **Temperature scaling + a tuned threshold.** Raw softmax from a fine-tuned ViT is badly overconfident (Run A measured T=2.37) and its best operating point was 0.34, not 0.5. Announcing "99.8% AI" about a holiday snap is worse than being quietly wrong. """ from __future__ import annotations import json import os import sys import gradio as gr import numpy as np import spaces # ZeroGPU: must be imported BEFORE torch import torch from huggingface_hub import hf_hub_download, snapshot_download from PIL import Image MODEL_REPO = os.environ.get("MODEL_REPO", "husseinelsaadi/aidetect-vit-b16") CODE_REPO = os.environ.get("CODE_REPO", "husseinelsaadi/aidetect-code") RUN = os.environ.get("RUN", "runC") # runC wins on unseen generators AND reals # The threshold in summary.json was fitted on IN-DISTRIBUTION validation, where real # images look familiar; it flagged ~1 in 4 genuine photos from unseen sources as AI. # 0.71 was refitted on half the OOD set and validated on the held-out half, cutting the # false-positive rate to ~9% for a modest cost in recall. For a live demo, wrongly # accusing a real photograph is the failure that matters. THRESHOLD_OVERRIDE = float(os.environ.get("THRESHOLD", "0.71")) TOKEN = os.environ.get("HF_TOKEN") sys.path.insert(0, snapshot_download(CODE_REPO, repo_type="dataset", token=TOKEN)) from preprocess import canonical_pil, whole_frame_tiles # noqa: E402 from transforms import CLIP_MEAN, CLIP_STD # noqa: E402 from models import build_model, load_checkpoint_into # noqa: E402 # ZeroGPU allocates a GPU only while a @spaces.GPU function is executing, so at import # time there is no CUDA device. Everything loads on CPU and moves to GPU inside the # decorated call. A free personal account may host up to 2 Gradio Spaces this way — # plain CPU Gradio Spaces now require a paid plan. DEVICE = "cpu" def _load(): ck = hf_hub_download(MODEL_REPO, f"{RUN}/checkpoints/best.pt", repo_type="model", token=TOKEN) blob = torch.load(ck, map_location="cpu", weights_only=False) cfg = blob.get("extra", {}).get("config", {}) model = build_model(cfg.get("model", "clip-vit-b16")) load_checkpoint_into(model, blob["model"]) model.eval() temperature, threshold = 1.0, 0.5 try: s = json.load(open(hf_hub_download(MODEL_REPO, f"{RUN}/summary.json", repo_type="model", token=TOKEN))) temperature = s.get("temperature", 1.0) threshold = s.get("threshold", 0.5) except Exception as e: print(f"no calibration found ({e}) — falling back to raw logits at 0.5") if THRESHOLD_OVERRIDE > 0: threshold = THRESHOLD_OVERRIDE return model, temperature, threshold MODEL, TEMPERATURE, THRESHOLD = _load() print(f"loaded {RUN} | T={TEMPERATURE:.3f} threshold={THRESHOLD:.3f}") # Attention rollout needs the per-layer attention matrices, and the SDPA kernel that # recent transformers selects by default does not return them — `output_attentions=True` # then yields None and the heatmap silently disappears. Eager attention is slower, but # this is five 224px crops, not a training loop. try: MODEL.backbone.config._attn_implementation = "eager" except Exception as e: # pragma: no cover print(f"could not force eager attention ({e}) — heatmap may be unavailable") _norm_mean = torch.tensor(CLIP_MEAN).view(3, 1, 1) _norm_std = torch.tensor(CLIP_STD).view(3, 1, 1) def _to_tensor(img: Image.Image) -> torch.Tensor: t = torch.from_numpy(np.array(img, dtype=np.float32) / 255.0).permute(2, 0, 1) return (t - _norm_mean) / _norm_std @torch.no_grad() def _attention_rollout(img: Image.Image, dev: str = "cpu") -> np.ndarray | None: """Attention rollout (Abnar & Zuidema): multiply per-layer attention across all blocks, with a residual term, to see which patches actually drove the decision.""" try: x = _to_tensor(img).unsqueeze(0).to(dev) out = MODEL.backbone(pixel_values=x, output_attentions=True) attns = out.attentions if not attns: return None result = torch.eye(attns[0].size(-1), device=dev) for a in attns: a = a.mean(dim=1)[0] # average the heads a = a + torch.eye(a.size(-1), device=dev) # residual connection a = a / a.sum(dim=-1, keepdim=True) result = a @ result mask = result[0, 1:] # CLS row, patch columns side = int(mask.numel() ** 0.5) mask = mask[: side * side].reshape(side, side).cpu().numpy() # np.ptp(), not mask.ptp() — the ndarray method was removed in NumPy 2.0 and # this whole function is wrapped in a try/except, so it failed invisibly. mask = (mask - mask.min()) / (np.ptp(mask) + 1e-8) return mask except Exception as e: print(f"rollout failed: {e}") return None def _overlay(img: Image.Image, mask: np.ndarray) -> Image.Image: import matplotlib.cm as cm m = np.array(Image.fromarray((mask * 255).astype(np.uint8)).resize(img.size, Image.BILINEAR)) heat = (cm.inferno(m / 255.0)[..., :3] * 255).astype(np.uint8) return Image.fromarray((0.55 * np.array(img.convert("RGB")) + 0.45 * heat) .astype(np.uint8)) def _banner(verdict: str, p_ai: float) -> str: """One unmissable line. A demo is read across a room, not squinted at.""" ai = verdict.startswith("AI") bg, fg = ("#3b0d0d", "#ff8f8f") if ai else ("#0d2e1a", "#7ee2a8") icon = "AI-GENERATED" if ai else "REAL PHOTOGRAPH" conf = p_ai if ai else 1 - p_ai return ( f"
" f"
{icon}
" f"
confidence {conf:.0%} · " f"calibrated probability of AI generation {p_ai:.1%}
") @spaces.GPU(duration=30) @torch.no_grad() def predict(image: Image.Image): if image is None: return ("", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)) dev = "cuda" if torch.cuda.is_available() else "cpu" MODEL.to(dev) # Windows spanning the WHOLE frame, each cropped at native resolution and then put # through the training pipeline. Averaged in LOGIT space (averaging probabilities # would let one saturated window dominate). See preprocess.whole_frame_tiles for # why centre-only cropping was wrong. crops = whole_frame_tiles(image) batch = torch.stack([_to_tensor(c) for c in crops]).to(dev) logits = MODEL(batch).float().cpu().numpy() / TEMPERATURE logit = float(np.mean(logits)) p_ai = float(1 / (1 + np.exp(-logit))) verdict = "AI-GENERATED" if p_ai >= THRESHOLD else "REAL (camera)" spread = float(np.std(1 / (1 + np.exp(-logits)))) heat = None mask = _attention_rollout(canonical_pil(image).crop((16, 16, 240, 240)), dev) if mask is not None: heat = _overlay(canonical_pil(image).crop((16, 16, 240, 240)), mask) agree = ("windows agree" if spread < 0.15 else "windows disagree — treat with caution") notes = ( f"| | |\n|---|---|\n" f"| Probability of AI generation | **{p_ai:.1%}** |\n" f"| Decision threshold | {THRESHOLD:.2f} — fitted on *unseen-generator* data |\n" f"| Calibration | temperature {TEMPERATURE:.2f} |\n" f"| Frame windows analysed | {len(crops)}, spread ±{spread:.1%} ({agree}) |\n\n" f"Evidence, not proof — detection is measurably weaker on generators the model " f"has never seen." ) return (_banner(verdict, p_ai), gr.update(value={"AI-generated": p_ai, "Real": 1 - p_ai}, visible=True), gr.update(value=heat, visible=heat is not None), gr.update(value=notes, visible=True)) CSS = """ /* One narrow column, centred. A detector has exactly one job and the page should look like it: upload, read the answer. Everything else is folded away. */ .gradio-container {max-width: 680px !important; margin: 0 auto !important;} #hdr {text-align: center; margin: 8px 0 18px 0;} #hdr h1 {font-size: 2.2rem; margin: 0 0 6px 0;} #hdr p {opacity: .7; margin: 0; font-size: .96rem; line-height: 1.5;} footer {visibility: hidden;} """ ABOUT = """ ### How it works The vision tower of `openai/clip-vit-base-patch16` with its last four transformer blocks fine-tuned (~28M trainable parameters). At inference, windows spanning the whole frame are cropped at native resolution, passed through the **identical** preprocessing used in training, and averaged in logit space — a single centre crop of a large photo throws away most of the evidence. ### Training data ~195,000 images, balanced 50/50, from **1,100+ distinct generators**: [OpenFake](https://huggingface.co/datasets/ComplexDataLab/OpenFake) (GPT-Image, nano-banana, Midjourney 6/7, Flux, Imagen, Ideogram, Recraft, SD 1.4→3.5, SDXL) and [Community Forensics](https://huggingface.co/datasets/OwensLab/CommunityForensics-Small) (thousands of community fine-tunes and LoRAs). Real images come from LAION, Pexels, COCO, FFHQ, VISION and Landscapes-HQ. Every image of both classes is decoded, cropped and re-encoded as JPEG q95 by one function with no branch on the label. Without that, a detector reaches 99% by learning "PNG means fake". Training on randomly permuted labels lands at 0.511 — chance — which is the evidence that it did not. ### Measured accuracy | Evaluation set | What it tests | Accuracy | AUROC | |---|---|---|---| | Held-out validation | seen generators | 0.976 | 0.997 | | **OpenFake `core/test`** | **unseen generators and unseen real sources** | **0.898** | **0.960** | | OpenFake `reddit/test` | in-the-wild Reddit uploads | 0.869 | 0.939 | | CIFAKE | Bird & Lotfi (2024), 32×32 | 0.720 | 0.815 | ### Limitations — please read before trusting a verdict - About **1 in 11 genuine photographs** from unfamiliar camera pipelines is still flagged as AI. Unfamiliar sensor and compression statistics resemble generation. - Accuracy drops on very low resolution images; heavy upscaling is its own shift. - Screenshots, memes and heavily edited photos sit between the two classes. - This is academic work, not a forensic authority. **Evidence, not proof.** --- CSC625 Deep Learning · Hussein El Saadi · Modern University for Business and Science (MUBS) · Summer 2026. Non-commercial academic use — the training data is CC-BY-NC. """ with gr.Blocks(title="AI Image Detector — Real Photo or AI-Generated?", theme=gr.themes.Soft(primary_hue="indigo"), css=CSS) as demo: gr.HTML( "

AI Image Detector

" "

Upload an image and find out whether a camera took it or a generative " "model made it.
Trained on 1,100+ generators — and measured on generators " "it has never seen.

") inp = gr.Image(type="pil", label="Your image", height=320, sources=["upload", "clipboard"]) btn = gr.Button("Analyse image", variant="primary", size="lg") banner = gr.HTML() label = gr.Label(num_top_classes=2, label="Probabilities", visible=False) heat = gr.Image(label="Where the model looked", visible=False, height=280) with gr.Accordion("Details", open=False, visible=False) as details: notes = gr.Markdown() with gr.Accordion("About this detector — method, data, accuracy, limitations", open=False): gr.Markdown(ABOUT) def _run(image): b, lab, h, n = predict(image) return b, lab, h, n, gr.update(visible=image is not None) # api_name is pinned: the endpoint is part of the public interface, and Gradio # otherwise names it after whichever function happens to be wired up. outputs = [banner, label, heat, notes, details] btn.click(_run, inputs=inp, outputs=outputs, api_name="predict") inp.upload(_run, inputs=inp, outputs=outputs, api_name=False) inp.clear(lambda: ("", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)), outputs=outputs, api_name=False) if __name__ == "__main__": demo.launch()