---
license: apache-2.0
library_name: timm
pipeline_tag: image-classification
tags:
  - image-classification
  - ai-detection
  - vit
datasets:
  - your-username/ai-generated-vs-real
metrics:
  - accuracy
  - f1
---

# AI Source Detector (ViT-Base)

Detects *and* classifies the source of AI-generated images into **five** classes  
(`stable_diffusion`, `midjourney`, `dalle`, `real`, `other_ai`).

## Model Details
* **Architecture:** ViT-Base Patch-16 × 224  
* **Parameters:** 86 M  
* **Fine-tuning epochs:** 10  
* **Optimizer:** AdamW (lr = 3e-5, wd = 0.01)  
* **Hardware:** 1× NVIDIA RTX 4090 (24 GB)

## Training Data
| Class | Images |
|-------|-------:|
| Stable Diffusion | 12 000 |
| Midjourney | 10 500 |
| DALL-E 3 | 9 400 |
| Real | 11 800 |
| Other AI | 8 200 |

Total ≈ 52 k images - 80 % train / 10 % val / 10 % test.

## Evaluation
| Metric | Top-1 | Macro F1 |
|--------|------:|---------:|
| Validation | 92.8 % | 0.928 |
| Test | 91.6 % | 0.914 |

<details>
<summary>Confusion Matrix (click to open)</summary>
<img src="confusion_matrix.png" width="480">
</details>

## Usage
```python
from transformers import ViTImageProcessor, ViTForImageClassification, pipeline
classifier = pipeline(
    task="image-classification",
    model="yaya36095/ai-source-detector",
    top_k=1
)
classifier("demo.jpg")
# → [{'label': 'stable_diffusion', 'score': 0.97}]