Ro-BustAI

Ro-BustAI is a robust AI-image detector combining spatial & frequency analysis to distinguish real from synthetic content, maintaining reliable detection even after real world image transformations

Comment

Ro-Bust AI: AIGC Image Detector

Problem Statement Analysis

The challenge requires more than simply distinguishing AI-generated images from authentic images. A successful prototype must also generalise to unseen data, achieve strong detection performance within limited compute, and remain robust after realistic image transformations.

We therefore identified three key requirements:

1. Generalisation:

The model must learn meaningful differences between real and AI-generated images rather than memorising training samples.

2. Effective Detection:

The model must achieve strong classification performance while remaining within the hackathon's <2B parameter and limited-compute constraints.

3. Robustness:

Detection performance should remain reliable when images undergo post-processing or redistribution.

Pain Point Identification

These requirements translate into three main pain points:

  1. Limited generalisation → Data quality and training strategy Reliable detection requires exposure to both real and AI-generated images, from multiple generators, with unseen data reserved for evaluation, and no single source of training data allowed to dominate.

  2. Detection performance → Feature extraction capability A single approach may not capture all relevant visual signals while remaining computationally efficient. The model therefore needs to extract strong and complementary image features.

  3. Performance degradation after transformation → Robustness Post-processing can alter the visual information used for detection without changing whether an image is real or AI-generated. The model must therefore be tested and strengthened against image transformations — including ones that change the file format itself, which we found could otherwise become a hidden shortcut the model exploits instead of learning genuine synthetic-vs-real cues.

These pain points directly informed our three solution components: a diverse, deliberately-balanced real-and-AI training dataset, a ResNet-18 + frequency-domain dual-branch detection approach, and a robustness-aware transformation pipeline.

Key Features of Proposed Solution

1. Diverse Multi-Source Training Dataset

To improve generalisation across different types of AI-generated images, our model combines three real and synthetic image datasets:

  • CIFAKE — provides labelled real and AI-generated images for baseline training.
  • SID_Set — introduces a larger and more diverse collection of real and synthetic images. Due to its ~140 GB size, the dataset is streamed and randomly sampled rather than fully downloaded.
  • Midjourney / DALL-E / Stable Diffusion Dataset — introduces images from modern generative models to address generalisation gaps identified during testing.

Rather than relying on each dataset's own predefined train/test split — which used three different, inconsistent ratios (CIFAKE's own Kaggle folder split, SID_Set's original split, and the Midjourney/DALL-E/SD set's own Hugging Face split) — we pool every image from each source and re-split it ourselves, stratified by class, at one consistent 70% train / 30% test ratio. A further 15% of the training portion is held back internally as a validation set, used only for early stopping and checkpoint selection, so the reported 70/30 train/test ratio itself is unaffected.

Because CIFAKE is numerically far larger than the other two sources, training uses a weighted random sampler so every dataset contributes equally to gradient updates regardless of its raw size. Without this, gradient updates would mostly reflect CIFAKE's older, lower-resolution generator, and the model would under-learn the modern-generator examples (SID_Set, Midjourney/DALL-E/SD) that matter most for generalising to unseen generators.

2. Dual-Branch Spatial + Frequency Detection

Our detector combines two complementary methods of analysing an image:

  • Spatial Branch — ResNet-18: A pretrained CNN extracts visual features such as textures, edges and image patterns that may distinguish authentic from AI-generated images.

  • Frequency Branch — FFT + CNN: The image is transformed into the frequency domain using a Fast Fourier Transform (FFT). A custom CNN then identifies frequency-level statistical patterns that may not be apparent from the image's pixels alone.

Features from both branches are combined before the final real/AI classification, allowing the model to use both what the image visually looks like and how its underlying frequency patterns behave.

3. Robustness-Aware Training & Evaluation

To reflect the post-processing scenarios specified in the problem statement, the model is To reflect the post-processing scenarios specified in the problem statement, the model is evaluated against the six official transformation categories: JPEG compression, Gaussian blur, resizing, Gaussian noise, colour jitter and center cropping.

Training additionally incorporates WEBP re-encoding, PNG round-tripping, rotation and horizontal flipping to improve resilience to other transformations encountered during testing.

A key refinement here: file-format re-encoding is applied as its own guaranteed augmentation stage on every training image, rather than being just one option among many in a single random pick. This matters because real-photo datasets are overwhelmingly JPEG-sourced while generator datasets are overwhelmingly saved as lossless PNG — under weaker augmentation, the model can shortcut-learn "has JPEG compression artifacts → real" instead of learning genuine synthetic-vs-real cues, a shortcut that breaks immediately on a real photo saved as WEBP/PNG or an AI image re-saved as JPEG. Forcing every image through a randomly-chosen format on every epoch decorrelates file format from label. We also apply EXIF-orientation correction to both training data and uploaded inference images, since phone photos are commonly stored sideways via an EXIF tag, which is its own source of misclassification if left uncorrected.

