Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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"<div style='background:{bg};color:{fg};border-radius:12px;padding:18px 22px;" | |
| f"border:1px solid {fg}33'>" | |
| f"<div style='font-size:1.6rem;font-weight:700;letter-spacing:.02em'>{icon}</div>" | |
| f"<div style='opacity:.85;margin-top:4px'>confidence {conf:.0%} · " | |
| f"calibrated probability of AI generation {p_ai:.1%}</div></div>") | |
| 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( | |
| "<div id='hdr'><h1>AI Image Detector</h1>" | |
| "<p>Upload an image and find out whether a camera took it or a generative " | |
| "model made it.<br>Trained on 1,100+ generators — and measured on generators " | |
| "it has never seen.</p></div>") | |
| 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() | |