Asteria — Robust AI Image Detection Under Real-World Transformations
Detecting AI-generated images is not enough: detection should remain useful after images are blurred, compressed, resized, or reposted.
Challenge: TikTok TechJam 2026, Problem 5 — Robust Detection of AI-Generated Images Under Real-World Transformations.
Code and English README: sunzk111/ai-image-detector · Workflow: English Colab notebook · Trained model and configuration: GitHub Release
Inspiration
An image rarely reaches a viewer in its original form. Messaging apps recompress it, platforms generate thumbnails, and users crop or edit it. A detector that works only on pristine images can therefore look convincing in a laboratory while failing in ordinary use.
Asteria addresses this gap with a compute-conscious, image-level detection prototype. Its emphasis is not a larger model or a single headline accuracy number, but a practical combination of diverse training data, transformation-aware learning, explicit decision thresholds, and detailed error reporting. False accusations against authentic images matter alongside missed synthetic images.
What it does
Asteria assigns an AI-generation score to an image. Its standalone predict.py accepts an unlabeled image directory, recursively processes supported images, and exports a JSON array with image_path and pred for each image. Here, pred is a continuous AI score, not a thresholded 0/1 label; a documented threshold can be applied separately when a binary Real/AI decision is needed.
The existing evaluate.py remains a separate, unchanged robustness-evaluation entry point. It tests clean images and 15 transformed conditions on labeled data, exports per-image scores, and reports accuracy, AUROC, F1, class recall, false positives, and false negatives. This separation supports both submission-time directory inference and reproducible benchmark analysis without changing the model or its training procedure.
The intended use is an auxiliary signal for content review and investigation—not conclusive proof of an image's origin and not an autonomous moderation decision. This is a single-image hackathon prototype, not a production platform, video detector, or audio detector.
How it was built
Model and preprocessing
The backbone is Hugging Face facebook/dinov2-base, fine-tuned for binary classification. The implemented model has approximately 87.4 million parameters, comfortably below the challenge's 2-billion-parameter limit.
Images are converted to RGB, resized with bicubic interpolation to 448 × 448, and normalized with mean [0.485, 0.456, 0.406] and standard deviation [0.229, 0.224, 0.225]. Alpha-channel images are composited onto a white background.
The representation concatenates the CLS token with the mean of the patch tokens. A compact head applies LayerNorm, a 512-unit linear layer, GELU, dropout 0.2, and a final linear layer producing one logit. A sigmoid produces the AI score:
RGB image → 448×448 preprocessing → DINOv2
→ [CLS token ; mean patch token] → MLP → logit z
→ score s = sigmoid(z)
label 0 = Real; label 1 = AI
AI decision if s >= threshold
Submission JSON pred = s, without applying the threshold
The score is useful for ranking, but it is not claimed to be a calibrated probability of real-world provenance. A score of 0.9 should not automatically be interpreted as a universally valid 90% probability.
Training data: diversity within a 100,000-image budget
Training only on CIFAKE risks learning source- or generator-specific shortcuts. Asteria therefore uses three sources while capping training at 100,000 images, balanced between 50,000 Real and 50,000 AI.
| Source | Real training | AI training | Additional internal validation |
|---|---|---|---|
| CIFAKE | 5,000 | 5,000 | 500 |
| SID_Set | 5,000 | 5,000 | 500 |
| WildFake | 40,000 | 40,000 | 3,000 |
| Total | 50,000 | 50,000 | 4,000 |
The WildFake real-image allocation is ImageNet 20,000; Church 6,000; AFHQ 6,000; FFHQ 4,000; and CelebA-HQ 4,000. The synthetic allocation is 8,000 images each from the DDIM, DDPM, ADM, Imagen, and VQDM archives under Diffusion_based.
CIFAKE and SID_Set each contribute 10% of training, rather than allowing the small-resolution CIFAKE source to dominate. SID_Set uses binary labels 0 and 1 only; label 2, representing tampering, is excluded from this full-image Real/AI task.
The preparation script explicitly excludes all WildFake COCO and DALL·E archives from training. The official demonstration subset is stored separately: 4,998 COCO val2017 real images and 8,843 DALL·E Advanced/DALLE3 images, totaling 13,841. These demonstration images were not used for gradient-based training or epoch checkpoint selection.
Preparation maintains resumable state, verifies source quotas, and removes exact decoded-RGB duplicates across training, internal validation, and the existing demonstration directory. The preparation record indexed all 13,841 demonstration images. This is an exact-duplicate safeguard, not a guarantee against resized, re-encoded, or semantic near-duplicates. Source exclusion also does not establish that every generator is absent from every other upstream dataset or from backbone pretraining.
Training for blur, JPEG compression, and combined degradation
Each training sample produces two views: the clean image and an on-the-fly degraded image. Both share the same label. The model processes the two views together and learns from:
L = BCE(clean_logit, label)
+ BCE(degraded_logit, label)
+ 0.2 × mean[1 − cosine_similarity(clean_feature, degraded_feature)]
The degraded-view classification loss directly trains recognition after information loss. Feature consistency encourages a stable representation between the original and its transformed version, while the clean-view loss preserves sensitivity to unmodified images. These are design motivations; a controlled ablation has not isolated each component's contribution.
The implementation includes the challenge's transformation ranges:
| Transformation | Implemented settings |
|---|---|
| JPEG | Quality 90, 70, 50, 30; re-encode and decode |
| Gaussian blur | Sigma/radius 0.5, 1.0, 2.0, using Pillow GaussianBlur |
| Resize | Downsample to 0.5× or 0.25×, then upscale |
| Gaussian noise | Standard deviation 0.02, 0.05, 0.10 on normalized RGB values |
| Color adjustment | Brightness, contrast, saturation within ±20% |
| Center crop | Keep 80% of each spatial dimension, then resize back |
Transformations are applied before model-input resizing, approximating edits that occur before an image reaches a detector. Real and AI images use the same transformation policy, reducing the risk that degradation itself becomes a class cue.
A curriculum increases difficulty using normalized epoch progress:
- Early stage, progress below 0.25: zero or one transformation, with mild settings such as JPEG q90 and blur sigma 0.5.
- Middle stage, progress from 0.25 to below 0.60: one or two transformations, using light/intermediate settings such as JPEG q90/q70 and blur sigma 0.5/1.0.
- Late stage, progress at least 0.60: one to three transformations, with the full ranges available, including JPEG q30 and blur sigma 2.0.
For the configured 10-epoch schedule, these stages correspond to epochs 1–3, 4–6, and 7–10. Degradation types are sampled without replacement per image. Thus, severe blur and JPEG artifacts receive explicit training coverage, including compound cases. However, the reported run does not implement extra blur-specific oversampling, a special JPEG-only loss, a frequency branch, or a deblurring module.
Optimization and engineering choices
| Setting | Current configuration |
|---|---|
| Epochs | 10 |
| Training batch size | 16 original images; 32 views after clean/degraded pairing |
| Gradient accumulation | 1 |
| Optimizer | AdamW; weight decay 0.05 |
| Backbone / classifier learning rate | 1e-5 / 1e-4 |
| Schedule | 10% warmup, followed by cosine decay |
| Backbone freezing | First epoch only |
| Gradient clipping | Maximum norm 1.0 |
| Evaluation batch size | 32 |
| Seed | 42 |
| Gradient checkpointing | Disabled |
The successful workflow ran in Google Colab on an NVIDIA A100-SXM4-40GB, using PyTorch 2.11.0+cu128 and BF16. Mixed precision and a fixed dataset budget keep the experiment feasible on one GPU. No multi-GPU or latency benchmark is claimed. Downloads use CPU/network resources rather than GPU compute.
Development tools include Colab/Jupyter, GitHub, and Codex-assisted implementation/documentation iteration. The stack uses Python, PyTorch, torchvision, Hugging Face Transformers and Datasets, NumPy, Pillow, scikit-learn, PyYAML, tqdm, kagglehub, modelscope-hub, and SQLite for resumable data preparation. Model and dataset hubs supply resources; the classifier does not call a hosted generative-model API to make each detection.
Threshold selection: training policy versus the final demonstration setting
Automatic internal calibration during training
The additional 4,000 internal validation images are class-balanced and disjoint from training. A fixed split based on sample IDs and seed 42 reserves 1,000 images for threshold calibration and 3,000 for epoch evaluation.
At each epoch, candidate decision boundaries are evaluated on the 1,000 calibration images to maximize balanced accuracy. Equally good choices prefer the threshold closest to the configured default, 0.5. That threshold is then evaluated on the other 3,000 images. Their accuracy determines the best checkpoint. Checkpoints retain decision_threshold and calibration metadata, and the split fingerprints are recorded for reproducibility.
This separates fitting the decision threshold from choosing the epoch on the same internal evaluation pass. It does not turn repeatedly used internal validation data into an untouched final test set.
Final demonstration operating point
Subsequent development exposed a mismatch between strong ranking performance and threshold-dependent classification accuracy on the demonstration distribution. Asteria therefore added per-image logits/scores, threshold comparison tables, and score-distribution diagnostics. Lowering a threshold makes more predictions AI-positive: it can reduce missed AI images but may increase false alarms on real images. Threshold selection changes the operating point, not the learned representation or the underlying score ordering.
The final report uses one fixed threshold, 0.000005 (5e-6), for all 16 conditions. It was selected through manual development-time comparison of demonstration results, rather than by the automatic internal calibration procedure. It is not asserted to be the globally optimal threshold or appropriate for every future model and dataset.
Evaluation disclosure: the COCO/DALL·E subset was not used to train model weights, but its results informed this manual threshold choice. The numbers below must therefore be interpreted as development-time demonstration results, not unbiased held-out test estimates. The organizers describe this subset as a reference benchmark that does not contribute to the final score. An independent, previously unused evaluation set is needed for a stronger generalization claim.
The evaluator computes sigmoid scores from FP32 logits after model inference and reports both probability-based and raw-logit AUROC. This makes saturation/precision effects visible without conflating high AUROC with high accuracy at a particular threshold. Newly trained checkpoints retain their own internally calibrated thresholds; they do not automatically inherit 5e-6.
Final robustness results
The supplied final robustness.csv contains 16 conditions, each evaluated on the same 13,841 images, using threshold 5e-6. F1 and recall treat AI as the positive class. Real FPR is the fraction of authentic images incorrectly marked as AI. Rates are shown as percentages.
| Condition | Accuracy | AUROC | F1 (AI) | AI recall | Real FPR |
|---|---|---|---|---|---|
| clean | 96.55% | 99.50% | 97.25% | 95.62% | 1.82% |
| jpeg_q90 | 96.77% | 99.65% | 97.42% | 95.53% | 1.04% |
| jpeg_q70 | 96.53% | 99.72% | 97.22% | 94.95% | 0.66% |
| jpeg_q50 | 96.40% | 99.73% | 97.11% | 94.67% | 0.54% |
| jpeg_q30 | 95.63% | 99.61% | 96.47% | 93.46% | 0.54% |
| blur_sigma0.5 | 95.79% | 99.12% | 96.65% | 94.98% | 2.78% |
| blur_sigma1 | 93.90% | 98.50% | 95.09% | 92.32% | 3.30% |
| blur_sigma2 | 91.03% | 97.13% | 92.72% | 89.38% | 6.04% |
| resize_0.5x | 96.00% | 99.41% | 96.78% | 94.27% | 0.94% |
| resize_0.25x | 93.87% | 98.33% | 95.09% | 92.91% | 4.44% |
| noise_sigma0.02 | 96.24% | 99.52% | 96.99% | 94.87% | 1.32% |
| noise_sigma0.05 | 95.86% | 99.53% | 96.67% | 94.02% | 0.88% |
| noise_sigma0.1 | 95.79% | 99.59% | 96.60% | 93.77% | 0.64% |
| color_jitter_minus0.2 | 93.87% | 99.08% | 95.02% | 91.42% | 1.78% |
| color_jitter_plus0.2 | 96.19% | 99.24% | 97.02% | 96.92% | 5.10% |
| center_crop_0.8 | 95.11% | 99.38% | 96.04% | 92.94% | 1.06% |
Summary: clean accuracy is 96.55%, the equally weighted mean accuracy over the 15 transformed conditions is 95.27%, and worst-condition accuracy is 91.03% under blur sigma 2.0. The transformed-condition mean excludes clean and is not a pooled AUROC or an official competition score. The transformed images are correlated versions of the same originals, not 15 independent datasets.
JPEG q30 reduces accuracy by 0.92 percentage points relative to clean, whereas blur sigma 2.0 reduces it by 5.51 points. Strong downscaling and negative color adjustment each cause approximately 2.7-point drops. These comparisons identify where the present system is less robust; they do not establish a causal improvement over an earlier model, since no controlled before/after ablation is included in this final CSV.
Error analysis and trade-offs
The confusion counts help distinguish false-alarm problems from missed-detection problems:
| Condition | Real → AI: false positives | AI → Real: false negatives | Interpretation |
|---|---|---|---|
| Clean | 91 / 4,998 | 387 / 8,843 | High ranking quality still leaves errors at the chosen operating point. |
| JPEG q30 | 27 / 4,998 | 578 / 8,843 | Missed AI images rise despite a low real-image false-alarm rate. |
| Blur sigma 2.0 | 302 / 4,998 | 939 / 8,843 | Both error types increase; this is the most important robustness gap. |
| Positive color adjustment | 255 / 4,998 | 272 / 8,843 | Better AI recall comes with more authentic-image false alarms. |
Clean AI precision is 98.94%, real recall is 98.18%, and balanced accuracy is 96.90%. These measures supplement raw accuracy because the demonstration set is not class-balanced.
Under strong blur, AUROC also falls, from 99.50% to 97.13%. This indicates reduced score separation, not just a poorly positioned threshold. Recovering robustness therefore requires better representations or training coverage; moving the threshold alone cannot recover information removed by blur.
A plausible mechanism is that blur suppresses texture and local evidence, while JPEG modifies fine-scale structure. This remains a hypothesis: the aggregate CSV cannot identify individual semantic failure modes or prove which visual cues the network uses. Per-image predictions are exported for subsequent inspection of representative false positives and false negatives; the aggregate report is not a substitute for viewing those actual examples.
Challenges, accomplishments, and lessons learned
- Balancing diversity and compute: a capped, balanced mixture broadened the training sources without making data volume unbounded. Generalization still needs independent evaluation; diversity alone is not proof of it.
- Handling downloads and ephemeral runtimes: selective archive access, resumable preparation state, manifests, and separate checkpoint backups made the Colab workflow more manageable. The trained checkpoint is now distributed as a GitHub Release asset, separately from source code; cloning the repository does not download the model or datasets.
- Separating ranking from decisions: high AUROC did not guarantee a useful default classification threshold. Explicit calibration and diagnostic exports made the difference observable.
- Training for degradation without discarding clean evidence: paired clean/degraded supervision made robustness an explicit training objective. However, the current results do not isolate the benefit of consistency loss from dataset or threshold changes.
- Reporting beyond the best number: the final evaluation retains all 15 transformation conditions, class-specific errors, and the weakest condition rather than reporting clean accuracy alone.
- Making the workflow accessible: the README, dataset guide, and Colab notebook are available in English, with verified release download links, checksum checks, separate training/evaluation paths, and instructions for the standalone JSON prediction interface.
What is next: targeted improvements, not claims about this run
- Prioritize severe blur in the degraded-view sampler. Increase exposure to sigma 1–2 cases and blur-plus-resize/compression combinations, while keeping the clean branch and identical policies for both labels. Reuse the current baseline as a controlled comparison.
- Continue JPEG q30 coverage with hard-example analysis. Its remaining misses justify inspection, but strong blur is the higher priority given the larger loss of accuracy and AUROC.
- Measure the effect with ablations. Compare the existing sampler, blur-focused sampling, and altered consistency weighting under the same data split and training budget; track clean accuracy, worst-condition performance, AI recall, and real FPR.
- Validate transfer independently. Add generator-disjoint and source-disjoint checks, stronger near-duplicate auditing, unseen transformation combinations, and real reposting pipelines.
- Improve deployment calibration and usability. Fit any future calibration on representative internal development data, freeze the operating point before independent testing, and measure memory, latency, and throughput. A review/abstention policy for uncertain cases is a future feature, not part of the current implementation.
Reproducing the current workflow
The public repository includes the model implementation, data pipeline, training loop, calibration module, unchanged robustness evaluator, and an English Colab notebook. The trained best.pt and its matching configuration are available from the model release. Datasets must be obtained separately from their upstream sources.
There are three distinct workflows: unlabeled directory → JSON, labeled robustness evaluation, and training a new model. Neither inference workflow requires downloading the training set or retraining the detector. Place the companion predict.py in the project root alongside model.py, augmentations.py, and utils.py before using the JSON interface.
1. Install and download the released model
For a fresh Linux or Colab setup:
git clone https://github.com/sunzk111/ai-image-detector.git
cd ai-image-detector
python -m pip install -r requirements.txt
python -m pip install kagglehub modelscope-hub
mkdir -p outputs/dinov2_mixed100k_v2
curl -fL --retry 3 \
https://github.com/sunzk111/ai-image-detector/releases/download/model/best.pt \
-o outputs/dinov2_mixed100k_v2/best.pt
curl -fL --retry 3 \
https://github.com/sunzk111/ai-image-detector/releases/download/model/config.mixed100k.yaml \
-o config.release.yaml
These download commands are intended for a fresh checkout: do not overwrite an existing experiment's weights. The release configuration is deliberately saved as config.release.yaml so that the repository's original training configuration, config.mixed100k.yaml, remains separate. Use current main for the updated English documentation and notebook; the release's automatic source-code archives correspond to its older tag snapshot.
The checkpoint is approximately 1 GB. Verify downloaded files with sha256sum; the expected SHA-256 values are:
best.pt f1e2d5470db116f59a64f65d8e3b7bccf2fd5fdec7c163023dbbd7b6d35b9508
config.release.yaml feb1765f6d18db1f0853f22669c7fdb8fd79204380e43774eb308138557849a4
The notebook performs these checks and refuses to replace a different existing checkpoint. Load only trusted checkpoints. Model initialization also obtains the upstream DINOv2 backbone from Hugging Face. Colab does not require .venv activation.
2. Submission inference: image directory → JSON
Run the standalone script on an arbitrary image directory; no labels, manifests, class folders, or demonstration data are required:
python predict.py --config config.release.yaml \
--checkpoint outputs/dinov2_mixed100k_v2/best.pt \
--input-dir /path/to/images --output-json predictions.json
The output schema is illustrated below; these example scores are placeholders, not measured predictions:
[
{"image_path": "example.jpg", "pred": 0.123},
{"image_path": "nested/another.png", "pred": 0.987}
]
image_pathis relative to the input directory and uses forward slashes. Recursive scanning supports JPG/JPEG, PNG, WEBP, BMP, and TIF/TIFF, with stable path ordering.predis the continuous AI score in[0, 1], computed with FP32 sigmoid after model inference. No decision threshold is applied to the JSON output. For a separate binary decision with this released checkpoint, compare the score with0.000005; do not replacepredwith a binary label.- The script reuses the checkpoint's model configuration and existing clean-image preprocessing. It does not retrain the model, alter weights, or modify
evaluate.py. --device cpu --num-workers 0supports CPU inference; lower--batch-sizeif memory is limited. The default worker count is 0.- Empty directories, unreadable images, nonfinite scores, and existing output files produce explicit errors. Symlinked image files are rejected, directory symlinks are not followed, and JSON is written after all images have been scored successfully.
The new export path passed eight local tests, including recursive path handling, preprocessing, continuous-score output, JSON writing, corrupt-image handling, and an end-to-end run with a small fixture model/checkpoint. This validates the export plumbing, not a new benchmark result: the released DINOv2 checkpoint has not yet been rerun through this new script on A100. The robustness results above remain those from the original evaluator.
3. Reproduce the released model's robustness benchmark
Obtain the COCO/DALL·E demonstration subset using the README or notebook. The release configuration expects data/wildfake_eval with immediate coco/ and DALLE/ class directories; nested image directories are supported. If your images are elsewhere, update data.validation_demo.path in config.release.yaml. Keep this benchmark separate from training data.
python evaluate.py --config config.release.yaml \
--checkpoint outputs/dinov2_mixed100k_v2/best.pt \
--source validation_demo
The configuration uses evaluation.threshold_source: checkpoint; this reads the released model's saved decision_threshold, 5e-6, rather than overriding it with a YAML number. Check the printed Decision threshold line. Appending --threshold 0.000005 explicitly fixes the same operating point for this report only and does not modify the checkpoint.
For a quick loading/pipeline check, append --conditions clean jpeg_q30 blur_sigma2 --max-samples 64; omit both options for the complete 13,841-image, 16-condition benchmark. A smoke check is not a complete benchmark result.
Each evaluation writes a separate run directory with robustness.csv/json, predictions.csv, threshold_sweep.csv, score_distributions.csv, and run_status.json. A full run also refreshes the top-level robustness reports beside the checkpoint. Saved predictions can be rescored without another GPU inference pass:
python evaluate.py \
--from-predictions outputs/dinov2_mixed100k_v2/evaluation_runs/<run_id>/predictions.csv \
--threshold 0.000005
Replace <run_id> with an existing evaluation directory. Exact reproduction of the reported numbers requires the corresponding released checkpoint, the same data and processing, and the recorded threshold—not just a fresh training run.
4. Train a new model separately
Skip this section if only using the released model. Prepare the demonstration images first for the exact-overlap check, then prepare the training set using the repository's original training template, not config.release.yaml:
python -u prepare_mixed_dataset.py --base-config config.mixed100k.yaml --workers 8
# After preparation, choose a fresh output_dir in config.mixed100k.yaml
# if the default directory already contains released weights or an older run.
python train.py --config config.mixed100k.yaml
Reusing the original config.mixed100k.yaml template preserves the training-time calibration policy and default tie-break threshold of 0.5. Preparation regenerates the training configuration and resets output_dir and training.resume, so set experiment-specific output/resume options after preparation. Re-running preparation can resume from its cache; --existing-only is available when all source quotas are already complete.
For a new model, use its training configuration and actual checkpoint path:
python evaluate.py --config config.mixed100k.yaml \
--checkpoint outputs/dinov2_mixed100k_v2/best.pt \
--source validation_demo
Update --checkpoint if you changed output_dir. Do not automatically append the released model's manual 5e-6 threshold: a new model should retain its own internally calibrated threshold. A new training run need not reproduce the exact saved model or the same optimal threshold; dependencies are not fully pinned, and GPU operations are not configured for strict determinism.
The English notebook defaults to evaluating the release and skips training unless TRAIN_FROM_SCRATCH is enabled. Its optional JSON-inference section uses the separate predict.py. The optional Drive backup includes code, weights, and reports but excludes data/; datasets require a separate backup before an ephemeral Colab runtime is discarded.
Resources and attribution
- Backbone: DINOv2 Base.
- Data: CIFAKE, SID_Set, and WildFake.
Datasets, pretrained weights, and third-party software remain subject to their respective permissions and licenses. Source provenance is acknowledged; this document does not claim that all underlying components were created from scratch.
Built With
- colab
- dinov2
- face
- hugging
- numpy
- pillow
- pytorch
- scikit-learn
- sqlite
- torchvision
- transformers
Log in or sign up for Devpost to join the conversation.