Upload 130 files

.gitattributes

@@ -119,3 +119,89 @@ camie-tagger-v2/images/training_monitor_overview.png filter=lfs diff=lfs merge=l
 camie-tagger-v2/images/training_monitor_predictions.png filter=lfs diff=lfs merge=lfs -text
 camie-tagger-v2/images/training_monitor_selection.png filter=lfs diff=lfs merge=lfs -text
 camie-tagger-v2/training/val_dataset.csv filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/1st_costume_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/animal_crossing_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/arknights_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/azur_lane_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/blue_archive_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/boku_no_hero_academia_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/casual_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/chainsaw_man_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/character_extended_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/company_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/cosplay_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/disgaea_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/disney_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/dragon_ball_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/dungeon_and_fighter_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/elsword_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/emblem_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/ensemble_stars!_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/fate_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/ff14_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/fire_emblem_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/flower_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/food_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/genshin_impact_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/girls_und_panzer_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/girls'_frontline_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/granblue_fantasy_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/honkai_impact_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/honkai_star_rail_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/housamo_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/idolmaster_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/jojo_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/kancolle_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/kemono_friends_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/kirby_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/league_of_legends_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/love_live!_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/madoka_magica_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/main_template.gif filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/manga_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/mega_man_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/meme_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/monster_girl_encyclopedia_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/naruto_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/new_year_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/nijisanji_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/nikke_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/omori_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/pokemon_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/precure_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/princess_connect!_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/punishing_gray_raven_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/ragnarok_online_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/sailor_moon_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/sao_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/sekaiju_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/senran_kagura_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/series_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/show_by_rock!!_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/skullgirls_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/sousou_no_frierem_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/splatoon_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/stand_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/street_fighter_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/style_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/symbol_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/tarot_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/tf2_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/umamusume_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/vocaloid_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/vtuber_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/warship_girls_r_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/weapon_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/wuthering_wave_template.png filter=lfs diff=lfs merge=lfs -text
+game/mosaics/templates/xenoblade_template.png filter=lfs diff=lfs merge=lfs -text
+images/app_screenshot.png filter=lfs diff=lfs merge=lfs -text
+images/collect_tags.PNG filter=lfs diff=lfs merge=lfs -text
+images/essence_tab.PNG filter=lfs diff=lfs merge=lfs -text
+images/lamp_essence.jpg filter=lfs diff=lfs merge=lfs -text
+images/library.PNG filter=lfs diff=lfs merge=lfs -text
+images/mosaics.PNG filter=lfs diff=lfs merge=lfs -text
+images/tag_results_example.png filter=lfs diff=lfs merge=lfs -text
+images/training_monitor_overview.png filter=lfs diff=lfs merge=lfs -text
+images/training_monitor_predictions.png filter=lfs diff=lfs merge=lfs -text
+images/training_monitor_selection.png filter=lfs diff=lfs merge=lfs -text
+training/val_dataset.csv filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,362 @@
 ---
 license: gpl-3.0
+datasets:
+- p1atdev/danbooru-2024
+language:
+- en
+pipeline_tag: image-classification
 ---
