Skip to content

Repository files navigation

TCSR-Net

A deep learning approach to anime character retrieval using gated multimodal fusion with robustness to incomplete and degraded user-provided tags.

Paper: Robust Multimodal Retrieval via Gated Contrastive Fusion: Adaptive Modality Weights for Incomplete Tags

Overview

This project implements a multimodal anime character retrieval system that combines visual and textual information through a learned gating mechanism. The key contribution is demonstrating that adaptive gated fusion preserves retrieval performance even when users provide sparse or incomplete tag information—a realistic scenario in production retrieval systems.

Key Results:

  • R@1 = 0.94 on validation set with full tags
  • Maintains R@1 = 0.67 when tags are 70% degraded
  • ResNet-50 backbone outperforms ViT-small under corruption
  • Gated fusion beats both single-modality baselines across all degradation levels

Hardware Requirements

Recommended (tested):

  • GPU: NVIDIA RTX 5070 Ti (16 GB VRAM)
  • CPU: AMD Ryzen 7 3700X
  • RAM: 32 GB DDR4
  • Storage: ~50 GB (dataset + models)

Minimum viable:

  • GPU: 12+ GB VRAM (NVIDIA RTX 3060 Ti or better)
  • CPU: 8-core modern processor
  • RAM: 16 GB
  • Storage: 50 GB SSD

Note: CPU inference is supported for the demo (presentation.py).

Installation

1. Clone/Setup Environment

# Create conda environment
conda create -n CBIR python=3.10
conda activate CBIR

# Install PyTorch with CUDA support
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

# Install additional dependencies
pip install --break-system-packages \
    transformers timm pillow pyyaml tqdm tensorboard \
    PyQt6

note: the model was traied using CUDA 13.X (nightly)

2. Project Structure (after running)

project_root/
├── data/
│   ├── gallery/              # Reference images
│   ├── distractors/          # Distractor images
│   └── annotations/
│       ├── gallery.csv       # Gallery metadata (image_id, tags, split)
│       ├── mapping.csv       # Query-to-source mapping (query_id, source_id, difficulty)
│       ├── reference_metadata.json
│       └── distractor_metadata.json
├── runs/                     # Checkpoints and logs
│   └── resnet50_gated_mixed_variable_dropout/
│       ├── best.pt
│       ├── args.json
│       └── eval_*.csv
├── architecture_diagrams/    # Generated architecture visualizations
├── model.py                  # Architecture definition
├── data_pipeline.py          # Data loading and preprocessing
├── data_loader.py            # PyTorch DataLoader classes
├── loss.py                   # Loss functions (InfoNCE)
├── train.py                  # Training loop
├── eval.py                   # Evaluation harness
├── presentation.py           # Qt6 interactive demo
└── README.md

Dataset Setup

Option 1: Manual Data

If you have the Danbooru and/or Safebooru data already scraped:

# Place images in the correct structure
data/
├── gallery/
│   ├── 11587418.jpg
│   ├── 11587419.jpg
│   └── ...
├── distractors/
│   └── [distractor images]
└── annotations/
    ├── gallery.csv
    └── mapping.csv

Option 2: Scrape Danbooru

You can scrape danbooru data using download_danbooru_subset.py

make sure that you have a danbooru account and have generated an API key

create a .env file on your project root with these content:

DANBOORU_LOGIN=<USERNAME>
DANBOORU_API_KEY=<API-KEY>

then run download_danbooru_subset.py

note: the data will not be the same as used on the training because there might be newer data when we run this

Option 3: Generate Data Pipeline

# Regenerate queries and splits with mixed difficulty
python data_pipeline.py --data-root ./data --difficulty mixed --drop-unreadable-distractors

Output:

  • data/annotations/mapping.csv — query→source mapping with difficulty levels
  • data/annotations/gallery.csv — gallery metadata (updated splits)
  • Mixed-difficulty queries across: easy, medium, hard, extreme (1:1:1:1 ratio)

Training

Basic Training (ResNet-50 + Gated Fusion)

python train.py \
  --data-root ./data \
  --backbone resnet50 \
  --fusion gated \
  --batch-size 256 \
  --epochs 15 \
  --eval-every 5 \
  --num-workers 8 \
  --tag-drop-min 0.0 \
  --tag-drop-max 0.95 \
  --modality-dropout 0.15 \
  --run-name resnet50_gated_mixed_variable_dropout

Key hyperparameters:

  • --batch-size 256: Larger batches (metric learning tolerates it, 2× faster)
  • --tag-drop-min/max 0.0 0.95: Tags randomly dropped 0–95% during training
  • --modality-dropout 0.15: 15% of training queries get all tags dropped (pure-image path)
  • --eval-every 5: Evaluate on validation set every 5 epochs