Training augmentations and evaluation transformations are kept separate, allowing robustness to be measured under fixed conditions rather than only on randomly augmented training data. The evaluation table also reports the WEBP and PNG conditions specifically, since earlier iterations of the pipeline applied those formats during training but never actually measured performance on them.

4. Development & Technical Stack

The prototype was developed using Google Colab (T4 GPU runtime) with PyTorch as the primary deep-learning framework.

Development tools: Google Colab (Jupyter notebook interface, T4 GPU runtime); Kaggle account (authentication for CIFAKE download); Hugging Face account (authentication for SID_Set/Midjourney+ dataset streaming).

Models used: ResNet-18 (ImageNet-pretrained, loaded via timm) as the spatial feature branch; a custom lightweight CNN operating on the log-magnitude FFT spectrum as the frequency-domain branch. The full dual-branch detector totals 11.3M parameters — well under the competition's 2B-parameter ceiling. No external LLM/API-based model is used for the detection task itself.

APIs used: Kaggle API (via kagglehub) to download CIFAKE; Hugging Face Hub API (datasets, huggingface_hub) to stream SID_Set and load the Midjourney/DALL-E/SD parquet files.

Libraries and frameworks: PyTorch and torchvision (model, training loop, data pipeline); timm (pretrained ResNet-18 backbone); Hugging Face datasets / huggingface_hub (dataset streaming and access); kagglehub (Kaggle dataset download); scikit-learn (roc_auc_score, accuracy_score); NumPy and Pillow (image processing and augmentation — JPEG/WEBP/PNG re-encoding, blur, noise, rotation, EXIF correction); Matplotlib (training-curve and robustness-heatmap visualisations).

Datasets and assets used: CIFAKE (Kaggle), SID_Set (Hugging Face, streamed), and the Midjourney/DALL-E/Stable Diffusion dataset (Hugging Face) — see Section 1 above for how they're combined and split.

Training uses cross-entropy loss with label smoothing and the Adam optimizer, with separate learning rates for the pretrained ResNet-18 branch (fine-tuned gently) and the newly-initialised frequency branch and classifier (trained faster). A cosine learning-rate schedule, gradient clipping, and validation-based early stopping (patience of 4 epochs, monitored on validation AUC) are used to prevent overfitting on later epochs. Only the checkpoint with the best validation AUC is kept as the final model, rather than whichever epoch training happened to stop on. All random seeds are fixed for reproducibility.

Robustness Evaluation Summary

(refer the media attached.)

Model Behaviour & Trade-offs

Our dual-branch design and robustness-first training involve several explicit trade-offs:

  • Compute vs. signal diversity — running two branches (spatial CNN + frequency CNN) costs more compute per image than a single-branch classifier, but captures artifacts that are only visible in one domain or the other; a purely spatial model would miss frequency-only tells, and vice versa.
  • Robustness vs. clean-image accuracy — always re-encoding every training image through a random file format, and stacking additional distortions on top, makes the training distribution deliberately harder than the clean test distribution. This is expected to cost a small amount of accuracy on clean, unmodified images in exchange for meaningfully better, more consistent accuracy across the WEBP/PNG/JPEG-mixed conditions the task specifically evaluates.
  • Cross-source balance vs. best-single-source performance — weighting SID_Set and the Midjourney/DALL-E/SD set equally against the much larger CIFAKE pool trades a small amount of CIFAKE-specific accuracy for better generalisation to the modern, photorealistic generators that matter more for real-world deployment.

Current Scope & Future Development

Within the hackathon timeframe and limited-compute environment, our solution focuses on demonstrating a robust proof of concept rather than a production-ready detection system.

Given more time, we would extend the solution by:

  • Expanding generator diversity to include newer and unseen generative models. With training now balanced across sources via the weighted sampler, additional data should translate directly into generalisation gains rather than being drowned out by the largest dataset, as it would have been in an earlier version of the pipeline.
  • Extending robustness testing across additional transformation combinations and severity levels to better simulate repeated real-world redistribution.
  • Optimising decision thresholds and calibration to further reduce false positives while maintaining strong AI-image detection.
  • Adding explainability mechanisms to indicate which spatial or frequency-domain characteristics contributed most strongly to a prediction.

These improvements build directly on the current dual-branch architecture and robustness pipeline, providing a clear pathway from the hackathon prototype toward a more generalisable real-world detection system.

Conclusion

Our solution addresses the challenge of AI-generated image detection by focusing on three key areas: generalisation, detection performance and robustness.

By combining diverse, deliberately-balanced real and AI-generated datasets, a dual-branch spatial and frequency-domain architecture, and robustness-aware training and evaluation, our prototype goes beyond clean-image classification to consider how detection performs after realistic image transformations — including transformations that change the file format itself.

Within the constraints of a hackathon-scale prototype, the solution demonstrates a practical and computationally efficient approach to robust AI-image detection, while providing a clear foundation for further improvements in robustness and explainability.

Built With

Share this project:

Updates