AI-Generated Image Detector — robust to real-world post-processing

Summary

Generative tools can now produce photorealistic images at scale, and the images that actually cause harm online are rarely perfect, as they've been screenshotted, re-compressed by a messaging app, cropped to a thumbnail, or run through a filter. A detector that only works on clean lab images is close to useless in that setting.

We built a prototype that classifies an image as real or AI-generated and is trained specifically to survive that redistribution pipeline. We train small models (all far under the 2B-parameter cap) with on-the-spot augmentation drawn from the exact transform family we're evaluated against, so the network learns cues that persist through compression and blur instead of weak high-frequency traces that vanish under any processing

Headline result, on the held-out demo benchmark (800 real COCO photos + 800 DALL·E images), decision threshold 0.5: 94.6% clean accuracy, 0.997 clean AUC, holding at ≥0.97 AUC and ≥89% accuracy under every post-processing transform except one — JPEG down to quality 30, Gaussian blur to σ=2, 4× downscale, sensor noise, ±20% colour jitter. The exception is an 80% centre-crop, where accuracy drops to 43%: the model leans on framing/border cues that an aggressive crop removes, and hardening that is the obvious next step.

In a separate ablation on our internal 25.5k-image test set, adding the robustness augmentation lifts worst-case accuracy under heavy transforms from 59% to 84% for under a 1-point cost on clean accuracy — evidence that the training pipeline, not model size, drives robustness. Both models are ~0.14M parameters.

How it addresses the problem statement

Requirement	What we did
Detect AIGC images	Binary classifier, outputs a calibrated P(AI) per image
Robustness to post-processing	Train-time augmentation matched to the eval grid (JPEG 90/70/50/30, blur σ 0.5/1/2, resize 0.5×/0.25×, noise σ 0.02/0.05/0.10, colour jitter ±20%, centre-crop 80%), applied as stacked 1-2 op chains because real redistribution stacks operations. Crop-based resizing instead of plain downsampling (SAFE, KDD 2025) to preserve high-frequency artifacts
Clear technical approach	Documented, single-purpose scripts: one `--backbone` flag switches between a from-scratch CNN, a pretrained EfficientNet-B0, and a Fourier-spectrum CNN
Evaluation strategy	Per-transform accuracy and AUC across the full severity grid, a clean /mean-transformed/worst-transformed 3-number summary, augmented-vs-baseline A/B, per-dataset breakdown
Trade-offs: robustness, generalisation, false positives	Temperature calibration + operating points at 1% / 5% / 10% FPR, a leave-one-dataset-out generalisation probe, a written error-analysis note on representative FPs/FNs
< 2B parameters, limited compute	From-scratch CNN is ~0.2M params, everything trains on a single Colab GPU or an Apple-Silicon laptop (MPS) in well under an hour
Deliverable inference script	`python -m src.infer --input_dir <dir> --output_json <path>` → JSON list of `{image_path, pred}`
Don't train on the reserved benchmark	The WildFake downloader hard-blocks `coco.zip` and `DALLE.zip`; a separate opt-in script pulls them into an isolated directory for demo/eval only

How we built it

Data acquisition:

We wrote streaming partial downloaders as all three source datasets were too large to download completely.