Training stats (RTX 5070 Ti):

  • ~80 seconds/epoch
  • Converges by epoch 8–10
  • 15 epochs ≈ 20 minutes total

Ablation Runs

ViT-Small Backbone:

python train.py --data-root ./data --backbone vit_small --fusion gated \
  --batch-size 256 --epochs 15 --eval-every 5 \
  --tag-drop-min 0.0 --tag-drop-max 0.95 --modality-dropout 0.15 \
  --run-name vit_small_gated_mixed_variable_dropout

Concat Fusion (Gate OFF):

python train.py --data-root ./data --backbone resnet50 --fusion concat \
  --batch-size 256 --epochs 15 --eval-every 5 \
  --tag-drop-min 0.0 --tag-drop-max 0.95 --modality-dropout 0.15 \
  --run-name resnet50_concat_mixed_variable_dropout

No Dropout Baseline:

python train.py --data-root ./data --backbone resnet50 --fusion gated \
  --batch-size 256 --epochs 15 --eval-every 5 \
  --tag-drop-min 0.4 --tag-drop-max 0.6 --modality-dropout 0.0 \
  --run-name resnet50_gated_no_dropout

Evaluation

Full Eval Harness (Per-Difficulty, Per-Mode)

# Evaluate on validation set with all modes (fusion, image_only, tag_only)
python eval.py \
  --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
  --data-root ./data \
  --split val \
  --modes fusion image_only tag_only \
  --batch-size 256 \
  --num-workers 8

Output:

  • Per-difficulty metrics: R@1, R@5, R@10, mAP
  • Per-mode comparison (fusion vs image_only vs tag_only)
  • CSV saved to runs/resnet50_gated_mixed_variable_dropout/eval_val.csv

Tag-Quality Robustness Sweep

Simulate users providing partial tags:

# Tag-drop 40% (user provides ~60% of tags)
python eval.py --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
  --data-root ./data --split val --modes fusion image_only --tag-drop 0.4

# Tag-drop 70% (user provides ~30% of tags)
python eval.py --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
  --data-root ./data --split val --modes fusion image_only --tag-drop 0.7

# No tags (pure image fallback)
python eval.py --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
  --data-root ./data --split val --modes fusion image_only --no-tags

Character-Tag Ablation

Remove character/copyright tags (e.g., kayoko_(blue_archive)), keep only descriptive tags:

python eval.py --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
  --data-root ./data --split val --modes tag_only --strip-character-tags

Batch Sweep (All Conditions)

# Run tag-quality sweep across all conditions (0.4, 0.7, 1.0 drop, no-tags)
for drop in 0.4 0.7 1.0; do
  python eval.py --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
    --data-root ./data --split val --modes fusion image_only --tag-drop $drop
done

python eval.py --run-dir ./runs/resnet50_gated_mixed_variable_dropout \
  --data-root ./data --split val --modes fusion image_only --no-tags

Interactive Demo (Qt6)

Build the Gallery Index (One-Time)

python presentation.py --build-index \
  --ckpt ./runs/resnet50_gated_mixed_variable_dropout/best.pt \
  --data-root ./data --split val \
  --index-cache ./demo_index.pt

Output: demo_index.pt (~100 MB for 974 items, encodings cached)

Launch the GUI

python presentation.py \
  --ckpt ./runs/resnet50_gated_mixed_variable_dropout/best.pt \
  --data-root ./data --index-cache ./demo_index.pt

How to use:

  1. Drag/drop an image or click to browse
  2. Type tags (comma-separated or natural language): 1girl, long_hair, blue_eyes
  3. Leave blank for image-only retrieval
  4. Click Search → top-5 matches with similarity scores
  5. Hover over any result → tooltip shows query gate weight and match metadata

Modes:

  • Full tags + image → fusion (gate adapts modality blend)
  • Image + blank tags → pure-image fallback (has_tags=False)
  • Natural language tags → automatically normalized and cleaned

Table 1: Per-Band Results (No Tag Drop, split=test)

Band Mode (1) ResNet gated (2) ViT gated (3) ResNet concat (4) ResNet gated, no-drop
overall fusion 0.9211 0.8842 0.8852 0.9013
image_only 0.5540 0.5629 0.5226 0.4818
tag_only 0.8857 0.8377 0.8365 0.8755
easy fusion 0.9571 0.9472 0.9439 0.9207
image_only 0.8573 0.8699 0.8336 0.8025
tag_only 0.8976 0.8402 0.8362 0.8844
medium fusion 0.9399 0.9247 0.9287 0.9148
image_only 0.6810 0.6942 0.6387 0.5865
tag_only 0.8844 0.8375 0.8329 0.8752
hard fusion 0.9102 0.8646 0.8600 0.8897
image_only 0.4122 0.4339 0.3877 0.3408
tag_only 0.8824 0.8309 0.8382 0.8725
extreme fusion 0.8771 0.8005 0.8085 0.8798
image_only 0.2655 0.2536 0.2305 0.1975
tag_only 0.8785 0.8421 0.8388 0.8699

