realphongha/danbooru-tag-query

DanbooruTagQuery

Lightweight multi-label anime image tagger using DINOv3 ViT backbone and a cross-attention tag query head, trained on Danbooru images.

• 11,424 tags (general + character + copyright)
• 448×448 input, ONNX-exported

GitHub repo
You can try our models here: realphongha/DanbooruTagQuery

Model description

Image → ViT backbone (DINOv3) → patch tokens → cross-attention → per-tag logits.

Each tag is a single learnable embedding vector (a "query"). An nn.MultiheadAttention layer computes attention between all tag queries and all ViT patch tokens; each query's attended feature is projected through a scalar linear classifier. No MLP, no positional encoding — the query itself encodes "what to look for". This design is modular: the head can be swapped without touching the dataset, training loop, or metrics.

Input image (448×448)
    │
    ▼
┌─────────────────────┐
│        DINOv3       │  pretrained ViT backbone
└─────────┬───────────┘
          │  tokens: (B, N_patches+5, D)
          ▼
┌─────────────────────┐
│   Tag Query Head    │
│                     │
│  tag_queries:       │  learned (num_tags, D)
│  (num_tags, D) ─────┼─→ cross-attention ──→ tag features (B, num_tags, D)
│                     │     queries attend to ViT patch tokens
│  classifier:        │
│  Linear(D→1) ───────┘  → logits (B, num_tags)
└─────────────────────┘
          │
          ▼
  sigmoid(logits) → per-tag probabilities

Evaluation

Evaluated against popular community taggers on the intersection evaluation subset (3,383 tags present in every model's vocabulary), drawn from a held-out danbooru2025 test set of images with post ID > 7220105 (avoids contamination with WD-SwinV2 training data). Each model ran at its native resolution; thresholds searched over [0.10, 0.15, …, 0.95].

Model	Params	Input	Latency	mAP	Macro F1	Micro F1	Best threshold
Ours (L/16) - trained from scratch	319.0M	448×448	36.6ms	0.5352	0.4775	0.6884	0.20
WD-eva02-large-tagger-v3	315.2M	448×448	50.3ms	0.4822	0.4344	0.6684	0.30
Ours (B/16) - trained from scratch	96.8M	448×448	24.9ms	0.4693	0.4195	0.6684	0.20
WD-SwinV2-tagger-v3	98.0M	448×448	35.8ms	0.4603	0.4140	0.6474	0.15
ML-Danbooru	68.9M	448×448	34.0ms	0.4023	0.3490	0.5952	0.60
JoyTag	91.5M	448×448	20.2ms	0.3783	0.3429	0.6179	0.35
DeepDanbooru (CNN)	161.0M	512×512	33.6ms	0.2100	0.1920	0.4692	0.15

Key results:

• Ours (L/16) leads all baselines by +0.053 mAP over WD-eva02 (the strongest comparable model) with 33% faster inference (36.6ms vs 50.3ms).
• Ours (B/16) outperforms all models in the <100M class (WD-SwinV2 +0.009 mAP, JoyTag +0.091 mAP) at 24.9ms latency.

Latency measured on a single NVIDIA RTX 5090 (PyTorch eager for local runs; ONNX Runtime for ONNX-exported models).

Intended uses

• Tagging anime / illustration images (general, character, and copyright tags)
• Generating caption / metadata text for image collections
• Danbooru-style workflows, e.g. injecting a tag preset into image-generation UIs
• Batch CLI tagging of local image folders

Note: this model tags adult (NSFW) content. Danbooru contains explicit imagery and its tag vocabulary reflects that. Do not use in contexts where such content is unacceptable.

Training data

• Metadata source for Danbooru posts/images: trojblue/danbooru2025-metadata
• Filters: deleted/banned/flagged/pending images removed, score > 10, file size 50 KB–20 MB, dimensions ≥ 224×224
• Tags: tag_string_general + tag_string_character + tag_string_copyright unified; tags with frequency ≥ 100; images with ≥ 2 tags
• Vocabulary: 11,424 tags; ~1M training images; last 100K images held out as validation
• Noise tags removed via data/ignored_tags.txt

Future plan

• Care more about the data (de-duplication, remove noise, etc.). Current data cleaning pipeline is too basic?
• Try KD + Transfer Learning from L16 => B16/S16 models.
• Another self-attention block for the head => more semantic for the tags.

Acknowledgements

• DINOv3 (Meta AI) — pretrained ViT backbones (licensed under the DINOv3 License)
• trojblue/danbooru2025-metadata — training data
• Baseline reimplementations live in src/compared_models/: DeepDanbooru, WD-tagger, ML-Danbooru, JoyTag
• timm for backbone loading and ONNX Runtime for deployment