CIFAKE (32×32, 120k images) via kagglehub: small enough to take whole.
SID_Set via Hugging Face datasets in streaming mode: iterate the shard stream, take the first N for training and the next M as a non-overlapping holdout. 3-way labels (real/fully-synthetic/tampered) collapse to REAL / FAKE.
WildFake (~1.3 TB, ModelScope): it ships as 6-50 GB per-generator zip archives with no streaming loader. We open an HTTP stream per archive, decode it with stream-unzip (which never needs the zip's central directory), and stop after the first N images. This uses a few hundred MB per generator instead of tens of GB. Nine generator families for FAKE (ADM, DDIM, DDPM, Imagen, VQDM, Midjourney, personalised SD, GAN, "other"), five sources for REAL. The two archives that make up the organisers' demo benchmark are blocklisted so they can never leak into training.

Balancing: CIFAKE alone is 100k images and would dominate the merge, so the trainer has --cap_per_source, --max_total, and --balance_classes to build an even, class-balanced training set.

Models: A single --backbone flag selects:

small_cnn : 6 conv layers, BatchNorm + dropout, ~0.2M params, from scratch at 32×32.
efficientnet_b0 : ImageNet-pretrained via timm, 224×224, optional --freeze_backbone for a minutes-per-epoch head-only fit.
freq_cnn : the same small CNN but fed a 3-channel log-magnitude Fourier spectrum (per RGB channel, fftshift, 64×64). GAN/diffusion up-sampling leaves periodic peaks in the spectrum that cameras don't produce.

Ensemble: src/ensemble.py averages any set of checkpoints. A spatial model and the spectrum model fail on disjoint inputs as the spatial model misses smooth low-detail real photos, and the spectrum model is blind to garbled text and warped geometry, so the average is flatter across the transform grid than either alone.

Calibration: src/calibrate.py fits a single temperature on a held-out set (min NLL via LBFGS) and derives decision thresholds that hit target false-positive rates, so the demo can report "recall at 5% FPR" instead of an arbitrary 0.5 cut.

Evaluation harness: eval_robustness (per-transform accuracy), eval_auc (per-transform AUC), robustness_summary (grouped-bar chart + the 3-number summary, auto-detecting whichever metric), eval_crossdataset (per-dataset clean/mean-transformed AUC + a shortcut-spread flag), and error_analysis (worst FP/FN saved as images + a written note + a "which images stay wrong across the most transforms" scan).

Robustness results:

Final model (frequency CNN) on the demo benchmark, 1,600 images, threshold 0.5:

	Clean	Mean over transforms	Worst transform
Accuracy	94.6%	91.6%	42.9% (centre-crop 80%)
AUC	0.997	0.950	0.385 (centre-crop 80%)

Excluding centre-crop, every transform holds at ≥88.7% accuracy / ≥0.972 AUC — JPEG q90–q30, blur σ 0.5–2.0, resize 0.5×/0.25×, noise σ 0.02–0.10 and ±20% colour jitter are all handled well (AUC ≈ 0.99 for most). The centre-crop collapse is real and we report it rather than dropping the row.

Augmentation ablation (spatial CNN, internal 25,550-image test set, identical data and split, only the training augmentation changes):

Model	Clean acc	Mean acc (transformed)	Worst acc (transformed)
No augmentation	91.6%	72.3%	58.7%
With augmentation	90.7%	87.4%	83.9%

~1 point of clean accuracy gains 25 points in worst-case robustness.

Limitations

Calibration: src/calibrate.py fits a temperature and reports operating points at fixed false-positive rates. The model is already near-calibrated (T = 0.92 on a 6,350-image holdout). At a 5% false-positive rate it recalls 70% of AI images (precision 95%); at 1% FPR, 41% recall (precision 98%). This makes the precision/recall trade-off explicit instead of hiding behind a fixed 0.5 cut.
Error analysis on the demo benchmark (threshold 0.5): false-positive rate 0.6% (5/800), false-negative rate 10.2% (82/800). The model is conservative as it almost never flags a real photo, but lets ~1 in 10 DALL·E images through. All 5 false positives sit at P(AI) 0.51-0.53, right on the boundary; there is no confident false positive, so a small threshold nudge clears them. The false negatives are the most photorealistic, texture-rich generations, including one with obviously garbled text that a person spots instantly and a spectrum-only model cannot, which is the argument for fusing a spatial branch.
Dataset shortcut, acknowledged: The sources carry a resolution tell (demo benchmark: COCO ~200px vs DALL·E 512px; training: SID fakes are always exactly 1024²). We mitigate with near-square random crops before resize and built a per-dataset probe (src/eval_crossdataset.py); a full leave-one-generator-out evaluation is the honest next step we ran out of time for.
Known limits: small input resolution discards fine detail; WildFake training uses only the leading slice of each generator archive; centre-crop robustness is poor; not tested against fully unseen generator families.

Development tools

Google Colab (GPU training)
VS Code + Jupyter (notebooks/training.ipynb for the end-to-end run)
Apple-Silicon MacBook for local development and evaluation — the code auto-selects CUDA → Apple MPS → CPU
GitHub with a branch-per-person workflow
Claude Code as a pair-programming assistant

Models / APIs

SmallCNN: custom 6-layer CNN, ~0.2M parameters, trained from scratch
EfficientNet-B0: ImageNet-pretrained, fine-tuned as a binary classifier (via timm)
freq_cnn: SmallCNN architecture over a log-magnitude Fourier spectrum
Ensemble: probability average of a spatial model and the frequency model
No third-party inference APIs: the detector is fully self-contained. External services used only for data acquisition: the ModelScope dataset file API and the Hugging Face datasets streaming API.

Libraries & frameworks

PyTorch / torchvision — model, training loop, data loading, transforms
timm — EfficientNet-B0 backbone
Hugging Face datasets — streaming partial download of SID_Set
stream-unzip + requests — streaming partial download of WildFake zip archives from ModelScope
kagglehub / kaggle CLI — CIFAKE download
scikit-learn — ROC-AUC, precision / recall / F1
NumPy — FFT spectra, noise, array ops
Pillow — image I/O and augmentation
matplotlib — robustness charts
tqdm — progress reporting

Datasets & assets

CIFAKE (Kaggle, birdy654/cifake-real-and-ai-generated-synthetic-images) — 100k train / 20k test, 32×32, balanced REAL/FAKE.
SID_Set (Hugging Face, saberzl/SID_Set) — streamed subset, full resolution, 3-way labels mapped to REAL/FAKE.
WildFake (ModelScope, hy2628982280/WildFake) — streamed per-generator subset across nine synthetic families and five real sources; the coco.zip and DALLE.zip archives are excluded from all training.
WildFake demo/validation subset (COCO real + DALL-E generated) — downloaded into an isolated directory and used only for demonstration and error analysis, per the challenge rules.
Self-created transformed test sets — every image passed through the JPEG / blur / resize / noise / colour-jitter / crop grid to measure per-transform performance.

What's next

Fuse the spatial and frequency branches into one end-to-end model instead of a late average.
Swap the 32×32 backbone for a frozen DINOv2 / CLIP feature extractor at higher resolution — strong cross-generator generalisation for very little compute.
Leave-one-generator-out evaluation for an honest out-of-distribution number.
Grad-CAM overlays to confirm the model attends to texture and boundaries rather than the subject.

Built with

python · pytorch · torchvision · timm · huggingface · scikit-learn · numpy · pillow · matplotlib · stream-unzip · google-colab · jupyter · modelscope · kaggle