jayesh20/ai-image-detector

CLIP-Based AI-Generated Image Detector

Model Description

This model is a binary image classifier that distinguishes real (natural) photographs from AI-generated images. It uses a frozen CLIP ViT-B/32 vision encoder (pretrained on LAION-2B) as a fixed feature extractor, with a lightweight multilayer perceptron classification head trained on top of the extracted image embeddings.

The model was developed and trained in a Kaggle notebook titled complete_Fakeddit_image. Despite the notebook name, the training data used is the tiny-genimage dataset rather than the Fakeddit dataset; this README describes the model as actually implemented and trained.

Model Details

• Base encoder: CLIP ViT-B-32, pretrained weights laion2b_s34b_b79k (loaded via open_clip)
• Encoder state: Frozen; no gradient updates applied to CLIP parameters during training
• Classification head: Fully connected network operating on 512-dimensional, L2-normalized CLIP image embeddings

Layer	Output Size	Normalization	Activation	Dropout
Linear	512	BatchNorm1d	GELU	0.4
Linear	256	BatchNorm1d	GELU	0.3
Linear	128	BatchNorm1d	GELU	0.2
Linear	2	-	-	-

• Output: Two logits corresponding to the classes real (label 0) and ai-generated (label 1)
• Framework: PyTorch, with open_clip for the CLIP backbone
• Input resolution: 224 x 224 pixels, RGB

Intended Use

The model is intended for research and experimentation in AI-generated image detection, such as:

• Screening images for likely synthetic origin
• Research on generalization of detectors across different generative model families
• Educational use in understanding CLIP-based transfer learning for detection tasks

Out of Scope Use

This model is not intended for:

• Legal, forensic, or high-stakes determinations of image authenticity without human review
• Detection of generative models or techniques not represented in the training distribution
• Use as a sole determinant of content moderation decisions

Training Data

The model was trained on the tiny-genimage dataset (source: yangsangtai/tiny-genimage), which pairs natural images with images produced by seven different generative model families.

• Total images: 35,000
• Class balance: 17,500 real images (label 0), 17,500 AI-generated images (label 1)
• Train / validation split: 28,000 / 7,000 images

Generators represented in the dataset, each contributing 5,000 images:

• imagenet_ai_0424_wukong
• imagenet_glide
• imagenet_ai_0419_biggan
• imagenet_ai_0419_vqdm
• imagenet_midjourney
• imagenet_ai_0424_sdv5
• imagenet_ai_0508_adm

Training Procedure

Preprocessing and Augmentation

Implemented using the albumentations library.

Training transforms:

• Resize to 224 x 224
• Horizontal flip (probability 0.5)
• Random brightness/contrast (probability 0.3)
• Normalization
• Conversion to tensor

Validation transforms:

• Resize to 224 x 224
• Normalization
• Conversion to tensor

Optimization

• Loss function: Cross-entropy loss
• Optimizer: AdamW, applied only to the classification head parameters
• Learning rate: 1e-4 (initial training phase), reduced to 1e-5 for a subsequent fine-tuning phase
• Learning rate schedule: Cosine annealing
• Batch size: 32
• Maximum epochs: 20, with early stopping (patience of 5 epochs, monitored on validation F1 score)
• Checkpointing: Best model saved whenever validation F1 improved; full training state (model, optimizer, scheduler, epoch, best F1) checkpointed for resumption
• Hardware: Single NVIDIA Tesla T4 GPU

Training was conducted in two stages within the notebook: an initial run to 18 epochs before early stopping triggered a checkpoint save, followed by a resumed run at a lower learning rate for 2 additional epochs.

Evaluation

Evaluation was performed on the held-out validation split (7,000 images) using the checkpoint with the best validation F1 score.

Final Reported Metrics

Metric	Score
Accuracy	0.9500
Precision	0.9474
Recall	0.9529
F1 Score	0.9501
ROC AUC	0.9894

A confusion matrix was also generated on the validation set to inspect class-wise performance; see the original notebook for the corresponding plot.

Usage

import torch
import torch.nn as nn
import torch.nn.functional as F
import open_clip
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load CLIP backbone
clip_model, _, _ = open_clip.create_model_and_transforms(
    "ViT-B-32",
    pretrained="laion2b_s34b_b79k"
)
clip_model = clip_model.to(DEVICE)
for param in clip_model.parameters():
    param.requires_grad = False

# Define classifier head
class CLIPBinaryClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.clip = clip_model
        self.classifier = nn.Sequential(
            nn.Linear(512, 512), nn.BatchNorm1d(512), nn.GELU(), nn.Dropout(0.4),
            nn.Linear(512, 256), nn.BatchNorm1d(256), nn.GELU(), nn.Dropout(0.3),
            nn.Linear(256, 128), nn.BatchNorm1d(128), nn.GELU(), nn.Dropout(0.2),
            nn.Linear(128, 2)
        )

    def forward(self, images):
        with torch.no_grad():
            features = self.clip.encode_image(images)
        features = F.normalize(features, dim=-1)
        return self.classifier(features)

# Load trained weights
model = CLIPBinaryClassifier().to(DEVICE)
model.load_state_dict(torch.load("best_clip_detector.pth", map_location=DEVICE))
model.eval()

# Preprocess an image
transform = A.Compose([
    A.Resize(224, 224),
    A.Normalize(),
    ToTensorV2()
])

image = np.array(Image.open("example.jpg").convert("RGB"))
image_tensor = transform(image=image)["image"].unsqueeze(0).to(DEVICE)

# Run inference
with torch.no_grad():
    logits = model(image_tensor)
    probs = torch.softmax(logits, dim=1)
    prediction = logits.argmax(1).item()

label_map = {0: "real", 1: "ai-generated"}
print(label_map[prediction], probs.cpu().numpy())

Limitations

• The training data is derived from ImageNet-based real images paired with a fixed set of seven generative model families; performance on generators, domains, or image types outside this distribution is not established.
• The dataset is referred to as "tiny-genimage," implying it is a reduced-scale subset of a larger dataset; results may not generalize to the full-scale version.
• The classification head was trained on frozen CLIP embeddings only; the underlying CLIP encoder was not fine-tuned, which may limit adaptation to subtle generation artifacts not captured by general-purpose CLIP features.
• No adversarial robustness testing was performed against images specifically crafted to evade detection.

Citation

If you use this model, please cite the underlying CLIP and dataset resources:

CLIP backbone: laion2b_s34b_b79k (OpenCLIP, LAION)
Dataset: tiny-genimage (yangsangtai/tiny-genimage)