File size: 5,758 Bytes
29271db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
---
license: mit
tags:
  - image-classification
  - ai-generated-image-detection
  - deepfake-detection
  - pytorch
  - clip
  - ensemble
library_name: pytorch
pipeline_tag: image-classification
---

# AI vs Real Image Detector β€” 3-Branch Ensemble
Full write-up, methodology, and deployment code: https://github.com/ammarsapru/ai-detector

Binary classifier that predicts whether an image is a **real photograph** or
**AI-generated**. It fuses three complementary views of an image:

- **CLIP ViT-L-14** (frozen) + trainable adapter β€” semantic plausibility (768-d)
- **EfficientNet-B3** (fine-tuned) β€” local texture / spatial artifacts (1536-d)
- **FFT-CNN** (custom, on the Fourier spectrum) β€” frequency-domain fingerprints
  left by generator upsampling (512-d)

The three feature vectors (2816-d) are concatenated and passed to a fusion MLP
(`2816 β†’ 512 β†’ 128 β†’ 2`).

- **Input:** RGB image, resized to **224Γ—224**
- **Output:** 2 logits β†’ softmax. **Label map: `0 = real`, `1 = ai`**
- **Params:** ~317.9M total, ~13.9M trainable

Full architecture write-up (per-branch design rationale): [MODEL_ARCHITECTURE.md](MODEL_ARCHITECTURE.md).

## How to use

This is a **custom `nn.Module`**, so `AutoModel` / `pipeline()` will not load it.
Download `model.py` + `best_model.pt` from this repo and:

```python
import torch, torch.nn.functional as F
import torchvision.transforms as T
from PIL import Image
from huggingface_hub import hf_hub_download
from model import load_detector  # model.py from this repo

repo = "YOUR-USERNAME/ai-detector"
ckpt = hf_hub_download(repo, "best_model.pt")
device = "cuda" if torch.cuda.is_available() else "cpu"
model, cfg = load_detector(ckpt, device=device)

tf = T.Compose([T.Resize((224, 224)), T.ToTensor()])
img = Image.open("example.jpg").convert("RGB")
with torch.no_grad():
    probs = F.softmax(model(tf(img).unsqueeze(0).to(device)), dim=1)[0]
print({"real": float(probs[0]), "ai": float(probs[1])})
```

## Training data

- **AI images:** Tiny-GenImage β€” 8 generators (Midjourney, SD v1.4/v1.5, SDXL,
  DALLΒ·E 2/3, VQDM, Wukong)
- **Real images:** Tiny-ImageNet
- Balanced 14,000 real / 14,000 AI; 85/15 train/val split; images resized to 224Γ—224.

## Results

In-distribution: **98.5% accuracy / 0.985 macro-F1** (98% real recall, 99% AI
recall), n=4,200, 3 epochs.

Out-of-distribution testing revealed an **asymmetric** generalization
profile β€” both classes have a genuine strength, with one narrow gap:

| Probe | Class | Condition | Result | Provenance |
|---|---|---|---|---|
| AIGC-Detection-Benchmark, standard (n=900, 17 gen.) | AI | full/variable native res, no resize trick | **94.0% recall** | recovered chat transcript |
| AIGC-Detection-Benchmark, standard (n=900, 17 gen.) | real | full/variable native res, no resize trick | **8.0% recall** | recovered chat transcript |
| Tiny-ImageNet `valid`, held-out (n=200) | real | native 64Γ—64 | **97.0% recall** | saved classification report |
| AIGC-Detection-Benchmark, forced to 64Γ—64 (n=900) | AI | forced downscale-then-upscale | 68.0% recall | saved classification report |
| AIGC-Detection-Benchmark, forced to 64Γ—64 (n=900) | real | forced downscale-then-upscale | 36.0% recall | saved classification report |

**AI-image detection already generalizes well to full-resolution external
images (94.0% recall, no resize trick needed).** **Real-image detection
performs extremely well within its native 64Γ—64 resolution domain (97.0%
recall on a held-out split never seen in training)** β€” essentially matching
in-distribution performance. The gap is narrow: that native-resolution
strength doesn't yet extend to full-resolution photos (8.0%). Forcing inputs
to 64Γ—64 before inference is **not a fix**: it partially helps real
(8.0%β†’36.0%) but hurts AI (94.0%β†’68.0%), a net loss. Full methodology,
mechanism, and per-generator breakdown: see `report.pdf` in this repo, or the
[technical report on GitHub](https://github.com/ammarsapru/ai-detector/blob/main/docs/REPORT.md).

![Training curves](curves.png)
![Confusion matrix, in-distribution validation](confusion_matrix.png)

## Limitations & intended use

- **Real-image detection does not generalize to full-resolution photos**
  (8.0% recall) β€” the central limitation. Works well on natively low-res real
  photos (97.0%), not on full-resolution ones. Traces to the real-image
  training source (Tiny-ImageNet) being natively 64Γ—64 β€” not an architectural
  flaw, since the AI branch does not show the same failure. See `report.pdf` Β§6.
- **Do not force inputs to 64Γ—64 before inference as a workaround** β€” it
  trades a partial real-recall gain (8.0%β†’36.0%) for a larger AI-recall loss
  (94.0%β†’68.0%), net negative.
- **One data point (8.0% real / 94.0% AI recall) is recovered from a
  preserved chat transcript, not the notebook's saved output** β€” the
  evaluation cell was edited in place to produce the forced-resize version,
  overwriting the original run's output. See `report.pdf` Β§5.1/Β§7.
- **Distribution note:** real and AI images came from *different* source
  datasets, so in-distribution numbers may be optimistic.
- **Cross-generator generalization has not been properly validated** β€” an
  intended train/test generator exclusion in the OOD benchmark code was not
  actually enforced (see `report.pdf` Β§5.3). Treat any apparent generalization
  in this report as provisional.
- Intended for research / educational use in synthetic-media detection, not as
  a sole authority for high-stakes decisions.

## Training config

AdamW (per-branch LRs: CLIP adapter 1e-5, EfficientNet 3e-5, FFT 1e-4,
head 1e-4), CosineAnnealingLR, grad-clip 1.0, weight decay 1e-4,
CrossEntropy loss, augmentation (flip, color jitter, grayscale, blur).