+
+# Camie Tagger v2
+
+An advanced deep learning model for automatically tagging anime/manga illustrations with relevant tags across multiple categories, achieving **67.3% micro F1 score** (50.6% macro F1 score using the macro optimized threshold preset) across 70,527 possible tags on a test set of 20,116 samples. Now with Vision Transformer backbone and significantly improved performance.
+
+## 🚀 What's New in v2
+
+### Major Performance Improvements
+- **Micro F1**: 58.1% → **67.3%** (+9.2 percentage points)
+- **Macro F1**: 31.5% → **50.6%** (+19.1 percentage points) 
+- **Model Size**: 424M → **143M parameters** (-66% reduction)
+- **Architecture**: Switched from EfficientNetV2-L to Vision Transformer (ViT) backbone
+- **Simplified Design**: Streamlined from dual-stage to single refined prediction model
+
+### Training Innovations
+- **Multi-Resolution Training**: Progressive scaling from 384px → 512px resolution
+- **IRFS (Instance-Aware Repeat Factor Sampling)**: Significant macro F1 improvements for rare tags
+- **Adaptive Training**: Models quickly adapt to resolution/distribution changes after initial pretraining
+
+*v2 demonstrates that Vision Transformers can achieve superior anime image tagging performance with fewer parameters and cleaner architecture.*
+
+## 🔑 Key Highlights
+
+- **Efficient Training**: Completed on just a single RTX 3060 GPU (12GB VRAM)  
+- **Fast Adaptation**: Models adapt to new resolutions/distributions within partial epochs after pretraining
+- **Comprehensive Coverage**: 70,527 tags across 7 categories (general, character, copyright, artist, meta, rating, year)
+- **Modern Architecture**: Vision Transformer backbone with cross-attention refinement
+- **User-Friendly Interface**: Easy-to-use application with customizable thresholds and tag collection game
+
+## ✨ Features
+
+- **Multi-category tagging system**: Handles general tags, characters, copyright (series), artists, meta information, and content ratings
+- **High performance**: 67.3% micro F1 score (50.6% macro F1) across 70,527 possible tags
+- **Windows compatibility**: Works on Windows without Flash Attention requirements
+- **Streamlit web interface**: User-friendly UI for uploading and analyzing images and a tag collection game
+- **Adjustable threshold profiles**: Micro, Macro, Balanced, Category-specific, High Precision, and High Recall profiles
+- **Fine-grained control**: Per-category threshold adjustments for precision-recall tradeoffs
+- **Safetensors and ONNX**: Original pickle files available in /models
+- **Vision Transformer Backbone**: Modern architecture with superior performance-to-parameter ratio
+
+## 📊 Performance Analysis
+
+### Complete v1 vs v2 Performance Comparison
+
+| CATEGORY | v1 Micro F1 | v2 Micro F1 | Micro Δ | v1 Macro F1 | v2 Macro F1 | Macro Δ |
+|----------|-------------|-------------|---------|-------------|-------------|---------|
+| **Overall** | 58.1% | **67.3%** | **+9.2pp** | 31.5% | **50.6%** | **+19.1pp** |
+| **Artist** | 47.4% | **70.0%** | **+22.6pp** | 29.8% | **64.4%** | **+34.6pp** |
+| **Character** | 74.6% | **83.4%** | **+8.8pp** | 47.8% | **64.5%** | **+16.7pp** |
+| **Copyright** | 76.3% | **86.6%** | **+10.3pp** | 37.7% | **53.1%** | **+15.4pp** |
+| **General** | 57.6% | **66.4%** | **+8.8pp** | 20.4% | **27.4%** | **+7.0pp** |
+| **Meta** | 55.7% | **61.2%** | **+5.5pp** | 14.4% | **19.2%** | **+4.8pp** |
+| **Rating** | 77.9% | **83.1%** | **+5.2pp** | 76.8% | **81.8%** | **+5.0pp** |
+| **Year** | 33.1% | **30.8%** | **-2.3pp** | 28.6% | **21.3%** | **-7.3pp** |
+
+*Both using the balanced preset.
+
+### Key Performance Insights
+
+The v2 model shows remarkable improvements across nearly all categories:
+
+- **Artist Recognition**: Massive +22.6pp micro F1 improvement, indicating much better artist identification
+- **Character Detection**: Strong +8.8pp micro F1 and +16.7pp macro F1 gains
+- **Copyright Recognition**: Excellent +10.3pp micro F1 improvement for series identification  
+- **General Tags**: Consistent +8.8pp micro F1 improvement for visual attributes
+- **Overall Macro F1**: Exceptional +19.1pp improvement shows much better rare tag recognition
+
+Only the year category shows slight regression, likely due to the reduced model complexity making temporal classification more challenging.
+
+### Detailed v2 Performance
+
+#### MACRO OPTIMIZED (Recommended)
+
+| CATEGORY | THRESHOLD | MICRO-F1 | MACRO-F1 |
+|----------|-----------|----------|----------|
+| **overall** | 0.492 | **60.9%** | **50.6%** |
+| artist | 0.492 | 62.3% | 66.1% |
+| character | 0.492 | 79.9% | 66.2% |
+| copyright | 0.492 | 81.8% | 56.2% |
+| general | 0.492 | 60.2% | 34.6% |
+| meta | 0.492 | 56.3% | 23.7% |
+| rating | 0.492 | 78.7% | 77.5% |
+| year | 0.492 | 37.2% | 32.6% |
+
+#### MICRO OPTIMIZED
+
+| CATEGORY | THRESHOLD | MICRO-F1 | MACRO-F1 |
+|----------|-----------|----------|----------|
+| **overall** | 0.614 | **67.3%** | **46.3%** |
+| artist | 0.614 | 70.0% | 64.4% |
+| character | 0.614 | 83.4% | 64.5% |
+| copyright | 0.614 | 86.6% | 53.1% |
+| general | 0.614 | 66.4% | 27.4% |
+| meta | 0.614 | 61.2% | 19.2% |
+| rating | 0.614 | 83.1% | 81.8% |
+| year | 0.614 | 30.8% | 21.3% |
+
+The model performs exceptionally well on character identification (83.4% F1 across 26,968 tags), copyright/series detection (86.6% F1 across 5,364 tags), and content rating classification (83.1% F1 across 4 tags).
+
+### Real-world Tag Accuracy
+
+The macro optimized threshold is recommended as many "false positives" according to the benchmark are actually correct tags missing from the Danbooru dataset. The model frequently identifies appropriate tags that weren't included in the original tagging, making perceived accuracy higher than formal metrics suggest.
+
+## 🧠 Architecture Overview
+
+### Vision Transformer Backbone
+- **Base Model**: Vision Transformer (ViT) with patch-based image processing
+- **Dual Output**: Patch feature map + CLS token for comprehensive image understanding
+- **Efficient Design**: 86.4M backbone parameters vs previous 214M+ classifier layers
+
+### Refined Prediction Pipeline
+1. **Feature Extraction**: ViT processes image into patch tokens and global CLS token
+2. **Global Pooling**: Combines mean-pooled patches with CLS token (dual-pool approach)  
+3. **Initial Predictions**: Shared weights between tag embeddings and classification layer
+4. **Candidate Selection**: Top-K tag selection based on initial confidence
+5. **Cross-Attention**: Tag embeddings attend to image patch features
+6. **Final Scoring**: Refined predictions for selected candidate tags
+
+### Key Improvements
+- **Shared Weights**: Tag embeddings directly used for initial classification
+- **Simplified Pipeline**: Single refined prediction stage (vs previous initial + refined)
+- **Native PyTorch**: Uses optimized MultiheadAttention instead of Flash Attention
+- **Custom Embeddings**: No dependency on external models like CLIP
+- **Gradient Checkpointing**: Memory-efficient training on consumer hardware
+
+## 🛠️ Training Details
+
+### Multi-Resolution Training Strategy
+
+The model was trained using an innovative multi-resolution approach:
+
+1. **Phase 1**: 3 epochs at 384px resolution with learning rate 1e-4
+2. **Phase 2**: IRFS (Instance-Aware Repeat Factor Sampling) - addresses long-tailed distribution imbalance  
+3. **Phase 3**: 512px resolution fine-tuning with learning rate 5e-5
+
+### Key Training Insights
+
+**Rapid Adaptation**: Once the model learns good general features during initial pretraining, it adapts to resolution changes and distribution shifts very quickly - often within a fraction of an epoch rather than requiring full retraining.
+
+**IRFS Benefits**: Instance-Aware Repeat Factor Sampling provided substantial macro F1 improvements by addressing the long-tailed distribution of anime tags, where instance counts vary dramatically between classes even with similar image counts.
+
+**Efficient Scaling**: The ViT architecture generalizes resolution and capacity changes to the entire dataset, making incremental training highly efficient.
+
+#### Training Data:
+- **Training subset**: 2,000,000 images
+- **Training duration**: 3+ epochs with multi-resolution scaling
+- **Final resolution**: 512x512 pixels
+
+## 🛠️ Requirements
+
+- **Python 3.11.9 specifically** (newer versions are incompatible)
+- PyTorch 1.10+
+- Streamlit
+- PIL/Pillow
+- NumPy
+- Flash Attention (note: doesn't work properly on Windows only needed for refined model which I'm not supporting that much anyway)
+
+## 🔧 Usage
+
+Setup the application and game by executing `setup.bat`. This installs the required virtual environment:
+
+- Upload your own images or select from example images
+- Choose different threshold profiles
+- Adjust category-specific thresholds
+- View predictions organized by category
+- Filter and sort tags based on confidence
+
+Use run_app.bat and run_game.bat.
+
+## 🎮 Tag Collector Game (Camie Collector)
+
+Introducing a Tagging game - a gamified approach to anime image tagging that helps you understand the performance and limits of the model. This was a shower thought gone to far! Lots of Project Moon references.
+
+### How to Play:
+1. Upload an image
+2. Scan for tags to discover them
+   ![Collect Tags Tab](images/collect_tags.PNG)
+3. Earn TagCoins for new discoveries
+4. Spend TagCoins on upgrades to lower the threshold
+   ![Upgrades Tab](images/upgrades.PNG)
+5. Lower thresholds reveal rarer tags!
+6. Collect sets of related tags for bonuses and reveal unique mosaics!
+   ![Mosaics Tab](images/mosaics.PNG)
+7. Visit the Library System to discover unique tags (not collect)
+   ![Library Tab](images/library.PNG)
+8. Use collected tags to either inspire new searches or generate essence
+9. Use Enkephalin to generate Tag Essences
+   ![Essence Tab](images/essence_tab.PNG)
+10. Use the Tag Essence Generator to collect the tag and related tags to it. Lamp Essence:
+    ![Lamp Essence](images/lamp_essence.jpg)
+
+## 🖥️ Web Interface Guide
+
+The interface is divided into three main sections:
+
+1. **Model Selection** (Sidebar):
+   - Choose between Full Model, Initial-only Model or ONNX accelerated (initial only)
+   - View model information and memory usage
+
+2. **Image Upload** (Left Panel):
+   - Upload your own images or select from examples
+   - View the selected image
+
+3. **Tagging Controls** (Right Panel):
+   - Select threshold profile
+   - Adjust thresholds for precision-recall and micro/macro tradeoff
+   - Configure display options
+   - View predictions organized by category
+
+### Display Options:
+
+- **Show all tags**: Display all tags including those below threshold
+- **Compact view**: Hide progress bars for cleaner display
+- **Minimum confidence**: Filter out low-confidence predictions
+- **Category selection**: Choose which categories to include in the summary
+
+### Interface Screenshots:
+
+![Application Interface](images/app_screenshot.png)
+
+![Tag Results Example](images/tag_results_example.png)
+
+## 🧠 Training Details
+
+### Dataset
+
+The model was trained on a carefully filtered subset of the [Danbooru 2024 dataset](https://huggingface.co/datasets/p1atdev/danbooru-2024), which contains a vast collection of anime/manga illustrations with comprehensive tagging.
+
+#### Filtering Process:
+
+The dataset was filtered with the following constraints:
+
+```python
+# Minimum tags per category required for each image
+min_tag_counts = {
+    'general': 25, 
+    'character': 1, 
+    'copyright': 1, 
+    'artist': 0, 
+    'meta': 0
+}
+
+# Minimum samples per tag required for tag to be included
+min_tag_samples = {
+    'general': 20, 
+    'character': 40, 
+    'copyright': 50, 
+    'artist': 200, 
+    'meta': 50
+}
+```
+
+This filtering process:
+1. First removed low-sample tags (tags with fewer occurrences than specified in `min_tag_samples`)
+2. Then removed images with insufficient tags per category (as specified in `min_tag_counts`)
+
+#### Training Data:
+
+- **Starting dataset size**: ~3,000,000 filtered images
+- **Training subset**: 2,000,000 images (due to storage and time constraints)
+- **Training duration**: 3.5 epochs
+
+#### Preprocessing:
+
+Images were preprocessed with minimal transformations:
+- Tensor normalization (scaled to 0-1 range)
+- ImageNet normalization.
+- Resized while maintaining original aspect ratio
+- No additional augmentations were applied
+
+#### Tag Categories:
+
+The model recognizes tags across these categories:
+- **General**: Visual elements, concepts, clothing, etc. (30,841 tags)
+- **Character**: Individual characters appearing in the image (26,968 tags)
+- **Copyright**: Source material (anime, manga, game) (5,364 tags)
+- **Artist**: Creator of the artwork (7,007 tags)
+- **Meta**: Meta information about the image (323 tags)
+- **Rating**: Content rating (4 tags)
+- **Year**: Year of upload (20 tags)
+
+All supported tags are stored in `model/metadata.json`, which maps tag IDs to their names and categories.
+
+### Training Notebooks
+
+The repository includes the main training notebook:
+
+1. **camie-tagger-v2.ipynb**:
+   - Main training notebook
+   - Dataset loading and preprocessing
+   - Model initialization
+   - Initial training loop with DeepSpeed integration
+   - Tag selection optimization
+   - Metric tracking and visualization
+
+### Training Monitor
+
+The project includes a real-time training monitor accessible via browser at `localhost:5000` during training:
+
+#### Performance Tips:
+
+⚠️ **Important**: For optimal training speed, keep VSCode minimized and the training monitor open in your browser. This can improve iteration speed by **3-5x** due to how the Windows/WSL graphics stack handles window focus and CUDA kernel execution.
+
+#### Monitor Features:
+
+The training monitor provides three main views:
+
+##### 1. Overview Tab:
+
+![Overview Tab](images/training_monitor_overview.png)
+
+- **Training Progress**: Real-time metrics including epoch, batch, speed, and time estimates
+- **Loss Chart**: Training and validation loss visualization
+- **F1 Scores**: Initial and refined F1 metrics for both training and validation
+
+##### 2. Predictions Tab:
+
+![Predictions Tab](images/training_monitor_predictions.png)
+
+- **Image Preview**: Shows the current sample being analyzed
+- **Prediction Controls**: Toggle between initial and refined predictions
+- **Tag Analysis**: 
+  - Color-coded tag results (correct, incorrect, missing)
+  - Confidence visualization with probability bars
+  - Category-based organization
+  - Filtering options for error analysis
+
+##### 3. Selection Analysis Tab:
+
+![Selection Analysis Tab](images/training_monitor_selection.png)
+
+- **Selection Metrics**: Statistics on tag selection quality
+  - Ground truth recall
+  - Average probability for ground truth vs. non-ground truth tags
+  - Unique tags selected
+- **Selection Graph**: Trends in selection quality over time
+- **Selected Tags Details**: Detailed view of model-selected tags with confidence scores
+
+The monitor provides invaluable insights into how the two-stage prediction model is performing, particularly how the tag selection process is working between the initial and refined prediction stages.
+
+### Training Notes:
+
+- Training notebooks may require WSL and 32GB+ of RAM to handle the dataset
+- With more computational resources, the model could be trained longer on the full dataset
+
+## 🙏 Acknowledgments
+
+- Claude Sonnet 3.5 and 4 for development assistance and architectural insights
+- [Vision Transformer](https://arxiv.org/abs/2010.11929) for the foundational architecture
+- [Danbooru](https://danbooru.donmai.us/) for the comprehensive tagged anime image dataset
+- [p1atdev](https://huggingface.co/p1atdev) for the processed Danbooru 2024 dataset
+- [IRFS paper](https://arxiv.org/abs/2305.08069) for Instance-Aware Repeat Factor Sampling methodology
+- PyTorch team for optimized attention implementations and gradient checkpointing
+- The open-source ML community for foundational tools and methods

app/app.py

@@ -0,0 +1,1050 @@
+#!/usr/bin/env python3
+"""
+Camie-Tagger-V2 Application
+A Streamlit web app for tagging images using an AI model.
+"""
+
+import streamlit as st
+import os
+import sys
+import traceback
+import tempfile
+import time
+import platform
+import subprocess
+import webbrowser
+import glob
+import numpy as np
+import matplotlib.pyplot as plt
+import io
+import base64
+import json
+from matplotlib.colors import LinearSegmentedColormap
+from PIL import Image
+from pathlib import Path
+
+# Add parent directory to path to allow importing from utils
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Import utilities
+from utils.image_processing import process_image, batch_process_images
+from utils.file_utils import save_tags_to_file, get_default_save_locations
+from utils.ui_components import display_progress_bar, show_example_images, display_batch_results
+from utils.onnx_processing import batch_process_images_onnx
+
+# Define the model directory
+MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+print(f"Using model directory: {MODEL_DIR}")
+
+# Define threshold profile descriptions and explanations
+threshold_profile_descriptions = {
+    "Micro Optimized": "Maximizes micro-averaged F1 score (best for dominant classes). Optimal for overall prediction quality.",
+    "Macro Optimized": "Maximizes macro-averaged F1 score (equal weight to all classes). Better for balanced performance across all tags.",
+    "Balanced": "Provides a trade-off between precision and recall with moderate thresholds. Good general-purpose setting.",
+    "Overall": "Uses a single threshold value across all categories. Simplest approach for consistent behavior.",
+    "Category-specific": "Uses different optimal thresholds for each category. Best for fine-tuning results."
+}
+
+threshold_profile_explanations = {
+    "Micro Optimized": """
+    ### Micro Optimized Profile
+    
+    **Technical definition**: Maximizes micro-averaged F1 score, which calculates metrics globally across all predictions.
+    
+    **When to use**: When you want the best overall accuracy, especially for common tags and dominant categories.
+    
+    **Effects**:
+    - Optimizes performance for the most frequent tags
+    - Gives more weight to categories with many examples (like 'character' and 'general')
+    - Provides higher precision in most common use cases
+    
+    **Performance from validation**:
+    - Micro F1: ~67.3%
+    - Macro F1: ~46.3%
+    - Threshold: ~0.614
+    """,
+    
+    "Macro Optimized": """
+    ### Macro Optimized Profile
+    
+    **Technical definition**: Maximizes macro-averaged F1 score, which gives equal weight to all categories regardless of size.
+    
+    **When to use**: When balanced performance across all categories is important, including rare tags.
+    
+    **Effects**:
+    - More balanced performance across all tag categories
+    - Better at detecting rare or unusual tags
+    - Generally has lower thresholds than micro-optimized
+    
+    **Performance from validation**:
+    - Micro F1: ~60.9%
+    - Macro F1: ~50.6%
+    - Threshold: ~0.492
+    """,
+    
+    "Balanced": """
+    ### Balanced Profile
+    
+    **Technical definition**: Same as Micro Optimized but provides a good reference point for manual adjustment.
+    
+    **When to use**: For general-purpose tagging when you don't have specific recall or precision requirements.
+    
+    **Effects**:
+    - Good middle ground between precision and recall
+    - Works well for most common use cases
+    - Default choice for most users
+    
+    **Performance from validation**:
+    - Micro F1: ~67.3%
+    - Macro F1: ~46.3%
+    - Threshold: ~0.614
+    """,
+    
+    "Overall": """
+    ### Overall Profile
+    
+    **Technical definition**: Uses a single threshold value across all categories.
+    
+    **When to use**: When you want consistent behavior across all categories and a simple approach.
+    
+    **Effects**:
+    - Consistent tagging threshold for all categories
+    - Simpler to understand than category-specific thresholds
+    - User-adjustable with a single slider
+    
+    **Default threshold value**: 0.5 (user-adjustable)
+    
+    **Note**: The threshold value is user-adjustable with the slider below.
+    """,
+    
+    "Category-specific": """
+    ### Category-specific Profile
+    
+    **Technical definition**: Uses different optimal thresholds for each category, allowing fine-tuning.
+    
+    **When to use**: When you want to customize tagging sensitivity for different categories.
+    
+    **Effects**:
+    - Each category has its own independent threshold
+    - Full control over category sensitivity
+    - Best for fine-tuning results when some categories need different treatment
+    
+    **Default threshold values**: Starts with balanced thresholds for each category
+    
+    **Note**: Use the category sliders below to adjust thresholds for individual categories.
+    """
+}
+
+def load_validation_results(results_path):
+    """Load validation results from JSON file"""
+    try:
+        with open(results_path, 'r') as f:
+            data = json.load(f)
+        return data
+    except Exception as e:
+        print(f"Error loading validation results: {e}")
+        return None
+
+def extract_thresholds_from_results(validation_data):
+    """Extract threshold information from validation results"""
+    if not validation_data or 'results' not in validation_data:
+        return {}
+    
+    thresholds = {
+        'overall': {},
+        'categories': {}
+    }
+    
+    # Process results to extract thresholds
+    for result in validation_data['results']:
+        category = result['CATEGORY'].lower()
+        profile = result['PROFILE'].lower().replace(' ', '_')
+        threshold = result['THRESHOLD']
+        micro_f1 = result['MICRO-F1']
+        macro_f1 = result['MACRO-F1']
+        
+        # Map profile names
+        if profile == 'micro_opt':
+            profile = 'micro_optimized'
+        elif profile == 'macro_opt':
+            profile = 'macro_optimized'
+        
+        threshold_info = {
+            'threshold': threshold,
+            'micro_f1': micro_f1,
+            'macro_f1': macro_f1
+        }
+        
+        if category == 'overall':
+            thresholds['overall'][profile] = threshold_info
+        else:
+            if category not in thresholds['categories']:
+                thresholds['categories'][category] = {}
+            thresholds['categories'][category][profile] = threshold_info
+    
+    return thresholds
+
+def load_model_and_metadata():
+    """Load model and metadata from available files"""
+    # Check for SafeTensors model
+    safetensors_path = os.path.join(MODEL_DIR, "camie-tagger-v2.safetensors")
+    safetensors_metadata_path = os.path.join(MODEL_DIR, "camie-tagger-v2-metadata.json")
+    
+    # Check for ONNX model
+    onnx_path = os.path.join(MODEL_DIR, "camie-tagger-v2.onnx")
+    
+    # Check for validation results
+    validation_results_path = os.path.join(MODEL_DIR, "full_validation_results.json")
+    
+    model_info = {
+        'safetensors_available': os.path.exists(safetensors_path) and os.path.exists(safetensors_metadata_path),
+        'onnx_available': os.path.exists(onnx_path) and os.path.exists(safetensors_metadata_path),
+        'validation_results_available': os.path.exists(validation_results_path)
+    }
+    
+    # Load metadata (same for both model types)
+    metadata = None
+    if os.path.exists(safetensors_metadata_path):
+        try:
+            with open(safetensors_metadata_path, 'r') as f:
+                metadata = json.load(f)
+        except Exception as e:
+            print(f"Error loading metadata: {e}")
+    
+    # Load validation results for thresholds
+    thresholds = {}
+    if model_info['validation_results_available']:
+        validation_data = load_validation_results(validation_results_path)
+        if validation_data:
+            thresholds = extract_thresholds_from_results(validation_data)
+    
+    # Add default thresholds if not available
+    if not thresholds:
+        thresholds = {
+            'overall': {
+                'balanced': {'threshold': 0.5, 'micro_f1': 0, 'macro_f1': 0},
+                'micro_optimized': {'threshold': 0.6, 'micro_f1': 0, 'macro_f1': 0},
+                'macro_optimized': {'threshold': 0.4, 'micro_f1': 0, 'macro_f1': 0}
+            },
+            'categories': {}
+        }
+    
+    return model_info, metadata, thresholds
+
+def load_safetensors_model(safetensors_path, metadata_path):
+    """Load SafeTensors model"""
+    try:
+        from safetensors.torch import load_file
+        import torch
+        
+        # Load metadata
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+        
+        # Import the model class (assuming it's available)
+        # You'll need to make sure the ImageTagger class is importable
+        from utils.model_loader import ImageTagger  # Update this import
+        
+        model_info = metadata['model_info']
+        dataset_info = metadata['dataset_info']
+        
+        # Recreate model architecture
+        model = ImageTagger(
+            total_tags=dataset_info['total_tags'],
+            dataset=None,
+            model_name=model_info['backbone'],
+            num_heads=model_info['num_attention_heads'],
+            dropout=0.0,
+            pretrained=False,
+            tag_context_size=model_info['tag_context_size'],
+            use_gradient_checkpointing=False,
+            img_size=model_info['img_size']
+        )
+        
+        # Load weights
+        state_dict = load_file(safetensors_path)
+        model.load_state_dict(state_dict)
+        model.eval()
+        
+        return model, metadata
+    except Exception as e:
+        raise Exception(f"Failed to load SafeTensors model: {e}")
+
+def get_profile_metrics(thresholds, profile_name):
+    """Extract metrics for the given profile from the thresholds dictionary"""
+    profile_key = None
+    
+    # Map UI-friendly names to internal keys
+    if profile_name == "Micro Optimized":
+        profile_key = "micro_optimized"
+    elif profile_name == "Macro Optimized":
+        profile_key = "macro_optimized"
+    elif profile_name == "Balanced":
+        profile_key = "balanced"
+    elif profile_name in ["Overall", "Category-specific"]:
+        profile_key = "macro_optimized"  # Use macro as default for these modes
+    
+    if profile_key and 'overall' in thresholds and profile_key in thresholds['overall']:
+        return thresholds['overall'][profile_key]
+    
+    return None
+
+def on_threshold_profile_change():
+    """Handle threshold profile changes"""
+    new_profile = st.session_state.threshold_profile
+    
+    if hasattr(st.session_state, 'thresholds') and hasattr(st.session_state, 'settings'):
+        # Initialize category thresholds if needed
+        if st.session_state.settings['active_category_thresholds'] is None:
+            st.session_state.settings['active_category_thresholds'] = {}
+        
+        current_thresholds = st.session_state.settings['active_category_thresholds']
+        
+        # Map profile names to keys
+        profile_key = None
+        if new_profile == "Micro Optimized":
+            profile_key = "micro_optimized"
+        elif new_profile == "Macro Optimized":
+            profile_key = "macro_optimized"
+        elif new_profile == "Balanced":
+            profile_key = "balanced"
+        
+        # Update thresholds based on profile
+        if profile_key and 'overall' in st.session_state.thresholds and profile_key in st.session_state.thresholds['overall']:
+            st.session_state.settings['active_threshold'] = st.session_state.thresholds['overall'][profile_key]['threshold']
+            
+            # Set category thresholds
+            for category in st.session_state.categories:
+                if category in st.session_state.thresholds['categories'] and profile_key in st.session_state.thresholds['categories'][category]:
+                    current_thresholds[category] = st.session_state.thresholds['categories'][category][profile_key]['threshold']
+                else:
+                    current_thresholds[category] = st.session_state.settings['active_threshold']
+        
+        elif new_profile == "Overall":
+            # Use balanced threshold for Overall profile
+            if 'overall' in st.session_state.thresholds and 'balanced' in st.session_state.thresholds['overall']:
+                st.session_state.settings['active_threshold'] = st.session_state.thresholds['overall']['balanced']['threshold']
+            else:
+                st.session_state.settings['active_threshold'] = 0.5
+            
+            # Clear category-specific overrides
+            st.session_state.settings['active_category_thresholds'] = {}
+        
+        elif new_profile == "Category-specific":
+            # Initialize with balanced thresholds
+            if 'overall' in st.session_state.thresholds and 'balanced' in st.session_state.thresholds['overall']:
+                st.session_state.settings['active_threshold'] = st.session_state.thresholds['overall']['balanced']['threshold']
+            else:
+                st.session_state.settings['active_threshold'] = 0.5
+            
+            # Initialize category thresholds
+            for category in st.session_state.categories:
+                if category in st.session_state.thresholds['categories'] and 'balanced' in st.session_state.thresholds['categories'][category]:
+                    current_thresholds[category] = st.session_state.thresholds['categories'][category]['balanced']['threshold']
+                else:
+                    current_thresholds[category] = st.session_state.settings['active_threshold']
+
+def apply_thresholds(all_probs, threshold_profile, active_threshold, active_category_thresholds, min_confidence, selected_categories):
+    """Apply thresholds to raw probabilities and return filtered tags"""
+    tags = {}
+    all_tags = []
+    
+    # Handle None case for active_category_thresholds
+    active_category_thresholds = active_category_thresholds or {}
+    
+    for category, cat_probs in all_probs.items():
+        # Get the appropriate threshold for this category
+        threshold = active_category_thresholds.get(category, active_threshold)
+        
+        # Filter tags above threshold
+        tags[category] = [(tag, prob) for tag, prob in cat_probs if prob >= threshold]
+        
+        # Add to all_tags if selected
+        if selected_categories.get(category, True):
+            for tag, prob in tags[category]:
+                all_tags.append(tag)
+    
+    return tags, all_tags
+
+def image_tagger_app():
+    """Main Streamlit application for image tagging."""
+    st.set_page_config(layout="wide", page_title="Camie Tagger", page_icon="🖼️")
+    
+    st.title("Camie-Tagger-v2 Interface")
+    st.markdown("---")
+
+    # Initialize settings
+    if 'settings' not in st.session_state:
+        st.session_state.settings = {
+            'show_all_tags': False,
+            'compact_view': True,
+            'min_confidence': 0.01,
+            'threshold_profile': "Macro",
+            'active_threshold': 0.5,
+            'active_category_thresholds': {},  # Initialize as empty dict, not None
+            'selected_categories': {},
+            'replace_underscores': False
+        }
+        st.session_state.show_profile_help = False
+
+    # Session state initialization for model
+    if 'model_loaded' not in st.session_state:
+        st.session_state.model_loaded = False
+        st.session_state.model = None
+        st.session_state.thresholds = None
+        st.session_state.metadata = None
+        st.session_state.model_type = "onnx"  # Default to ONNX
+
+    # Sidebar for model selection and information
+    with st.sidebar:
+        # Support information
+        st.subheader("💡 Notes")
+        
+        st.markdown("""
+        This tagger was trained on a subset of the available data due to hardware limitations.
+        
+        A more comprehensive model trained on the full 3+ million image dataset would provide:
+        - More recent characters and tags.
+        - Improved accuracy.
+        
+        If you find this tool useful and would like to support future development:
+        """)
+        
+        # Add Buy Me a Coffee button with Star of the City-like glow effect
+        st.markdown("""
+        <style>
+        @keyframes coffee-button-glow {
+            0% { box-shadow: 0 0 5px #FFD700; }
+            50% { box-shadow: 0 0 15px #FFD700; }
+            100% { box-shadow: 0 0 5px #FFD700; }
+        }
+        
+        .coffee-button {
+            display: inline-block;
+            animation: coffee-button-glow 2s infinite;
+            border-radius: 5px;
+            transition: transform 0.3s ease;
+        }
+        
+        .coffee-button:hover {
+            transform: scale(1.05);
+        }
+        </style>
+        
+        <a href="https://ko-fi.com/camais" target="_blank" class="coffee-button">
+            <img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" 
+                alt="Buy Me A Coffee" 
+                style="height: 45px; width: 162px; border-radius: 5px;" />
+        </a>
+        """, unsafe_allow_html=True)
+        
+        st.markdown("""
+        Your support helps with:
+        - GPU costs for training
+        - Storage for larger datasets
+        - Development of new features
+        - Future projects
+        
+        Thank you! 🙏
+                    
+        Full Details: https://huggingface.co/Camais03/camie-tagger
+        """)
+
+        st.header("Model Selection")
+        
+        # Load model information
+        model_info, metadata, thresholds = load_model_and_metadata()
+        
+        # Determine available model options
+        model_options = []
+        if model_info['onnx_available']:
+            model_options.append("ONNX (Recommended)")
+        if model_info['safetensors_available']:
+            model_options.append("SafeTensors (PyTorch)")
+        
+        if not model_options:
+            st.error("No model files found!")
+            st.info(f"Looking for models in: {MODEL_DIR}")
+            st.info("Expected files:")
+            st.info("- camie-tagger-v2.onnx")
+            st.info("- camie-tagger-v2.safetensors")
+            st.info("- camie-tagger-v2-metadata.json")
+            st.stop()
+        
+        # Model type selection
+        default_index = 0 if model_info['onnx_available'] else 0
+        model_type = st.radio(
+            "Select Model Type:",
+            model_options,
+            index=default_index,
+            help="ONNX: Optimized for speed and compatibility\nSafeTensors: Native PyTorch format"
+        )
+        
+        # Convert selection to internal model type
+        if model_type == "ONNX (Recommended)":
+            selected_model_type = "onnx"
+        else:
+            selected_model_type = "safetensors"
+        
+        # If model type changed, reload
+        if selected_model_type != st.session_state.model_type:
+            st.session_state.model_loaded = False
+            st.session_state.model_type = selected_model_type
+        
+        # Reload button
+        if st.button("Reload Model") and st.session_state.model_loaded:
+            st.session_state.model_loaded = False
+            st.info("Reloading model...")
+
+    # Try to load the model
+    if not st.session_state.model_loaded:
+        try:
+            with st.spinner(f"Loading {st.session_state.model_type.upper()} model..."):
+                if st.session_state.model_type == "onnx":
+                    # Load ONNX model
+                    import onnxruntime as ort
+                    
+                    onnx_path = os.path.join(MODEL_DIR, "camie-tagger-v2.onnx")
+                    
+                    # Check ONNX providers
+                    providers = ort.get_available_providers()
+                    gpu_available = any('CUDA' in provider for provider in providers)
+                    
+                    # Create ONNX session
+                    session = ort.InferenceSession(onnx_path, providers=providers)
+                    
+                    st.session_state.model = session
+                    st.session_state.device = f"ONNX Runtime ({'GPU' if gpu_available else 'CPU'})"
+                    st.session_state.param_dtype = "float32"
+                    
+                else:
+                    # Load SafeTensors model
+                    safetensors_path = os.path.join(MODEL_DIR, "camie-tagger-v2.safetensors")
+                    metadata_path = os.path.join(MODEL_DIR, "camie-tagger-v2-metadata.json")
+                    
+                    model, loaded_metadata = load_safetensors_model(safetensors_path, metadata_path)
+                    
+                    st.session_state.model = model
+                    device = next(model.parameters()).device
+                    param_dtype = next(model.parameters()).dtype
+                    st.session_state.device = device
+                    st.session_state.param_dtype = param_dtype
+                    metadata = loaded_metadata  # Use loaded metadata instead
+                
+                # Store common info
+                st.session_state.thresholds = thresholds
+                st.session_state.metadata = metadata
+                st.session_state.model_loaded = True
+                
+                # Get categories
+                if metadata and 'dataset_info' in metadata:
+                    tag_mapping = metadata['dataset_info']['tag_mapping']
+                    categories = list(set(tag_mapping['tag_to_category'].values()))
+                    st.session_state.categories = categories
+                    
+                    # Initialize selected categories
+                    if not st.session_state.settings['selected_categories']:
+                        st.session_state.settings['selected_categories'] = {cat: True for cat in categories}
+                
+                # Set initial threshold from validation results
+                if 'overall' in thresholds and 'balanced' in thresholds['overall']:
+                    st.session_state.settings['active_threshold'] = thresholds['overall']['macro_optimized']['threshold']
+                
+        except Exception as e:
+            st.error(f"Error loading model: {str(e)}")
+            st.code(traceback.format_exc())
+            st.stop()
+
+    # Display model information in sidebar
+    with st.sidebar:
+        st.header("Model Information")
+        if st.session_state.model_loaded:
+            if st.session_state.model_type == "onnx":
+                st.success("Using ONNX Model")
+            else:
+                st.success("Using SafeTensors Model")
+            
+            st.write(f"Device: {st.session_state.device}")
+            st.write(f"Precision: {st.session_state.param_dtype}")
+            
+            if st.session_state.metadata:
+                if 'dataset_info' in st.session_state.metadata:
+                    total_tags = st.session_state.metadata['dataset_info']['total_tags']
+                    st.write(f"Total tags: {total_tags}")
+                elif 'total_tags' in st.session_state.metadata:
+                    st.write(f"Total tags: {st.session_state.metadata['total_tags']}")
+            
+            # Show categories
+            with st.expander("Available Categories"):
+                for category in sorted(st.session_state.categories):
+                    st.write(f"- {category.capitalize()}")
+            
+            # About section
+            with st.expander("About this app"):
+                st.write("""
+                This app uses a trained image tagging model to analyze and tag images.
+                
+                **Model Options**:
+                - **ONNX (Recommended)**: Optimized for inference speed with broad compatibility
+                - **SafeTensors**: Native PyTorch format for advanced users
+                
+                **Features**:
+                - Upload or process images in batches
+                - Multiple threshold profiles based on validation results
+                - Category-specific threshold adjustment
+                - Export tags in various formats
+                - Fast inference with GPU acceleration (when available)
+                
+                **Threshold Profiles**:
+                - **Micro Optimized**: Best overall F1 score (67.3% micro F1)
+                - **Macro Optimized**: Balanced across categories (50.6% macro F1)
+                - **Balanced**: Good general-purpose setting
+                - **Overall**: Single adjustable threshold
+                - **Category-specific**: Fine-tune each category individually
+                """)
+
+    # Main content area - Image upload and processing
+    col1, col2 = st.columns([1, 1.5])
+    
+    with col1:
+        st.header("Image")
+        
+        upload_tab, batch_tab = st.tabs(["Upload Image", "Batch Processing"])
+        
+        image_path = None
+        
+        with upload_tab:
+            uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+            
+            if uploaded_file:
+                # Create temporary file
+                with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
+                    tmp_file.write(uploaded_file.getvalue())
+                    image_path = tmp_file.name
+                
+                st.session_state.original_filename = uploaded_file.name
+                
+                # Display image
+                image = Image.open(uploaded_file)
+                st.image(image, use_container_width=True)
+        
+        with batch_tab:
+            st.subheader("Batch Process Images")
+            
+            # Folder selection
+            batch_folder = st.text_input("Enter folder path containing images:", "")
+            
+            # Save options
+            save_options = st.radio(
+                "Where to save tag files:",
+                ["Same folder as images", "Custom location", "Default save folder"],
+                index=0
+            )
+            
+            # Batch size control
+            st.subheader("Performance Options")
+            batch_size = st.number_input("Batch size", min_value=1, max_value=32, value=4,
+                                       help="Higher values may improve speed but use more memory")
+            
+            # Category limits
+            enable_category_limits = st.checkbox("Limit tags per category in batch output", value=False)
+            
+            if enable_category_limits and hasattr(st.session_state, 'categories'):
+                if 'category_limits' not in st.session_state:
+                    st.session_state.category_limits = {}
+                
+                st.markdown("**Limit Values:** -1 = no limit, 0 = exclude, N = top N tags")
+                
+                limit_cols = st.columns(2)
+                for i, category in enumerate(sorted(st.session_state.categories)):
+                    col_idx = i % 2
+                    with limit_cols[col_idx]:
+                        current_limit = st.session_state.category_limits.get(category, -1)
+                        new_limit = st.number_input(
+                            f"{category.capitalize()}:",
+                            value=current_limit,
+                            min_value=-1,
+                            step=1,
+                            key=f"limit_{category}"
+                        )
+                        st.session_state.category_limits[category] = new_limit
+            
+            # Process batch button
+            if batch_folder and os.path.isdir(batch_folder):
+                image_files = []
+                for ext in ['*.jpg', '*.jpeg', '*.png']:
+                    image_files.extend(glob.glob(os.path.join(batch_folder, ext)))
+                    image_files.extend(glob.glob(os.path.join(batch_folder, ext.upper())))
+                
+                if image_files:
+                    st.write(f"Found {len(image_files)} images")
+                    
+                    if st.button("🔄 Process All Images", type="primary"):
+                        if not st.session_state.model_loaded:
+                            st.error("Model not loaded")
+                        else:
+                            with st.spinner("Processing images..."):
+                                progress_bar = st.progress(0)
+                                status_text = st.empty()
+                                
+                                def update_progress(current, total, image_path):
+                                    progress = current / total if total > 0 else 0
+                                    progress_bar.progress(progress)
+                                    status_text.text(f"Processing {current}/{total}: {os.path.basename(image_path) if image_path else 'Complete'}")
+                                
+                                # Determine save directory
+                                if save_options == "Same folder as images":
+                                    save_dir = batch_folder
+                                elif save_options == "Custom location":
+                                    save_dir = st.text_input("Custom save directory:", batch_folder)
+                                else:
+                                    save_dir = os.path.join(os.path.dirname(__file__), "saved_tags")
+                                    os.makedirs(save_dir, exist_ok=True)
+                                
+                                # Get current settings
+                                category_limits = st.session_state.category_limits if enable_category_limits else None
+                                
+                                # Process based on model type
+                                if st.session_state.model_type == "onnx":
+                                    batch_results = batch_process_images_onnx(
+                                        folder_path=batch_folder,
+                                        model_path=os.path.join(MODEL_DIR, "camie-tagger-v2.onnx"),
+                                        metadata_path=os.path.join(MODEL_DIR, "camie-tagger-v2-metadata.json"),
+                                        threshold_profile=st.session_state.settings['threshold_profile'],
+                                        active_threshold=st.session_state.settings['active_threshold'],
+                                        active_category_thresholds=st.session_state.settings['active_category_thresholds'],
+                                        save_dir=save_dir,
+                                        progress_callback=update_progress,
+                                        min_confidence=st.session_state.settings['min_confidence'],
+                                        batch_size=batch_size,
+                                        category_limits=category_limits
+                                    )
+                                else:
+                                    # SafeTensors processing (would need to implement)
+                                    st.error("SafeTensors batch processing not implemented yet")
+                                    batch_results = None
+                                
+                                if batch_results:
+                                    display_batch_results(batch_results)
+
+    # Column 2: Controls and Results
+    with col2:
+        st.header("Tagging Controls")
+        
+        # Threshold profile selection
+        all_profiles = [
+            "Micro Optimized",
+            "Macro Optimized", 
+            "Balanced",
+            "Overall",
+            "Category-specific"
+        ]
+        
+        profile_col1, profile_col2 = st.columns([3, 1])
+        
+        with profile_col1:
+            threshold_profile = st.selectbox(
+                "Select threshold profile",
+                options=all_profiles,
+                index=1,  # Default to Macro
+                key="threshold_profile",
+                on_change=on_threshold_profile_change
+            )
+        
+        with profile_col2:
+            if st.button("ℹ️ Help", key="profile_help"):
+                st.session_state.show_profile_help = not st.session_state.get('show_profile_help', False)
+        
+        # Show profile help
+        if st.session_state.get('show_profile_help', False):
+            st.markdown(threshold_profile_explanations[threshold_profile])
+        else:
+            st.info(threshold_profile_descriptions[threshold_profile])
+        
+        # Show profile metrics if available
+        if st.session_state.model_loaded:
+            metrics = get_profile_metrics(st.session_state.thresholds, threshold_profile)
+            
+            if metrics:
+                metrics_cols = st.columns(3)
+                
+                with metrics_cols[0]:
+                    st.metric("Threshold", f"{metrics['threshold']:.3f}")
+                
+                with metrics_cols[1]:
+                    st.metric("Micro F1", f"{metrics['micro_f1']:.1f}%")
+                
+                with metrics_cols[2]:
+                    st.metric("Macro F1", f"{metrics['macro_f1']:.1f}%")
+        
+        # Threshold controls based on profile
+        if st.session_state.model_loaded:
+            active_threshold = st.session_state.settings.get('active_threshold', 0.5)
+            active_category_thresholds = st.session_state.settings.get('active_category_thresholds', {})
+            
+            if threshold_profile in ["Micro Optimized", "Macro Optimized", "Balanced"]:
+                # Show reference threshold (disabled)
+                st.slider(
+                    "Threshold (from validation)", 
+                    min_value=0.01, 
+                    max_value=1.0, 
+                    value=float(active_threshold),
+                    step=0.01,
+                    disabled=True,
+                    help="This threshold is optimized from validation results"
+                )
+                
+            elif threshold_profile == "Overall":
+                # Adjustable overall threshold
+                active_threshold = st.slider(
+                    "Overall threshold", 
+                    min_value=0.01, 
+                    max_value=1.0, 
+                    value=float(active_threshold),
+                    step=0.01
+                )
+                st.session_state.settings['active_threshold'] = active_threshold
+                
+            elif threshold_profile == "Category-specific":
+                # Show reference overall threshold
+                st.slider(
+                    "Overall threshold (reference)", 
+                    min_value=0.01, 
+                    max_value=1.0, 
+                    value=float(active_threshold),
+                    step=0.01,
+                    disabled=True
+                )
+                
+                st.write("Adjust thresholds for individual categories:")
+                
+                # Category sliders
+                slider_cols = st.columns(2)
+                
+                if not active_category_thresholds:
+                    active_category_thresholds = {}
+                
+                for i, category in enumerate(sorted(st.session_state.categories)):
+                    col_idx = i % 2
+                    with slider_cols[col_idx]:
+                        default_val = active_category_thresholds.get(category, active_threshold)
+                        new_threshold = st.slider(
+                            f"{category.capitalize()}", 
+                            min_value=0.01, 
+                            max_value=1.0, 
+                            value=float(default_val),
+                            step=0.01,
+                            key=f"slider_{category}"
+                        )
+                        active_category_thresholds[category] = new_threshold
+                
+                st.session_state.settings['active_category_thresholds'] = active_category_thresholds
+        
+        # Display options
+        with st.expander("Display Options", expanded=False):
+            col1, col2 = st.columns(2)
+            with col1:
+                show_all_tags = st.checkbox("Show all tags (including below threshold)", 
+                                        value=st.session_state.settings['show_all_tags'])
+                compact_view = st.checkbox("Compact view (hide progress bars)", 
+                                        value=st.session_state.settings['compact_view'])
+                replace_underscores = st.checkbox("Replace underscores with spaces", 
+                                            value=st.session_state.settings.get('replace_underscores', False))
+            
+            with col2:
+                min_confidence = st.slider("Minimum confidence to display", 0.0, 0.5, 
+                                        st.session_state.settings['min_confidence'], 0.01)
+            
+            # Update settings
+            st.session_state.settings.update({
+                'show_all_tags': show_all_tags,
+                'compact_view': compact_view,
+                'min_confidence': min_confidence,
+                'replace_underscores': replace_underscores
+            })
+            
+            # Category selection
+            st.write("Categories to include in 'All Tags' section:")
+            
+            category_cols = st.columns(3)
+            selected_categories = {}
+            
+            if hasattr(st.session_state, 'categories'):
+                for i, category in enumerate(sorted(st.session_state.categories)):
+                    col_idx = i % 3
+                    with category_cols[col_idx]:
+                        default_val = st.session_state.settings['selected_categories'].get(category, True)
+                        selected_categories[category] = st.checkbox(
+                            f"{category.capitalize()}", 
+                            value=default_val,
+                            key=f"cat_select_{category}"
+                        )
+                
+                st.session_state.settings['selected_categories'] = selected_categories
+        
+        # Run tagging button
+        if image_path and st.button("Run Tagging"):
+            if not st.session_state.model_loaded:
+                st.error("Model not loaded")
+            else:
+                with st.spinner("Analyzing image..."):
+                    try:
+                        # Process image based on model type
+                        if st.session_state.model_type == "onnx":
+                            from utils.onnx_processing import process_single_image_onnx
+                            
+                            result = process_single_image_onnx(
+                                image_path=image_path,
+                                model_path=os.path.join(MODEL_DIR, "camie-tagger-v2.onnx"),
+                                metadata=st.session_state.metadata,
+                                threshold_profile=threshold_profile,
+                                active_threshold=st.session_state.settings['active_threshold'],
+                                active_category_thresholds=st.session_state.settings.get('active_category_thresholds', {}),
+                                min_confidence=st.session_state.settings['min_confidence']
+                            )
+                        else:
+                            # SafeTensors processing
+                            result = process_image(
+                                image_path=image_path,
+                                model=st.session_state.model,
+                                thresholds=st.session_state.thresholds,
+                                metadata=st.session_state.metadata,
+                                threshold_profile=threshold_profile,
+                                active_threshold=st.session_state.settings['active_threshold'],
+                                active_category_thresholds=st.session_state.settings.get('active_category_thresholds', {}),
+                                min_confidence=st.session_state.settings['min_confidence']
+                            )
+                        
+                        if result['success']:
+                            st.session_state.all_probs = result['all_probs']
+                            st.session_state.tags = result['tags']
+                            st.session_state.all_tags = result['all_tags']
+                            st.success("Analysis completed!")
+                        else:
+                            st.error(f"Analysis failed: {result.get('error', 'Unknown error')}")
+                    
+                    except Exception as e:
+                        st.error(f"Error during analysis: {str(e)}")
+                        st.code(traceback.format_exc())
+        
+        # Display results
+        if image_path and hasattr(st.session_state, 'all_probs'):
+            st.header("Predictions")
+            
+            # Apply current thresholds
+            filtered_tags, current_all_tags = apply_thresholds(
+                st.session_state.all_probs,
+                threshold_profile,
+                st.session_state.settings['active_threshold'],
+                st.session_state.settings.get('active_category_thresholds', {}),
+                st.session_state.settings['min_confidence'],
+                st.session_state.settings['selected_categories']
+            )
+            
+            all_tags = []
+            
+            # Display by category
+            for category in sorted(st.session_state.all_probs.keys()):
+                all_tags_in_category = st.session_state.all_probs.get(category, [])
+                filtered_tags_in_category = filtered_tags.get(category, [])
+                
+                if all_tags_in_category:
+                    expander_label = f"{category.capitalize()} ({len(filtered_tags_in_category)} tags)"
+                    
+                    with st.expander(expander_label, expanded=True):
+                        # Get threshold for this category (handle None case)
+                        active_category_thresholds = st.session_state.settings.get('active_category_thresholds') or {}
+                        threshold = active_category_thresholds.get(category, st.session_state.settings['active_threshold'])
+                        
+                        # Determine tags to display
+                        if st.session_state.settings['show_all_tags']:
+                            tags_to_display = all_tags_in_category
+                        else:
+                            tags_to_display = [(tag, prob) for tag, prob in all_tags_in_category if prob >= threshold]
+                        
+                        if not tags_to_display:
+                            st.info(f"No tags above {st.session_state.settings['min_confidence']:.2f} confidence")
+                            continue
+                        
+                        # Display tags
+                        if st.session_state.settings['compact_view']:
+                            # Compact view
+                            tag_list = []
+                            replace_underscores = st.session_state.settings.get('replace_underscores', False)
+                            
+                            for tag, prob in tags_to_display:
+                                percentage = int(prob * 100)
+                                display_tag = tag.replace('_', ' ') if replace_underscores else tag
+                                tag_list.append(f"{display_tag} ({percentage}%)")
+                                
+                                if prob >= threshold and st.session_state.settings['selected_categories'].get(category, True):
+                                    all_tags.append(tag)
+                            
+                            st.markdown(", ".join(tag_list))
+                        else:
+                            # Expanded view with progress bars
+                            for tag, prob in tags_to_display:
+                                replace_underscores = st.session_state.settings.get('replace_underscores', False)
+                                display_tag = tag.replace('_', ' ') if replace_underscores else tag
+                                
+                                if prob >= threshold and st.session_state.settings['selected_categories'].get(category, True):
+                                    all_tags.append(tag)
+                                    tag_display = f"**{display_tag}**"
+                                else:
+                                    tag_display = display_tag
+                                
+                                st.write(tag_display)
+                                st.markdown(display_progress_bar(prob), unsafe_allow_html=True)
+            
+            # All tags summary
+            st.markdown("---")
+            st.subheader(f"All Tags ({len(all_tags)} total)")
+            if all_tags:
+                replace_underscores = st.session_state.settings.get('replace_underscores', False)
+                if replace_underscores:
+                    display_tags = [tag.replace('_', ' ') for tag in all_tags]
+                    st.write(", ".join(display_tags))
+                else:
+                    st.write(", ".join(all_tags))
+            else:
+                st.info("No tags detected above the threshold.")
+            
+            # Save tags section
+            st.markdown("---")
+            st.subheader("Save Tags")
+            
+            if 'custom_folders' not in st.session_state:
+                st.session_state.custom_folders = get_default_save_locations()
+            
+            selected_folder = st.selectbox(
+                "Select save location:",
+                options=st.session_state.custom_folders,
+                format_func=lambda x: os.path.basename(x) if os.path.basename(x) else x
+            )
+            
+            if st.button("💾 Save to Selected Location"):
+                try:
+                    original_filename = st.session_state.original_filename if hasattr(st.session_state, 'original_filename') else None
+                    
+                    saved_path = save_tags_to_file(
+                        image_path=image_path,
+                        all_tags=all_tags,
+                        original_filename=original_filename,
+                        custom_dir=selected_folder,
+                        overwrite=True
+                    )
+                    
+                    st.success(f"Tags saved to: {os.path.basename(saved_path)}")
+                    st.info(f"Full path: {saved_path}")
+                    
+                    # Show file preview
+                    with st.expander("File Contents", expanded=True):
+                        with open(saved_path, 'r', encoding='utf-8') as f:
+                            content = f.read()
+                        st.code(content, language='text')
+                
+                except Exception as e:
+                    st.error(f"Error saving tags: {str(e)}")
+
+if __name__ == "__main__":
+    image_tagger_app()

app/run_app.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+"""
+Launcher script for the Image Tagger application.
+"""
+import os
+import sys
+import subprocess
+import webbrowser
+import time
+from pathlib import Path
+
+def run_app():
+    """Run the Streamlit app"""
+    # Check if app.py exists
+    app_path = "app.py"
+    if not os.path.exists(app_path):
+        print(f"Error: {app_path} not found")
+        return False
+    
+    # Get parent directory path (where venv is located)
+    parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    
+    # Get the path to streamlit in the virtual environment (in parent directory)
+    if sys.platform == "win32":
+        streamlit_path = os.path.join(parent_dir, "venv", "Scripts", "streamlit.exe")
+    else:
+        streamlit_path = os.path.join(parent_dir, "venv", "bin", "streamlit")
+    
+    if not os.path.exists(streamlit_path):
+        print(f"Error: Streamlit not found at {streamlit_path}")
+        print("Make sure you've run setup.py first to create the virtual environment")
+        return False
+    
+    print("=" * 60)
+    print("      Starting Image Tagger Application")
+    print("=" * 60)
+    
+    print("\nLaunching the web interface...")
+    
+    # Create a directory for example images if it doesn't exist
+    examples_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
+    os.makedirs(examples_dir, exist_ok=True)
+    
+    # Check if there are example images
+    example_files = [f for f in os.listdir(examples_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    if not example_files:
+        print("\nTip: Add some example images to the 'examples' directory for testing")
+    
+    # Run Streamlit app - using streamlit's built-in browser opening
+    # This avoids the double browser opening issue
+    try:
+        command = [streamlit_path, "run", app_path]
+        subprocess.run(command, check=True)
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"Error running the app: {e}")
+        return False
+    except KeyboardInterrupt:
+        print("\nApplication stopped by user")
+        return True
+
+if __name__ == "__main__":
+    success = run_app()
+    sys.exit(0 if success else 1)

app/utils/__init__.py

@@ -0,0 +1 @@
+# Make utils a proper Python package

app/utils/file_utils.py

@@ -0,0 +1,127 @@
+"""
+File utilities for Image Tagger application.
+"""
+
+import os
+import time
+
+def save_tags_to_file(image_path, all_tags, original_filename=None, custom_dir=None, overwrite=False):
+    """
+    Save tags to a text file in a dedicated 'saved_tags' folder or custom directory.
+    
+    Args:
+        image_path: Path to the original image
+        all_tags: List of all tags to save
+        original_filename: Original filename if uploaded through Streamlit
+        custom_dir: Custom directory to save tags to (if None, uses 'saved_tags' folder)
+    
+    Returns:
+        Path to the saved file
+    """
+    # Determine the save directory
+    if custom_dir and os.path.isdir(custom_dir):
+        save_dir = custom_dir
+    else:
+        # Create a dedicated folder for saved tags in the app's root directory
+        app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        save_dir = os.path.join(app_dir, "saved_tags")
+    
+    # Ensure the directory exists
+    os.makedirs(save_dir, exist_ok=True)
+    
+    # Determine the filename
+    if original_filename:
+        # For uploaded files, use original filename
+        base_name = os.path.splitext(original_filename)[0]
+    else:
+        # For non-uploaded files, use the image path
+        base_name = os.path.splitext(os.path.basename(image_path))[0]
+    
+    # Create the output path
+    output_path = os.path.join(save_dir, f"{base_name}.txt")
+    
+    # If overwrite is False and file exists, add a timestamp to avoid overwriting
+    if not overwrite and os.path.exists(output_path):
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        output_path = os.path.join(save_dir, f"{base_name}_{timestamp}.txt")
+    
+    # Write the tags to file
+    with open(output_path, 'w', encoding='utf-8') as f:
+        if all_tags:
+            # Add comma after each tag including the last one
+            tag_text = ", ".join(all_tags) + ","
+            f.write(tag_text)
+    
+    return output_path
+
+def get_default_save_locations():
+    """
+    Get default save locations for tag files.
+    
+    Returns:
+        List of default save locations
+    """
+    # App directory
+    app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    save_dir = os.path.join(app_dir, "saved_tags")
+    
+    # Common user directories
+    desktop_dir = os.path.expanduser("~/Desktop")
+    download_dir = os.path.expanduser("~/Downloads")
+    documents_dir = os.path.expanduser("~/Documents")
+    
+    # List of default save locations
+    save_locations = [
+        save_dir,
+        desktop_dir,
+        download_dir,
+        documents_dir,
+    ]
+    
+    # Ensure directories exist
+    for folder in save_locations:
+        os.makedirs(folder, exist_ok=True)
+    
+    return save_locations
+
+def apply_category_limits(result, category_limits):
+    """
+    Apply category limits to a result dictionary.
+    
+    Args:
+        result: Result dictionary containing tags and all_tags
+        category_limits: Dictionary mapping categories to their tag limits
+                         (0 = exclude category, -1 = no limit/include all)
+        
+    Returns:
+        Updated result dictionary with limits applied
+    """
+    if not category_limits or not result['success']:
+        return result
+    
+    # Get the filtered tags
+    filtered_tags = result['tags']
+    
+    # Apply limits to each category
+    for category, cat_tags in list(filtered_tags.items()):
+        # Get limit for this category, default to -1 (no limit)
+        limit = category_limits.get(category, -1)
+        
+        if limit == 0:
+            # Exclude this category entirely
+            del filtered_tags[category]
+        elif limit > 0 and len(cat_tags) > limit:
+            # Limit to top N tags for this category
+            filtered_tags[category] = cat_tags[:limit]
+    
+    # Regenerate all_tags list after applying limits
+    all_tags = []
+    for category, cat_tags in filtered_tags.items():
+        for tag, _ in cat_tags:
+            all_tags.append(tag)
+    
+    # Update the result with limited tags
+    result['tags'] = filtered_tags
+    result['all_tags'] = all_tags
+    
+    return result

app/utils/image_processing.py

@@ -0,0 +1,511 @@
+"""
+Image processing functions for the Image Tagger application.
+"""
+
+import os
+import traceback
+import glob
+
+
+def process_image(image_path, model, thresholds, metadata, threshold_profile, active_threshold, active_category_thresholds, min_confidence=0.1):
+    """
+    Process a single image and return the tags.
+    
+    Args:
+        image_path: Path to the image
+        model: The image tagger model
+        thresholds: Thresholds dictionary
+        metadata: Metadata dictionary
+        threshold_profile: Selected threshold profile
+        active_threshold: Overall threshold value
+        active_category_thresholds: Category-specific thresholds
+        min_confidence: Minimum confidence to include in results
+        
+    Returns:
+        Dictionary with tags, all probabilities, and other info
+    """
+    try:
+        # Run inference directly using the model's predict method
+        if threshold_profile in ["Category-specific", "High Precision", "High Recall"]:
+            results = model.predict(
+                image_path=image_path,
+                category_thresholds=active_category_thresholds
+            )
+        else:
+            results = model.predict(
+                image_path=image_path,
+                threshold=active_threshold
+            )
+        
+        # Extract and organize all probabilities
+        all_probs = {}
+        probs = results['refined_probabilities'][0]  # Remove batch dimension
+        
+        for idx in range(len(probs)):
+            prob_value = probs[idx].item()
+            if prob_value >= min_confidence:
+                tag, category = model.dataset.get_tag_info(idx)
+                
+                if category not in all_probs:
+                    all_probs[category] = []
+                
+                all_probs[category].append((tag, prob_value))
+        
+        # Sort tags by probability within each category
+        for category in all_probs:
+            all_probs[category] = sorted(
+                all_probs[category], 
+                key=lambda x: x[1], 
+                reverse=True
+            )
+        
+        # Get the filtered tags based on the selected threshold
+        tags = {}
+        for category, cat_tags in all_probs.items():
+            threshold = active_category_thresholds.get(category, active_threshold) if active_category_thresholds else active_threshold
+            tags[category] = [(tag, prob) for tag, prob in cat_tags if prob >= threshold]
+        
+        # Create a flat list of all tags above threshold
+        all_tags = []
+        for category, cat_tags in tags.items():
+            for tag, _ in cat_tags:
+                all_tags.append(tag)
+        
+        return {
+            'tags': tags,
+            'all_probs': all_probs,
+            'all_tags': all_tags,
+            'success': True
+        }
+    
+    except Exception as e:
+        print(f"Error processing {image_path}: {str(e)}")
+        traceback.print_exc()
+        return {
+            'tags': {},
+            'all_probs': {},
+            'all_tags': [],
+            'success': False,
+            'error': str(e)
+        }
+
+def apply_category_limits(result, category_limits):
+    """
+    Apply category limits to a result dictionary.
+    
+    Args:
+        result: Result dictionary containing tags and all_tags
+        category_limits: Dictionary mapping categories to their tag limits
+                         (0 = exclude category, -1 = no limit/include all)
+        
+    Returns:
+        Updated result dictionary with limits applied
+    """
+    if not category_limits or not result['success']:
+        return result
+    
+    # Get the filtered tags
+    filtered_tags = result['tags']
+    
+    # Apply limits to each category
+    for category, cat_tags in list(filtered_tags.items()):
+        # Get limit for this category, default to -1 (no limit)
+        limit = category_limits.get(category, -1)
+        
+        if limit == 0:
+            # Exclude this category entirely
+            del filtered_tags[category]
+        elif limit > 0 and len(cat_tags) > limit:
+            # Limit to top N tags for this category
+            filtered_tags[category] = cat_tags[:limit]
+    
+    # Regenerate all_tags list after applying limits
+    all_tags = []
+    for category, cat_tags in filtered_tags.items():
+        for tag, _ in cat_tags:
+            all_tags.append(tag)
+    
+    # Update the result with limited tags
+    result['tags'] = filtered_tags
+    result['all_tags'] = all_tags
+    
+    return result
+
+def batch_process_images(folder_path, model, thresholds, metadata, threshold_profile, active_threshold, 
+                        active_category_thresholds, save_dir=None, progress_callback=None, 
+                        min_confidence=0.1, batch_size=1, category_limits=None):
+    """
+    Process all images in a folder with optional batching for improved performance.
+    
+    Args:
+        folder_path: Path to folder containing images
+        model: The image tagger model
+        thresholds: Thresholds dictionary
+        metadata: Metadata dictionary
+        threshold_profile: Selected threshold profile
+        active_threshold: Overall threshold value
+        active_category_thresholds: Category-specific thresholds
+        save_dir: Directory to save tag files (if None uses default)
+        progress_callback: Optional callback for progress updates
+        min_confidence: Minimum confidence threshold
+        batch_size: Number of images to process at once (default: 1)
+        category_limits: Dictionary mapping categories to their tag limits (0 = unlimited)
+        
+    Returns:
+        Dictionary with results for each image
+    """
+    from .file_utils import save_tags_to_file  # Import here to avoid circular imports
+    import torch
+    from PIL import Image
+    import time
+    
+    print(f"Starting batch processing on {folder_path} with batch size {batch_size}")
+    start_time = time.time()
+    
+    # Find all image files in the folder
+    image_extensions = ['*.jpg', '*.jpeg', '*.png']
+    image_files = []
+
+    for ext in image_extensions:
+        image_files.extend(glob.glob(os.path.join(folder_path, ext)))
+        image_files.extend(glob.glob(os.path.join(folder_path, ext.upper())))
+
+    # Use a set to remove duplicate files (Windows filesystems are case-insensitive)
+    if os.name == 'nt':  # Windows
+        # Use lowercase paths for comparison on Windows
+        unique_paths = set()
+        unique_files = []
+        for file_path in image_files:
+            normalized_path = os.path.normpath(file_path).lower()
+            if normalized_path not in unique_paths:
+                unique_paths.add(normalized_path)
+                unique_files.append(file_path)
+        image_files = unique_files
+    
+    # Sort files for consistent processing order
+    image_files.sort()
+    
+    if not image_files:
+        return {
+            'success': False,
+            'error': f"No images found in {folder_path}",
+            'results': {}
+        }
+    
+    print(f"Found {len(image_files)} images to process")
+    
+    # Use the provided save directory or create a default one
+    if save_dir is None:
+        app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        save_dir = os.path.join(app_dir, "saved_tags")
+    
+    # Ensure the directory exists
+    os.makedirs(save_dir, exist_ok=True)
+    
+    # Process images in batches
+    results = {}
+    total_images = len(image_files)
+    processed = 0
+    
+    # Process in batches
+    for i in range(0, total_images, batch_size):
+        batch_start = time.time()
+        # Get current batch of images
+        batch_files = image_files[i:i+batch_size]
+        batch_size_actual = len(batch_files)
+        
+        print(f"Processing batch {i//batch_size + 1}/{(total_images + batch_size - 1)//batch_size}: {batch_size_actual} images")
+        
+        if batch_size > 1:
+            # True batch processing for multiple images at once
+            try:
+                # Using batch processing if batch_size > 1
+                batch_results = process_image_batch(
+                    image_paths=batch_files,
+                    model=model,
+                    thresholds=thresholds,
+                    metadata=metadata,
+                    threshold_profile=threshold_profile,
+                    active_threshold=active_threshold,
+                    active_category_thresholds=active_category_thresholds,
+                    min_confidence=min_confidence
+                )
+                
+                # Process and save results for each image in the batch
+                for j, image_path in enumerate(batch_files):
+                    # Update progress if callback provided
+                    if progress_callback:
+                        progress_callback(processed + j, total_images, image_path)
+                    
+                    if j < len(batch_results):
+                        result = batch_results[j]
+                        
+                        # Apply category limits if specified
+                        if category_limits and result['success']:
+                            # Use the apply_category_limits function instead of the inline code
+                            result = apply_category_limits(result, category_limits)
+                            
+                            # Debug print if you want
+                            print(f"Applied limits for {os.path.basename(image_path)}, remaining tags: {len(result['all_tags'])}")
+                        
+                        # Save the tags to a file
+                        if result['success']:
+                            output_path = save_tags_to_file(
+                                image_path=image_path,
+                                all_tags=result['all_tags'],
+                                custom_dir=save_dir,
+                                overwrite=True
+                            )
+                            result['output_path'] = str(output_path)
+                        
+                        # Store the result
+                        results[image_path] = result
+                    else:
+                        # Handle case where batch processing returned fewer results than expected
+                        results[image_path] = {
+                            'success': False,
+                            'error': 'Batch processing error: missing result',
+                            'all_tags': []
+                        }
+                
+            except Exception as e:
+                print(f"Batch processing error: {str(e)}")
+                traceback.print_exc()
+                
+                # Fall back to processing images one by one in this batch
+                for j, image_path in enumerate(batch_files):
+                    if progress_callback:
+                        progress_callback(processed + j, total_images, image_path)
+                    
+                    result = process_image(
+                        image_path=image_path,
+                        model=model,
+                        thresholds=thresholds,
+                        metadata=metadata,
+                        threshold_profile=threshold_profile,
+                        active_threshold=active_threshold,
+                        active_category_thresholds=active_category_thresholds,
+                        min_confidence=min_confidence
+                    )
+                    
+                    # Apply category limits if specified
+                    if category_limits and result['success']:
+                        # Use the apply_category_limits function
+                        result = apply_category_limits(result, category_limits)
+                    
+                    if result['success']:
+                        output_path = save_tags_to_file(
+                            image_path=image_path,
+                            all_tags=result['all_tags'],
+                            custom_dir=save_dir,
+                            overwrite=True
+                        )
+                        result['output_path'] = str(output_path)
+                    
+                    results[image_path] = result
+        else:
+            # Process one by one if batch_size is 1
+            for j, image_path in enumerate(batch_files):
+                if progress_callback:
+                    progress_callback(processed + j, total_images, image_path)
+                
+                result = process_image(
+                    image_path=image_path,
+                    model=model,
+                    thresholds=thresholds,
+                    metadata=metadata,
+                    threshold_profile=threshold_profile,
+                    active_threshold=active_threshold,
+                    active_category_thresholds=active_category_thresholds,
+                    min_confidence=min_confidence
+                )
+                
+                # Apply category limits if specified
+                if category_limits and result['success']:
+                    # Use the apply_category_limits function
+                    result = apply_category_limits(result, category_limits)
+                
+                if result['success']:
+                    output_path = save_tags_to_file(
+                        image_path=image_path,
+                        all_tags=result['all_tags'],
+                        custom_dir=save_dir,
+                        overwrite=True
+                    )
+                    result['output_path'] = str(output_path)
+                
+                results[image_path] = result
+        
+        # Update processed count
+        processed += batch_size_actual
+        
+        # Calculate batch timing
+        batch_end = time.time()
+        batch_time = batch_end - batch_start
+        print(f"Batch processed in {batch_time:.2f} seconds ({batch_time/batch_size_actual:.2f} seconds per image)")
+        
+    # Final progress update
+    if progress_callback:
+        progress_callback(total_images, total_images, None)
+    
+    end_time = time.time()
+    total_time = end_time - start_time
+    print(f"Batch processing finished. Total time: {total_time:.2f} seconds, Average: {total_time/total_images:.2f} seconds per image")
+    
+    return {
+        'success': True,
+        'total': total_images,
+        'processed': len(results),
+        'results': results,
+        'save_dir': save_dir,
+        'time_elapsed': end_time - start_time
+    }
+
+def process_image_batch(image_paths, model, thresholds, metadata, threshold_profile, active_threshold, active_category_thresholds, min_confidence=0.1):
+    """
+    Process a batch of images at once.
+    
+    Args:
+        image_paths: List of paths to the images
+        model: The image tagger model
+        thresholds: Thresholds dictionary
+        metadata: Metadata dictionary
+        threshold_profile: Selected threshold profile
+        active_threshold: Overall threshold value
+        active_category_thresholds: Category-specific thresholds
+        min_confidence: Minimum confidence to include in results
+        
+    Returns:
+        List of dictionaries with tags, all probabilities, and other info for each image
+    """
+    try:
+        import torch
+        from PIL import Image
+        import torchvision.transforms as transforms
+        
+        # Identify the model type we're using for better error handling
+        model_type = model.__class__.__name__
+        print(f"Running batch processing with model type: {model_type}")
+        
+        # Prepare the transformation for the images
+        transform = transforms.Compose([
+            transforms.Resize((512, 512)),  # Adjust based on your model's expected input
+            transforms.ToTensor(),
+        ])
+        
+        # Get model information
+        device = next(model.parameters()).device
+        dtype = next(model.parameters()).dtype
+        print(f"Model is using device: {device}, dtype: {dtype}")
+        
+        # Load and preprocess all images
+        batch_tensor = []
+        valid_images = []
+        
+        for img_path in image_paths:
+            try:
+                img = Image.open(img_path).convert('RGB')
+                img_tensor = transform(img)
+                img_tensor = img_tensor.to(device=device, dtype=dtype)
+                batch_tensor.append(img_tensor)
+                valid_images.append(img_path)
+            except Exception as e:
+                print(f"Error loading image {img_path}: {str(e)}")
+        
+        if not batch_tensor:
+            return []
+        
+        # Stack all tensors into a single batch
+        batch_input = torch.stack(batch_tensor)
+        
+        # Process entire batch at once
+        with torch.no_grad():
+            try:
+                # Forward pass on the whole batch
+                output = model(batch_input)
+                
+                # Handle tuple output format
+                if isinstance(output, tuple):
+                    probs_batch = torch.sigmoid(output[1])
+                else:
+                    probs_batch = torch.sigmoid(output)
+                
+                # Process each image's results
+                results = []
+                for i, img_path in enumerate(valid_images):
+                    probs = probs_batch[i].unsqueeze(0)  # Add batch dimension back
+                    
+                    # Extract and organize all probabilities
+                    all_probs = {}
+                    for idx in range(probs.size(1)):
+                        prob_value = probs[0, idx].item()
+                        if prob_value >= min_confidence:
+                            tag, category = model.dataset.get_tag_info(idx)
+                            
+                            if category not in all_probs:
+                                all_probs[category] = []
+                            
+                            all_probs[category].append((tag, prob_value))
+                    
+                    # Sort tags by probability
+                    for category in all_probs:
+                        all_probs[category] = sorted(all_probs[category], key=lambda x: x[1], reverse=True)
+                    
+                    # Get filtered tags
+                    tags = {}
+                    for category, cat_tags in all_probs.items():
+                        threshold = active_category_thresholds.get(category, active_threshold) if active_category_thresholds else active_threshold
+                        tags[category] = [(tag, prob) for tag, prob in cat_tags if prob >= threshold]
+                    
+                    # Create a flat list of all tags above threshold
+                    all_tags = []
+                    for category, cat_tags in tags.items():
+                        for tag, _ in cat_tags:
+                            all_tags.append(tag)
+                    
+                    results.append({
+                        'tags': tags,
+                        'all_probs': all_probs,
+                        'all_tags': all_tags,
+                        'success': True
+                    })
+                
+                return results
+                
+            except RuntimeError as e:
+                # If we encounter CUDA out of memory or another runtime error, 
+                # fall back to processing one by one
+                print(f"Error in batch processing: {str(e)}")
+                print("Falling back to one-by-one processing...")
+                
+                # Process one by one as fallback
+                results = []
+                for i, (img_tensor, img_path) in enumerate(zip(batch_tensor, valid_images)):
+                    try:
+                        input_tensor = img_tensor.unsqueeze(0)
+                        output = model(input_tensor)
+                        
+                        if isinstance(output, tuple):
+                            probs = torch.sigmoid(output[1])
+                        else:
+                            probs = torch.sigmoid(output)
+                        
+                        # Same post-processing as before...
+                        # [Code omitted for brevity]
+                        
+                    except Exception as e:
+                        print(f"Error processing image {img_path}: {str(e)}")
+                        results.append({
+                            'tags': {},
+                            'all_probs': {},
+                            'all_tags': [],
+                            'success': False,
+                            'error': str(e)
+                        })
+                
+                return results
+                
+    except Exception as e:
+        print(f"Error in batch processing: {str(e)}")
+        import traceback
+        traceback.print_exc()

app/utils/model_loader.py

@@ -0,0 +1,379 @@
+import torch
+import torch.nn as nn
+from torch.nn import GroupNorm, LayerNorm
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+import timm
+
+class ViTWrapper(nn.Module):
+    """Wrapper to make ViT compatible with feature extraction for ImageTagger"""
+    def __init__(self, vit_model):
+        super().__init__()
+        self.vit = vit_model
+        self.out_indices = (-1,)  # mimic timm.features_only
+        
+        # Get patch size and embedding dim from the model
+        self.patch_size = vit_model.patch_embed.patch_size[0]
+        self.embed_dim = vit_model.embed_dim
+        
+    def forward(self, x):
+        B = x.size(0)
+
+        # ➊ patch tokens
+        x = self.vit.patch_embed(x)                       # (B, N, C)
+
+        # ➋ prepend CLS
+        cls_tok = self.vit.cls_token.expand(B, -1, -1)    # (B, 1, C)
+        x = torch.cat((cls_tok, x), dim=1)                # (B, 1+N, C)
+
+        # ➌ add positional encodings (full, incl. CLS)
+        if self.vit.pos_embed is not None:
+            x = x + self.vit.pos_embed[:, : x.size(1), :]
+
+        x = self.vit.pos_drop(x)
+
+        for blk in self.vit.blocks:
+            x = blk(x)
+
+        x = self.vit.norm(x)                              # (B, 1+N, C)
+
+        # ➍ split back out
+        cls_final   = x[:, 0]              # (B, C)
+        patch_tokens = x[:, 1:]            # (B, N, C)
+
+        # ➎ reshape patches to (B, C, H, W)
+        B, N, C = patch_tokens.shape
+        h = w = int(N ** 0.5)              # square assumption
+        patch_features = patch_tokens.permute(0, 2, 1).reshape(B, C, h, w)
+
+        # Return **both**: (patch map, CLS)
+        return patch_features, cls_final
+    
+    def set_grad_checkpointing(self, enable=True):
+        """Enable gradient checkpointing if supported"""
+        if hasattr(self.vit, 'set_grad_checkpointing'):
+            self.vit.set_grad_checkpointing(enable)
+            return True
+        return False
+
+class ImageTagger(nn.Module):
+    """
+    ImageTagger with Vision Transformer backbone
+    """
+    def __init__(self, total_tags, dataset, model_name='vit_base_patch16_224', 
+                 num_heads=16, dropout=0.1, pretrained=True, tag_context_size=256,
+                 use_gradient_checkpointing=False, img_size=224):
+        super().__init__()
+        
+        # Store checkpointing config
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.model_name = model_name
+        self.img_size = img_size
+        
+        # Debug and stats flags
+        self._flags = {
+            'debug': False,
+            'model_stats': True
+        }
+        
+        # Core model config
+        self.dataset = dataset
+        self.tag_context_size = tag_context_size
+        self.total_tags = total_tags
+        
+        print(f"🏗️ Building ImageTagger with ViT backbone and {total_tags} tags")
+        print(f"   Backbone: {model_name}")
+        print(f"   Image size: {img_size}x{img_size}")
+        print(f"   Tag context size: {tag_context_size}")
+        print(f"   Gradient checkpointing: {use_gradient_checkpointing}")
+        print(f"   🎯 Custom embeddings, PyTorch native attention, no ground truth inclusion")
+        
+        # 1. Vision Transformer Backbone
+        print("📦 Loading Vision Transformer backbone...")
+        self._load_vit_backbone()
+        
+        # Get backbone dimensions by running a test forward pass
+        self._determine_backbone_dimensions()
+        
+        self.embedding_dim = self.backbone.embed_dim
+
+        # 2. Custom Tag Embeddings (no CLIP)
+        print("🎯 Using custom tag embeddings (no CLIP)")
+        self.tag_embedding = nn.Embedding(total_tags, self.embedding_dim)
+        
+        # 3. Shared weights approach - tag bias for initial predictions
+        print("🔗 Using shared weights between initial head and tag embeddings")
+        self.tag_bias = nn.Parameter(torch.zeros(total_tags))
+
+     
+        # 4. Image token extraction (for attention AND global pooling)
+        self.image_token_proj = nn.Identity()
+        
+        # 5. Tags-as-queries cross-attention (using PyTorch's optimized implementation)
+        self.cross_attention = nn.MultiheadAttention(
+            embed_dim=self.embedding_dim, 
+            num_heads=num_heads, 
+            dropout=dropout,
+            batch_first=True  # Use (batch, seq, feature) format
+        )
+        self.cross_norm = nn.LayerNorm(self.embedding_dim)
+        
+        # Initialize weights
+        self._init_weights()
+        
+        # Enable gradient checkpointing
+        if self.use_gradient_checkpointing:
+            self._enable_gradient_checkpointing()
+        
+        print(f"✅ ImageTagger with ViT initialized!")
+        self._print_parameter_count()
+    
+    def _load_vit_backbone(self):
+        """Load Vision Transformer model from timm"""
+        print(f"   Loading from timm: {self.model_name}")
+        
+        # Load the ViT model (not features_only, we want the full model for token extraction)
+        vit_model = timm.create_model(
+            self.model_name, 
+            pretrained=True,
+            img_size=self.img_size,
+            num_classes=0  # Remove classification head
+        )
+        
+        # Wrap it in our compatibility layer
+        self.backbone = ViTWrapper(vit_model)
+        
+        print(f"   ✅ ViT loaded successfully")
+        print(f"   Patch size: {self.backbone.patch_size}x{self.backbone.patch_size}")
+        print(f"   Embed dim: {self.backbone.embed_dim}")
+    
+    def _determine_backbone_dimensions(self):
+        """Determine backbone output dimensions"""
+        print("   🔍 Determining backbone dimensions...")
+        
+        with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
+            # Create a dummy input 
+            dummy_input = torch.randn(1, 3, self.img_size, self.img_size)
+            
+            # Get features
+            backbone_features, cls_dummy = self.backbone(dummy_input)
+            feature_tensor = backbone_features
+            
+            self.backbone_dim = feature_tensor.shape[1]
+            self.feature_map_size = feature_tensor.shape[2]
+            
+        print(f"   Backbone output: {self.backbone_dim}D, {self.feature_map_size}x{self.feature_map_size} spatial")
+        print(f"   Total patch tokens: {self.feature_map_size * self.feature_map_size}")
+    
+    def _enable_gradient_checkpointing(self):
+        """Enable gradient checkpointing for memory efficiency"""
+        print("🔄 Enabling gradient checkpointing...")
+        
+        # Enable checkpointing for ViT backbone
+        if self.backbone.set_grad_checkpointing(True):
+            print("   ✅ ViT backbone checkpointing enabled")
+        else:
+            print("   ⚠️ ViT backbone doesn't support built-in checkpointing, will checkpoint manually")
+    
+    def _checkpoint_backbone(self, x):
+        """Wrapper for backbone with gradient checkpointing"""
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(self.backbone, x, use_reentrant=False)
+        else:
+            return self.backbone(x)
+    
+    def _checkpoint_image_proj(self, x):
+        """Wrapper for image projection with gradient checkpointing"""
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(self.image_token_proj, x, use_reentrant=False)
+        else:
+            return self.image_token_proj(x)
+    
+    def _checkpoint_cross_attention(self, query, key, value):
+        """Wrapper for cross attention with gradient checkpointing"""
+        def _attention_forward(q, k, v):
+            attended_features, _ = self.cross_attention(query=q, key=k, value=v)
+            return self.cross_norm(attended_features)
+        
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(_attention_forward, query, key, value, use_reentrant=False)
+        else:
+            return _attention_forward(query, key, value)
+    
+    def _checkpoint_candidate_selection(self, initial_logits):
+        """Wrapper for candidate selection with gradient checkpointing"""
+        def _candidate_forward(logits):
+            return self._get_candidate_tags(logits)
+        
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(_candidate_forward, initial_logits, use_reentrant=False)
+        else:
+            return _candidate_forward(initial_logits)
+    
+    def _checkpoint_final_scoring(self, attended_features, candidate_indices):
+        """Wrapper for final scoring with gradient checkpointing"""
+        def _scoring_forward(features, indices):
+            emb = self.tag_embedding(indices)
+            # BF16 in, BF16 out
+            return (features * emb).sum(dim=-1)
+        
+        if self.use_gradient_checkpointing and self.training:
+            return checkpoint.checkpoint(_scoring_forward, attended_features, candidate_indices, use_reentrant=False)
+        else:
+            return _scoring_forward(attended_features, candidate_indices)
+   
+    def _init_weights(self):
+        """Initialize weights for new modules"""
+        def _init_layer(layer):
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.zeros_(layer.bias)
+            elif isinstance(layer, nn.Conv2d):
+                nn.init.kaiming_normal_(layer.weight, mode='fan_out', nonlinearity='relu')
+                if layer.bias is not None:
+                    nn.init.zeros_(layer.bias)
+            elif isinstance(layer, nn.Embedding):
+                nn.init.normal_(layer.weight, mean=0, std=0.02)
+        
+        # Initialize new components
+        self.image_token_proj.apply(_init_layer)
+        
+        # Initialize tag embeddings with normal distribution
+        nn.init.normal_(self.tag_embedding.weight, mean=0, std=0.02)
+        
+        # Initialize tag bias
+        nn.init.zeros_(self.tag_bias)
+    
+    def _print_parameter_count(self):
+        """Print parameter statistics"""
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        backbone_params = sum(p.numel() for p in self.backbone.parameters())
+        
+        print(f"📊 Parameter Statistics:")
+        print(f"   Total parameters: {total_params/1e6:.1f}M")
+        print(f"   Trainable parameters: {trainable_params/1e6:.1f}M")
+        print(f"   Frozen parameters: {(total_params-trainable_params)/1e6:.1f}M")
+        print(f"   Backbone parameters: {backbone_params/1e6:.1f}M")
+        
+        if self.use_gradient_checkpointing:
+            print(f"   🔄 Gradient checkpointing enabled for memory efficiency")
+    
+    @property
+    def debug(self):
+        return self._flags['debug']
+    
+    @property
+    def model_stats(self):
+        return self._flags['model_stats']
+    
+    def _get_candidate_tags(self, initial_logits, target_tags=None, hard_negatives=None):
+        """Select candidate tags - no ground truth inclusion"""
+        batch_size = initial_logits.size(0)
+        
+        # Simply select top K candidates based on initial predictions
+        top_probs, top_indices = torch.topk(
+            torch.sigmoid(initial_logits), 
+            k=min(self.tag_context_size, self.total_tags),
+            dim=1, largest=True, sorted=True
+        )
+        
+        return top_indices
+    
+    def _analyze_predictions(self, predictions, tag_indices):
+        """Analyze prediction patterns"""
+        if not self.model_stats:
+            return {}
+        
+        if torch._dynamo.is_compiling():
+            return {}
+        
+        with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
+            probs = torch.sigmoid(predictions)
+            relevant_probs = torch.gather(probs, 1, tag_indices)
+            
+            return {
+                'prediction_confidence': relevant_probs.mean().item(),
+                'prediction_entropy': -(relevant_probs * torch.log(relevant_probs + 1e-9)).mean().item(),
+                'high_confidence_ratio': (relevant_probs > 0.7).float().mean().item(),
+                'above_threshold_ratio': (relevant_probs > 0.5).float().mean().item(),
+            }
+    
+    def forward(self, x, targets=None, hard_negatives=None):
+        """
+        Forward pass with ViT backbone, CLS token support and gradient-checkpointing.
+        All arithmetic tensors stay in the backbone’s dtype (BF16 under autocast,
+        FP32 otherwise).  Anything that must mix dtypes is cast to match.
+        """
+        batch_size  = x.size(0)
+        model_stats = {} if self.model_stats else {}
+
+        # ------------------------------------------------------------------
+        # 1. Backbone  →  patch map + CLS token
+        # ------------------------------------------------------------------
+        patch_map, cls_token = self._checkpoint_backbone(x)         # patch_map: [B, C, H, W]
+                                                                    # cls_token: [B, C]
+
+        # ------------------------------------------------------------------
+        # 2. Tokens  →  global image vector
+        # ------------------------------------------------------------------
+        image_tokens_4d = self._checkpoint_image_proj(patch_map)    # [B, C, H, W]
+        image_tokens    = image_tokens_4d.flatten(2).transpose(1, 2)  # [B, N, C]
+
+        # “Dual-pool”: mean-pool patches ⊕ CLS
+        global_features = 0.5 * (image_tokens.mean(dim=1, dtype=image_tokens.dtype) + cls_token)  # [B, C]
+
+        compute_dtype = global_features.dtype                       # BF16 or FP32
+
+        # ------------------------------------------------------------------
+        # 3. Initial logits  (shared weights)
+        # ------------------------------------------------------------------
+        tag_weights = self.tag_embedding.weight.to(compute_dtype)   # [T, C]
+        tag_bias    = self.tag_bias.to(compute_dtype)               # [T]
+
+        initial_logits = global_features @ tag_weights.t() + tag_bias   # [B, T]
+        initial_logits = initial_logits.to(compute_dtype)               # keep dtype uniform
+        initial_preds  = initial_logits                                 # alias
+
+        # ------------------------------------------------------------------
+        # 4. Candidate set
+        # ------------------------------------------------------------------
+        candidate_indices = self._checkpoint_candidate_selection(initial_logits)  # [B, K]
+
+        tag_embeddings   = self.tag_embedding(candidate_indices).to(compute_dtype)  # [B, K, C]
+
+        attended_features = self._checkpoint_cross_attention(       # [B, K, C]
+            tag_embeddings, image_tokens, image_tokens
+        )
+
+        # ------------------------------------------------------------------
+        # 5. Score candidates  &  scatter back
+        # ------------------------------------------------------------------
+        candidate_logits = self._checkpoint_final_scoring(attended_features, candidate_indices)  # [B, K]
+
+        # --- align dtypes so scatter never throws ---
+        if candidate_logits.dtype != initial_logits.dtype:
+            candidate_logits = candidate_logits.to(initial_logits.dtype)
+
+        refined_logits = initial_logits.clone()
+        refined_logits.scatter_(1, candidate_indices, candidate_logits)
+        refined_preds = refined_logits
+
+        # ------------------------------------------------------------------
+        # 6. Optional stats
+        # ------------------------------------------------------------------
+        if self.model_stats and targets is not None and not torch._dynamo.is_compiling():
+            model_stats['initial_prediction_stats'] = self._analyze_predictions(initial_preds,
+                                                                                candidate_indices)
+            model_stats['refined_prediction_stats'] = self._analyze_predictions(refined_preds,
+                                                                                candidate_indices)
+
+        return {
+            'initial_predictions': initial_preds,
+            'refined_predictions': refined_preds,
+            'selected_candidates': candidate_indices,
+            'model_stats': model_stats
+        }
+    
+    def predict

app/utils/onnx_processing.py

@@ -0,0 +1,729 @@
+"""
+ONNX-based batch image processing for the Image Tagger application.
+Updated with proper ImageNet normalization and new metadata format.
+"""
+
+import os
+import json
+import time
+import traceback
+import numpy as np
+import glob
+import onnxruntime as ort
+from PIL import Image
+import torchvision.transforms as transforms
+from concurrent.futures import ThreadPoolExecutor
+
+def preprocess_image(image_path, image_size=512):
+    """
+    Process an image for ImageTagger inference with proper ImageNet normalization
+    """
+    if not os.path.exists(image_path):
+        raise ValueError(f"Image not found at path: {image_path}")
+    
+    # ImageNet normalization - CRITICAL for your model
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+    
+    try:
+        with Image.open(image_path) as img:
+            # Convert RGBA or Palette images to RGB
+            if img.mode in ('RGBA', 'P'):
+                img = img.convert('RGB')
+            
+            # Get original dimensions
+            width, height = img.size
+            aspect_ratio = width / height
+            
+            # Calculate new dimensions to maintain aspect ratio
+            if aspect_ratio > 1:
+                new_width = image_size
+                new_height = int(new_width / aspect_ratio)
+            else:
+                new_height = image_size
+                new_width = int(new_height * aspect_ratio)
+            
+            # Resize with LANCZOS filter
+            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+            
+            # Create new image with padding (use ImageNet mean for padding)
+            # Using RGB values close to ImageNet mean: (0.485*255, 0.456*255, 0.406*255)
+            pad_color = (124, 116, 104)
+            new_image = Image.new('RGB', (image_size, image_size), pad_color)
+            paste_x = (image_size - new_width) // 2
+            paste_y = (image_size - new_height) // 2
+            new_image.paste(img, (paste_x, paste_y))
+            
+            # Apply transforms (including ImageNet normalization)
+            img_tensor = transform(new_image)
+            return img_tensor.numpy()
+            
+    except Exception as e:
+        raise Exception(f"Error processing {image_path}: {str(e)}")
+
+def process_single_image_onnx(image_path, model_path, metadata, threshold_profile="Overall", 
+                           active_threshold=0.35, active_category_thresholds=None, 
+                           min_confidence=0.1):
+    """
+    Process a single image using ONNX model with new metadata format
+    
+    Args:
+        image_path: Path to the image file
+        model_path: Path to the ONNX model file
+        metadata: Model metadata dictionary
+        threshold_profile: The threshold profile being used
+        active_threshold: Overall threshold value
+        active_category_thresholds: Category-specific thresholds
+        min_confidence: Minimum confidence to include in results
+        
+    Returns:
+        Dictionary with tags and probabilities
+    """
+    try:
+        # Create ONNX tagger for this image (or reuse an existing one)
+        if hasattr(process_single_image_onnx, 'tagger'):
+            tagger = process_single_image_onnx.tagger
+        else:
+            # Create new tagger
+            tagger = ONNXImageTagger(model_path, metadata)
+            # Cache it for future calls
+            process_single_image_onnx.tagger = tagger
+        
+        # Preprocess the image
+        start_time = time.time()
+        img_array = preprocess_image(image_path)
+        
+        # Run inference
+        results = tagger.predict_batch(
+            [img_array], 
+            threshold=active_threshold,
+            category_thresholds=active_category_thresholds,
+            min_confidence=min_confidence
+        )
+        inference_time = time.time() - start_time
+        
+        if results:
+            result = results[0]
+            result['inference_time'] = inference_time
+            result['success'] = True
+            return result
+        else:
+            return {
+                'success': False,
+                'error': 'Failed to process image',
+                'all_tags': [],
+                'all_probs': {},
+                'tags': {}
+            }
+    
+    except Exception as e:
+        print(f"Error in process_single_image_onnx: {str(e)}")
+        traceback.print_exc()
+        return {
+            'success': False,
+            'error': str(e),
+            'all_tags': [],
+            'all_probs': {},
+            'tags': {}
+        }
+
+def preprocess_images_parallel(image_paths, image_size=512, max_workers=8):
+    """Process multiple images in parallel"""
+    processed_images = []
+    valid_paths = []
+    
+    # Define a worker function
+    def process_single_image(path):
+        try:
+            return preprocess_image(path, image_size), path
+        except Exception as e:
+            print(f"Error processing {path}: {str(e)}")
+            return None, path
+    
+    # Process images in parallel
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        results = list(executor.map(process_single_image, image_paths))
+    
+    # Filter results
+    for img_array, path in results:
+        if img_array is not None:
+            processed_images.append(img_array)
+            valid_paths.append(path)
+    
+    return processed_images, valid_paths
+
+def apply_category_limits(result, category_limits):
+    """
+    Apply category limits to a result dictionary.
+    
+    Args:
+        result: Result dictionary containing tags and all_tags
+        category_limits: Dictionary mapping categories to their tag limits
+                         (0 = exclude category, -1 = no limit/include all)
+        
+    Returns:
+        Updated result dictionary with limits applied
+    """
+    if not category_limits or not result['success']:
+        return result
+    
+    # Get the filtered tags
+    filtered_tags = result['tags']
+    
+    # Apply limits to each category
+    for category, cat_tags in list(filtered_tags.items()):
+        # Get limit for this category, default to -1 (no limit)
+        limit = category_limits.get(category, -1)
+        
+        if limit == 0:
+            # Exclude this category entirely
+            del filtered_tags[category]
+        elif limit > 0 and len(cat_tags) > limit:
+            # Limit to top N tags for this category
+            filtered_tags[category] = cat_tags[:limit]
+    
+    # Regenerate all_tags list after applying limits
+    all_tags = []
+    for category, cat_tags in filtered_tags.items():
+        for tag, _ in cat_tags:
+            all_tags.append(tag)
+    
+    # Update the result with limited tags
+    result['tags'] = filtered_tags
+    result['all_tags'] = all_tags
+    
+    return result
+
+class ONNXImageTagger:
+    """ONNX-based image tagger for fast batch inference with updated metadata format"""
+    
+    def __init__(self, model_path, metadata):
+        # Load model
+        self.model_path = model_path
+        try:
+            self.session = ort.InferenceSession(
+                model_path, 
+                providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
+            )
+            print(f"Using providers: {self.session.get_providers()}")
+        except Exception as e:
+            print(f"CUDA not available, using CPU: {e}")
+            self.session = ort.InferenceSession(
+                model_path, 
+                providers=['CPUExecutionProvider']
+            )
+            print(f"Using providers: {self.session.get_providers()}")
+        
+        # Store metadata (passed as dict, not loaded from file)
+        self.metadata = metadata
+        
+        # Extract tag mappings from new metadata structure
+        if 'dataset_info' in metadata:
+            # New metadata format
+            self.tag_mapping = metadata['dataset_info']['tag_mapping']
+            self.idx_to_tag = self.tag_mapping['idx_to_tag']
+            self.tag_to_category = self.tag_mapping['tag_to_category']
+            self.total_tags = metadata['dataset_info']['total_tags']
+        else:
+            # Fallback for older format
+            self.idx_to_tag = metadata.get('idx_to_tag', {})
+            self.tag_to_category = metadata.get('tag_to_category', {})
+            self.total_tags = metadata.get('total_tags', len(self.idx_to_tag))
+        
+        # Get input name
+        self.input_name = self.session.get_inputs()[0].name
+        print(f"Model loaded successfully. Input name: {self.input_name}")
+        print(f"Total tags: {self.total_tags}, Categories: {len(set(self.tag_to_category.values()))}")
+        
+    def predict_batch(self, image_arrays, threshold=0.5, category_thresholds=None, min_confidence=0.1):
+        """Run batch inference on preprocessed image arrays"""
+        # Stack arrays into batch
+        batch_input = np.stack(image_arrays)
+        
+        # Run inference
+        start_time = time.time()
+        outputs = self.session.run(None, {self.input_name: batch_input})
+        inference_time = time.time() - start_time
+        print(f"Batch inference completed in {inference_time:.4f} seconds ({inference_time/len(image_arrays):.4f} s/image)")
+        
+        # Process outputs - handle both single and multi-output models
+        if len(outputs) >= 2:
+            # Multi-output model (initial_predictions, refined_predictions, selected_candidates)
+            initial_logits = outputs[0]
+            refined_logits = outputs[1]
+            # Use refined predictions as main output
+            main_logits = refined_logits
+            print(f"Using refined predictions (shape: {refined_logits.shape})")
+        else:
+            # Single output model
+            main_logits = outputs[0]
+            print(f"Using single output (shape: {main_logits.shape})")
+        
+        # Apply sigmoid to get probabilities
+        main_probs = 1.0 / (1.0 + np.exp(-main_logits))
+        
+        # Process results for each image in batch
+        batch_results = []
+        
+        for i in range(main_probs.shape[0]):
+            probs = main_probs[i]
+            
+            # Extract and organize all probabilities
+            all_probs = {}
+            for idx in range(probs.shape[0]):
+                prob_value = float(probs[idx])
+                if prob_value >= min_confidence:
+                    idx_str = str(idx)
+                    tag_name = self.idx_to_tag.get(idx_str, f"unknown-{idx}")
+                    category = self.tag_to_category.get(tag_name, "general")
+                    
+                    if category not in all_probs:
+                        all_probs[category] = []
+                    
+                    all_probs[category].append((tag_name, prob_value))
+            
+            # Sort tags by probability within each category
+            for category in all_probs:
+                all_probs[category] = sorted(
+                    all_probs[category], 
+                    key=lambda x: x[1], 
+                    reverse=True
+                )
+            
+            # Get the filtered tags based on the selected threshold
+            tags = {}
+            for category, cat_tags in all_probs.items():
+                # Use category-specific threshold if available
+                if category_thresholds and category in category_thresholds:
+                    cat_threshold = category_thresholds[category]
+                else:
+                    cat_threshold = threshold
+                    
+                tags[category] = [(tag, prob) for tag, prob in cat_tags if prob >= cat_threshold]
+            
+            # Create a flat list of all tags above threshold
+            all_tags = []
+            for category, cat_tags in tags.items():
+                for tag, _ in cat_tags:
+                    all_tags.append(tag)
+            
+            batch_results.append({
+                'tags': tags,
+                'all_probs': all_probs,
+                'all_tags': all_tags,
+                'success': True
+            })
+            
+        return batch_results
+
+def batch_process_images_onnx(folder_path, model_path, metadata_path, threshold_profile, 
+                            active_threshold, active_category_thresholds, save_dir=None, 
+                            progress_callback=None, min_confidence=0.1, batch_size=16,
+                            category_limits=None):
+    """
+    Process all images in a folder using the ONNX model with new metadata format.
+    
+    Args:
+        folder_path: Path to folder containing images
+        model_path: Path to the ONNX model file
+        metadata_path: Path to the model metadata file
+        threshold_profile: Selected threshold profile
+        active_threshold: Overall threshold value
+        active_category_thresholds: Category-specific thresholds
+        save_dir: Directory to save tag files (if None uses default)
+        progress_callback: Optional callback for progress updates
+        min_confidence: Minimum confidence threshold
+        batch_size: Number of images to process at once
+        category_limits: Dictionary mapping categories to their tag limits
+        
+    Returns:
+        Dictionary with results for each image
+    """
+    from utils.file_utils import save_tags_to_file  # Import here to avoid circular imports
+    
+    # Find all image files in the folder
+    image_extensions = ['*.jpg', '*.jpeg', '*.png']
+    image_files = []
+    
+    for ext in image_extensions:
+        image_files.extend(glob.glob(os.path.join(folder_path, ext)))
+        image_files.extend(glob.glob(os.path.join(folder_path, ext.upper())))
+    
+    # Remove duplicates (Windows case-insensitive filesystems)
+    if os.name == 'nt':  # Windows
+        unique_paths = set()
+        unique_files = []
+        for file_path in image_files:
+            normalized_path = os.path.normpath(file_path).lower()
+            if normalized_path not in unique_paths:
+                unique_paths.add(normalized_path)
+                unique_files.append(file_path)
+        image_files = unique_files
+    
+    if not image_files:
+        return {
+            'success': False,
+            'error': f"No images found in {folder_path}",
+            'results': {}
+        }
+    
+    # Use the provided save directory or create a default one
+    if save_dir is None:
+        app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        save_dir = os.path.join(app_dir, "saved_tags")
+    
+    # Ensure the directory exists
+    os.makedirs(save_dir, exist_ok=True)
+    
+    # Load metadata
+    try:
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+    except Exception as e:
+        return {
+            'success': False,
+            'error': f"Failed to load metadata: {e}",
+            'results': {}
+        }
+    
+    # Create ONNX tagger
+    try:
+        tagger = ONNXImageTagger(model_path, metadata)
+    except Exception as e:
+        return {
+            'success': False,
+            'error': f"Failed to load model: {e}",
+            'results': {}
+        }
+    
+    # Process images in batches
+    results = {}
+    total_images = len(image_files)
+    processed = 0
+    
+    start_time = time.time()
+    
+    # Process in batches
+    for i in range(0, total_images, batch_size):
+        batch_start = time.time()
+        
+        # Get current batch of images
+        batch_files = image_files[i:i+batch_size]
+        batch_size_actual = len(batch_files)
+        
+        # Update progress if callback provided
+        if progress_callback:
+            progress_callback(processed, total_images, batch_files[0] if batch_files else None)
+        
+        print(f"Processing batch {i//batch_size + 1}/{(total_images + batch_size - 1)//batch_size}: {batch_size_actual} images")
+        
+        try:
+            # Preprocess images in parallel
+            processed_images, valid_paths = preprocess_images_parallel(batch_files)
+            
+            if processed_images:
+                # Run batch prediction
+                batch_results = tagger.predict_batch(
+                    processed_images, 
+                    threshold=active_threshold,
+                    category_thresholds=active_category_thresholds,
+                    min_confidence=min_confidence
+                )
+                
+                # Process results for each image
+                for j, (image_path, result) in enumerate(zip(valid_paths, batch_results)):
+                    # Update progress if callback provided
+                    if progress_callback:
+                        progress_callback(processed + j, total_images, image_path)
+                    
+                    # Apply category limits if specified
+                    if category_limits and result['success']:
+                        print(f"Applying limits to {os.path.basename(image_path)}: {len(result['all_tags'])} → ", end="")
+                        result = apply_category_limits(result, category_limits)
+                        print(f"{len(result['all_tags'])} tags")
+                    
+                    # Save the tags to a file
+                    if result['success']:
+                        try:
+                            output_path = save_tags_to_file(
+                                image_path=image_path,
+                                all_tags=result['all_tags'],
+                                custom_dir=save_dir,
+                                overwrite=True
+                            )
+                            result['output_path'] = str(output_path)
+                        except Exception as e:
+                            print(f"Error saving tags for {image_path}: {e}")
+                            result['save_error'] = str(e)
+                    
+                    # Store the result
+                    results[image_path] = result
+            
+            processed += batch_size_actual
+            
+            # Calculate batch timing
+            batch_end = time.time()
+            batch_time = batch_end - batch_start
+            print(f"Batch processed in {batch_time:.2f} seconds ({batch_time/batch_size_actual:.2f} seconds per image)")
+            
+        except Exception as e:
+            print(f"Error processing batch: {str(e)}")
+            traceback.print_exc()
+            
+            # Process failed images one by one as fallback
+            for j, image_path in enumerate(batch_files):
+                try:
+                    # Update progress if callback provided
+                    if progress_callback:
+                        progress_callback(processed + j, total_images, image_path)
+                    
+                    # Preprocess single image
+                    img_array = preprocess_image(image_path)
+                    
+                    # Run inference on single image
+                    single_results = tagger.predict_batch(
+                        [img_array], 
+                        threshold=active_threshold,
+                        category_thresholds=active_category_thresholds,
+                        min_confidence=min_confidence
+                    )
+                    
+                    if single_results:
+                        result = single_results[0]
+                        
+                        # Apply category limits if specified
+                        if category_limits and result['success']:
+                            result = apply_category_limits(result, category_limits)
+                        
+                        # Save the tags to a file
+                        if result['success']:
+                            try:
+                                output_path = save_tags_to_file(
+                                    image_path=image_path,
+                                    all_tags=result['all_tags'],
+                                    custom_dir=save_dir,
+                                    overwrite=True
+                                )
+                                result['output_path'] = str(output_path)
+                            except Exception as e:
+                                print(f"Error saving tags for {image_path}: {e}")
+                                result['save_error'] = str(e)
+                        
+                        results[image_path] = result
+                    else:
+                        results[image_path] = {
+                            'success': False,
+                            'error': 'Failed to process image',
+                            'all_tags': []
+                        }
+                        
+                except Exception as img_e:
+                    print(f"Error processing single image {image_path}: {str(img_e)}")
+                    results[image_path] = {
+                        'success': False,
+                        'error': str(img_e),
+                        'all_tags': []
+                    }
+            
+            processed += batch_size_actual
+    
+    # Final progress update
+    if progress_callback:
+        progress_callback(total_images, total_images, None)
+    
+    end_time = time.time()
+    total_time = end_time - start_time
+    print(f"Batch processing finished. Total time: {total_time:.2f} seconds, Average: {total_time/total_images:.2f} seconds per image")
+    
+    return {
+        'success': True,
+        'total': total_images,
+        'processed': len(results),
+        'results': results,
+        'save_dir': save_dir,
+        'time_elapsed': end_time - start_time
+    }
+
+def test_onnx_imagetagger(model_path, metadata_path, image_path, threshold=0.5, top_k=256):
+    """
+    Test ImageTagger ONNX model with proper handling of all outputs and new metadata format
+    
+    Args:
+        model_path: Path to ONNX model file
+        metadata_path: Path to metadata JSON file  
+        image_path: Path to test image
+        threshold: Confidence threshold for predictions
+        top_k: Maximum number of predictions to show
+    """
+    import onnxruntime as ort
+    import numpy as np
+    import json
+    import time
+    from collections import defaultdict
+    
+    print(f"Loading ImageTagger ONNX model from {model_path}")
+    
+    # Load metadata with proper error handling
+    try:
+        with open(metadata_path, 'r') as f:
+            metadata = json.load(f)
+    except Exception as e:
+        raise ValueError(f"Failed to load metadata: {e}")
+    
+    # Extract tag mappings from new metadata structure
+    try:
+        if 'dataset_info' in metadata:
+            # New metadata format
+            dataset_info = metadata['dataset_info']
+            tag_mapping = dataset_info['tag_mapping']
+            idx_to_tag = tag_mapping['idx_to_tag']
+            tag_to_category = tag_mapping['tag_to_category']
+            total_tags = dataset_info['total_tags']
+        else:
+            # Fallback for older format
+            idx_to_tag = metadata.get('idx_to_tag', {})
+            tag_to_category = metadata.get('tag_to_category', {})
+            total_tags = metadata.get('total_tags', len(idx_to_tag))
+        
+        print(f"Model info: {total_tags} tags, {len(set(tag_to_category.values()))} categories")
+        
+    except KeyError as e:
+        raise ValueError(f"Invalid metadata structure, missing key: {e}")
+    
+    # Initialize ONNX session with robust provider handling
+    providers = []
+    if ort.get_device() == 'GPU':
+        providers.append('CUDAExecutionProvider')
+    providers.append('CPUExecutionProvider')
+    
+    try:
+        session = ort.InferenceSession(model_path, providers=providers)
+        active_provider = session.get_providers()[0]
+        print(f"Using provider: {active_provider}")
+        
+        # Print model info
+        inputs = session.get_inputs()
+        outputs = session.get_outputs()
+        print(f"Model inputs: {len(inputs)}")
+        print(f"Model outputs: {len(outputs)}")
+        for i, output in enumerate(outputs):
+            print(f"  Output {i}: {output.name} {output.shape}")
+            
+    except Exception as e:
+        raise RuntimeError(f"Failed to create ONNX session: {e}")
+    
+    # Preprocess image
+    print(f"Processing image: {image_path}")
+    try:
+        # Get image size from metadata
+        img_size = metadata.get('model_info', {}).get('img_size', 512)
+        img_tensor = preprocess_image(image_path, image_size=img_size)
+        img_numpy = img_tensor[np.newaxis, :]  # Add batch dimension
+        print(f"Input shape: {img_numpy.shape}, dtype: {img_numpy.dtype}")
+        
+    except Exception as e:
+        raise ValueError(f"Image preprocessing failed: {e}")
+    
+    # Run inference
+    input_name = session.get_inputs()[0].name
+    print("Running inference...")
+    
+    start_time = time.time()
+    try:
+        outputs = session.run(None, {input_name: img_numpy})
+        inference_time = time.time() - start_time
+        print(f"Inference completed in {inference_time:.4f} seconds")
+        
+    except Exception as e:
+        raise RuntimeError(f"Inference failed: {e}")
+    
+    # Handle outputs properly
+    if len(outputs) >= 2:
+        initial_logits = outputs[0]
+        refined_logits = outputs[1] 
+        selected_candidates = outputs[2] if len(outputs) > 2 else None
+        
+        # Use refined predictions as main output
+        main_logits = refined_logits
+        print(f"Using refined predictions (shape: {refined_logits.shape})")
+        
+    else:
+        # Fallback to single output
+        main_logits = outputs[0]
+        print(f"Using single output (shape: {main_logits.shape})")
+    
+    # Apply sigmoid to get probabilities
+    main_probs = 1.0 / (1.0 + np.exp(-main_logits))
+    
+    # Apply threshold and get predictions
+    predictions_mask = (main_probs >= threshold)
+    indices = np.where(predictions_mask[0])[0]
+    
+    if len(indices) == 0:
+        print(f"No predictions above threshold {threshold}")
+        # Show top 5 regardless of threshold
+        top_indices = np.argsort(main_probs[0])[-5:][::-1]
+        print("Top 5 predictions:")
+        for idx in top_indices:
+            idx_str = str(idx)
+            tag_name = idx_to_tag.get(idx_str, f"unknown-{idx}")
+            prob = float(main_probs[0, idx])
+            print(f"  {tag_name}: {prob:.3f}")
+        return {}
+    
+    # Group by category  
+    tags_by_category = defaultdict(list)
+    
+    for idx in indices:
+        idx_str = str(idx)
+        tag_name = idx_to_tag.get(idx_str, f"unknown-{idx}")
+        category = tag_to_category.get(tag_name, "general")
+        prob = float(main_probs[0, idx])
+        
+        tags_by_category[category].append((tag_name, prob))
+    
+    # Sort by probability within each category
+    for category in tags_by_category:
+        tags_by_category[category] = sorted(
+            tags_by_category[category], 
+            key=lambda x: x[1], 
+            reverse=True
+        )[:top_k]  # Limit per category
+    
+    # Print results
+    total_predictions = sum(len(tags) for tags in tags_by_category.values())
+    print(f"\nPredicted tags (threshold: {threshold}): {total_predictions} total")
+    
+    # Category order for consistent display
+    category_order = ['general', 'character', 'copyright', 'artist', 'meta', 'year', 'rating']
+    
+    for category in category_order:
+        if category in tags_by_category:
+            tags = tags_by_category[category]
+            print(f"\n{category.upper()} ({len(tags)}):")
+            for tag, prob in tags:
+                print(f"  {tag}: {prob:.3f}")
+    
+    # Show any other categories not in standard order
+    for category in sorted(tags_by_category.keys()):
+        if category not in category_order:
+            tags = tags_by_category[category]
+            print(f"\n{category.upper()} ({len(tags)}):")
+            for tag, prob in tags:
+                print(f"  {tag}: {prob:.3f}")
+    
+    # Performance stats
+    print(f"\nPerformance:")
+    print(f"  Inference time: {inference_time:.4f}s")
+    print(f"  Provider: {active_provider}")
+    print(f"  Max confidence: {main_probs.max():.3f}")
+    if total_predictions > 0:
+        avg_conf = np.mean([prob for tags in tags_by_category.values() for _, prob in tags])
+        print(f"  Average confidence: {avg_conf:.3f}")
+    
+    return dict(tags_by_category)

app/utils/ui_components.py

@@ -0,0 +1,137 @@
+"""
+UI components for the Image Tagger application.
+"""
+
+import os
+import streamlit as st
+from PIL import Image
+
+
+def display_progress_bar(prob):
+    """
+    Create an HTML progress bar for displaying probability.
+    
+    Args:
+        prob: Probability value between 0 and 1
+        
+    Returns:
+        HTML string for the progress bar
+    """
+    # Convert probability to percentage
+    percentage = int(prob * 100)
+    
+    # Choose color based on confidence level
+    if prob >= 0.8:
+        color = "green"
+    elif prob >= 0.5:
+        color = "orange"
+    else:
+        color = "red"
+    
+    # Return HTML for a styled progress bar
+    return f"""
+    <div style="margin-bottom: 5px; display: flex; align-items: center;">
+        <div style="flex-grow: 1; background-color: #f0f0f0; border-radius: 3px; height: 8px; position: relative;">
+            <div style="position: absolute; width: {percentage}%; background-color: {color}; height: 8px; border-radius: 3px;"></div>
+        </div>
+        <div style="margin-left: 8px; min-width: 40px; text-align: right; font-size: 0.9em;">{percentage}%</div>
+    </div>
+    """
+
+
+def show_example_images(examples_dir):
+    """
+    Display example images from a directory.
+    
+    Args:
+        examples_dir: Directory containing example images
+        
+    Returns:
+        Selected image path or None
+    """
+    selected_image = None
+    
+    if os.path.exists(examples_dir):
+        example_files = [f for f in os.listdir(examples_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+        
+        if example_files:
+            st.write("Select an example image:")
+            
+            # Create a 2-column layout for examples
+            example_cols = st.columns(2)
+            
+            for i, example_file in enumerate(example_files):
+                col_idx = i % 2
+                with example_cols[col_idx]:
+                    example_path = os.path.join(examples_dir, example_file)
+                    
+                    # Display thumbnail
+                    try:
+                        img = Image.open(example_path)
+                        st.image(img, width=150, caption=example_file)
+                        
+                        # Button to select this example
+                        if st.button(f"Use", key=f"example_{i}"):
+                            selected_image = example_path
+                            st.session_state.original_filename = example_file
+                            
+                            # Display full image
+                            st.image(img, use_container_width=True)
+                            st.success(f"Example '{example_file}' selected!")
+                    except Exception as e:
+                        st.error(f"Error loading {example_file}: {str(e)}")
+        else:
+            st.info("No example images found.")
+            st.write("Add some JPG or PNG images to the 'examples' directory.")
+    else:
+        st.info("Examples directory not found.")
+        st.write("Create an 'examples' directory and add some JPG or PNG images.")
+        
+    return selected_image
+
+
+def display_batch_results(batch_results):
+    """
+    Display batch processing results.
+    
+    Args:
+        batch_results: Dictionary with batch processing results
+    """
+    if batch_results['success']:
+        st.success(f"✅ Processed {batch_results['processed']} of {batch_results['total']} images")
+        
+        # Show details in an expander
+        with st.expander("Batch Processing Results", expanded=True):
+            # Count successes and failures
+            successes = sum(1 for r in batch_results['results'].values() if r['success'])
+            failures = batch_results['total'] - successes
+            
+            st.write(f"- Successfully tagged: {successes}")
+            st.write(f"- Failed to process: {failures}")
+            
+            if failures > 0:
+                # Show errors
+                st.write("### Processing Errors")
+                for img_path, result in batch_results['results'].items():
+                    if not result['success']:
+                        st.write(f"- **{os.path.basename(img_path)}**: {result.get('error', 'Unknown error')}")
+            
+            # Show the location of the output files
+            if successes > 0:
+                st.write("### Output Files")
+                st.write(f"Tag files have been saved to the 'saved_tags' folder.")
+                
+                # Show the first few as examples
+                st.write("Example outputs:")
+                sample_results = [(path, res) for path, res in batch_results['results'].items() if res['success']][:3]
+                for img_path, result in sample_results:
+                    output_path = result.get('output_path', '')
+                    if output_path and os.path.exists(output_path):
+                        st.write(f"- **{os.path.basename(output_path)}**")
+                        
+                        # Show file contents in a collapsible code block
+                        with open(output_path, 'r', encoding='utf-8') as f:
+                            content = f.read()
+                        st.code(content, language='text')
+    else:
+        st.error(f"Batch processing failed: {batch_results.get('error', 'Unknown error')}")

camie-tagger-v2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab0aaf253e3d546090001bec9bebc776c354ab6800f442ab9167af87b4a953ac
+size 788983561

camie-tagger-v2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81b1f1caaf84afc2e91204e1d01163487b90720f797af8b64eea331bf16b3f25
+size 572151796

config.json

@@ -0,0 +1,21 @@
+{
+  "model_name": "camie-tagger",
+  "model_type": "pytorch",
+  "models": [
+    {
+      "name": "camie-tagger-v2-safetensors",
+      "path": "camie-tagger-v2.safetensors",
+      "description": "Safetensors"
+    },
+    {
+      "name": "camie-tagger-v2-onnx",
+      "path": "camie-tagger-v2.onnx",
+      "description": "ONNX"
+    }
+  ],
+  "default_model": "camie-tagger-v2-onnx",
+  "description": "A custom tagging model",
+  "version": "1.0.0",
+  "tags": ["tagger", "classification"],
+  "license": "gpl-3.0"
+}

config.yaml

@@ -0,0 +1,26 @@
+model_info:
+  name: camie-tagger
+  version: 1.0.0
+  type: pytorch  
+  
+models:
+  - name: camie-tagger-v2-safetensors
+    path: camie-tagger-v2.safetensors
+    description: Safetensors
+
+  - name: camie-tagger-v2-onnx
+    path: camie-tagger-v2.onnx
+    description: ONNX
+    
+default_model: camie-tagger-v2-onnx
+  
+framework:
+  name: pytorch
+  version: "1.X"
+  
+metadata:
+  task: classification
+  tags:
+    - tagger
+    - classification
+  license: GPL-3.0

full_validation_results.json

@@ -0,0 +1,178 @@
+{
+  "metadata": {
+    "total_samples": 20116,
+    "num_profiles": 3,
+    "num_categories": 8,
+    "evaluation_date": "2025-08-30T19:03:30.091334"
+  },
+  "results": [
+    {
+      "CATEGORY": "OVERALL",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 67.34597711325426,
+      "MACRO-F1": 46.29084825934598
+    },
+    {
+      "CATEGORY": "YEAR",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 30.752530027663965,
+      "MACRO-F1": 21.295489610362065
+    },
+    {
+      "CATEGORY": "RATING",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 83.067447729746,
+      "MACRO-F1": 81.83200483395994
+    },
+    {
+      "CATEGORY": "GENERAL",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 66.41444399192685,
+      "MACRO-F1": 27.40425414261813
+    },
+    {
+      "CATEGORY": "CHARACTER",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 83.42532813054274,
+      "MACRO-F1": 64.54679672170435
+    },
+    {
+      "CATEGORY": "COPYRIGHT",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 86.57996129008917,
+      "MACRO-F1": 53.128368093914744
+    },
+    {
+      "CATEGORY": "ARTIST",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 70.03317039106145,
+      "MACRO-F1": 64.35697070900147
+    },
+    {
+      "CATEGORY": "META",
+      "PROFILE": "MICRO OPT",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 61.17673478553783,
+      "MACRO-F1": 19.199087616327382
+    },
+    {
+      "CATEGORY": "OVERALL",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 60.90569893586309,
+      "MACRO-F1": 50.61684357922967
+    },
+    {
+      "CATEGORY": "YEAR",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 37.21798791030943,
+      "MACRO-F1": 32.62783637958324
+    },
+    {
+      "CATEGORY": "RATING",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 78.70415252867991,
+      "MACRO-F1": 77.4818397805115
+    },
+    {
+      "CATEGORY": "GENERAL",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 60.23192293959081,
+      "MACRO-F1": 34.578761327417006
+    },
+    {
+      "CATEGORY": "CHARACTER",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 79.91802201844503,
+      "MACRO-F1": 66.17825085829298
+    },
+    {
+      "CATEGORY": "COPYRIGHT",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 81.84935439943561,
+      "MACRO-F1": 56.18607804162489
+    },
+    {
+      "CATEGORY": "ARTIST",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 62.32690082490389,
+      "MACRO-F1": 66.10039817274198
+    },
+    {
+      "CATEGORY": "META",
+      "PROFILE": "MACRO OPT",
+      "THRESHOLD": 0.49183673469387756,
+      "MICRO-F1": 56.30398796087284,
+      "MACRO-F1": 23.673363465036253
+    },
+    {
+      "CATEGORY": "OVERALL",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 67.34597711325426,
+      "MACRO-F1": 46.29084825934598
+    },
+    {
+      "CATEGORY": "YEAR",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 30.752530027663965,
+      "MACRO-F1": 21.295489610362065
+    },
+    {
+      "CATEGORY": "RATING",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 83.067447729746,
+      "MACRO-F1": 81.83200483395994
+    },
+    {
+      "CATEGORY": "GENERAL",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 66.41444399192685,
+      "MACRO-F1": 27.40425414261813
+    },
+    {
+      "CATEGORY": "CHARACTER",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 83.42532813054274,
+      "MACRO-F1": 64.54679672170435
+    },
+    {
+      "CATEGORY": "COPYRIGHT",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 86.57996129008917,
+      "MACRO-F1": 53.128368093914744
+    },
+    {
+      "CATEGORY": "ARTIST",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 70.03317039106145,
+      "MACRO-F1": 64.35697070900147
+    },
+    {
+      "CATEGORY": "META",
+      "PROFILE": "BALANCED",
+      "THRESHOLD": 0.6142857142857143,
+      "MICRO-F1": 61.17673478553783,
+      "MACRO-F1": 19.199087616327382
+    }
+  ]
+}

game/dev_tools.py

@@ -0,0 +1,580 @@
+#!/usr/bin/env python3
+"""
+Developer Tools for Tag Collector Game
+A hidden panel with tools for testing and debugging game features.
+"""
+
+import streamlit as st
+import random
+import time
+from game_constants import (
+    TAG_CURRENCY_NAME, 
+    RARITY_LEVELS, 
+    ACHIEVEMENTS,
+)
+
+from tag_categories import (
+    TAG_CATEGORIES,
+    TAG_DETECTOR_UPGRADES,
+    PROGRESSION_ACHIEVEMENTS
+)
+
+def display_dev_tools():
+    """
+    Display the developer tools interface
+    This should be hidden in production builds or behind a developer toggle
+    """
+    st.title("🛠️ Developer Tools")
+    st.warning("These tools are for testing and debugging only. They can break game balance!")
+    
+    # Create tabs for different tool categories
+    resource_tab, tag_tab, progression_tab, mosaic_tab, reset_tab = st.tabs([
+        "Resources", "Tag Management", "Progression", "Mosaic Tools", "Reset Tools"
+    ])
+    
+    with resource_tab:
+        display_resource_tools()
+    
+    with tag_tab:
+        display_tag_tools()
+    
+    with progression_tab:
+        display_progression_tools()
+    
+    with mosaic_tab:
+        display_mosaic_tools()
+    
+    with reset_tab:
+        display_reset_tools()
+    
+def display_resource_tools():
+    """Display tools for managing game resources"""
+    st.subheader("Currency and Resources")
+    
+    # Add TagCoins
+    col1, col2 = st.columns([3, 1])
+    with col1:
+        amount = st.number_input("Amount of TagCoins to add:", min_value=0, max_value=1000000, value=1000, step=100)
+    with col2:
+        if st.button("Add TagCoins", key="add_currency"):
+            st.session_state.tag_currency += amount
+            st.session_state.game_stats["total_currency_earned"] += amount
+            st.success(f"Added {amount} {TAG_CURRENCY_NAME}!")
+    
+    # Set threshold directly
+    col1, col2 = st.columns([3, 1])
+    with col1:
+        threshold = st.slider("Set threshold value:", min_value=0.0, max_value=1.0, value=st.session_state.threshold, step=0.01)
+    with col2:
+        if st.button("Set Threshold", key="set_threshold"):
+            st.session_state.threshold = threshold
+            st.success(f"Set threshold to {threshold:.2f}")
+    
+    # Add tag power bonuses
+    col1, col2 = st.columns([3, 1])
+    with col1:
+        power = st.number_input("Add tag power bonus:", min_value=0.0, max_value=0.1, value=0.01, step=0.001, format="%.3f")
+    with col2:
+        if st.button("Add Power", key="add_power"):
+            if not hasattr(st.session_state, 'tag_power_bonus'):
+                st.session_state.tag_power_bonus = 0
+            st.session_state.tag_power_bonus += power
+            st.success(f"Added {power:.3f} tag power!")
+
+def display_tag_tools():
+    """Display tools for managing tags"""
+    st.subheader("Tag Management")
+    
+    # Add specific tag
+    with st.expander("Add Specific Tag", expanded=True):
+        col1, col2, col3 = st.columns([4, 2, 1])
+        
+        with col1:
+            tag_name = st.text_input("Tag name:", value="custom_tag")
+        
+        with col2:
+            rarities = list(RARITY_LEVELS.keys())
+            rarity = st.selectbox("Rarity:", rarities)
+        
+        with col3:
+            # Get categories from session state or fallback to general
+            categories = ["general", "character", "copyright", "meta", "rating", "artist", "year"]
+            category = st.selectbox("Category:", categories)
+        
+        if st.button("Add Tag", key="add_specific_tag"):
+            # Check if tag already exists
+            is_new = tag_name not in st.session_state.collected_tags
+            
+            # Add tag to collection
+            st.session_state.collected_tags[tag_name] = {
+                "count": 1,
+                "rarity": rarity,
+                "category": category,
+                "discovery_time": time.strftime("%Y-%m-%d %H:%M:%S")
+            }
+            
+            # Show confirmation
+            if is_new:
+                st.success(f"Added new tag '{tag_name}' ({rarity}, {category})")
+            else:
+                st.session_state.collected_tags[tag_name]["count"] += 1
+                st.info(f"Incremented count for existing tag '{tag_name}'")
+    
+    # Generate random tags
+    with st.expander("Generate Random Tags", expanded=False):
+        col1, col2 = st.columns([3, 1])
+        
+        with col1:
+            num_tags = st.number_input("Number of random tags to generate:", min_value=1, max_value=1000, value=10)
+            
+            # Options for distribution
+            advanced = st.checkbox("Advanced options")
+            if advanced:
+                st.write("Rarity distribution:")
+                common_pct = st.slider("Common tags %:", 0, 100, 70)
+                uncommon_pct = st.slider("Uncommon tags %:", 0, 100, 20)
+                rare_pct = st.slider("Rare tags %:", 0, 100, 8)
+                super_rare_pct = st.slider("Super rare tags %:", 0, 100, 2)
+                
+                # Ensure total is 100%
+                total = common_pct + uncommon_pct + rare_pct + super_rare_pct
+                if total != 100:
+                    st.warning(f"Distribution totals {total}%, should be 100%")
+        
+        with col2:
+            generate_button = st.button("Generate", key="generate_random_tags")
+        
+        if generate_button:
+            generated_count = 0
+            
+            # Determine distribution of rarities
+            if advanced and total == 100:
+                # Custom distribution
+                rarity_weights = {
+                    "Whispered Word": common_pct / 100,
+                    "Common Canard": uncommon_pct / 100 * 0.6,
+                    "Urban Footnote": uncommon_pct / 100 * 0.4,
+                    "Urban Myth": rare_pct / 100 * 0.5,
+                    "Urban Legend": rare_pct / 100 * 0.5,
+                    "Urban Nightmare": super_rare_pct / 100 * 0.8,
+                    "Impuritas Civitas": super_rare_pct / 100 * 0.2
+                }
+            else:
+                # Default distribution
+                rarity_weights = {
+                    "Whispered Word": 0.70,
+                    "Common Canard": 0.15,
+                    "Urban Footnote": 0.08,
+                    "Urban Myth": 0.04,
+                    "Urban Legend": 0.02,
+                    "Urban Nightmare": 0.008,
+                    "Impuritas Civitas": 0.002
+                }
+            
+            # Generate the tags
+            for i in range(num_tags):
+                # Create a random tag name if we don't have metadata
+                tag_name = f"random_tag_{int(time.time() % 10000)}_{i}"
+                
+                # Determine rarity
+                rarity = random.choices(
+                    list(rarity_weights.keys()),
+                    weights=list(rarity_weights.values()),
+                    k=1
+                )[0]
+                
+                # Determine category
+                categories = list(TAG_CATEGORIES.keys())
+                category = random.choice(categories)
+                
+                # Add to collection
+                timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
+                
+                # Check if this is a new tag
+                is_new = tag_name not in st.session_state.collected_tags
+                
+                if is_new:
+                    st.session_state.collected_tags[tag_name] = {
+                        "count": 1,
+                        "rarity": rarity,
+                        "category": category,
+                        "discovery_time": timestamp
+                    }
+                    generated_count += 1
+                else:
+                    # Increment count if already exists
+                    st.session_state.collected_tags[tag_name]["count"] += 1
+            
+            # Show confirmation
+            st.success(f"Generated {generated_count} new random tags!")
+
+def display_progression_tools():
+    """Display tools for managing progression"""
+    st.subheader("Progression System Tools")
+    
+    # Unlock categories
+    with st.expander("Unlock Categories", expanded=True):
+        st.write("Select categories to unlock:")
+        
+        # Get currently unlocked categories
+        unlocked = []
+        if hasattr(st.session_state, 'unlocked_tag_categories'):
+            unlocked = st.session_state.unlocked_tag_categories
+        
+        # Display each category with a checkbox
+        category_checkboxes = {}
+        for category, info in TAG_CATEGORIES.items():
+            # Skip default unlocked
+            if info["unlocked_by_default"]:
+                continue
+                
+            # Check if already unlocked
+            is_unlocked = category in unlocked
+            category_checkboxes[category] = st.checkbox(
+                f"{info['name']} ({info['cost']} {TAG_CURRENCY_NAME})",
+                value=is_unlocked,
+                key=f"cat_{category}"
+            )
+        
+        # Button to apply changes
+        if st.button("Apply Category Changes", key="apply_categories"):
+            # Initialize if needed
+            if not hasattr(st.session_state, 'unlocked_tag_categories'):
+                st.session_state.unlocked_tag_categories = []
+                
+                # Add default unlocked
+                for cat, info in TAG_CATEGORIES.items():
+                    if info["unlocked_by_default"]:
+                        st.session_state.unlocked_tag_categories.append(cat)
+            
+            # Update unlocked categories
+            for category, checked in category_checkboxes.items():
+                # If checked but not unlocked, add it
+                if checked and category not in st.session_state.unlocked_tag_categories:
+                    st.session_state.unlocked_tag_categories.append(category)
+                    st.success(f"Unlocked {TAG_CATEGORIES[category]['name']}!")
+                
+                # If unchecked but unlocked, remove it
+                elif not checked and category in st.session_state.unlocked_tag_categories:
+                    st.session_state.unlocked_tag_categories.remove(category)
+                    st.info(f"Locked {TAG_CATEGORIES[category]['name']}")
+    
+    # Upgrade detector level
+    with st.expander("Set Detector Level", expanded=False):
+        # Get current level
+        current_level = 0
+        if hasattr(st.session_state, 'detector_level'):
+            current_level = st.session_state.detector_level
+            
+        # Display slider for detector level
+        new_level = st.slider(
+            "Detector Level:",
+            min_value=0,
+            max_value=len(TAG_DETECTOR_UPGRADES) - 1,
+            value=current_level
+        )
+        
+        # Show info about selected level
+        upgrade = TAG_DETECTOR_UPGRADES[new_level]
+        max_tags = upgrade["max_tags"]
+        if max_tags == 0:
+            st.write(f"Selected: {upgrade['name']} (Unlimited tags)")
+        else:
+            st.write(f"Selected: {upgrade['name']} ({max_tags} tags)")
+        
+        # Button to apply changes
+        if st.button("Set Detector Level", key="set_detector_level"):
+            st.session_state.detector_level = new_level
+            st.success(f"Set detector level to {new_level} ({upgrade['name']})")
+    
+    # Unlock achievements
+    with st.expander("Manage Achievements", expanded=False):
+        # Combine standard and progression achievements
+        all_achievements = {**ACHIEVEMENTS, **PROGRESSION_ACHIEVEMENTS}
+        
+        # Initialize achievements if needed
+        if not hasattr(st.session_state, 'achievements'):
+            st.session_state.achievements = set()
+        
+        # Create tabs for unlocked and locked
+        unlocked_tab, locked_tab = st.tabs(["Unlocked", "Locked"])
+        
+        with unlocked_tab:
+            st.write("Currently unlocked achievements:")
+            
+            # Show unlocked achievements with option to remove
+            for achievement_id in sorted(st.session_state.achievements):
+                if achievement_id in all_achievements:
+                    col1, col2 = st.columns([3, 1])
+                    
+                    with col1:
+                        achievement = all_achievements[achievement_id]
+                        st.write(f"**{achievement['name']}**: {achievement['description']}")
+                    
+                    with col2:
+                        if st.button("Remove", key=f"remove_{achievement_id}"):
+                            st.session_state.achievements.remove(achievement_id)
+                            st.info(f"Removed achievement: {achievement['name']}")
+                            st.rerun()
+        
+        with locked_tab:
+            st.write("Currently locked achievements:")
+            
+            # Show locked achievements with option to add
+            locked_achievements = [a for a in all_achievements if a not in st.session_state.achievements]
+            
+            for achievement_id in sorted(locked_achievements):
+                col1, col2 = st.columns([3, 1])
+                
+                with col1:
+                    achievement = all_achievements[achievement_id]
+                    st.write(f"**{achievement['name']}**: {achievement['description']}")
+                
+                with col2:
+                    if st.button("Unlock", key=f"unlock_{achievement_id}"):
+                        st.session_state.achievements.add(achievement_id)
+                        
+                        # Apply rewards if applicable
+                        if "reward" in achievement:
+                            from scan_handler import apply_achievement_reward
+                            apply_achievement_reward(achievement_id, achievement["reward"])
+                            
+                        st.success(f"Unlocked achievement: {achievement['name']}")
+                        st.rerun()
+
+def display_mosaic_tools():
+    """Display tools for managing the tag mosaic"""
+    st.subheader("Tag Mosaic Tools")
+    
+    # Check if mosaic exists
+    has_mosaic = hasattr(st.session_state, 'tag_mosaic')
+    
+    if not has_mosaic:
+        st.warning("Tag Mosaic not initialized yet. Visit the Tag Collection tab first.")
+        return
+    
+    # Fill random portions of the mosaic
+    with st.expander("Fill Random Portions", expanded=True):
+        col1, col2 = st.columns([3, 1])
+        
+        with col1:
+            fill_percentage = st.slider(
+                "Percentage to fill:",
+                min_value=0,
+                max_value=100,
+                value=10,
+                step=1
+            )
+            
+            # Options for distribution
+            st.write("Fill with tags of rarity:")
+            fill_rarities = {}
+            for rarity in RARITY_LEVELS:
+                fill_rarities[rarity] = st.checkbox(rarity, value=True, key=f"fill_{rarity}")
+        
+        with col2:
+            fill_button = st.button("Fill Mosaic", key="fill_mosaic")
+        
+        if fill_button:
+            # Get the mosaic from session state
+            mosaic = st.session_state.tag_mosaic
+            
+            # Calculate how many cells to fill
+            total_cells = mosaic.total_cells
+            existing_filled = len(mosaic.filled_cells)
+            target_filled = int(total_cells * fill_percentage / 100)
+            cells_to_add = max(0, target_filled - existing_filled)
+            
+            # Get active rarities
+            active_rarities = [r for r, checked in fill_rarities.items() if checked]
+            if not active_rarities:
+                st.error("Select at least one rarity to fill with")
+                return
+            
+            # Create artificial tags and add them
+            added_count = 0
+            added_tags = {}
+            
+            # Generate random positions
+            all_positions = [(x, y) for x in range(mosaic.grid_width) for y in range(mosaic.grid_height)]
+            # Remove already filled positions
+            available_positions = [pos for pos in all_positions if pos not in mosaic.filled_cells]
+            
+            # If we need more than available, just use what's available
+            cells_to_add = min(cells_to_add, len(available_positions))
+            
+            # Randomly select positions
+            selected_positions = random.sample(available_positions, cells_to_add)
+            
+            # Create tags for each position
+            for pos in selected_positions:
+                x, y = pos
+                
+                # Create a tag name 
+                tag_name = f"mosaic_fill_{x}_{y}_{int(time.time() % 10000)}"
+                
+                # Select a random rarity from active rarities
+                rarity = random.choice(active_rarities)
+                
+                # Add to tags dictionary (this won't be saved to session_state)
+                added_tags[tag_name] = {
+                    "count": 1,
+                    "rarity": rarity,
+                    "category": "general"
+                }
+                
+                added_count += 1
+            
+            # Update the mosaic (this does save to disk)
+            if added_count > 0:
+                mosaic.update_with_tags(added_tags)
+                st.success(f"Added {added_count} random cells to the mosaic!")
+                
+                # Show updated stats
+                stats = mosaic.get_stats()
+                st.write(f"New completion: {stats['completion_percentage']:.2f}%")
+                st.write(f"Emerging pattern: {stats['completion_pattern']}")
+                
+                # Show image
+                mosaic_img = mosaic.get_image(show_highlights=True)
+                st.image(mosaic_img, caption="Updated Mosaic", width=400)
+            else:
+                st.info("No new cells added. Mosaic may already be filled to the requested level.")
+    
+    # Reset mosaic without affecting collection
+    with st.expander("Reset Mosaic", expanded=False):
+        if st.button("Reset Mosaic", key="reset_mosaic"):
+            # Confirm
+            confirm = st.checkbox("I understand this will clear the mosaic visualization (not your collection)")
+            
+            if confirm:
+                # Get the mosaic from session state
+                mosaic = st.session_state.tag_mosaic
+                
+                # Reset the mosaic by creating a new one
+                from tag_mosaic import TagMosaic
+                st.session_state.tag_mosaic = TagMosaic()
+                
+                # Delete the mosaic save file
+                import os
+                if os.path.exists("tag_mosaic.png"):
+                    try:
+                        os.remove("tag_mosaic.png")
+                    except Exception as e:
+                        st.error(f"Error removing mosaic file: {e}")
+                
+                st.success("Mosaic has been reset!")
+
+def display_reset_tools():
+    """Display tools for resetting the game"""
+    st.subheader("Reset Tools")
+    st.warning("These tools will reset parts of your game progress. Use with caution!")
+    
+    # Reset currency
+    with st.expander("Reset Currency", expanded=False):
+        col1, col2 = st.columns([3, 1])
+        
+        with col1:
+            new_amount = st.number_input("Set currency to:", min_value=0, value=0)
+            
+        with col2:
+            if st.button("Reset Currency", key="reset_currency"):
+                st.session_state.tag_currency = new_amount
+                st.success(f"Reset currency to {new_amount} {TAG_CURRENCY_NAME}")
+    
+    # Reset collection
+    with st.expander("Reset Collection", expanded=False):
+        st.write("This will remove all collected tags or specific rarities.")
+        
+        # Options to keep certain rarities
+        st.write("Keep tags with these rarities:")
+        keep_rarities = {}
+        for rarity in RARITY_LEVELS:
+            keep_rarities[rarity] = st.checkbox(rarity, value=False, key=f"keep_{rarity}")
+        
+        if st.button("Reset Collection", key="reset_collection"):
+            # Confirm
+            confirm = st.checkbox("I understand this will delete collected tags")
+            
+            if confirm:
+                # Get rarities to keep
+                rarities_to_keep = [r for r, checked in keep_rarities.items() if checked]
+                
+                # If keeping some rarities, filter the collection
+                if rarities_to_keep:
+                    # Create a new collection with only the kept rarities
+                    kept_tags = {}
+                    for tag, info in st.session_state.collected_tags.items():
+                        if info.get("rarity") in rarities_to_keep:
+                            kept_tags[tag] = info
+                    
+                    # Replace the collection
+                    removed_count = len(st.session_state.collected_tags) - len(kept_tags)
+                    st.session_state.collected_tags = kept_tags
+                    st.success(f"Removed {removed_count} tags. Kept {len(kept_tags)} tags with rarities: {', '.join(rarities_to_keep)}")
+                else:
+                    # Remove all tags
+                    removed_count = len(st.session_state.collected_tags)
+                    st.session_state.collected_tags = {}
+                    st.success(f"Removed all {removed_count} tags from your collection")
+    
+    # Reset complete game
+    with st.expander("Reset ENTIRE Game", expanded=False):
+        st.error("This will reset ALL game progress including collection, currency, achievements, and upgrades.")
+        
+        if st.button("Reset EVERYTHING", key="reset_everything"):
+            # Double confirm
+            confirm1 = st.checkbox("I understand ALL progress will be lost")
+            confirm2 = st.checkbox("This cannot be undone")
+            
+            if confirm1 and confirm2:
+                # Reset everything
+                st.session_state.threshold = 0.25  # Default starting threshold
+                st.session_state.tag_currency = 0
+                st.session_state.collected_tags = {}
+                st.session_state.purchased_upgrades = []
+                st.session_state.achievements = set()
+                st.session_state.tag_history = []
+                st.session_state.current_scan = None
+                st.session_state.game_stats = {
+                    "images_processed": 0,
+                    "total_tags_found": 0,
+                    "total_currency_earned": 0,
+                    "currency_spent": 0
+                }
+                
+                # Reset progression
+                if hasattr(st.session_state, 'unlocked_tag_categories'):
+                    st.session_state.unlocked_tag_categories = []
+                    
+                    # Add default unlocked categories
+                    for cat, info in TAG_CATEGORIES.items():
+                        if info["unlocked_by_default"]:
+                            st.session_state.unlocked_tag_categories.append(cat)
+                
+                if hasattr(st.session_state, 'detector_level'):
+                    st.session_state.detector_level = 0
+                
+                if hasattr(st.session_state, 'tag_power_bonus'):
+                    st.session_state.tag_power_bonus = 0
+                
+                if hasattr(st.session_state, 'coin_multiplier'):
+                    st.session_state.coin_multiplier = 1.0
+                
+                if hasattr(st.session_state, 'essence_generator_count'):
+                    st.session_state.essence_generator_count = 0
+                
+                # Reset mosaic
+                import os
+                if os.path.exists("tag_mosaic.png"):
+                    try:
+                        os.remove("tag_mosaic.png")
+                    except Exception as e:
+                        st.error(f"Error removing mosaic file: {e}")
+                
+                if hasattr(st.session_state, 'tag_mosaic'):
+                    from tag_mosaic import TagMosaic
+                    st.session_state.tag_mosaic = TagMosaic()
+                
+                st.success("Game completely reset to initial state!")
+                st.info("Refresh the page to see changes take effect")

game/game_constants.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python3
+"""
+Game Constants for Tag Collector Game - Updated with new tag rarity names
+This file contains shared constants used by both the main game and the library system.
+"""
+
+# Game currency names
+TAG_CURRENCY_NAME = "TagCoins"
+ENKEPHALIN_CURRENCY_NAME = "Enkephalin"
+ENKEPHALIN_ICON = "💧"
+
+STARTING_THRESHOLD = 0.55
+MIN_THRESHOLD = 0.1
+
+# Tag animations and theme settings
+TAG_ANIMATIONS = {
+    "Star of the City": {
+        "css_class": "star-of-city",
+        "animation": """
+        @keyframes glowing {
+            0% { box-shadow: 0 0 5px #FFD700; }
+            50% { box-shadow: 0 0 20px #FFD700; }
+            100% { box-shadow: 0 0 5px #FFD700; }
+        }
+        .star-of-city {
+            background-color: rgba(255, 215, 0, 0.2);
+            padding: 8px;
+            border-radius: 5px;
+            border: 2px solid gold;
+            animation: glowing 2s infinite;
+        }
+        """
+    },
+    "Impuritas Civitas": {
+        "css_class": "impuritas-civitas",
+        "animation": """
+        @keyframes rainbow-border {
+            0% { border-color: red; }
+            14% { border-color: orange; }
+            28% { border-color: yellow; }
+            42% { border-color: green; }
+            57% { border-color: blue; }
+            71% { border-color: indigo; }
+            85% { border-color: violet; }
+            100% { border-color: red; }
+        }
+        
+        @keyframes rainbow-text {
+            0% { color: red; }
+            14% { color: orange; }
+            28% { color: yellow; }
+            42% { color: green; }
+            57% { color: blue; }
+            71% { color: indigo; }
+            85% { color: violet; }
+            100% { color: red; }
+        }
+        
+        @keyframes rainbow-bg {
+            0% { background-color: rgba(255,0,0,0.1); }
+            14% { background-color: rgba(255,165,0,0.1); }
+            28% { background-color: rgba(255,255,0,0.1); }
+            42% { background-color: rgba(0,128,0,0.1); }
+            57% { background-color: rgba(0,0,255,0.1); }
+            71% { background-color: rgba(75,0,130,0.1); }
+            85% { background-color: rgba(238,130,238,0.1); }
+            100% { background-color: rgba(255,0,0,0.1); }
+        }
+        
+        .impuritas-civitas {
+            background-color: rgba(0, 0, 0, 0.1);
+            padding: 10px;
+            border-radius: 5px;
+            border: 3px solid red;
+            animation: rainbow-border 4s linear infinite, rainbow-bg 4s linear infinite;
+        }
+        
+        .impuritas-text {
+            font-weight: bold;
+            animation: rainbow-text 4s linear infinite;
+        }
+        """
+    }
+}
+
+# Rarity levels with appropriate colors (updated to match new rarity tiers)
+RARITY_LEVELS = {
+    "Canard": {"color": "#AAAAAA", "value": 1},                # Gray
+    "Urban Myth": {"color": "#5D9C59", "value": 5},            # Green
+    "Urban Legend": {"color": "#2196F3", "value": 10},          # Blue
+    "Urban Plague": {"color": "#9C27B0", "value": 25},         # Purple
+    "Urban Nightmare": {"color": "#FF9800", "value": 50},      # Orange
+    "Star of the City": {"color": "#FFEB3B", "value": 250},    # Yellow/Gold
+    "Impuritas Civitas": {"color": "#F44336", "value": 1000}    # Red
+}
+
+# Essence generation costs in enkephalin
+ESSENCE_COSTS = {
+    "Canard": 10,                # Common tags
+    "Urban Myth": 30,            # Uncommon tags
+    "Urban Legend": 75,          # Rare tags
+    "Urban Plague": 150,         # Very rare tags
+    "Urban Nightmare": 300,      # Extremely rare tags
+    "Star of the City": 600,     # Nearly mythical tags
+    "Impuritas Civitas": 1200    # Legendary tags
+}
+
+# Tag power system
+TAG_POWER_BONUSES = {
+    "Canard": {"coin_multiplier": 0, "enkephalin_reward": 0},
+    "Urban Myth": {"coin_multiplier": 0, "enkephalin_reward": 0},
+    "Urban Legend": {"coin_multiplier": 0, "enkephalin_reward": 1},
+    "Urban Plague": {"coin_multiplier": 0.001, "enkephalin_reward": 3},
+    "Urban Nightmare": {"coin_multiplier": 0.0025, "enkephalin_reward": 5},
+    "Star of the City": {"coin_multiplier": 0.005, "enkephalin_reward": 10},
+    "Impuritas Civitas": {"coin_multiplier": 0.01, "enkephalin_reward": 25}
+}
+
+THRESHOLD_UPGRADES = [
+    {
+        "name": "Pattern Recognition Module", 
+        "threshold_setting": 0.48367345,  # High precision threshold
+        "cost": 300, 
+        "description": "Basic algorithm focused on high-precision identification. Reduces false positives but may miss some tags."
+    },
+    {
+        "name": "Neural Network Enhancement", 
+        "threshold_setting": 0.40000000, 
+        "cost": 500, 
+        "description": "Improved tag detection using multi-layer perceptrons. Offers good precision with moderate recall."
+    },
+    {
+        "name": "Deep Learning Framework", 
+        "threshold_setting": 0.35000000, 
+        "cost": 1000, 
+        "description": "Advanced algorithms that learn from previous scans. Provides better balance between precision and recall."
+    },
+    {
+        "name": "Quantum Probability Engine", 
+        "threshold_setting": 0.32857141,  # Balanced optimal F1 score threshold
+        "cost": 2500, 
+        "description": "Leverages quantum uncertainty for optimal detection balance. Perfect calibration point for F1 score."
+    },
+    {
+        "name": "Recursive Self-Improvement", 
+        "threshold_setting": 0.31224489,  # Weighted F1 threshold
+        "cost": 7500, 
+        "description": "Scanner enhances its own detection capabilities. Optimized for weighted tag discovery."
+    },
+    {
+        "name": "Consciousness Emulation", 
+        "threshold_setting": 0.25000000, 
+        "cost": 15000, 
+        "description": "Scanner develops intuition-like abilities. Favors higher recall while maintaining reasonable precision."
+    },
+    {
+        "name": "Technological Singularity", 
+        "threshold_setting": 0.20612246,  # High recall threshold
+        "cost": 50000, 
+        "description": "The scanner transcends conventional limitations. Maximizes tag discovery at the cost of some precision."
+    }
+]
+
+# Achievements
+ACHIEVEMENTS = {
+    # Collection achievements
+    "tag_collector_beginner": {"name": "Novice Archivist", "requirement": 50, "description": "Collect 50 different tags", "reward": {"coin_bonus": 0.01}},
+    "tag_collector_expert": {"name": "Senior Cataloger", "requirement": 250, "description": "Collect 250 different tags", "reward": {"coin_bonus": 0.01}},
+    "tag_collector_master": {"name": "Master Librarian", "requirement": 500, "description": "Collect 500 different tags", "reward": {"coin_bonus": 0.01}},
+    "tag_master": {"name": "Grand Archivist", "requirement": 1000, "description": "Collect 1000 different tags", "reward": {"coin_bonus": 0.01}},
+    
+    # Rarity achievements
+    "legendary_hunter": {"name": "Impuritas Seeker", "requirement": 1, "description": "Find your first Impuritas Civitas tag", "reward": {"coin_bonus": 0.01, "enkephalin": 50}},
+    "multi_legendary": {"name": "Forbidden Collection", "requirement": 5, "description": "Collect 5 Impuritas Civitas tags", "reward": {"coin_bonus": 0.01, "enkephalin": 100}},
+    "canard_collector": {"name": "Canard Chronicler", "requirement": 30, "description": "Collect 30 Canard tags", "reward": {"coin_bonus": 0.01}},
+    "urban_myth_collector": {"name": "Myth Curator", "requirement": 15, "description": "Collect 15 Urban Myth tags", "reward": {"coin_bonus": 0.01}},
+    "urban_legend_collector": {"name": "Legend Preserver", "requirement": 10, "description": "Collect 10 Urban Legend tags", "reward": {"coin_bonus": 0.01}},
+    "urban_plague_collector": {"name": "Plague Archivist", "requirement": 5, "description": "Collect 5 Urban Plague tags", "reward": {"coin_bonus": 0.01}},
+    "urban_nightmare_collector": {"name": "Nightmare Keeper", "requirement": 5, "description": "Collect 5 Urban Nightmare tags", "reward": {"coin_bonus": 0.01}},
+    "star_collector": {"name": "Star Collector", "requirement": 3, "description": "Collect 3 Star of the City tags", "reward": {"coin_bonus": 0.01, "enkephalin": 30}},
+    "impuritas_collector": {"name": "Impuritas Scholar", "requirement": 3, "description": "Collect 3 Impuritas Civitas tags", "reward": {"coin_bonus": 0.01, "enkephalin": 75}},
+    
+    # Progress achievements
+    "perfect_scanner": {"name": "Omniscient Observer", "description": "Reach the minimum threshold", "reward": {"coin_bonus": 0.01}},
+    "optimal_threshold": {"name": "Perfect Calibration", "description": "Reach the optimal F1 score threshold of 0.328", "reward": {"coin_bonus": 0.01}},
+    "collection_milestone_100": {"name": "Century Collector", "requirement": 100, "description": "Collect 100 different tags", "reward": {"tagcoins": 100, "coin_bonus": 0.01}},
+    "collection_milestone_1000": {"name": "Millennium Collector", "requirement": 1000, "description": "Collect 1000 different tags", "reward": {"tagcoins": 1000, "coin_bonus": 0.01}},
+    "collection_milestone_5000": {"name": "Epic Collector", "requirement": 5000, "description": "Collect 5000 different tags", "reward": {"tagcoins": 5000, "coin_bonus": 0.01}},
+    
+    # Essence & library achievements
+    "essence_creator": {"name": "Essence Creator", "requirement": 5, "description": "Generate 5 tag essences", "reward": {"essence_cost_reduction": 0.2, "coin_bonus": 0.01}},
+    "tag_explorer": {"name": "Tag Explorer", "requirement": 20, "description": "Explore all library tiers", "reward": {"library_cost_reduction": 0.15, "coin_bonus": 0.01}},
+    "enkephalin_master": {"name": "Enkephalin Master", "requirement": 5000, "description": "Generate 5000 Enkephalin", "reward": {"essence_cost_reduction": 0.25, "coin_bonus": 0.01}},
+    "sacrifice_devotee": {"name": "Sacrifice Devotee", "requirement": 100, "description": "Sacrifice 100 tags", "reward": {"enkephalin_bonus": 0.2, "coin_bonus": 0.01}},
+    
+    # New achievements
+    "category_explorer": {"name": "Category Explorer", "requirement": 10, "description": "Collect tags from 10 different categories", "reward": {"coin_bonus": 0.01}},
+    "series_collector": {"name": "Series Collector", "requirement": 3, "description": "Complete 3 series mosaics", "reward": {"coin_bonus": 0.01, "enkephalin": 25}},
+    "rapid_tagger": {"name": "Rapid Tagger", "requirement": 100, "description": "Scan 100 images", "reward": {"coin_bonus": 0.01}},
+    "enkephalin_harvester": {"name": "Enkephalin Harvester", "requirement": 1000, "description": "Generate 1000 Enkephalin", "reward": {"enkephalin_bonus": 0.1, "coin_bonus": 0.01}},
+    "library_scholar": {"name": "Library Scholar", "requirement": 50, "description": "Extract 50 tags from the library", "reward": {"library_cost_reduction": 0.1, "coin_bonus": 0.01}},
+    "rarity_hunter": {"name": "Rarity Hunter", "description": "Find tags of all rarity levels", "reward": {"coin_bonus": 0.02}},
+    "essence_master": {"name": "Essence Master", "requirement": 25, "description": "Generate 25 tag essences", "reward": {"essence_cost_reduction": 0.15, "coin_bonus": 0.01}},
+    "legendary_librarian": {"name": "Legendary Librarian", "description": "Extract an Impuritas Civitas tag from the library", "reward": {"library_cost_reduction": 0.2, "coin_bonus": 0.01, "enkephalin": 100}}
+}

game/library_system.py

@@ -0,0 +1,2010 @@
+"""
+Redesigned Library System with Instant Discovery and Cooldown for Tag Collector Game
+"""
+
+import streamlit as st
+import random
+import time
+import math
+import pandas as pd
+import datetime
+from game_constants import (
+    TAG_CURRENCY_NAME, 
+    RARITY_LEVELS, 
+    ENKEPHALIN_CURRENCY_NAME, 
+    ENKEPHALIN_ICON,
+    TAG_POWER_BONUSES
+)
+from essence_generator import display_essence_generator
+from tag_categories import (
+    TAG_CATEGORIES, 
+    get_collection_power_level
+)
+import tag_storage
+
+# Define library properties
+LIBRARY_INFO = {
+    "name": "The Library",
+    "description": "A vast repository of knowledge where tags are discovered through patient exploration and research.",
+    "color": "#4A148C",  # Deep purple
+    "rarities_available": ["Canard", "Urban Myth", "Urban Legend", "Urban Plague", "Urban Nightmare", "Star of the City", "Impuritas Civitas"],
+    "odds_multiplier": 2.0
+}
+
+# Define library floors with their unlocking requirements and rarity boosts
+LIBRARY_FLOORS = [
+    {
+        "name": "Floor of General Works",
+        "description": "The foundation of knowledge. Contains basic tags with limited rarity.",
+        "required_tags": 0,  # Available from the start
+        "rarity_boost": 0.0,
+        "color": "#8D99AE",  # Light blue-gray
+        "unlocked": True,  # Always unlocked
+        "rarities": ["Canard", "Urban Myth", "Urban Legend"],
+        "odds_multiplier": 1.0  # Base odds multiplier
+    },
+    {
+        "name": "Floor of History",
+        "description": "Archives of past knowledge. Offers more access to uncommon tags.",
+        "required_tags": 25,  # Unlocked after collecting 25 tags
+        "rarity_boost": 0.2,
+        "color": "#457B9D",  # Moderate blue
+        "rarities": ["Canard", "Urban Myth", "Urban Legend", "Urban Plague"],
+        "odds_multiplier": 1.2
+    },
+    {
+        "name": "Floor of Technological Sciences",
+        "description": "Repository of technical knowledge. Access to rare tags begins here.",
+        "required_tags": 75,  # Unlocked after collecting 75 tags
+        "rarity_boost": 0.4,
+        "color": "#2B9348",  # Green
+        "rarities": ["Canard", "Urban Myth", "Urban Legend", "Urban Plague"],
+        "odds_multiplier": 1.5
+    },
+    {
+        "name": "Floor of Literature",
+        "description": "A vast collection of narrative concepts. Higher chance of rare discoveries.",
+        "required_tags": 150,  # Unlocked after collecting 150 tags
+        "rarity_boost": 0.6,
+        "color": "#6A0572",  # Purple
+        "rarities": ["Canard", "Urban Myth", "Urban Legend", "Urban Plague", "Urban Nightmare"],
+        "odds_multiplier": 1.8
+    },
+    {
+        "name": "Floor of Art",
+        "description": "The realm of aesthetic concepts. First access to Urban Nightmare tags.",
+        "required_tags": 250,  # Unlocked after collecting 250 tags
+        "rarity_boost": 0.8,
+        "color": "#D90429",  # Red
+        "rarities": ["Canard", "Urban Myth", "Urban Legend", "Urban Plague", "Urban Nightmare"],
+        "odds_multiplier": 2.2
+    },
+    {
+        "name": "Floor of Natural Sciences",
+        "description": "Where empirical knowledge is cataloged. Significant chance of very rare tags.",
+        "required_tags": 500,  # Unlocked after collecting 400 tags
+        "rarity_boost": 1.0,
+        "color": "#1A759F",  # Deep blue
+        "rarities": ["Urban Myth", "Urban Legend", "Urban Plague", "Urban Nightmare", "Star of the City"],
+        "odds_multiplier": 2.5
+    },
+    {
+        "name": "Floor of Language",
+        "description": "The domain of linguistic concepts. First glimpse of Star of the City tags.",
+        "required_tags": 1000,  # Unlocked after collecting 600 tags
+        "rarity_boost": 1.2,
+        "color": "#FF8C00",  # Orange
+        "rarities": ["Urban Myth", "Urban Legend", "Urban Plague", "Urban Nightmare", "Star of the City"],
+        "odds_multiplier": 3.0
+    },
+    {
+        "name": "Floor of Social Sciences",
+        "description": "Complex social patterns and abstractions. Notable chance of exceptional rarities.",
+        "required_tags": 2000,  # Unlocked after collecting 1000 tags
+        "rarity_boost": 1.4,
+        "color": "#76B041",  # Brighter green
+        "rarities": ["Urban Legend", "Urban Plague", "Urban Nightmare", "Star of the City"],
+        "odds_multiplier": 3.5
+    },
+    {
+        "name": "Floor of Philosophy",
+        "description": "The realm of profound thought. First access to the rarest 'Impuritas Civitas' tags.",
+        "required_tags": 5000,  # Unlocked after collecting 1500 tags
+        "rarity_boost": 1.6,
+        "color": "#7209B7",  # Deep purple
+        "rarities": ["Urban Legend", "Urban Plague", "Urban Nightmare", "Star of the City", "Impuritas Civitas"],
+        "odds_multiplier": 5.0
+    },
+    {
+        "name": "Floor of Religion",
+        "description": "The ultimate repository of the most profound conceptual territories.",
+        "required_tags": 10000,
+        "rarity_boost": 2.0,
+        "color": "#FFBD00",  # Gold
+        "rarities": ["Urban Plague", "Urban Nightmare", "Star of the City", "Impuritas Civitas"],
+        "odds_multiplier": 10.0
+    }
+]
+
+def start_instant_expedition():
+    """
+    Start an instant expedition with a cooldown before next expedition.
+    
+    Returns:
+        List of discoveries or None if on cooldown
+    """
+    # Check if we're on cooldown
+    current_time = time.time()
+    
+    if hasattr(st.session_state, 'last_expedition_time'):
+        elapsed_time = current_time - st.session_state.last_expedition_time
+        cooldown_duration = calculate_expedition_duration()
+        
+        if elapsed_time < cooldown_duration:
+            # Still on cooldown
+            time_remaining = cooldown_duration - elapsed_time
+            minutes, seconds = divmod(int(time_remaining), 60)
+            st.error(f"Expedition on cooldown. {minutes:02d}:{seconds:02d} remaining.")
+            return None
+    
+    # Generate instant discoveries
+    discoveries = generate_expedition_discoveries()
+    
+    # Set the cooldown timer
+    st.session_state.last_expedition_time = current_time
+    
+    # Save state
+    tag_storage.save_game(st.session_state)
+    
+    # Preserve the current tab
+    if 'library_tab_index' not in st.session_state:
+        st.session_state.library_tab_index = 0
+    
+    return discoveries
+
+def generate_expedition_discoveries():
+    """
+    Generate expedition discoveries instantly.
+    
+    Returns:
+        List of discovered tags and their info
+    """
+    # Get current library floor
+    current_floor = None
+    collection_size = len(st.session_state.collected_tags) if hasattr(st.session_state, 'collected_tags') else 0
+    
+    if hasattr(st.session_state, 'library_floors'):
+        # Find the highest unlocked floor
+        for floor in reversed(st.session_state.library_floors):
+            if collection_size >= floor["required_tags"]:
+                current_floor = floor
+                break
+    
+    # Default to first floor if we couldn't find one
+    if not current_floor:
+        current_floor = st.session_state.library_floors[0] if hasattr(st.session_state, 'library_floors') else {
+            "name": "Archival Records",
+            "rarities": ["Canard", "Urban Myth"],
+            "rarity_boost": 0.0
+        }
+    
+    # Calculate rarity odds for discoveries
+    rarity_odds = calculate_rarity_odds()
+    
+    # Calculate capacity from upgrades
+    tags_capacity = calculate_expedition_capacity()
+    
+    # Generate discoveries
+    discoveries = []
+    for _ in range(tags_capacity):
+        # Select a rarity based on calculated odds
+        rarities = list(rarity_odds.keys())
+        weights = list(rarity_odds.values())
+        selected_rarity = random.choices(rarities, weights=weights, k=1)[0]
+        
+        # Now select a random tag with this rarity that hasn't been discovered yet
+        possible_tags = []
+        
+        # Check if we have tag metadata with rarity info
+        if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
+            # Find all tags of the selected rarity
+            for tag, tag_info in st.session_state.tag_rarity_metadata.items():
+                # Skip if already discovered
+                if tag in st.session_state.discovered_tags:
+                    continue
+                    
+                # Handle both formats - new (dict with rarity) and old (just rarity string)
+                if isinstance(tag_info, dict) and "rarity" in tag_info:
+                    if tag_info["rarity"] == selected_rarity:
+                        possible_tags.append(tag)
+                elif tag_info == selected_rarity:
+                    possible_tags.append(tag)
+        
+        # If no undiscovered tags found in the selected rarity, fallback to already discovered tags
+        if not possible_tags:
+            if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
+                for tag, tag_info in st.session_state.tag_rarity_metadata.items():
+                    # Skip if in different rarity
+                    tag_rarity = tag_info.get("rarity", tag_info) if isinstance(tag_info, dict) else tag_info
+                    if tag_rarity != selected_rarity:
+                        continue
+                    
+                    possible_tags.append(tag)
+        
+        # If still no tags found (or no metadata), create a fallback
+        if not possible_tags:
+            # If we have the model's full tag list, use it
+            if hasattr(st.session_state, 'metadata') and 'idx_to_tag' in st.session_state.metadata:
+                all_tags = list(st.session_state.metadata['idx_to_tag'].values())
+                # Just pick a random tag and assign the selected rarity
+                possible_tags = random.sample(all_tags, min(20, len(all_tags)))
+            else:
+                # Complete fallback - use some generic tags
+                possible_tags = ["portrait", "landscape", "digital_art", "anime", "realistic", 
+                                "fantasy", "sci-fi", "city", "nature", "character"]
+        
+        # If we found possible tags, select one randomly
+        if possible_tags:
+            selected_tag = random.choice(possible_tags)
+            
+            # Get category from metadata if available
+            category = "unknown"
+            if hasattr(st.session_state, 'metadata') and 'tag_to_category' in st.session_state.metadata:
+                if selected_tag in st.session_state.metadata['tag_to_category']:
+                    category = st.session_state.metadata['tag_to_category'][selected_tag]
+            
+            # Use the enhanced tag storage function to add the discovered tag
+            is_new = tag_storage.add_discovered_tag(
+                tag=selected_tag,
+                rarity=selected_rarity,
+                session_state=st.session_state,
+                library_floor=current_floor["name"],
+                category=category  # Pass the category we found
+            )
+            
+            # Record for library growth
+            st.session_state.library_growth["total_discoveries"] += 1
+            st.session_state.library_growth["last_discovery_time"] = time.time()
+            
+            # Create timestamp for display
+            timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
+            
+            # Add to results
+            discoveries.append({
+                "tag": selected_tag,
+                "rarity": selected_rarity,
+                "is_new": is_new,
+                "timestamp": timestamp,
+                "library": current_floor["name"]
+            })
+    
+    # Save the game state after discoveries
+    tag_storage.save_library_state(session_state=st.session_state)
+    tag_storage.save_game(st.session_state)
+    
+    return discoveries
+
+def update_discovered_tag_categories():
+    """Update categories of discovered tags from metadata if they're unknown"""
+    if not hasattr(st.session_state, 'discovered_tags') or not st.session_state.discovered_tags:
+        return 0
+        
+    updated_count = 0
+    
+    # First try from metadata.tag_to_category
+    if hasattr(st.session_state, 'metadata') and 'tag_to_category' in st.session_state.metadata:
+        tag_to_category = st.session_state.metadata['tag_to_category']
+        
+        for tag, info in st.session_state.discovered_tags.items():
+            if info.get('category', 'unknown') == 'unknown' and tag in tag_to_category:
+                info['category'] = tag_to_category[tag]
+                updated_count += 1
+    
+    # Then try from tag_rarity_metadata
+    if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
+        for tag, info in st.session_state.discovered_tags.items():
+            if info.get('category', 'unknown') == 'unknown' and tag in st.session_state.tag_rarity_metadata:
+                tag_metadata = st.session_state.tag_rarity_metadata[tag]
+                if isinstance(tag_metadata, dict) and "category" in tag_metadata:
+                    info['category'] = tag_metadata["category"]
+                    updated_count += 1
+    
+    if updated_count > 0:
+        print(f"Updated categories for {updated_count} discovered tags")
+        tag_storage.save_library_state(session_state=st.session_state)
+    
+    return updated_count
+
+def calculate_expedition_duration():
+    """
+    Calculate the duration of cooldown after an expedition based on upgrades.
+    
+    Returns:
+        Duration in seconds for the cooldown
+    """
+    base_duration = 10 # Default to 10 seconds
+    
+    # Apply speed upgrades if they exist
+    speed_level = 1
+    if hasattr(st.session_state, 'library_upgrades'):
+        speed_level = st.session_state.library_upgrades.get("speed", 1)
+    
+    # Each speed level reduces duration by 10% (multiplicative)
+    duration_multiplier = 0.9 ** (speed_level - 1)
+    
+    # Calculate final duration (minimum 1 second)
+    duration = max(1, base_duration * duration_multiplier)
+    
+    return duration
+
+def calculate_expedition_capacity():
+    """
+    Calculate how many tags can be discovered in one expedition.
+    
+    Returns:
+        Number of tags that can be discovered
+    """
+    base_capacity = 1  # Default to 1 discovery per expedition
+    
+    # Apply capacity upgrades if they exist
+    capacity_level = 1
+    if hasattr(st.session_state, 'library_upgrades'):
+        capacity_level = st.session_state.library_upgrades.get("capacity", 1)
+    
+    # Each capacity level increases discoveries by 1
+    capacity = base_capacity + (capacity_level - 1)
+    
+    return capacity
+
+def calculate_rarity_odds():
+    """
+    Calculate rarity odds based on library floor level and upgrades.
+    
+    Returns:
+        Dictionary of {rarity: probability} for available rarities
+    """
+    # Get current library floor
+    current_floor = None
+    if hasattr(st.session_state, 'library_floors'):
+        collection_size = len(st.session_state.collected_tags) if hasattr(st.session_state, 'collected_tags') else 0
+        
+        # Find the highest unlocked floor
+        for floor in reversed(st.session_state.library_floors):
+            if collection_size >= floor["required_tags"]:
+                current_floor = floor
+                break
+    
+    # Default to first floor if we couldn't find one
+    if not current_floor:
+        current_floor = st.session_state.library_floors[0] if hasattr(st.session_state, 'library_floors') else {
+            "rarities": ["Canard", "Urban Myth"],
+            "rarity_boost": 0.0,
+            "odds_multiplier": 1.0
+        }
+    
+    # Get available rarities from current floor
+    available_rarities = current_floor.get("rarities", ["Canard", "Urban Myth"])
+    odds_multiplier = current_floor.get("odds_multiplier", 1.0)
+    
+    # Base weights for each rarity
+    base_weights = {
+        "Canard": 70,
+        "Urban Myth": 20,
+        "Urban Legend": 7,
+        "Urban Plague": 2,
+        "Urban Nightmare": 1,
+        "Star of the City": 0.1,
+        "Impuritas Civitas": 0.01
+    }
+    
+    # Apply floor's rarity boost
+    floor_rarity_boost = current_floor.get("rarity_boost", 0.0)
+    
+    # Apply rarity upgrades if they exist
+    rarity_level = 1
+    if hasattr(st.session_state, 'library_upgrades'):
+        rarity_level = st.session_state.library_upgrades.get("rarity", 1)
+    
+    # Calculate boost based on rarity level
+    upgrade_rarity_boost = (rarity_level - 1) * 0.2  # Each level gives 20% more chance for rare tags
+    
+    # Combine boosts
+    total_boost = floor_rarity_boost + upgrade_rarity_boost
+    
+    # Adjust weights based on rarity boost
+    adjusted_weights = {}
+    for rarity in available_rarities:
+        if rarity == "Canard":
+            # Reduce common tag odds as rarity level increases
+            adjusted_weights[rarity] = base_weights[rarity] * (1.0 - total_boost * 0.7)
+        elif rarity == "Urban Myth":
+            # Slight reduction for uncommon as rarity level increases
+            adjusted_weights[rarity] = base_weights[rarity] * (1.0 - total_boost * 0.3)
+        else:
+            # Increase rare tag odds as rarity level increases
+            rarity_index = list(RARITY_LEVELS.keys()).index(rarity)
+            # Higher rarities get larger boosts
+            boost_factor = 1.0 + (total_boost * odds_multiplier * (rarity_index + 1))
+            adjusted_weights[rarity] = base_weights[rarity] * boost_factor
+    
+    # Normalize weights
+    total = sum(adjusted_weights.values())
+    normalized_weights = {r: w/total for r, w in adjusted_weights.items()}
+    
+    return normalized_weights
+
+def format_time_remaining(seconds):
+    """
+    Format seconds into a human-readable time remaining format.
+    
+    Args:
+        seconds: Seconds remaining
+        
+    Returns:
+        String with formatted time
+    """
+    if seconds < 60:
+        return f"{int(seconds)} seconds"
+    elif seconds < 3600:
+        minutes = seconds / 60
+        return f"{int(minutes)} minutes"
+    else:
+        hours = seconds / 3600
+        minutes = (seconds % 3600) / 60
+        if minutes > 0:
+            return f"{int(hours)} hours, {int(minutes)} minutes"
+        else:
+            return f"{int(hours)} hours"
+
+def display_cooldown_timer():
+    """Display a countdown timer until the next expedition is available"""
+    # Check if on cooldown
+    current_time = time.time()
+    cooldown_remaining = 0
+    cooldown_duration = calculate_expedition_duration()
+    
+    if hasattr(st.session_state, 'last_expedition_time'):
+        elapsed_time = current_time - st.session_state.last_expedition_time
+        if elapsed_time < cooldown_duration:
+            cooldown_remaining = cooldown_duration - elapsed_time
+    
+    # If on cooldown, show timer
+    if cooldown_remaining > 0:
+        minutes, seconds = divmod(int(cooldown_remaining), 60)
+        
+        # Create a timer display with dark-mode styling
+        st.markdown("""
+        <div style="background-color: rgba(255, 152, 0, 0.15); 
+                   border: 1px solid #FF9800; 
+                   border-radius: 5px; 
+                   padding: 10px; 
+                   text-align: center;
+                   margin-bottom: 15px;
+                   color: #ffffff;">
+            <p style="margin: 0; font-weight: bold;">⏱️ Next expedition available in:</p>
+            <p style="font-size: 1.2em; margin: 5px 0;">{:02d}:{:02d}</p>
+        </div>
+        """.format(minutes, seconds), unsafe_allow_html=True)
+        
+        # Add refresh button for the timer
+        if st.button("🔄 Refresh Timer", key="refresh_timer"):
+            st.rerun()
+        
+        return True  # Still on cooldown
+    
+    return False  # Not on cooldown
+
+def display_library_exploration_interface():
+    """Display the unified interface for library exploration using Streamlit elements."""
+    # Tag collection progress
+    tag_count = len(st.session_state.collected_tags) if hasattr(st.session_state, 'collected_tags') else 0
+    
+    # Check if we have tags to start exploring
+    if not hasattr(st.session_state, 'collected_tags') or not st.session_state.collected_tags:
+        st.warning("Start scanning images to collect tags first. The library will grow as you collect more tags!")
+        return
+    
+    # Get current library floor
+    current_floor = None
+    if hasattr(st.session_state, 'library_floors'):
+        # Find the highest unlocked floor
+        for floor in reversed(st.session_state.library_floors):
+            if tag_count >= floor["required_tags"]:
+                current_floor = floor
+                break
+    
+    # Default to first floor if we couldn't find one
+    if not current_floor:
+        current_floor = st.session_state.library_floors[0] if hasattr(st.session_state, 'library_floors') else {
+            "name": "Floor of General Works",
+            "description": "The foundational level of knowledge.",
+            "color": "#607D8B",
+            "rarities": ["Canard", "Urban Myth"]
+        }
+    
+    # Library growth progress
+    total_discoveries = st.session_state.library_growth["total_discoveries"]
+    
+    # Create container with colored border for current floor
+    floor_container = st.container()
+    with floor_container:
+        # Use a stylized container with dark mode theme
+        st.markdown(f"""
+        <div style="border-left: 5px solid {current_floor['color']}; 
+                    border-radius: 5px;
+                    background-color: rgba({int(current_floor['color'][1:3], 16)}, 
+                                           {int(current_floor['color'][3:5], 16)}, 
+                                           {int(current_floor['color'][5:7], 16)}, 0.15);
+                    padding: 15px 10px 10px 15px;
+                    margin-bottom: 15px;
+                    color: #ffffff;">
+            <h3 style="margin-top: 0; color: {current_floor['color']};">{current_floor['name']}</h3>
+            <p>{current_floor['description']}</p>
+            <p>Total Discoveries: <strong>{total_discoveries}</strong></p>
+        </div>
+        """, unsafe_allow_html=True)
+    
+    # Create a nice divider for dark theme
+    st.markdown("<hr style='margin: 20px 0; border: 0; height: 1px; background-image: linear-gradient(to right, rgba(255, 255, 255, 0), rgba(255, 255, 255, 0.2), rgba(255, 255, 255, 0));'>", unsafe_allow_html=True)
+    
+    # Display expedition details
+    st.subheader("Expedition Details")
+    
+    # Calculate capacity
+    capacity = calculate_expedition_capacity()
+    
+    # Two columns for expedition stats
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        # Expedition duration/timer
+        cooldown_duration = calculate_expedition_duration()
+        st.write(f"📊 Cooldown: {format_time_remaining(cooldown_duration)}")
+        st.write(f"🔍 Tag Discoveries: {capacity} per expedition")
+    
+    with col2:
+        # Calculate and display rarity odds with Streamlit elements
+        rarity_odds = calculate_rarity_odds()
+        available_rarities = current_floor.get("rarities", ["Canard", "Urban Myth"])
+        
+        # Display rarity chances with dark theme styling
+        for rarity in available_rarities:
+            if rarity in rarity_odds:
+                color = RARITY_LEVELS[rarity]["color"]
+                percentage = rarity_odds[rarity]*100
+                
+                # Custom styling based on rarity
+                if rarity == "Impuritas Civitas":
+                    st.markdown(f"""
+                    <div style="display: flex; align-items: center; margin-bottom: 5px;">
+                        <span style="animation: rainbow-text 4s linear infinite; font-weight: bold; width: 140px;">{rarity}:</span>
+                        <div style="flex-grow: 1; background-color: #2c2c2c; border-radius: 5px; height: 10px;">
+                            <div style="width: {min(percentage*5, 100)}%; background: linear-gradient(to right, red, orange, yellow, green, blue, indigo, violet); height: 10px; border-radius: 5px;"></div>
+                        </div>
+                        <span style="margin-left: 10px;">{percentage:.1f}%</span>
+                    </div>
+                    """, unsafe_allow_html=True)
+                elif rarity == "Star of the City":
+                    st.markdown(f"""
+                    <div style="display: flex; align-items: center; margin-bottom: 5px;">
+                        <span style="color:{color}; text-shadow: 0 0 3px gold; font-weight: bold; width: 140px;">{rarity}:</span>
+                        <div style="flex-grow: 1; background-color: #2c2c2c; border-radius: 5px; height: 10px;">
+                            <div style="width: {min(percentage*5, 100)}%; background-color: {color}; box-shadow: 0 0 5px gold; height: 10px; border-radius: 5px;"></div>
+                        </div>
+                        <span style="margin-left: 10px;">{percentage:.1f}%</span>
+                    </div>
+                    """, unsafe_allow_html=True)
+                elif rarity == "Urban Nightmare":
+                    st.markdown(f"""
+                    <div style="display: flex; align-items: center; margin-bottom: 5px;">
+                        <span style="color:{color}; text-shadow: 0 0 1px #FF5722; font-weight: bold; width: 140px;">{rarity}:</span>
+                        <div style="flex-grow: 1; background-color: #2c2c2c; border-radius: 5px; height: 10px;">
+                            <div style="width: {min(percentage*5, 100)}%; background-color: {color}; animation: pulse-bar 3s infinite; height: 10px; border-radius: 5px;"></div>
+                        </div>
+                        <span style="margin-left: 10px;">{percentage:.1f}%</span>
+                    </div>
+                    """, unsafe_allow_html=True)
+                else:
+                    st.markdown(f"""
+                    <div style="display: flex; align-items: center; margin-bottom: 5px;">
+                        <span style="color:{color}; font-weight: bold; width: 140px;">{rarity}:</span>
+                        <div style="flex-grow: 1; background-color: #2c2c2c; border-radius: 5px; height: 10px;">
+                            <div style="width: {min(percentage*5, 100)}%; background-color: {color}; height: 10px; border-radius: 5px;"></div>
+                        </div>
+                        <span style="margin-left: 10px;">{percentage:.1f}%</span>
+                    </div>
+                    """, unsafe_allow_html=True)
+    
+    # Check for cooldown and display timer if needed
+    on_cooldown = display_cooldown_timer()
+    
+    # Add button to start expedition regardless of cooldown status
+    # The button will always be displayed, but if on cooldown, expedition won't start
+    if st.button("🚀 Start Expedition", key="start_expedition", use_container_width=True, disabled=on_cooldown):
+        if not on_cooldown:
+            discoveries = start_instant_expedition()
+            if discoveries:
+                # Store discovery results for display
+                st.session_state.expedition_results = discoveries
+                # Show success message
+                st.success(f"Expedition completed! Discovered {len(discoveries)} new tags!")
+                # Show balloons for celebration
+                st.balloons()
+                # Display the results
+                display_expedition_results(discoveries)
+                # Save state
+                tag_storage.save_game(st.session_state)
+        else:
+            # This should not be reached due to disabled button, but just in case
+            st.error("Expedition on cooldown. Please wait until the timer expires.")
+    
+    # Display library upgrades
+    display_library_upgrades()
+
+def display_expedition_results(results):
+    """Display results from completed expeditions using Streamlit elements with enhanced dark-mode visuals."""
+    st.subheader("Expedition Discoveries")
+    
+    # Add animations CSS for dark theme
+    st.markdown("""
+    <style>
+    @keyframes rainbow-text {
+        0% { color: red; }
+        14% { color: orange; }
+        28% { color: yellow; }
+        42% { color: green; }
+        57% { color: blue; }
+        71% { color: indigo; }
+        85% { color: violet; }
+        100% { color: red; }
+    }
+    
+    @keyframes rainbow-border {
+        0% { border-color: red; }
+        14% { border-color: orange; }
+        28% { border-color: yellow; }
+        42% { border-color: green; }
+        57% { border-color: blue; }
+        71% { border-color: indigo; }
+        85% { border-color: violet; }
+        100% { border-color: red; }
+    }
+    
+    @keyframes star-glow {
+        0% { box-shadow: 0 0 5px #FFD700; }
+        50% { box-shadow: 0 0 15px #FFD700; }
+        100% { box-shadow: 0 0 5px #FFD700; }
+    }
+    
+    @keyframes nightmare-pulse {
+        0% { border-color: #FF9800; }
+        50% { border-color: #FF5722; }
+        100% { border-color: #FF9800; }
+    }
+    
+    @keyframes pulse-bar {
+        0% { opacity: 0.8; }
+        50% { opacity: 1; }
+        100% { opacity: 0.8; }
+    }
+    
+    .expedition-tag-impuritas {
+        animation: rainbow-text 4s linear infinite;
+        font-weight: bold;
+    }
+    
+    .expedition-card-impuritas {
+        background-color: rgba(255, 0, 0, 0.15);
+        border-radius: 8px;
+        border: 3px solid red;
+        padding: 12px;
+        animation: rainbow-border 4s linear infinite;
+        color: #ffffff;
+    }
+    
+    .expedition-card-star {
+        background-color: rgba(255, 215, 0, 0.15);
+        border-radius: 8px;
+        border: 2px solid gold;
+        padding: 12px;
+        animation: star-glow 2s infinite;
+        color: #ffffff;
+    }
+    
+    .expedition-card-nightmare {
+        background-color: rgba(255, 152, 0, 0.15);
+        border-radius: 8px;
+        border: 2px solid #FF9800;
+        padding: 12px;
+        animation: nightmare-pulse 3s infinite;
+        color: #ffffff;
+    }
+    
+    .expedition-card-plague {
+        background-color: rgba(156, 39, 176, 0.12);
+        border-radius: 8px;
+        border: 1px solid #9C27B0;
+        padding: 12px;
+        box-shadow: 0 0 3px #9C27B0;
+        color: #ffffff;
+    }
+    
+    .expedition-card-legend {
+        background-color: rgba(33, 150, 243, 0.15);
+        border-radius: 8px;
+        border: 1px solid #2196F3;
+        padding: 12px;
+        color: #ffffff;
+    }
+    
+    .expedition-card-myth {
+        background-color: rgba(76, 175, 80, 0.15);
+        border-radius: 8px;
+        border: 1px solid #4CAF50;
+        padding: 12px;
+        color: #ffffff;
+    }
+    
+    .expedition-card-canard {
+        background-color: rgba(170, 170, 170, 0.15);
+        border-radius: 8px;
+        border: 1px solid #AAAAAA;
+        padding: 12px;
+        color: #ffffff;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    # Group by rarity first for better organization
+    results_by_rarity = {}
+    for result in results:
+        rarity = result["rarity"]
+        if rarity not in results_by_rarity:
+            results_by_rarity[rarity] = []
+        results_by_rarity[rarity].append(result)
+    
+    # Get ordered rarities (rarest first)
+    ordered_rarities = list(RARITY_LEVELS.keys())
+    ordered_rarities.reverse()  # Reverse to display rarest first
+    
+    # Display rare discoveries first
+    for rarity in ordered_rarities:
+        if rarity not in results_by_rarity:
+            continue
+            
+        rarity_results = results_by_rarity[rarity]
+        if not rarity_results:
+            continue
+            
+        color = RARITY_LEVELS[rarity]["color"]
+        
+        # Special styling for rare discoveries
+        if rarity == "Impuritas Civitas":
+            st.markdown(f"""
+            <div style="background-color: rgba(255, 0, 0, 0.15); border-radius: 10px; padding: 15px; margin-bottom: 20px; border: 2px solid red; animation: rainbow-border 4s linear infinite; color: #ffffff;">
+                <h3 style="margin-top: 0; animation: rainbow-text 4s linear infinite;">✨ EXTRAORDINARY DISCOVERY!</h3>
+                <p>You found {len(rarity_results)} Impuritas Civitas tag(s)!</p>
+            </div>
+            """, unsafe_allow_html=True)
+            st.balloons()  # Add celebration effect
+        elif rarity == "Star of the City":
+            st.markdown(f"""
+            <div style="background-color: rgba(255, 215, 0, 0.15); border-radius: 10px; padding: 15px; margin-bottom: 20px; border: 2px solid gold; animation: star-glow 2s infinite; color: #ffffff;">
+                <h3 style="margin-top: 0; color: {color}; text-shadow: 0 0 3px gold;">🌟 EXCEPTIONAL DISCOVERY!</h3>
+                <p>You found {len(rarity_results)} Star of the City tag(s)!</p>
+            </div>
+            """, unsafe_allow_html=True)
+        elif rarity == "Urban Nightmare":
+            st.markdown(f"""
+            <div style="background-color: rgba(255, 152, 0, 0.15); border-radius: 10px; padding: 15px; margin-bottom: 20px; border: 2px solid #FF9800; animation: nightmare-pulse 3s infinite; color: #ffffff;">
+                <h3 style="margin-top: 0; color: {color}; text-shadow: 0 0 1px #FF5722;">👑 RARE DISCOVERY!</h3>
+                <p>You found {len(rarity_results)} Urban Nightmare tag(s)!</p>
+            </div>
+            """, unsafe_allow_html=True)
+        elif rarity == "Urban Plague":
+            st.markdown(f"""
+            <div style="background-color: rgba(156, 39, 176, 0.12); border-radius: 10px; padding: 15px; margin-bottom: 20px; border: 1px solid #9C27B0; box-shadow: 0 0 3px #9C27B0; color: #ffffff;">
+                <h3 style="margin-top: 0; color: {color}; text-shadow: 0 0 1px #9C27B0;">⚔️ UNCOMMON DISCOVERY!</h3>
+                <p>You found {len(rarity_results)} Urban Plague tag(s)!</p>
+            </div>
+            """, unsafe_allow_html=True)
+        else:
+            st.markdown(f"### {rarity} ({len(rarity_results)} discoveries)")
+        
+        # Display tags in this rarity
+        cols = st.columns(3)
+        for i, result in enumerate(rarity_results):
+            col_idx = i % 3
+            with cols[col_idx]:
+                tag = result["tag"]
+                floor_name = result.get("library", "Library")
+                
+                # Get the appropriate card class based on rarity
+                rarity_class = rarity.lower().replace(' ', '-')
+                card_class = f"expedition-card-{rarity_class}"
+                
+                # Create styled card for each tag
+                tag_html = f"""<div class="{card_class}">"""
+                
+                # Special styling for the tag name based on rarity
+                if rarity == "Impuritas Civitas":
+                    tag_html += f"""<p style="font-size: 1.1em; margin-bottom: 5px;"><span class="expedition-tag-impuritas">✨ {tag}</span></p>"""
+                elif rarity == "Star of the City":
+                    tag_html += f"""<p style="font-size: 1.1em; margin-bottom: 5px;"><span style="color: {color}; text-shadow: 0 0 3px gold; font-weight: bold;">🌟 {tag}</span></p>"""
+                elif rarity == "Urban Nightmare":
+                    tag_html += f"""<p style="font-size: 1.1em; margin-bottom: 5px;"><span style="color: {color}; text-shadow: 0 0 1px #FF5722; font-weight: bold;">👑 {tag}</span></p>"""
+                elif rarity == "Urban Plague":
+                    tag_html += f"""<p style="font-size: 1.1em; margin-bottom: 5px;"><span style="color: {color}; text-shadow: 0 0 1px #9C27B0; font-weight: bold;">⚔️ {tag}</span></p>"""
+                else:
+                    tag_html += f"""<p style="font-size: 1.1em; margin-bottom: 5px;"><span style="color: {color}; font-weight: bold;">{tag}</span></p>"""
+                
+                # Mark as new if it is
+                is_new = result.get("is_new", False)
+                new_badge = """<span style="background-color: #4CAF50; color: white; padding: 2px 6px; border-radius: 10px; font-size: 0.7em; margin-left: 5px;">NEW</span>""" if is_new else ""
+                
+                # Add other tag details
+                tag_html += f"""
+                    <p style="margin: 0; font-size: 0.9em;">Found in: {floor_name} {new_badge}</p>
+                    <p style="margin: 5px 0 0 0; font-size: 0.9em;">Rarity: <span style="color: {color};">{rarity}</span></p>
+                </div>
+                """
+                
+                st.markdown(tag_html, unsafe_allow_html=True)
+                
+        # Add separator between rarity groups
+        st.markdown("<hr style='margin: 20px 0; border: 0; height: 1px; background-image: linear-gradient(to right, rgba(255, 255, 255, 0), rgba(255, 255, 255, 0.2), rgba(255, 255, 255, 0));'>", unsafe_allow_html=True)
+
+def display_library_building():
+    """Display a visual representation of the library building with all floors."""
+    st.subheader("The Great Library Building")
+    
+    # Get collection size
+    collection_size = len(st.session_state.collected_tags) if hasattr(st.session_state, 'collected_tags') else 0
+    
+    # Determine current floor
+    current_floor_index = 0
+    for i, floor in enumerate(st.session_state.library_floors):
+        if collection_size >= floor["required_tags"]:
+            current_floor_index = i
+    
+    # Create a visual representation of the library building
+    total_floors = len(st.session_state.library_floors)
+    
+    # Enhanced CSS for the library building with dark theme
+    st.markdown("""
+    <style>
+    @keyframes floor-glow {
+        0% { box-shadow: 0 0 5px rgba(13, 110, 253, 0.5); }
+        50% { box-shadow: 0 0 15px rgba(13, 110, 253, 0.8); }
+        100% { box-shadow: 0 0 5px rgba(13, 110, 253, 0.5); }
+    }
+    
+    .library-roof {
+        background: linear-gradient(90deg, #8D6E63, #A1887F);
+        height: 35px;
+        width: 90%;
+        margin: 0 auto;
+        border-radius: 8px 8px 0 0;
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        color: white;
+        font-weight: bold;
+        box-shadow: 0 -2px 10px rgba(0,0,0,0.4);
+    }
+    
+    .library-floor {
+        height: 65px;
+        width: 80%;
+        margin: 0 auto;
+        border: 1px solid #444;
+        display: flex;
+        align-items: center;
+        padding: 0 20px;
+        position: relative;
+        transition: all 0.3s ease;
+        color: #ffffff;
+        background-color: #2c2c2c;
+    }
+    
+    .library-floor:hover {
+        transform: translateX(10px);
+    }
+    
+    .library-floor.current {
+        box-shadow: 0 0 10px rgba(13, 110, 253, 0.5);
+        z-index: 2;
+        animation: floor-glow 2s infinite;
+        border-left: 5px solid #0d6efd;
+    }
+    
+    .library-floor.locked {
+        background-color: #1e1e1e;
+        color: #777;
+        filter: grayscale(50%);
+    }
+    
+    .library-floor-number {
+        position: absolute;
+        left: -30px;
+        width: 25px;
+        height: 25px;
+        background-color: #0d6efd;
+        color: white;
+        border-radius: 50%;
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        font-weight: bold;
+    }
+    
+    .library-floor.locked .library-floor-number {
+        background-color: #555;
+    }
+    
+    .library-entrance {
+        background: linear-gradient(90deg, #5D4037, #795548);
+        height: 45px;
+        width: 35%;
+        margin: 0 auto;
+        border-radius: 10px 10px 0 0;
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        color: white;
+        font-weight: bold;
+        box-shadow: 0 -2px 10px rgba(0,0,0,0.4);
+    }
+    
+    .library-floor-details {
+        flex: 1;
+    }
+    
+    .library-floor-name {
+        font-weight: bold;
+        margin: 0;
+    }
+    
+    .library-floor-description {
+        font-size: 0.85em;
+        margin: 3px 0 0 0;
+        opacity: 0.9;
+    }
+    
+    .library-floor-status {
+        display: flex;
+        align-items: center;
+        font-weight: bold;
+    }
+    
+    .library-floor-rarities {
+        font-size: 0.8em;
+        margin-top: 4px;
+    }
+    
+    .rarity-dot {
+        display: inline-block;
+        width: 8px;
+        height: 8px;
+        border-radius: 50%;
+        margin-right: 3px;
+    }
+    
+    /* Special animations for rarer floors */
+    .library-floor.star {
+        background-color: rgba(255, 215, 0, 0.15);
+    }
+    
+    .library-floor.impuritas {
+        background-color: rgba(255, 0, 0, 0.15);
+    }
+    
+    @keyframes rainbow-border {
+        0% { border-color: red; }
+        14% { border-color: orange; }
+        28% { border-color: yellow; }
+        42% { border-color: green; }
+        57% { border-color: blue; }
+        71% { border-color: indigo; }
+        85% { border-color: violet; }
+        100% { border-color: red; }
+    }
+    
+    .rainbow-border {
+        animation: rainbow-border 4s linear infinite;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    # Roof
+    st.markdown('<div class="library-roof">🏛️ The Great Library</div>', unsafe_allow_html=True)
+    
+    # Display floors from top (highest) to bottom
+    for i in reversed(range(total_floors)):
+        floor = st.session_state.library_floors[i]
+        is_current = i == current_floor_index
+        is_unlocked = collection_size >= floor["required_tags"]
+        
+        # Style based on floor status
+        floor_class = "library-floor"
+        if is_current:
+            floor_class += " current"
+        if not is_unlocked:
+            floor_class += " locked"
+            
+        # Add special classes for highest floors
+        if i >= 8 and is_unlocked:  # Impuritas Civitas level floors
+            floor_class += " impuritas"
+        elif i >= 6 and is_unlocked:  # Star of the City level floors
+            floor_class += " star"
+            
+        # Determine rarity dots HTML
+        rarity_dots = ""
+        for rarity in floor.get("rarities", []):
+            color = RARITY_LEVELS[rarity]["color"]
+            rarity_dots += f'<span class="rarity-dot" style="background-color:{color};"></span>'
+            
+        # Floor style based on color
+        if is_unlocked:
+            floor_style = f"background-color: rgba({int(floor['color'][1:3], 16)}, {int(floor['color'][3:5], 16)}, {int(floor['color'][5:7], 16)}, 0.25);"
+        else:
+            floor_style = ""
+            
+        # Special border animation for highest floor
+        border_class = ""
+        if i == 9 and is_unlocked:  # Top floor
+            border_class = "rainbow-border"
+            
+        # Display the floor
+        floor_content = f"""
+        <div class="{floor_class} {border_class}" style="{floor_style}">
+            <span class="library-floor-number">{i+1}</span>
+            <div class="library-floor-details">
+                <p class="library-floor-name">{floor['name']}</p>
+                <p class="library-floor-description">{floor['description'] if is_unlocked else 'Locked'}</p>
+                <div class="library-floor-rarities">{rarity_dots}</div>
+            </div>
+            <div class="library-floor-status">
+                {"🔓" if is_unlocked else "🔒"} {floor['required_tags']} tags
+            </div>
+        </div>
+        """
+        st.markdown(floor_content, unsafe_allow_html=True)
+    
+    # Entrance
+    st.markdown('<div class="library-entrance">📚 Entrance</div>', unsafe_allow_html=True)
+    
+    # Floor details expander
+    with st.expander("Floor Details", expanded=False):
+        # Create a table with styled rarities for dark theme
+        st.markdown("""
+        <style>
+        .floor-details-table {
+            width: 100%;
+            border-collapse: collapse;
+            color: #ffffff;
+        }
+        
+        .floor-details-table th {
+            background-color: #333;
+            padding: 8px;
+            text-align: left;
+            border: 1px solid #444;
+        }
+        
+        .floor-details-table td {
+            padding: 8px;
+            border: 1px solid #444;
+        }
+        
+        .floor-details-table tr:nth-child(even) {
+            background-color: rgba(255,255,255,0.03);
+        }
+        
+        .floor-details-table tr:nth-child(odd) {
+            background-color: rgba(0,0,0,0.2);
+        }
+        
+        .floor-details-table tr:hover {
+            background-color: rgba(13, 110, 253, 0.1);
+        }
+        
+        .current-floor {
+            background-color: rgba(13, 110, 253, 0.15) !important;
+        }
+        </style>
+        
+        <table class="floor-details-table">
+            <tr>
+                <th>Floor</th>
+                <th>Name</th>
+                <th>Status</th>
+                <th>Req. Tags</th>
+                <th>Rarities</th>
+                <th>Rarity Boost</th>
+            </tr>
+        """, unsafe_allow_html=True)
+        
+        # Add each floor to the table
+        for i, floor in enumerate(st.session_state.library_floors):
+            is_unlocked = collection_size >= floor["required_tags"]
+            is_current = i == current_floor_index
+            
+            # Format rarities with colors
+            rarity_text = ""
+            for rarity in floor.get("rarities", []):
+                color = RARITY_LEVELS[rarity]["color"]
+                
+                # Special styling based on rarity
+                if rarity == "Impuritas Civitas":
+                    rarity_text += f"<span style='animation: rainbow-text 4s linear infinite;'>{rarity}</span>, "
+                elif rarity == "Star of the City":
+                    rarity_text += f"<span style='color:{color}; text-shadow: 0 0 3px gold;'>{rarity}</span>, "
+                elif rarity == "Urban Nightmare":
+                    rarity_text += f"<span style='color:{color}; text-shadow: 0 0 1px #FF5722;'>{rarity}</span>, "
+                elif rarity == "Urban Plague":
+                    rarity_text += f"<span style='color:{color}; text-shadow: 0 0 1px #9C27B0;'>{rarity}</span>, "
+                else:
+                    rarity_text += f"<span style='color:{color};'>{rarity}</span>, "
+            
+            # Current floor class
+            row_class = "current-floor" if is_current else ""
+            
+            # Add the floor row
+            st.markdown(f"""
+            <tr class="{row_class}">
+                <td>{i+1}</td>
+                <td>{floor["name"]}</td>
+                <td>{"🔓 Unlocked" if is_unlocked else "🔒 Locked"}</td>
+                <td>{floor["required_tags"]}</td>
+                <td>{rarity_text[:-2] if rarity_text else ""}</td>
+                <td>+{int(floor.get('rarity_boost', 0) * 100)}%</td>
+            </tr>
+            """, unsafe_allow_html=True)
+            
+        # Close the table
+        st.markdown("</table>", unsafe_allow_html=True)
+
+def add_discovered_tag(tag, rarity, library_floor=None):
+    """
+    Add a tag to the discovered tags with enriched metadata
+    
+    Args:
+        tag: The tag name
+        rarity: The tag rarity level
+        library_floor: The library floor where it was discovered (optional)
+    
+    Returns:
+        bool: True if it's a new discovery, False if already discovered
+    """
+    is_new = tag not in st.session_state.discovered_tags
+    
+    # Get current time
+    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
+    
+    # Get tag category if metadata is available
+    category = "unknown"
+    if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
+        if tag in st.session_state.tag_rarity_metadata:
+            tag_info = st.session_state.tag_rarity_metadata[tag]
+            if isinstance(tag_info, dict) and "category" in tag_info:
+                category = tag_info["category"]
+    
+    # Get or update tag info
+    if is_new:
+        tag_info = {
+            "rarity": rarity,
+            "discovery_time": timestamp,
+            "category": category,
+            "discovery_count": 1,
+            "last_seen": timestamp
+        }
+        
+        if library_floor:
+            tag_info["library_floor"] = library_floor
+            
+        st.session_state.discovered_tags[tag] = tag_info
+        
+        # Track exploration of library tiers
+        if 'explored_library_tiers' not in st.session_state:
+            st.session_state.explored_library_tiers = set()
+        
+        if library_floor:
+            st.session_state.explored_library_tiers.add(library_floor)
+    else:
+        # Update existing tag
+        tag_info = st.session_state.discovered_tags[tag]
+        tag_info["discovery_count"] = tag_info.get("discovery_count", 1) + 1
+        tag_info["last_seen"] = timestamp
+        
+        # Only update library floor if provided
+        if library_floor:
+            tag_info["library_floor"] = library_floor
+            
+            # Track exploration
+            if 'explored_library_tiers' not in st.session_state:
+                st.session_state.explored_library_tiers = set()
+            
+            st.session_state.explored_library_tiers.add(library_floor)
+    
+    # Save state after updating
+    tag_storage.save_library_state(session_state=st.session_state)
+    
+    return is_new
+
+def calculate_upgrade_cost(upgrade_type, current_level):
+    base_cost = {
+        "speed": 50,
+        "capacity": 100,
+        "rarity": 150
+    }
+    # Cost increases with each level
+    return int(base_cost[upgrade_type] * (current_level * 1.5))
+    
+def purchase_library_upgrade(upgrade_type):
+    """
+    Purchase a library upgrade
+    
+    Args:
+        upgrade_type: The type of upgrade ("speed", "capacity", or "rarity")
+        
+    Returns:
+        bool: True if purchased successfully, False otherwise
+    """
+    # Get current level and calculate cost
+    current_level = st.session_state.library_upgrades.get(upgrade_type, 1)
+    cost = calculate_upgrade_cost(upgrade_type, current_level)
+    
+    # Check if player can afford it
+    if st.session_state.tag_currency < cost:
+        return False
+    
+    # Apply purchase
+    st.session_state.tag_currency -= cost
+    st.session_state.library_upgrades[upgrade_type] = current_level + 1
+    
+    # Update stats
+    if hasattr(st.session_state, 'game_stats'):
+        if "currency_spent" not in st.session_state.game_stats:
+            st.session_state.game_stats["currency_spent"] = 0
+        st.session_state.game_stats["currency_spent"] += cost
+    
+    # Save state
+    tag_storage.save_library_state(session_state=st.session_state)
+    tag_storage.save_game(st.session_state)
+    
+    return True
+
+def display_library_upgrades():
+    """Display and manage upgrades for the library using Streamlit elements with enhanced visuals."""
+    st.subheader("Library Upgrades")
+    
+    # Add styling for upgrade cards with dark theme
+    st.markdown("""
+    <style>
+    .upgrade-card {
+        border: 1px solid #444;
+        border-radius: 10px;
+        padding: 15px;
+        margin-bottom: 20px;
+        background-color: #222;
+        color: #ffffff;
+        transition: transform 0.2s, box-shadow 0.2s;
+    }
+    
+    .upgrade-card:hover {
+        transform: translateY(-2px);
+        box-shadow: 0 4px 8px rgba(0,0,0,0.3);
+    }
+    
+    .upgrade-title {
+        font-size: 1.2em;
+        font-weight: bold;
+        margin-bottom: 10px;
+        color: #ffffff;
+    }
+    
+    .upgrade-level {
+        display: inline-block;
+        background-color: #0d6efd;
+        color: white;
+        padding: 3px 8px;
+        border-radius: 10px;
+        font-size: 0.8em;
+        margin-bottom: 10px;
+    }
+    
+    .upgrade-stat {
+        display: flex;
+        align-items: center;
+        margin-bottom: 5px;
+    }
+    
+    .upgrade-stat-label {
+        width: 100px;
+        font-size: 0.9em;
+        color: #adb5bd;
+    }
+    
+    .upgrade-stat-value {
+        font-weight: bold;
+    }
+    
+    .upgrade-cost {
+        margin-top: 10px;
+        font-weight: bold;
+        color: #6610f2;
+    }
+    
+    .level-bar {
+        height: 6px;
+        background-color: #333;
+        border-radius: 3px;
+        margin-bottom: 10px;
+        overflow: hidden;
+    }
+    
+    .level-progress {
+        height: 100%;
+        background-color: #0d6efd;
+        border-radius: 3px;
+    }
+    
+    @keyframes pulse-button {
+        0% { transform: scale(1); }
+        50% { transform: scale(1.05); }
+        100% { transform: scale(1); }
+    }
+    
+    .pulse-button {
+        animation: pulse-button 2s infinite;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    st.write("Improve your expeditions with these upgrades:")
+    
+    # Create columns for each upgrade type
+    col1, col2, col3 = st.columns(3)
+    
+    # Get current upgrade levels
+    upgrades = st.session_state.library_upgrades
+    
+    # Speed upgrade
+    with col1:
+        speed_level = upgrades.get("speed", 1)
+        speed_cost = calculate_upgrade_cost("speed", speed_level)
+        
+        # Calculate max level for progress bar
+        max_level = 10
+        progress_percentage = min(100, (speed_level / max_level) * 100)
+        
+        # Create an upgrade card with progress bar in dark theme
+        st.markdown(f"""
+        <div class="upgrade-card">
+            <div class="upgrade-title">⏱️ Speed Upgrade</div>
+            <div class="upgrade-level">Level {speed_level}</div>
+            <div class="level-bar">
+                <div class="level-progress" style="width: {progress_percentage}%;"></div>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Effect:</span>
+                <span class="upgrade-stat-value">Reduces cooldown time</span>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Current:</span>
+                <span class="upgrade-stat-value">{format_time_remaining(calculate_expedition_duration())}</span>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Next Level:</span>
+                <span class="upgrade-stat-value">{format_time_remaining(calculate_expedition_duration() * 0.9)}</span>
+            </div>
+            <div class="upgrade-cost" style="color: #9D4EDD;">Cost: {speed_cost} {TAG_CURRENCY_NAME}</div>
+        </div>
+        """, unsafe_allow_html=True)
+        
+        # Upgrade button
+        can_afford = st.session_state.tag_currency >= speed_cost
+        button_class = "pulse-button" if can_afford else ""
+                
+        if st.button(f"Upgrade Speed", key=f"upgrade_speed", disabled=not can_afford, use_container_width=True):
+            if purchase_library_upgrade("speed"):
+                st.success(f"Speed upgraded to level {speed_level + 1}!")
+                st.rerun()
+            else:
+                st.error(f"Not enough {TAG_CURRENCY_NAME}. Need {speed_cost}.")
+    
+    # Capacity upgrade
+    with col2:
+        capacity_level = upgrades.get("capacity", 1)
+        capacity_cost = calculate_upgrade_cost("capacity", capacity_level)
+        
+        # Calculate max level for progress bar
+        max_level = 10
+        progress_percentage = min(100, (capacity_level / max_level) * 100)
+        
+        # Create an upgrade card with progress bar in dark theme
+        st.markdown(f"""
+        <div class="upgrade-card">
+            <div class="upgrade-title">🔍 Capacity Upgrade</div>
+            <div class="upgrade-level">Level {capacity_level}</div>
+            <div class="level-bar">
+                <div class="level-progress" style="width: {progress_percentage}%;"></div>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Effect:</span>
+                <span class="upgrade-stat-value">Increases tags discovered</span>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Current:</span>
+                <span class="upgrade-stat-value">{calculate_expedition_capacity()} tags</span>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Next Level:</span>
+                <span class="upgrade-stat-value">{calculate_expedition_capacity() + 1} tags</span>
+            </div>
+            <div class="upgrade-cost" style="color: #9D4EDD;">Cost: {capacity_cost} {TAG_CURRENCY_NAME}</div>
+        </div>
+        """, unsafe_allow_html=True)
+        
+        # Upgrade button
+        can_afford = st.session_state.tag_currency >= capacity_cost
+        button_class = "pulse-button" if can_afford else ""
+        
+        if st.button(f"Upgrade Capacity", key=f"upgrade_capacity", disabled=not can_afford, use_container_width=True):
+            if purchase_library_upgrade("capacity"):
+                st.success(f"Capacity upgraded to level {capacity_level + 1}!")
+                st.rerun()
+            else:
+                st.error(f"Not enough {TAG_CURRENCY_NAME}. Need {capacity_cost}.")
+    
+    # Rarity upgrade
+    with col3:
+        rarity_level = upgrades.get("rarity", 1)
+        rarity_cost = calculate_upgrade_cost("rarity", rarity_level)
+        
+        # Calculate max level for progress bar
+        max_level = 10
+        progress_percentage = min(100, (rarity_level / max_level) * 100)
+        
+        # Create an upgrade card with progress bar in dark theme
+        st.markdown(f"""
+        <div class="upgrade-card">
+            <div class="upgrade-title">💎 Rarity Upgrade</div>
+            <div class="upgrade-level">Level {rarity_level}</div>
+            <div class="level-bar">
+                <div class="level-progress" style="width: {progress_percentage}%;"></div>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Effect:</span>
+                <span class="upgrade-stat-value">Improves rare tag chance</span>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Current:</span>
+                <span class="upgrade-stat-value">+{(rarity_level - 1) * 20}% boost</span>
+            </div>
+            <div class="upgrade-stat">
+                <span class="upgrade-stat-label">Next Level:</span>
+                <span class="upgrade-stat-value">+{rarity_level * 20}% boost</span>
+            </div>
+            <div class="upgrade-cost" style="color: #9D4EDD;">Cost: {rarity_cost} {TAG_CURRENCY_NAME}</div>
+        </div>
+        """, unsafe_allow_html=True)
+        
+        # Upgrade button
+        can_afford = st.session_state.tag_currency >= rarity_cost
+        button_class = "pulse-button" if can_afford else ""
+        
+        if st.button(f"Upgrade Rarity", key=f"upgrade_rarity", disabled=not can_afford, use_container_width=True):
+            if purchase_library_upgrade("rarity"):
+                st.success(f"Rapacity upgraded to level {rarity_level + 1}!")
+                st.rerun()
+            else:
+                st.error(f"Not enough {TAG_CURRENCY_NAME}. Need {rarity_cost}.")
+                
+    # Add a styled info box about library growth with dark theme
+    st.markdown("""
+    <div style="background-color: rgba(13, 110, 253, 0.15); 
+               border-left: 4px solid #0d6efd; 
+               border-radius: 4px; 
+               padding: 15px; 
+               margin-top: 20px;
+               color: #ffffff;">
+        <h4 style="margin-top: 0; color: #6495ED;">📚 Library Growth</h4>
+        <p style="margin-bottom: 0;">
+            Your library will grow as you collect more tags. Each floor of the library unlocks new rarities and 
+            improves your chances of finding rare tags. Continue collecting tags to unlock deeper levels of the library!
+        </p>
+    </div>
+    """, unsafe_allow_html=True)
+
+def initialize_library_system():
+    """Initialize the library system state in session state if not already present."""
+    if 'library_system_initialized' not in st.session_state:
+        st.session_state.library_system_initialized = True
+        
+        # Try to load from storage first
+        library_state = tag_storage.load_library_state(st.session_state)
+        
+        if library_state:
+            # We already have the state loaded into session_state by the load function
+            print("Library system loaded from storage.")
+        else:
+            # Initialize with defaults
+            st.session_state.discovered_tags = {}  # {tag_name: {"rarity": str, "discovery_time": timestamp, "category": str}}
+            st.session_state.library_exploration_history = []  # List of recent library explorations
+            
+            # Initialize enkephalin if not present
+            if 'enkephalin' not in st.session_state:
+                st.session_state.enkephalin = 0
+            
+            # For the library interface
+            st.session_state.expedition_results = []  # Results from completed expeditions
+            
+            # Library growth system
+            st.session_state.library_growth = {
+                "total_discoveries": 0,
+                "last_discovery_time": time.time()
+            }
+            
+            # Upgrade system for library
+            st.session_state.library_upgrades = {
+                "speed": 1,      # Expedition speed (reduces cooldown time)
+                "capacity": 1,   # Tags discovered per expedition
+                "rarity": 1      # Rare tag chance
+            }
+            
+            # Set of explored library tiers
+            st.session_state.explored_library_tiers = set()
+            
+            print("Library system initialized with defaults.")
+
+    # Store library floors in session state if not already there
+    if 'library_floors' not in st.session_state:
+        st.session_state.library_floors = LIBRARY_FLOORS
+    
+    # Update categories for any "unknown" category tags
+    update_discovered_tag_categories()
+        
+    # Add CSS animations for styling
+    st.markdown("""
+    <style>
+    /* Star of the City animation */
+    @keyframes star-glow {
+        0% { text-shadow: 0 0 5px #FFD700; }
+        50% { text-shadow: 0 0 15px #FFD700; }
+        100% { text-shadow: 0 0 5px #FFD700; }
+    }
+    
+    /* Impuritas Civitas animation */
+    @keyframes rainbow-text {
+        0% { color: red; }
+        14% { color: orange; }
+        28% { color: yellow; }
+        42% { color: green; }
+        57% { color: blue; }
+        71% { color: indigo; }
+        85% { color: violet; }
+        100% { color: red; }
+    }
+    
+    @keyframes rainbow-border {
+        0% { border-color: red; }
+        14% { border-color: orange; }
+        28% { border-color: yellow; }
+        42% { border-color: green; }
+        57% { border-color: blue; }
+        71% { border-color: indigo; }
+        85% { border-color: violet; }
+        100% { border-color: red; }
+    }
+    
+    /* Urban Nightmare animation */
+    @keyframes nightmare-pulse {
+        0% { border-color: #FF9800; }
+        50% { border-color: #FF5722; }
+        100% { border-color: #FF9800; }
+    }
+    
+    /* Urban Plague subtle effect */
+    .glow-purple {
+        text-shadow: 0 0 3px #9C27B0;
+    }
+    
+    /* Apply the animations to specific rarity classes */
+    .star-of-city {
+        animation: star-glow 2s infinite;
+        font-weight: bold;
+    }
+    
+    .impuritas-civitas {
+        animation: rainbow-text 4s linear infinite;
+        font-weight: bold;
+    }
+    
+    .urban-nightmare {
+        animation: nightmare-pulse 3s infinite;
+        font-weight: bold;
+    }
+    
+    .urban-plague {
+        text-shadow: 0 0 3px #9C27B0;
+        font-weight: bold;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+
+def display_library_extraction():
+    """Display the library exploration interface."""
+    initialize_library_system()
+    
+    st.title(f"Welcome to {LIBRARY_INFO['name']}")
+    st.markdown(f"""
+    <div style="background-color: rgba(74, 20, 140, 0.15); 
+               border-radius: 10px; 
+               padding: 15px; 
+               margin-bottom: 20px; 
+               border-left: 5px solid {LIBRARY_INFO['color']};
+               color: #ffffff;">
+        <p style="margin: 0;">{LIBRARY_INFO['description']}</p>
+    </div>
+    """, unsafe_allow_html=True)
+    
+    # Create tabs with enhanced styling for dark theme
+    st.markdown("""
+    <style>
+    /* Custom styling for tabs */
+    .stTabs [data-baseweb="tab-list"] {
+        gap: 2px;
+    }
+    
+    .stTabs [data-baseweb="tab"] {
+        border-radius: 5px 5px 0 0;
+        padding: 10px 16px;
+        font-weight: 600;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    # Store current tab index in session state if not present
+    if 'library_tab_index' not in st.session_state:
+        st.session_state.library_tab_index = 0
+    
+    # Create expanded tabs - removed the essence tab
+    explore_tab, discovered_tab, building_tab = st.tabs([
+        "📚 Library Exploration", 
+        "🔍 Discovered Tags",
+        "🏛️ Library Building"
+    ])
+    
+    with explore_tab:
+        st.session_state.library_tab_index = 0
+        display_library_exploration_interface()
+    
+    with discovered_tab:
+        st.session_state.library_tab_index = 1
+        st.subheader("Your Discovered Tags")
+        st.write("These are tags you've discovered through the library system. They differ from your collected tags, which are obtained from scanning images.")
+        
+        # Display discovered tags using our new function
+        display_discovered_tags()
+    
+    with building_tab:
+        st.session_state.library_tab_index = 2
+        display_library_building()
+
+def display_discovered_tags():
+    """Display the user's discovered tags with the same visual style as the tag collection"""
+    # Show total unique discovered tags
+    if not hasattr(st.session_state, 'discovered_tags') or not st.session_state.discovered_tags:
+        st.info("Explore the library to discover new tags!")
+        return
+        
+    unique_tags = len(st.session_state.discovered_tags)
+    st.write(f"You have discovered {unique_tags} unique tags.")
+    
+    # Count tags by rarity
+    rarity_counts = {}
+    for tag_info in st.session_state.discovered_tags.values():
+        rarity = tag_info.get("rarity", "Unknown")
+        if rarity not in rarity_counts:
+            rarity_counts[rarity] = 0
+        rarity_counts[rarity] += 1
+    
+    # Only display rarity categories that have tags
+    active_rarities = {r: c for r, c in rarity_counts.items() if c > 0}
+    
+    # If there are active rarities to display
+    if active_rarities:
+        display_discovered_rarity_distribution(active_rarities)
+    
+    # Add a sorting option
+    sort_options = ["Category (rarest first)", "Rarity", "Discovery Time"]
+    selected_sort = st.selectbox("Sort tags by:", sort_options, key="discovered_tags_sort")
+    
+    # Group tags by the selected method
+    if selected_sort == "Category (rarest first)":
+        # Group tags by category
+        categories = {}
+        for tag, info in st.session_state.discovered_tags.items():
+            category = info.get("category", "unknown")
+            if category not in categories:
+                categories[category] = []
+            categories[category].append((tag, info))
+        
+        # Display tags by category in expanders
+        for category, tags in sorted(categories.items()):
+            # Get rarity order for sorting
+            rarity_order = list(RARITY_LEVELS.keys())
+            
+            # Sort tags by rarity (rarest first)
+            def get_rarity_index(tag_tuple):
+                tag, info = tag_tuple
+                rarity = info.get("rarity", "Unknown")
+                if rarity in rarity_order:
+                    return len(rarity_order) - rarity_order.index(rarity)
+                return 0
+            
+            sorted_tags = sorted(tags, key=get_rarity_index, reverse=True)
+            
+            # Check if category has any rare tags
+            has_rare_tags = any(info.get("rarity") in ["Impuritas Civitas", "Star of the City"] 
+                               for _, info in sorted_tags)
+            
+            # Get category info if available
+            category_display = category.capitalize()
+            if category in TAG_CATEGORIES:
+                category_info = TAG_CATEGORIES[category]
+                category_icon = category_info.get("icon", "")
+                category_color = category_info.get("color", "#888888")
+                category_display = f"<span style='color:{category_color};'>{category_icon} {category.capitalize()}</span>"
+            
+            # Create header with information about rare tags if present
+            header = f"{category_display} ({len(tags)} tags)"
+            if has_rare_tags:
+                header += " ✨ Contains rare tags!"
+                
+            # Display the header and expander
+            st.markdown(header, unsafe_allow_html=True)
+            with st.expander("Show/Hide", expanded=has_rare_tags):
+                # Group by rarity within category
+                rarity_groups = {}
+                for tag, info in sorted_tags:
+                    rarity = info.get("rarity", "Unknown")
+                    if rarity not in rarity_groups:
+                        rarity_groups[rarity] = []
+                    rarity_groups[rarity].append((tag, info))
+                
+                # Display each rarity group in order (rarest first)
+                for rarity in reversed(rarity_order):
+                    if rarity in rarity_groups:
+                        tags_in_rarity = rarity_groups[rarity]
+                        if tags_in_rarity:
+                            color = RARITY_LEVELS[rarity]["color"]
+                            
+                            # Special styling for rare rarities
+                            if rarity == "Impuritas Civitas":
+                                rarity_style = f"animation:rainbow-text 4s linear infinite;font-weight:bold;"
+                            elif rarity == "Star of the City":
+                                rarity_style = f"color:{color};text-shadow:0 0 3px gold;font-weight:bold;"
+                            elif rarity == "Urban Nightmare":
+                                rarity_style = f"color:{color};text-shadow:0 0 1px #FF5722;font-weight:bold;"
+                            else:
+                                rarity_style = f"color:{color};font-weight:bold;"
+                            
+                            st.markdown(f"<span style='{rarity_style}'>{rarity.capitalize()}</span> ({len(tags_in_rarity)} tags)", unsafe_allow_html=True)
+                            display_discovered_tag_grid(tags_in_rarity)
+                            st.markdown("---")
+                
+    elif selected_sort == "Rarity":
+        # Group tags by rarity level
+        rarity_groups = {}
+        for tag, info in st.session_state.discovered_tags.items():
+            rarity = info.get("rarity", "Unknown")
+            if rarity not in rarity_groups:
+                rarity_groups[rarity] = []
+            rarity_groups[rarity].append((tag, info))
+        
+        # Get ordered rarities (rarest first)
+        ordered_rarities = list(RARITY_LEVELS.keys())
+        ordered_rarities.reverse()  # Reverse to show rarest first
+        
+        # Display tags by rarity
+        for rarity in ordered_rarities:
+            if rarity in rarity_groups:
+                tags = rarity_groups[rarity]
+                color = RARITY_LEVELS[rarity]["color"]
+                
+                # Add special styling for rare rarities
+                rarity_html = f"<span style='color:{color};font-weight:bold;'>{rarity.capitalize()}</span>"
+                if rarity == "Impuritas Civitas":
+                    rarity_html = f"<span style='animation:rainbow-text 4s linear infinite;font-weight:bold;'>{rarity.capitalize()}</span>"
+                elif rarity == "Star of the City":
+                    rarity_html = f"<span style='color:{color};text-shadow:0 0 3px gold;font-weight:bold;'>{rarity.capitalize()}</span>"
+                elif rarity == "Urban Nightmare":
+                    rarity_html = f"<span style='color:{color};text-shadow:0 0 1px #FF5722;font-weight:bold;'>{rarity.capitalize()}</span>"
+                
+                # First create the title with HTML, then use it in the expander
+                st.markdown(f"### {rarity_html} ({len(tags)} tags)", unsafe_allow_html=True)
+                with st.expander("Show/Hide", expanded=rarity in ["Impuritas Civitas", "Star of the City"]):
+                    # Group by category within rarity
+                    category_groups = {}
+                    for tag, info in tags:
+                        category = info.get("category", "unknown")
+                        if category not in category_groups:
+                            category_groups[category] = []
+                        category_groups[category].append((tag, info))
+                    
+                    # Display each category within this rarity level
+                    for category, category_tags in sorted(category_groups.items()):
+                        # Get category info if available
+                        category_display = category.capitalize()
+                        if category in TAG_CATEGORIES:
+                            category_info = TAG_CATEGORIES[category]
+                            category_icon = category_info.get("icon", "")
+                            category_color = category_info.get("color", "#888888")
+                            category_display = f"<span style='color:{category_color};'>{category_icon} {category.capitalize()}</span>"
+                        
+                        st.markdown(f"#### {category_display} ({len(category_tags)} tags)", unsafe_allow_html=True)
+                        display_discovered_tag_grid(category_tags)
+                        st.markdown("---")
+    
+    elif selected_sort == "Discovery Time":
+        # Sort all tags by discovery time (newest first)
+        sorted_tags = []
+        for tag, info in st.session_state.discovered_tags.items():
+            discovery_time = info.get("discovery_time", "")
+            sorted_tags.append((tag, info, discovery_time))
+        
+        sorted_tags.sort(key=lambda x: x[2], reverse=True)  # Sort by time, newest first
+        
+        # Group by date
+        date_groups = {}
+        for tag, info, time_str in sorted_tags:
+            # Extract just the date part if timestamp has date and time
+            date = time_str.split()[0] if " " in time_str else time_str
+            
+            if date not in date_groups:
+                date_groups[date] = []
+            date_groups[date].append((tag, info))
+        
+        # Display tags grouped by discovery date
+        for date, tags in date_groups.items():
+            date_display = date if date else "Unknown date"
+            st.markdown(f"### Discovered on {date_display} ({len(tags)} tags)")
+            
+            with st.expander("Show/Hide", expanded=date == list(date_groups.keys())[0]):  # Expand most recent by default
+                display_discovered_tag_grid(tags)
+                st.markdown("---")
+
+def display_discovered_rarity_distribution(active_rarities):
+    """Display distribution of discovered tags by rarity with themed animations"""
+    # Add the necessary CSS for animations
+    st.markdown("""
+    <style>
+    @keyframes grid-glow {
+        0% { text-shadow: 0 0 2px gold; }
+        50% { text-shadow: 0 0 6px gold; }
+        100% { text-shadow: 0 0 2px gold; }
+    }
+    
+    @keyframes grid-rainbow {
+        0% { color: red; }
+        14% { color: orange; }
+        28% { color: yellow; }
+        42% { color: green; }
+        57% { color: blue; }
+        71% { color: indigo; }
+        85% { color: violet; }
+        100% { color: red; }
+    }
+    
+    @keyframes grid-pulse {
+        0% { opacity: 0.8; }
+        50% { opacity: 1; }
+        100% { opacity: 0.8; }
+    }
+    
+    .grid-star {
+        text-shadow: 0 0 3px gold;
+        animation: grid-glow 2s infinite;
+    }
+    
+    .grid-impuritas {
+        animation: grid-rainbow 4s linear infinite;
+    }
+    
+    .grid-nightmare {
+        text-shadow: 0 0 1px #FF5722;
+        animation: grid-pulse 3s infinite;
+    }
+    
+    .grid-plague {
+        text-shadow: 0 0 1px #9C27B0;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+    
+    rarity_cols = st.columns(len(active_rarities))
+    for i, (rarity, count) in enumerate(active_rarities.items()):
+        with rarity_cols[i]:
+            # Get color with fallback
+            color = RARITY_LEVELS.get(rarity, {}).get("color", "#888888")
+            
+            # Apply special styling based on rarity
+            style = f"color:{color};font-weight:bold;"
+            class_name = ""
+            
+            if rarity == "Impuritas Civitas":
+                class_name = "grid-impuritas"
+            elif rarity == "Star of the City":
+                class_name = "grid-star"
+            elif rarity == "Urban Nightmare":
+                class_name = "grid-nightmare"
+            elif rarity == "Urban Plague":
+                class_name = "grid-plague"
+            
+            if class_name:
+                st.markdown(
+                    f"<div style='text-align:center;'><span class='{class_name}' style='font-weight:bold;'>{rarity.capitalize()}</span><br>{count}</div>",
+                    unsafe_allow_html=True
+                )
+            else:
+                st.markdown(
+                    f"<div style='text-align:center;'><span style='{style}'>{rarity.capitalize()}</span><br>{count}</div>",
+                    unsafe_allow_html=True
+                )
+
+def display_discovered_tag_grid(tags):
+    """Display discovered tags in a grid layout with discovery information"""
+    # Create a grid layout for tags
+    cols = st.columns(3)
+    for i, (tag, info) in enumerate(sorted(tags)):
+        col_idx = i % 3
+        with cols[col_idx]:
+            rarity = info.get("rarity", "Unknown")
+            discovery_time = info.get("discovery_time", "")
+            library_floor = info.get("library_floor", "")
+            discovery_count = info.get("discovery_count", 1)
+            
+            color = RARITY_LEVELS.get(rarity, {}).get("color", "#888888")
+            
+            # Get sample count if available
+            sample_count = None
+            if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
+                if tag in st.session_state.tag_rarity_metadata:
+                    tag_info = st.session_state.tag_rarity_metadata[tag]
+                    if isinstance(tag_info, dict) and "sample_count" in tag_info:
+                        sample_count = tag_info["sample_count"]
+            
+            # Format sample count
+            sample_display = ""
+            if sample_count is not None:
+                if sample_count >= 1000000:
+                    sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count/1000000:.1f}M)</span>"
+                elif sample_count >= 1000:
+                    sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count/1000:.1f}K)</span>"
+                else:
+                    sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count})</span>"
+            
+            # Apply special styling for rare tags
+            tag_html = tag
+            if rarity == "Impuritas Civitas":
+                tag_html = f"<span style='animation: rainbow-text 4s linear infinite;'>{tag}</span>"
+            elif rarity == "Star of the City":
+                tag_html = f"<span style='text-shadow: 0 0 3px gold;'>{tag}</span>"
+            elif rarity == "Urban Nightmare":
+                tag_html = f"<span style='text-shadow: 0 0 1px #FF9800;'>{tag}</span>"
+            
+            # Display tag with rarity badge and discovery info
+            st.markdown(
+                f"{tag_html} <span style='background-color:{color};color:white;padding:2px 6px;border-radius:10px;font-size:0.8em;'>{rarity.capitalize()}</span> {sample_display}",
+                unsafe_allow_html=True
+            )
+            
+            # Show discovery details
+            if library_floor:
+                st.markdown(f"<span style='font-size:0.85em;'>Found in: {library_floor}</span>", unsafe_allow_html=True)
+            
+            if discovery_count > 1:
+                st.markdown(f"<span style='font-size:0.85em;'>Seen {discovery_count} times</span>", unsafe_allow_html=True)