Table 2: Tag-Drop Robustness Sweep (Fusion R@1)

Model drop 0.0 (full tags) drop 0.4 (60% kept) drop 0.7 (30% kept) drop 1.0 (no tags) image-only floor
(1) ResNet gated 0.9211 0.8350 0.6640 0.5505 0.5540
(2) ViT gated 0.8842 0.8078 0.6433 0.5431 0.5629
(3) ResNet concat 0.8852 0.7782 0.5380 0.5073 0.5226
(4) ResNet, no-drop 0.9013 0.5769 0.1986 0.4361 0.4818

Key findings from final eval:

  1. ResNet + gated wins cleanly — highest fusion R@1 across all drops (0.0 through 1.0)
  2. Robustness ranking: ResNet gated > ViT gated > ResNet concat > ResNet no-drop
  3. The crossover: No-drop starts strong (0.9013 at drop 0.0) but collapses to 0.1986 at drop 0.7 — proves dropout is non-negotiable
  4. Concat weakness: Steady decline 0.8852 → 0.5380 at drop 0.7, then holds; gated goes 0.9211 → 0.6640 (better throughout)
  5. At drop 1.0: All models converge toward image-only floor, as expected

Model Architecture

Backbone: ResNet-50 (frozen, ImageNet pretrained)
Text Encoder: all-MiniLM-L6-v2 (frozen)
Fusion: Gated (learned adaptive blend per embedding dimension)
Embedding Dimension: 512
Total Parameters: 48.7M (2.56M trainable, 46.2M frozen)

Loss: Symmetric InfoNCE (τ = 0.07) with same-source false-negative masking

Project Structure Details

Core Training Files

  • model.py — Architecture definition (ResNet/ViT + MLP heads + gate + fusion)
  • train.py — Main training loop (AdamW + cosine LR schedule)
  • loss.py — InfoNCE loss with same-source masking
  • data_loader.py — QueryPairDataset, GalleryDataset, tag corruption
  • data_pipeline.py — Data scraping, splits, difficulty assignment, preprocessing

Evaluation & Inference

  • eval.py — Per-difficulty, per-mode evaluation harness with tag-drop sweep
  • presentation.py — Qt6 GUI for interactive retrieval demo

Data Files

  • gallery.csv — image_id, image_path, tags, tags_text, split, is_distractor
  • mapping.csv — query_id, source_id, query_image_path, source_image_path, difficulty, tags
  • reference_metadata.json — scraped metadata from Danbooru
  • distractor_metadata.json — scraped distractor metadata

Troubleshooting

CUDA Out of Memory

Reduce batch size:

python train.py ... --batch-size 128  # or 64

Slow Training (CPU-bound data loading)

Increase workers:

python train.py ... --num-workers 16  # or adjust to CPU cores

MiniLM not downloading

Ensure internet connection and HuggingFace cache writable:

export HF_HOME=~/.cache/huggingface

Citation

If you use this code or results, please cite:

@article{naufal2024robust,
  title={Robust Multimodal Retrieval via Gated Contrastive Fusion: 
         Adaptive Modality Weights for Incomplete Tags},
  author={M. N. Musyaffa, F. D. C. Willard},
  year={2026},
  note={Deep Learning Final Project, Universitas Gadjah Mada}
}

Limitations & Future Work

Known Limitations

  1. Domain shift: ResNet-50 pretrained on ImageNet. Performance degrades on Safebooru (out-of-distribution) and screenshots.
  2. Gallery scale: Evaluated on ~974 items (validation). Efficiency not tested at 100k+.
  3. Frozen encoders: Transfer learning only. Fine-tuning on anime domain likely improves performance.
  4. User study: Evaluation uses synthetic corruption and catalogued tags. Real user queries not tested.

Future Directions

  1. Domain-robust encoding: Fine-tune backbone on anime or use anime-pretrained CLIP.
  2. Tag-quality gating: Gate should key on tag reliability, not just presence.
  3. Scale & indexing: Test FAISS/HNSW at 100k+ gallery size.
  4. Real user queries: Collect free-text, misspelled, natural-language tag queries.
  5. Broader backbone comparison: ConvNeXt, DINOv2, anime-tuned encoders.

License

MIT

About

Deep Learning Course (MIE213401) Universitas Gadjah Mada, Final Project

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Contributors

Languages