|
|
--- |
|
|
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}] |
|
|
|
