File size: 4,946 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 124 125 126 | """
Model definition for the AI-vs-Real image detector.
This MUST ship alongside best_model.pt — the checkpoint only stores weights
(a state_dict), so the class definitions here are required to reconstruct the
network before loading those weights.
Architecture: 3-branch ensemble
- CLIP ViT-L-14 (frozen) + trainable adapter -> 768 (semantic)
- EfficientNet-B3 (fine-tuned) -> 1536 (texture)
- FFT-CNN (custom, on Fourier spectrum) -> 512 (frequency)
concatenated -> fusion MLP -> 2 logits (0=real, 1=ai)
"""
import torch
import torch.nn as nn
import torchvision.models as tv_models
import open_clip
class FFTBranch(nn.Module):
"""Detects frequency-domain fingerprints left by generator upsampling."""
def __init__(self, out_dim=512):
super().__init__()
self.cnn = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.ReLU(True), nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True), nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(True), nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(True),
nn.AdaptiveAvgPool2d(1), nn.Flatten(),
)
self.proj = nn.Sequential(nn.Linear(256, out_dim), nn.LayerNorm(out_dim), nn.GELU())
def fft(self, x):
gray = 0.299 * x[:, 0] + 0.587 * x[:, 1] + 0.114 * x[:, 2]
mag = torch.log1p(torch.abs(torch.fft.fft2(gray)))
mag = torch.fft.fftshift(mag, dim=(-2, -1))
B = mag.shape[0]
mn = mag.view(B, -1).min(1).values.view(B, 1, 1)
mx = mag.view(B, -1).max(1).values.view(B, 1, 1)
return ((mag - mn) / (mx - mn + 1e-8)).unsqueeze(1)
def forward(self, x):
return self.proj(self.cnn(self.fft(x)))
class CLIPBranch(nn.Module):
"""Frozen CLIP visual encoder + small trainable adapter (semantic view)."""
def __init__(self, model_name="ViT-L-14", pretrained="openai"):
super().__init__()
clip_model, _, _ = open_clip.create_model_and_transforms(model_name, pretrained=pretrained)
self.visual = clip_model.visual
self.visual.eval()
for p in self.visual.parameters():
p.requires_grad = False
self.register_buffer("mean", torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
self.register_buffer("std", torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
d = 768
self.adapter = nn.Sequential(
nn.Linear(d, d), nn.LayerNorm(d), nn.GELU(), nn.Dropout(0.1),
nn.Linear(d, d), nn.LayerNorm(d),
)
def forward(self, x):
with torch.no_grad():
f = self.visual((x - self.mean) / self.std)
return self.adapter(f)
class EfficientNetBranch(nn.Module):
"""Fine-tuned EfficientNet-B3 (local texture / spatial artifacts)."""
def __init__(self, dropout=0.4):
super().__init__()
m = tv_models.efficientnet_b3(weights=tv_models.EfficientNet_B3_Weights.IMAGENET1K_V1)
self.features = m.features
self.avgpool = m.avgpool
self.drop = nn.Dropout(dropout)
self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
def forward(self, x):
x = (x - self.mean) / self.std
return self.drop(torch.flatten(self.avgpool(self.features(x)), 1))
class ArtifactDetector(nn.Module):
"""Full 3-branch ensemble with a fusion head."""
def __init__(self, cfg):
super().__init__()
self.clip = CLIPBranch(cfg["clip_model"], cfg["clip_pretrain"])
self.effnet = EfficientNetBranch(cfg["effnet_dropout"])
self.fft = FFTBranch(cfg["fft_out_dim"])
total = 768 + 1536 + cfg["fft_out_dim"] # 2816
self.fusion = nn.Sequential(
nn.Linear(total, 512), nn.BatchNorm1d(512), nn.GELU(), nn.Dropout(0.4),
nn.Linear(512, 128), nn.BatchNorm1d(128), nn.GELU(), nn.Dropout(0.2),
nn.Linear(128, 2),
)
def forward(self, x):
return self.fusion(torch.cat([self.clip(x), self.effnet(x), self.fft(x)], dim=1))
# Default config — used as a fallback if the checkpoint does not embed its own.
DEFAULT_CFG = {
"clip_model": "ViT-L-14",
"clip_pretrain": "openai",
"effnet_dropout": 0.4,
"fft_out_dim": 512,
"image_size": 224,
}
def load_detector(checkpoint_path, device="cpu"):
"""Rebuild the model from its embedded cfg and load trained weights."""
ckpt = torch.load(checkpoint_path, map_location=device)
cfg = ckpt.get("cfg", DEFAULT_CFG)
model = ArtifactDetector(cfg).to(device)
model.load_state_dict(ckpt["model_state"])
model.eval()
return model, cfg
|