#!/usr/bin/env python3 """ Image Tagger 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 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.model_loader import load_exported_model, is_windows, check_flash_attention 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__))), "model") 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.", "High Precision": "Uses higher thresholds to prioritize accuracy over recall. Fewer but more confident predictions.", "High Recall": "Uses lower thresholds to capture more potential tags at the expense of accuracy. More comprehensive tagging.", "Overall": "Uses a single threshold value across all categories. Simplest approach for consistent behavior.", "Weighted": "Uses thresholds weighted by category importance. Better balance for tags that matter most.", "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 **Threshold value**: Approximately 0.33 (optimized on validation data) **Performance metrics**: - Micro F1: ~0.62 - Macro F1: ~0.35 - Precision: ~0.63 - Recall: ~0.60 """, "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 **Threshold value**: Approximately 0.19-0.21 (optimized on validation data) **Performance metrics**: - Micro F1: ~0.49 - Macro F1: ~0.41 - Precision: ~0.37 - Recall: ~0.53 """, "Balanced": """ ### Balanced Profile **Technical definition**: Provides a compromise between precision and recall with moderate thresholds. **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 **Threshold value**: Approximately 0.26 (optimized on validation data) **Performance metrics**: - Micro F1: ~0.59 - Macro F1: ~0.39 - Precision: ~0.51 - Recall: ~0.70 """, "High Precision": """ ### High Precision Profile **Technical definition**: Uses higher thresholds to prioritize precision (correctness) over recall (coverage). **When to use**: When you need high confidence in the tags that are returned and prefer to miss tags rather than include incorrect ones. **Effects**: - Much higher precision (84-97% of returned tags are correct) - Lower recall (only captures 35-60% of relevant tags) - Returns fewer tags overall, but with higher confidence **Threshold value**: 0.50 (optimized for precision on validation data) **Performance metrics**: - Micro F1: ~0.50 - Macro F1: ~0.22 - Precision: ~0.84 - Recall: ~0.35 """, "High Recall": """ ### High Recall Profile **Technical definition**: Uses lower thresholds to prioritize recall (coverage) over precision (correctness). **When to use**: When you want to capture as many potential tags as possible and don't mind some incorrect suggestions. **Effects**: - Very high recall (captures 90%+ of relevant tags) - Much lower precision (only 18-49% of returned tags may be correct) - Returns many more tags, including less confident ones **Threshold value**: 0.10-0.12 (optimized for recall on validation data) **Performance metrics**: - Micro F1: ~0.30 - Macro F1: ~0.35 - Precision: ~0.18 - Recall: ~0.90 """, "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.35 (uses "balanced" threshold by default) **Note**: The threshold value is user-adjustable with the slider below. """, "Weighted": """ ### Weighted Profile **Technical definition**: Uses thresholds weighted by category importance. **When to use**: When you want different sensitivity for different categories based on their importance. **Effects**: - More important categories (like character and copyright) get optimized thresholds - Less important categories get adjusted thresholds based on their contribution - Better balance for the tags that matter most **Default threshold values**: Varies by category (based on importance weighting) **Note**: This uses pre-calculated optimal thresholds that can't be adjusted directly. """, "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 get_profile_metrics(thresholds, profile_name, model_type="refined"): """ Extract metrics for the given profile from the thresholds dictionary Args: thresholds: The thresholds dictionary profile_name: Name of the profile (micro_opt, macro_opt, balanced, etc.) model_type: 'initial' or 'refined' Returns: Dictionary with metrics or None if not found """ profile_key = None # Map UI-friendly names to internal keys if profile_name == "Micro Optimized": profile_key = "micro_opt" elif profile_name == "Macro Optimized": profile_key = "macro_opt" elif profile_name == "Balanced": profile_key = "balanced" elif profile_name == "High Precision": profile_key = "high_precision" elif profile_name == "High Recall": profile_key = "high_recall" # For overall/weighted/category-specific, we're using the balanced profile metrics elif profile_name in ["Overall", "Weighted", "Category-specific"]: profile_key = "balanced" # Handle the new JSON structure with "initial" and "refined" top-level keys if "initial" in thresholds and "refined" in thresholds: # Get the appropriate model type model_type_key = model_type # Use the passed in model_type # Make sure the model_type_key is valid if model_type_key not in thresholds: model_type_key = "refined" if "refined" in thresholds else "initial" # Check if the profile exists if "overall" in thresholds[model_type_key] and profile_key in thresholds[model_type_key]["overall"]: return thresholds[model_type_key]["overall"][profile_key] else: # Fallback to the old structure if "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 to ensure smooth transitions between modes and preserve user customizations """ # Get the new profile new_profile = st.session_state.threshold_profile # Initialize active_threshold and active_category_thresholds based on profile if hasattr(st.session_state, 'thresholds') and hasattr(st.session_state, 'settings'): # Create category thresholds if they don't exist if st.session_state.settings['active_category_thresholds'] is None: st.session_state.settings['active_category_thresholds'] = {} # Get existing thresholds current_thresholds = st.session_state.settings['active_category_thresholds'] # Get model type for accessing thresholds - handle the new JSON structure if "initial" in st.session_state.thresholds and "refined" in st.session_state.thresholds: model_type_key = "refined" if hasattr(st.session_state, 'model_type') and st.session_state.model_type == "full" else "initial" # Make sure the model_type_key is valid if model_type_key not in st.session_state.thresholds: model_type_key = "refined" if "refined" in st.session_state.thresholds else "initial" else: # Use None for old structure model_type_key = None # Map profile display names to internal keys profile_key = None if new_profile == "Micro Optimized": profile_key = "micro_opt" elif new_profile == "Macro Optimized": profile_key = "macro_opt" elif new_profile == "Balanced": profile_key = "balanced" elif new_profile == "High Precision": profile_key = "high_precision" elif new_profile == "High Recall": profile_key = "high_recall" # For specialized profiles, update thresholds from the thresholds dictionary if profile_key: # Set overall threshold based on JSON structure if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and profile_key in st.session_state.thresholds[model_type_key]["overall"]: st.session_state.settings['active_threshold'] = st.session_state.thresholds[model_type_key]["overall"][profile_key]["threshold"] else: # Old structure if "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 based on JSON structure for category in st.session_state.categories: if model_type_key is not None: # New structure if "categories" in st.session_state.thresholds[model_type_key] and category in st.session_state.thresholds[model_type_key]["categories"]: if profile_key in st.session_state.thresholds[model_type_key]["categories"][category]: current_thresholds[category] = st.session_state.thresholds[model_type_key]["categories"][category][profile_key]["threshold"] else: # Fallback to overall threshold if profile not found for this category current_thresholds[category] = st.session_state.settings['active_threshold'] else: # Old structure if "categories" in st.session_state.thresholds and category in st.session_state.thresholds["categories"]: if 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'] # For "Overall" profile, reset to use just the overall threshold elif new_profile == "Overall": # Use the balanced threshold for Overall profile based on JSON structure if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and "balanced" in st.session_state.thresholds[model_type_key]["overall"]: st.session_state.settings['active_threshold'] = st.session_state.thresholds[model_type_key]["overall"]["balanced"]["threshold"] else: # Old structure 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"] # Clear any category-specific overrides st.session_state.settings['active_category_thresholds'] = {} # For "Weighted" profile, use weighted thresholds elif new_profile == "Weighted": # Use the balanced threshold as base for Weighted profile based on JSON structure if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and "balanced" in st.session_state.thresholds[model_type_key]["overall"]: st.session_state.settings['active_threshold'] = st.session_state.thresholds[model_type_key]["overall"]["balanced"]["threshold"] else: # Old structure 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"] # Get weighted thresholds if they exist, otherwise use balanced if model_type_key is not None: # New structure if "weighted" in st.session_state.thresholds[model_type_key]: weighted_thresholds = st.session_state.thresholds[model_type_key]["weighted"] for category in st.session_state.categories: if category in weighted_thresholds: current_thresholds[category] = weighted_thresholds[category] else: # Fallback to balanced threshold for categories not in weighted if "categories" in st.session_state.thresholds[model_type_key] and category in st.session_state.thresholds[model_type_key]["categories"]: if "balanced" in st.session_state.thresholds[model_type_key]["categories"][category]: current_thresholds[category] = st.session_state.thresholds[model_type_key]["categories"][category]["balanced"]["threshold"] else: current_thresholds[category] = st.session_state.settings['active_threshold'] else: current_thresholds[category] = st.session_state.settings['active_threshold'] else: # Old structure if "weighted" in st.session_state.thresholds: weighted_thresholds = st.session_state.thresholds["weighted"] for category in st.session_state.categories: if category in weighted_thresholds: current_thresholds[category] = weighted_thresholds[category] else: # Fallback to balanced threshold if "categories" in st.session_state.thresholds and category in st.session_state.thresholds["categories"]: if "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'] else: current_thresholds[category] = st.session_state.settings['active_threshold'] # For "Category-specific", initialize with balanced thresholds elif new_profile == "Category-specific": # Use the balanced threshold as base for Category-specific based on JSON structure if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and "balanced" in st.session_state.thresholds[model_type_key]["overall"]: st.session_state.settings['active_threshold'] = st.session_state.thresholds[model_type_key]["overall"]["balanced"]["threshold"] else: # Old structure 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"] # Initialize with balanced thresholds for each category for category in st.session_state.categories: if model_type_key is not None: # New structure if "categories" in st.session_state.thresholds[model_type_key] and category in st.session_state.thresholds[model_type_key]["categories"]: if "balanced" in st.session_state.thresholds[model_type_key]["categories"][category]: current_thresholds[category] = st.session_state.thresholds[model_type_key]["categories"][category]["balanced"]["threshold"] else: current_thresholds[category] = st.session_state.settings['active_threshold'] else: current_thresholds[category] = st.session_state.settings['active_threshold'] else: # Old structure if "categories" in st.session_state.thresholds and category in st.session_state.thresholds["categories"]: if "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'] else: current_thresholds[category] = st.session_state.settings['active_threshold'] def create_micro_macro_comparison(): """ Creates a visual explanation of micro vs macro optimization Returns: HTML for the visualization """ html = """ """ return html 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 Args: all_probs: Dictionary with all probabilities organized by category threshold_profile: Current threshold profile active_threshold: Overall threshold value active_category_thresholds: Dictionary of category-specific thresholds min_confidence: Minimum confidence to include selected_categories: Dictionary of selected categories Returns: tags: Dictionary of filtered tags above threshold by category all_tags: List of all tags above threshold """ # Apply thresholds to each category tags = {} all_tags = [] for category, cat_probs in all_probs.items(): # Get the appropriate threshold for this category # For Overall and Weighted profiles, respect per-category overrides if available threshold = active_category_thresholds.get(category, active_threshold) if active_category_thresholds else 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("Image Tagging Interface") st.markdown("---") # Check platform and Flash Attention windows_system = is_windows() flash_attn_installed = check_flash_attention() if 'settings' not in st.session_state: st.session_state.settings = { 'show_all_tags': False, 'compact_view': True, 'min_confidence': 0.01, 'threshold_profile': "Balanced", 'active_threshold': 0.35, 'active_category_thresholds': None, 'selected_categories': {}, 'replace_underscores': False # Added new setting } # Initialize show_profile_help state st.session_state.show_profile_help = False # Define default threshold values (initialized here to avoid errors) default_threshold_values = { 'overall': 0.35, 'weighted': 0.35, 'category_thresholds': {}, 'high_precision_thresholds': {}, 'high_recall_thresholds': {} } # 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 # Check if ONNX model files exist onnx_model_path = os.path.join(os.path.dirname(MODEL_DIR), "model_initial.onnx") onnx_metadata_path = os.path.join(os.path.dirname(MODEL_DIR), "model_initial_metadata.json") onnx_available = os.path.exists(onnx_model_path) and os.path.exists(onnx_metadata_path) # Default to ONNX if available, otherwise fallback to initial_only on Windows or full elsewhere if onnx_available: st.session_state.model_type = "onnx" else: st.session_state.model_type = "initial_only" if windows_system else "full" # Sidebar for model selection and information with st.sidebar: st.header("Model Selection") # Check if ONNX model files exist onnx_model_path = os.path.join(os.path.dirname(MODEL_DIR), "model_initial.onnx") onnx_metadata_path = os.path.join(os.path.dirname(MODEL_DIR), "model_initial_metadata.json") onnx_available = os.path.exists(onnx_model_path) and os.path.exists(onnx_metadata_path) # Define model options, including ONNX if available model_options = [ "Refined (Tag Embeddings)", "Initial (Base Model)" ] # Add ONNX option if available if onnx_available: model_options.append("ONNX Accelerated (Fastest)") # Determine the default index for model selection if st.session_state.model_type == "onnx" and onnx_available: default_index = 2 # ONNX is first priority when available elif windows_system or st.session_state.model_type == "initial_only": default_index = 1 # Initial Only is second priority, default on Windows else: default_index = 0 # Full model is last priority # Model type selection with radio buttons model_type = st.radio( "Select Model Type:", model_options, index=min(default_index, len(model_options)-1), # Ensure index is valid help=""" Full Model: Uses both initial and refined predictions for highest accuracy (requires more VRAM) Initial Only: Uses only the initial classifier, reducing VRAM usage at slight quality cost ONNX Accelerated: Optimized for inference speed, best for batch processing (if available) """ ) # Convert selection to internal model type if model_type == "Full Model (Best Quality)": selected_model_type = "full" elif model_type == "ONNX Accelerated (Fastest)": selected_model_type = "onnx" # Store ONNX paths in session state for later use st.session_state.onnx_model_path = onnx_model_path st.session_state.onnx_metadata_path = onnx_metadata_path else: # "Initial Only (Lower VRAM)" selected_model_type = "initial_only" # If the model type has changed, we need to reload the model if selected_model_type != st.session_state.model_type: st.session_state.model_loaded = False st.session_state.model_type = selected_model_type # Add Windows warning if relevant if windows_system and selected_model_type=="full": st.warning(""" ### Windows Compatibility Note The refined model requires Flash Attention which is difficult to install on Windows. For Windows users, I recommend using the "Initial Only" or ONNX Accelerated model which: - Does not require Flash Attention - Uses less memory - Provides very close to full prediction quality (check performance notes on HF) """) # Add a button to reload the model with current settings 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 when not loaded if not st.session_state.model_loaded: try: with st.spinner(f"Loading {st.session_state.model_type} model..."): if st.session_state.model_type == "onnx": # Load ONNX model and metadata import json import onnxruntime as ort try: # Check ONNX providers (for info display) providers = ort.get_available_providers() gpu_available = any('GPU' in provider for provider in providers) # Store provider info in session state st.session_state.onnx_providers = providers st.session_state.onnx_gpu_available = gpu_available # Load metadata with open(st.session_state.onnx_metadata_path, 'r') as f: metadata = json.load(f) # Load thresholds from a separate file or use defaults thresholds_path = os.path.join(MODEL_DIR, "thresholds.json") if os.path.exists(thresholds_path): with open(thresholds_path, 'r') as f: thresholds = json.load(f) else: # If no thresholds file, extract from metadata if available if 'thresholds' in metadata: thresholds = metadata['thresholds'] else: # Use default thresholds if not available thresholds = { 'overall': {'balanced': {'threshold': 0.35}}, 'weighted': {'f1': {'threshold': 0.35}}, 'categories': {} } # Build category thresholds if necessary if 'tag_to_category' in metadata: categories = set(metadata['tag_to_category'].values()) thresholds['categories'] = { cat: { 'balanced': {'threshold': 0.35}, 'high_precision': {'threshold': 0.45}, 'high_recall': {'threshold': 0.25} } for cat in categories } # Set device info for display device = "ONNX Runtime" + (" (GPU)" if gpu_available else " (CPU)") param_dtype = "float32" # Most ONNX models use float32 # Store information in session state st.session_state.model = None # No PyTorch model for ONNX st.session_state.device = device st.session_state.param_dtype = param_dtype st.session_state.thresholds = thresholds st.session_state.metadata = metadata st.session_state.model_loaded = True # Get categories from metadata categories = list(set(metadata['tag_to_category'].values())) st.session_state.categories = categories # Initialize selected categories if needed if not st.session_state.settings['selected_categories']: st.session_state.settings['selected_categories'] = {cat: True for cat in categories} except Exception as e: st.error(f"Error loading ONNX model: {str(e)}") st.info(f"Make sure the ONNX model and metadata files exist at: {st.session_state.onnx_model_path} and {st.session_state.onnx_metadata_path}") st.code(traceback.format_exc()) st.stop() else: # Load PyTorch model as before model, thresholds, metadata = load_exported_model( MODEL_DIR, model_type=st.session_state.model_type ) # Extract device and precision info device = next(model.parameters()).device param_dtype = next(model.parameters()).dtype # Store model in session state for PyTorch models st.session_state.model = model # Common code for all model types # Get available categories categories = list(set(metadata['tag_to_category'].values())) # Initialize selected categories (all selected by default) if not st.session_state.settings['selected_categories']: st.session_state.settings['selected_categories'] = {cat: True for cat in categories} # Store common info in session state st.session_state.device = device st.session_state.param_dtype = param_dtype st.session_state.thresholds = thresholds st.session_state.metadata = metadata st.session_state.model_loaded = True st.session_state.categories = categories # Debug: Print loaded thresholds to verify they're loaded correctly print("Loaded thresholds:", thresholds) if "initial" in thresholds and "refined" in thresholds: # Choose which model type to use as default model_type_key = "refined" if st.session_state.model_type == "full" else "initial" # Set overall threshold from the balanced profile if "overall" in thresholds[model_type_key] and "balanced" in thresholds[model_type_key]["overall"]: default_threshold_values['overall'] = thresholds[model_type_key]["overall"]["balanced"]["threshold"] # Get weighted threshold if available if "weighted" in thresholds[model_type_key] and "f1" in thresholds[model_type_key]["weighted"]: default_threshold_values['weighted'] = thresholds[model_type_key]["weighted"]["f1"]["threshold"] # Set category thresholds if "categories" in thresholds[model_type_key]: default_threshold_values['category_thresholds'] = { cat: opt['balanced']['threshold'] for cat, opt in thresholds[model_type_key]["categories"].items() } # Set high precision and high recall thresholds default_threshold_values['high_precision_thresholds'] = { cat: opt['high_precision']['threshold'] for cat, opt in thresholds[model_type_key]["categories"].items() } default_threshold_values['high_recall_thresholds'] = { cat: opt['high_recall']['threshold'] for cat, opt in thresholds[model_type_key]["categories"].items() } else: # Fallback to the old structure for backward compatibility if "overall" in thresholds and "balanced" in thresholds["overall"]: default_threshold_values['overall'] = thresholds["overall"]["balanced"]["threshold"] # Get weighted threshold if available if "weighted" in thresholds and "f1" in thresholds["weighted"]: default_threshold_values['weighted'] = thresholds["weighted"]["f1"]["threshold"] # Set category thresholds if "categories" in thresholds: default_threshold_values['category_thresholds'] = { cat: opt['balanced']['threshold'] for cat, opt in thresholds["categories"].items() } # Set high precision and high recall thresholds default_threshold_values['high_precision_thresholds'] = { cat: opt['high_precision']['threshold'] for cat, opt in thresholds["categories"].items() } default_threshold_values['high_recall_thresholds'] = { cat: opt['high_recall']['threshold'] for cat, opt in thresholds["categories"].items() } # Update session state with current threshold values # This part is crucial - store the values in session state st.session_state.default_threshold_values = default_threshold_values # Update active threshold with default values if st.session_state.settings['threshold_profile'] == "Overall": st.session_state.settings['active_threshold'] = default_threshold_values['overall'] elif st.session_state.settings['threshold_profile'] == "Weighted": st.session_state.settings['active_threshold'] = default_threshold_values['weighted'] except Exception as e: st.error(f"Error loading model: {str(e)}") st.info(f"Looking for model in: {os.path.abspath(MODEL_DIR)}") # Check for specific files if st.session_state.model_type == "initial_only": expected_model_paths = [ os.path.join(MODEL_DIR, "model_initial_only.pt"), os.path.join(MODEL_DIR, "model_initial.pt") ] if not any(os.path.exists(p) for p in expected_model_paths): st.error(f"Initial-only model file not found. Checked: {', '.join(expected_model_paths)}") st.info("Make sure you've exported both model types.") else: expected_model_paths = [ os.path.join(MODEL_DIR, "model_refined.pt"), os.path.join(MODEL_DIR, "model.pt"), os.path.join(MODEL_DIR, "model_full.pt") ] if not any(os.path.exists(p) for p in expected_model_paths): st.error(f"Full model file not found. Checked: {', '.join(expected_model_paths)}") st.code(traceback.format_exc()) st.stop() with st.sidebar: st.header("Model Information") if st.session_state.model_loaded: # Show model type if st.session_state.model_type == "onnx": st.success("Using ONNX Accelerated Model") # Show ONNX-specific info if hasattr(st.session_state, 'onnx_gpu_available') and st.session_state.onnx_gpu_available: st.write("Acceleration: GPU available") else: st.write("Acceleration: CPU only") elif st.session_state.model_type == "full": st.success("Using Full Model (Best Quality)") # Show Flash Attention info for full model if not flash_attn_installed and is_windows(): st.warning("Note: Flash Attention not available on Windows") else: st.success("Using Initial-Only Model (Lower VRAM)") # Show common model info st.write(f"Device: {st.session_state.device}") st.write(f"Precision: {st.session_state.param_dtype}") st.write(f"Total tags: {st.session_state.metadata['total_tags']}") # Show categories in an expander with st.expander("Available Categories"): for category in sorted(st.session_state.categories): st.write(f"- {category.capitalize()}") # Add an expander with information with st.expander("About this app"): st.write(""" This app uses a trained image tagging model to analyze and tag images. **Model Options**: - **ONNX Accelerated (Fastest)**: Optimized for inference speed with minimal VRAM usage, ideal for batch processing - **Refined Model (Tag Embeddings)**: Higher quality predictions using both initial and refined layers (uses more VRAM) - **Initial Model (Base model)**: Reduced VRAM usage with slightly lower accuracy (good for systems with limited resources) **Platform Notes**: - **Windows Users**: ONNX Accelerated model is recommended for best performance - **CUDA Support**: GPU acceleration is available for ONNX models if CUDA 12.x and cuDNN are installed - **Linux Users**: The Refined Model with Flash Attention provides the best quality results **Features**: - Upload or select an image - Process multiple images in batch mode with customizable batch size - Choose from different threshold profiles - Adjust category-specific thresholds - View predictions organized by category - Limit results to top N tags within each category - Save tags to text files in various locations - Export tags with consistent formatting for external use - Fast batch processing **Threshold profiles**: - **Micro Optimized**: Optimizes micro-averaged F1 score (best for common tags) - **Macro Optimized**: Optimizes macro-averaged F1 score (better for rare tags) - **Balanced**: Provides a good balance of precision and recall - **High Precision**: Prioritizes accuracy over recall - **High Recall**: Captures more potential tags but may be less accurate """) with st.sidebar: # Add separator for visual clarity st.markdown("---") # Support information st.subheader("💡 Notes") st.markdown(""" This tagger was trained on a subset of the available data and for limited epochs due to hardware limitations. A more comprehensive model trained on the full 3+ million image dataset and many more epochs 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(""" Buy Me A Coffee """, 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 """) with st.sidebar: # Add game link at the bottom of the sidebar st.markdown("---") st.subheader("Try the Tag Collector Game!") st.write("Test your tagging skills in our gamified version of the image tagger!") if st.button("🎮 Launch Tag Collector Game", type="primary"): # Get the current port to determine the game URL current_port = os.environ.get("STREAMLIT_SERVER_PORT", "8501") # The game should run on a different port (8502 if this is 8501, or vice versa) game_port = "8502" if current_port == "8501" else "8501" # Check if the game file exists game_path = os.path.join(os.path.dirname(__file__), "tag_collector_game.py") if os.path.exists(game_path): # Launch the game in a new process try: # Determine streamlit path if sys.platform == "win32": streamlit_path = os.path.join("venv", "Scripts", "streamlit.exe") else: streamlit_path = os.path.join("venv", "bin", "streamlit") if not os.path.exists(streamlit_path): streamlit_path = "streamlit" # Fallback to global streamlit # Build command to run game on different port command = [streamlit_path, "run", game_path, "--server.port", game_port] # Launch in background if sys.platform == "win32": subprocess.Popen(command, shell=True, creationflags=subprocess.CREATE_NEW_CONSOLE) else: subprocess.Popen(command) # Open in browser game_url = f"http://localhost:{game_port}" webbrowser.open(game_url) st.success(f"Launching Tag Collector Game!") except Exception as e: st.error(f"Failed to launch game: {str(e)}") else: st.error(f"Game file not found: {game_path}") st.info("Make sure tag_collector_game.py is in the same directory as this app.") # Main content area col1, col2 = st.columns([1, 1.5]) # Column 1: Image upload and display with col1: st.header("Image") # Add tabs for Upload, Examples, and Batch Processing 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 a temporary file with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: tmp_file.write(uploaded_file.getvalue()) image_path = tmp_file.name # Store the original filename for saving tags st.session_state.original_filename = uploaded_file.name # Display the image image = Image.open(uploaded_file) st.image(image, use_container_width=True) with batch_tab: st.subheader("Batch Process Images") st.write("Process multiple images from a folder and save tags to text files.") # Folder selection batch_folder = st.text_input("Enter folder path containing images:", "") if st.button("Browse Folder..."): # This is a dummy button since Streamlit doesn't have a native folder picker st.info("Please type the folder path manually in the text input above.") # Save location options save_options = st.radio( "Where to save tag files:", ["Same folder as images", "Custom location", "Default save folder"], index=0 ) # Add batch size control st.subheader("Performance Options") batch_size_text = st.text_input( "Batch size (images processed at once)", value="4", help="Higher values may improve processing speed but use more memory. Recommended: 4-16" ) # Convert to integer with error handling try: batch_size = int(batch_size_text) if batch_size < 1: st.warning("Batch size must be at least 1. Using batch size of 1.") batch_size = 1 except ValueError: st.warning("Please enter a valid number for batch size. Using default batch size of 4.") batch_size = 4 # Performance note if batch_size > 8: st.info(f"Using larger batch size ({batch_size}). If you encounter memory issues, try reducing this value.") st.write("Set tag limits per category for batch processing:") # Create a toggle for enabling category limits enable_category_limits = st.checkbox("Limit tags per category in batch output", value=False) # Initialize the category limit dictionary if it doesn't exist if 'category_limits' not in st.session_state: st.session_state.category_limits = {} if enable_category_limits: # Create a two-column layout for more compact display limit_cols = st.columns(2) # Add an explanation about the values st.markdown(""" **Limit Values:** * **-1** = No limit (include all tags) * **0** = Exclude category entirely * **N** (positive number) = Include only top N tags """) if hasattr(st.session_state, 'categories'): # Create text inputs for each category for i, category in enumerate(sorted(st.session_state.categories)): col_idx = i % 2 # Alternate between columns with limit_cols[col_idx]: # Get current limit value (default to -1 for unlimited) current_limit = st.session_state.category_limits.get(category, -1) # Add a text input for this category limit_text = st.text_input( f"{category.capitalize()} (top N):", value=str(current_limit), key=f"limit_{category}", help="-1 = no limit, 0 = exclude, N = top N tags" ) # Convert to integer with error handling try: limit = int(limit_text) if limit < -1: st.warning(f"Limit for {category} must be -1 or greater. Using -1 (unlimited).") limit = -1 except ValueError: st.warning(f"Invalid limit for {category}. Using -1 (unlimited).") limit = -1 # Display a clear indicator of what this setting means if limit == -1: st.caption(f"✅ Including all {category} tags") elif limit == 0: st.caption(f"❌ Excluding all {category} tags") else: st.caption(f"⚙ī¸ Including top {limit} {category} tags") # Store the limit in session state st.session_state.category_limits[category] = limit else: st.info("Categories will be available after loading a model.") else: # Clear any existing limits if disabled st.session_state.category_limits = {} custom_save_dir = None if save_options == "Custom location": # Allow selecting a custom save location if 'custom_folders' not in st.session_state: st.session_state.custom_folders = get_default_save_locations() custom_save_dir = 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 ) # Allow adding a new folder new_folder = st.text_input("Or enter a new folder path:", key="batch_new_folder") if st.button("Add Folder", key="batch_add_folder") and new_folder: if os.path.isdir(new_folder): if new_folder not in st.session_state.custom_folders: st.session_state.custom_folders.append(new_folder) st.success(f"Added folder: {new_folder}") st.rerun() else: st.info("This folder is already in the list.") else: try: # Try to create the folder if it doesn't exist os.makedirs(new_folder, exist_ok=True) st.session_state.custom_folders.append(new_folder) st.success(f"Created and added folder: {new_folder}") st.rerun() except Exception as e: st.error(f"Could not create folder: {str(e)}") # Check if folder exists and count images if batch_folder and os.path.isdir(batch_folder): # Count image files image_extensions = ['*.jpg', '*.jpeg', '*.png'] image_files = [] for ext in image_extensions: image_files.extend(glob.glob(os.path.join(batch_folder, ext))) image_files.extend(glob.glob(os.path.join(batch_folder, 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 total_images = len(image_files) st.write(f"Found {total_images} image files in the folder.") # Show first few images as thumbnails (more compact layout) if image_files: st.write("Sample images:") num_preview = min(8, len(image_files)) thumbnail_cols = st.columns(4) for i, img_path in enumerate(image_files[:num_preview]): with thumbnail_cols[i % 4]: try: img = Image.open(img_path) # Make thumbnails smaller for more compact view st.image(img, width=80, caption=os.path.basename(img_path)) except: st.write(f"Error loading {os.path.basename(img_path)}") # Determine save directory based on user selection if save_options == "Same folder as images": save_dir = batch_folder elif save_options == "Custom location": save_dir = custom_save_dir else: # Default save folder app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) save_dir = os.path.join(app_dir, "saved_tags") os.makedirs(save_dir, exist_ok=True) # Add a more prominent button for batch processing st.markdown("---") process_col1, process_col2 = st.columns([3, 1]) with process_col1: st.write(f"Ready to process {total_images} images") st.write(f"Tags will be saved to: **{save_dir}**") with process_col2: # Check if model is loaded process_button_disabled = not st.session_state.model_loaded # Add a prominent batch processing button if st.button("🔄 Process All Images", key="process_batch_btn", use_container_width=True, disabled=process_button_disabled, type="primary"): if not st.session_state.model_loaded: st.error("Model not loaded. Please check the model settings.") else: with st.spinner("Processing images..."): # Create progress bar progress_bar = st.progress(0) status_text = st.empty() # Define progress callback def update_progress(current, total, image_path): if total > 0: progress = min(current / total, 1.0) progress_bar.progress(progress) if image_path: status_text.text(f"Processing {current}/{total}: {os.path.basename(image_path)}") else: status_text.text(f"Completed processing {current}/{total} images") # Get settings for batch processing curr_threshold_profile = st.session_state.settings['threshold_profile'] curr_active_threshold = st.session_state.settings['active_threshold'] curr_active_category_thresholds = st.session_state.settings['active_category_thresholds'] curr_min_confidence = st.session_state.settings['min_confidence'] # Get category limits if enabled curr_category_limits = None if 'category_limits' in st.session_state and enable_category_limits: curr_category_limits = st.session_state.category_limits # Print debugging info in a more structured way if curr_category_limits: st.write("Category limit settings:") # Group by type for cleaner display excluded = [] limited = [] unlimited = [] for cat, limit in sorted(curr_category_limits.items()): if limit == 0: excluded.append(cat) elif limit > 0: limited.append(f"{cat}: top {limit}") else: # limit == -1 unlimited.append(cat) # Show in a more structured way if excluded: st.write("❌ Excluded categories: " + ", ".join(excluded)) if limited: st.write("⚙ī¸ Limited categories: " + ", ".join(limited)) if unlimited: st.write("✅ Unlimited categories: " + ", ".join(unlimited)) if not excluded and not limited: st.write("No limits set (all categories included)") # Then when you call the batch processing functions, pass this parameter: if st.session_state.model_type == "onnx": # Use ONNX batch processing batch_results = batch_process_images_onnx( folder_path=batch_folder, model_path=st.session_state.onnx_model_path, metadata_path=st.session_state.onnx_metadata_path, threshold_profile=curr_threshold_profile, active_threshold=curr_active_threshold, active_category_thresholds=curr_active_category_thresholds, save_dir=save_dir, progress_callback=update_progress, min_confidence=curr_min_confidence, batch_size=batch_size, category_limits=curr_category_limits ) else: # Use standard PyTorch processing batch_results = batch_process_images( folder_path=batch_folder, model=st.session_state.model, thresholds=st.session_state.thresholds, metadata=st.session_state.metadata, threshold_profile=curr_threshold_profile, active_threshold=curr_active_threshold, active_category_thresholds=curr_active_category_thresholds, save_dir=save_dir, progress_callback=update_progress, min_confidence=curr_min_confidence, batch_size=batch_size, category_limits=st.session_state.category_limits if enable_category_limits else None ) # Display results display_batch_results(batch_results) # If model is not loaded, show a hint if not st.session_state.model_loaded: st.warning("Please load a model before processing images.") else: st.warning("No image files found in the selected folder.") elif batch_folder: st.error(f"Folder not found: {batch_folder}") # Column 2: Threshold controls and predictions with col2: st.header("Tagging Controls") # Define all available profiles all_profiles = [ "Micro Optimized", "Macro Optimized", "Balanced", "High Precision", "High Recall", "Overall", "Weighted", "Category-specific" ] # Define a default index based on what's most similar to the current setting default_index = 2 # Default to "Balanced" if "threshold_profile" in st.session_state.settings: # If there's an existing selection, try to match it existing_profile = st.session_state.settings['threshold_profile'] if existing_profile in all_profiles: default_index = all_profiles.index(existing_profile) # Map old profile names to new ones elif existing_profile == "Overall": default_index = all_profiles.index("Overall") elif existing_profile == "Weighted": default_index = all_profiles.index("Weighted") elif existing_profile == "Category-specific": default_index = all_profiles.index("Category-specific") elif existing_profile == "High Precision": default_index = all_profiles.index("High Precision") elif existing_profile == "High Recall": default_index = all_profiles.index("High Recall") # Create the profile selection UI with help button profile_col1, profile_col2 = st.columns([3, 1]) with profile_col1: # Create the profile dropdown threshold_profile = st.selectbox( "Select threshold profile", options=all_profiles, index=default_index, key="threshold_profile", on_change=on_threshold_profile_change ) with profile_col2: # Add a help button that expands to show detailed information if st.button("ℹī¸ Help", key="profile_help"): st.session_state.show_profile_help = not st.session_state.get('show_profile_help', False) # Display the help text if the button was clicked if st.session_state.get('show_profile_help', False): st.markdown(threshold_profile_explanations[threshold_profile]) else: # Just show the description st.info(threshold_profile_descriptions[threshold_profile]) # Get profile metrics from thresholds if st.session_state.model_loaded: # Try to get metrics for the selected profile model_type = "refined" if st.session_state.model_type == "full" else "initial" metrics = get_profile_metrics(st.session_state.thresholds, threshold_profile, model_type) if metrics: # Create metrics display metrics_cols = st.columns(3) with metrics_cols[0]: # Display threshold threshold_value = metrics.get("threshold", 0.35) st.metric("Threshold", f"{threshold_value:.3f}") with metrics_cols[1]: # Display micro F1 micro_f1 = metrics.get("micro_f1", metrics.get("micro_precision", 0)) st.metric("Micro F1", f"{micro_f1:.3f}" if micro_f1 else "N/A") # Show precision if available precision = metrics.get("precision", metrics.get("micro_precision", 0)) if precision: st.metric("Precision", f"{precision:.3f}") with metrics_cols[2]: # Display macro F1 macro_f1 = metrics.get("macro_f1", 0) st.metric("Macro F1", f"{macro_f1:.3f}" if macro_f1 else "N/A") # Show recall if available recall = metrics.get("recall", metrics.get("micro_recall", 0)) if recall: st.metric("Recall", f"{recall:.3f}") # Initialize thresholds based on the selected profile active_threshold = None active_category_thresholds = {} # Handle the new JSON structure if st.session_state.model_loaded: # Determine which model type key to use if "initial" in st.session_state.thresholds and "refined" in st.session_state.thresholds: model_type_key = "refined" if st.session_state.model_type == "full" else "initial" # Make sure the model_type_key is valid if model_type_key not in st.session_state.thresholds: model_type_key = "refined" if "refined" in st.session_state.thresholds else "initial" else: model_type_key = None # Old structure - access thresholds directly # Map profile display names to internal keys profile_key = None if threshold_profile == "Micro Optimized": profile_key = "micro_opt" elif threshold_profile == "Macro Optimized": profile_key = "macro_opt" elif threshold_profile == "Balanced": profile_key = "balanced" elif threshold_profile == "High Precision": profile_key = "high_precision" elif threshold_profile == "High Recall": profile_key = "high_recall" # For specialized profiles, get thresholds from the thresholds dictionary if profile_key: # Get overall threshold if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and profile_key in st.session_state.thresholds[model_type_key]["overall"]: active_threshold = st.session_state.thresholds[model_type_key]["overall"][profile_key]["threshold"] else: # Old structure if "overall" in st.session_state.thresholds and profile_key in st.session_state.thresholds["overall"]: active_threshold = st.session_state.thresholds["overall"][profile_key]["threshold"] # Get category thresholds for category in st.session_state.categories: if model_type_key is not None: # New structure if "categories" in st.session_state.thresholds[model_type_key] and category in st.session_state.thresholds[model_type_key]["categories"]: if profile_key in st.session_state.thresholds[model_type_key]["categories"][category]: active_category_thresholds[category] = st.session_state.thresholds[model_type_key]["categories"][category][profile_key]["threshold"] else: # Fallback to overall threshold if profile not found for this category active_category_thresholds[category] = active_threshold else: active_category_thresholds[category] = active_threshold else: # Old structure if "categories" in st.session_state.thresholds and category in st.session_state.thresholds["categories"]: if profile_key in st.session_state.thresholds["categories"][category]: active_category_thresholds[category] = st.session_state.thresholds["categories"][category][profile_key]["threshold"] else: active_category_thresholds[category] = active_threshold else: active_category_thresholds[category] = active_threshold # Show informational text for these profiles st.info(f"The '{threshold_profile}' profile uses pre-optimized thresholds.") # Show disabled slider for overall threshold (for informational purposes) st.slider( "Overall threshold (reference)", min_value=0.01, max_value=1.0, value=float(active_threshold), step=0.01, disabled=True ) elif threshold_profile == "Overall" and st.session_state.model_loaded: # Use the balanced threshold for Overall profile if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and "balanced" in st.session_state.thresholds[model_type_key]["overall"]: active_threshold = st.session_state.thresholds[model_type_key]["overall"]["balanced"]["threshold"] else: # Old structure if "overall" in st.session_state.thresholds and "balanced" in st.session_state.thresholds["overall"]: active_threshold = st.session_state.thresholds["overall"]["balanced"]["threshold"] # Show slider for adjusting the overall threshold active_threshold = st.slider( "Overall threshold", min_value=0.01, max_value=1.0, value=float(active_threshold), step=0.01 ) elif threshold_profile == "Weighted" and st.session_state.model_loaded: # Use the balanced threshold as base for Weighted profile if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and "balanced" in st.session_state.thresholds[model_type_key]["overall"]: active_threshold = st.session_state.thresholds[model_type_key]["overall"]["balanced"]["threshold"] else: # Old structure if "overall" in st.session_state.thresholds and "balanced" in st.session_state.thresholds["overall"]: active_threshold = st.session_state.thresholds["overall"]["balanced"]["threshold"] # Show disabled slider for overall threshold (for informational purposes) st.slider( "Overall threshold (reference)", min_value=0.01, max_value=1.0, value=float(active_threshold), step=0.01, disabled=True ) st.info("The 'Weighted' profile uses different optimized thresholds for each category.") # Get weighted thresholds if they exist, otherwise use balanced if model_type_key is not None: # New structure if "weighted" in st.session_state.thresholds[model_type_key]: weighted_thresholds = st.session_state.thresholds[model_type_key]["weighted"] for category in st.session_state.categories: if category in weighted_thresholds: active_category_thresholds[category] = weighted_thresholds[category] else: # Fallback to balanced threshold if "categories" in st.session_state.thresholds[model_type_key] and category in st.session_state.thresholds[model_type_key]["categories"]: if "balanced" in st.session_state.thresholds[model_type_key]["categories"][category]: active_category_thresholds[category] = st.session_state.thresholds[model_type_key]["categories"][category]["balanced"]["threshold"] else: active_category_thresholds[category] = active_threshold else: active_category_thresholds[category] = active_threshold else: # Old structure if "weighted" in st.session_state.thresholds: weighted_thresholds = st.session_state.thresholds["weighted"] for category in st.session_state.categories: if category in weighted_thresholds: active_category_thresholds[category] = weighted_thresholds[category] else: # Fallback to balanced threshold if "categories" in st.session_state.thresholds and category in st.session_state.thresholds["categories"]: if "balanced" in st.session_state.thresholds["categories"][category]: active_category_thresholds[category] = st.session_state.thresholds["categories"][category]["balanced"]["threshold"] else: active_category_thresholds[category] = active_threshold else: active_category_thresholds[category] = active_threshold elif threshold_profile == "Category-specific" and st.session_state.model_loaded: # Use the balanced threshold as base for Category-specific if model_type_key is not None: # New structure if "overall" in st.session_state.thresholds[model_type_key] and "balanced" in st.session_state.thresholds[model_type_key]["overall"]: active_threshold = st.session_state.thresholds[model_type_key]["overall"]["balanced"]["threshold"] else: # Old structure if "overall" in st.session_state.thresholds and "balanced" in st.session_state.thresholds["overall"]: active_threshold = st.session_state.thresholds["overall"]["balanced"]["threshold"] # Show disabled slider for overall threshold (for informational purposes) 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:") # Create two columns for better layout of sliders slider_cols = st.columns(2) # Initialize with balanced thresholds for i, category in enumerate(sorted(st.session_state.categories)): # Get the balanced threshold for this category category_threshold = None if model_type_key is not None: # New structure if "categories" in st.session_state.thresholds[model_type_key] and category in st.session_state.thresholds[model_type_key]["categories"]: if "balanced" in st.session_state.thresholds[model_type_key]["categories"][category]: category_threshold = st.session_state.thresholds[model_type_key]["categories"][category]["balanced"]["threshold"] else: category_threshold = active_threshold else: category_threshold = active_threshold else: # Old structure if "categories" in st.session_state.thresholds and category in st.session_state.thresholds["categories"]: if "balanced" in st.session_state.thresholds["categories"][category]: category_threshold = st.session_state.thresholds["categories"][category]["balanced"]["threshold"] else: category_threshold = active_threshold else: category_threshold = active_threshold # Add slider to appropriate column col_idx = i % 2 # Alternate between columns with slider_cols[col_idx]: active_category_thresholds[category] = st.slider( f"{category.capitalize()}", min_value=0.01, max_value=1.0, value=float(category_threshold), step=0.01, key=f"slider_{category}" ) # Update session state with the thresholds if active_threshold is not None: st.session_state.settings['active_threshold'] = active_threshold if active_category_thresholds: st.session_state.settings['active_category_thresholds'] = active_category_thresholds # Add threshold profile details expander with st.expander("Threshold Profile Details"): # Add tabs for the visualizations if st.session_state.model_loaded: threshold_tabs = st.tabs(["About Metrics"]) with threshold_tabs[0]: st.markdown(""" ### Understanding Performance Metrics **F1 Score** is the harmonic mean of precision and recall: `2 * (precision * recall) / (precision + recall)` **Micro F1** calculates metrics globally by considering each example/prediction pair. This gives more weight to categories with more examples. **Macro F1** calculates F1 separately for each category and then takes the average. This treats all categories equally regardless of their size. """) st.markdown(create_micro_macro_comparison(), unsafe_allow_html=True) st.markdown(""" ### Other Metrics **Precision** measures how many of the predicted tags are correct: `true_positives / (true_positives + false_positives)` **Recall** measures how many of the relevant tags are captured: `true_positives / (true_positives + false_negatives)` ### The Precision-Recall Tradeoff There's an inherent tradeoff between precision and recall: - Higher threshold → Higher precision, Lower recall - Lower threshold → Lower precision, Higher recall The best threshold depends on your specific use case: - **Prefer Precision**: When false positives are costly (e.g., you want only accurate tags) - **Prefer Recall**: When false negatives are costly (e.g., you don't want to miss any potentially relevant tags) - **Balanced**: When both types of errors are equally important """) else: st.info("Load a model to see detailed threshold information.") # Display options display_options = st.expander("Display Options", expanded=False) with display_options: # Tag display options 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']) # Add the new checkbox for replacing underscores with spaces 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 session state with display options st.session_state.settings['show_all_tags'] = show_all_tags st.session_state.settings['compact_view'] = compact_view st.session_state.settings['min_confidence'] = min_confidence st.session_state.settings['replace_underscores'] = replace_underscores # Category selection for the "All Tags" section st.write("Categories to include in 'All Tags' section:") # Create a multi-column layout for category checkboxes category_cols = st.columns(3) selected_categories = {} # If categories exist in session state, create checkboxes for each if hasattr(st.session_state, 'categories'): for i, category in enumerate(sorted(st.session_state.categories)): col_idx = i % 3 # Distribute across 3 columns with category_cols[col_idx]: # Use previously selected value or default to True 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}" ) # Update session state with selected categories st.session_state.settings['selected_categories'] = selected_categories if st.session_state.model_loaded: if st.session_state.model_type == "onnx": model_type_display = "ONNX Accelerated Model" elif st.session_state.model_type == "full": model_type_display = "Full Model" else: model_type_display = "Initial-Only Model (Lower VRAM)" st.info(f"Using: {model_type_display}") # Run inference button for single image if image_path and st.button("Run Tagging"): if not st.session_state.model_loaded: st.error("Model not loaded. Please check the model settings.") else: with st.spinner("Analyzing image..."): try: inference_start = time.time() # Different processing based on model type if st.session_state.model_type == "onnx": # Use the appropriate function for ONNX inference from utils.onnx_processing import process_single_image_onnx # Run ONNX inference result = process_single_image_onnx( image_path=image_path, model_path=st.session_state.onnx_model_path, metadata=st.session_state.metadata, threshold_profile=threshold_profile, active_threshold=active_threshold, active_category_thresholds=active_category_thresholds, min_confidence=min_confidence ) else: # Run standard PyTorch inference 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=active_threshold, active_category_thresholds=active_category_thresholds, min_confidence=min_confidence ) inference_time = time.time() - inference_start if result['success']: # Store results in session state st.session_state.all_probs = result['all_probs'] st.session_state.tags = result['tags'] st.session_state.all_tags = result['all_tags'] st.success(f"Analysis completed in {inference_time:.2f} seconds") else: st.error(f"Inference failed: {result.get('error', 'Unknown error')}") except Exception as e: st.error(f"Inference error: {str(e)}") st.code(traceback.format_exc()) # Display the predictions if available if image_path and hasattr(st.session_state, 'all_probs'): st.header("Predictions") # Apply current thresholds to stored probabilities filtered_tags, current_all_tags = apply_thresholds( st.session_state.all_probs, threshold_profile, active_threshold, active_category_thresholds, min_confidence, st.session_state.settings['selected_categories'] ) # Store the updated results back to session state st.session_state.tags = filtered_tags st.session_state.all_tags = current_all_tags # Create an empty list to collect all tags that pass the thresholds # We'll rebuild this list as we process each category all_tags = [] for category in sorted(st.session_state.all_probs.keys()): # Get all tags for this category and the filtered ones all_tags_in_category = st.session_state.all_probs.get(category, []) filtered_tags_in_category = filtered_tags.get(category, []) # Only show categories with tags if all_tags_in_category: # Get the appropriate threshold for this category if threshold_profile in ["Overall", "Weighted"]: threshold = active_threshold else: threshold = active_category_thresholds.get(category, active_threshold) # Create expander with count information expander_label = f"{category.capitalize()} ({len(filtered_tags_in_category)} tags)" with st.expander(expander_label, expanded=True): # Add threshold control specific to this category threshold_row = st.columns([1, 2]) with threshold_row[0]: # Display the current threshold pass with threshold_row[1]: # Always show category slider regardless of threshold profile cat_slider_key = f"cat_threshold_{category}" # For the initial value, respect the threshold profile if threshold_profile in ["Overall", "Weighted"]: # In these modes, start with the global threshold value current_cat_threshold = active_threshold else: # In other modes, use the category-specific threshold current_cat_threshold = active_category_thresholds.get(category, active_threshold) # Add a slider for this specific category new_threshold = st.slider( f"Adjust {category.capitalize()} threshold:", min_value=0.01, max_value=1.0, value=float(current_cat_threshold), step=0.01, key=cat_slider_key, disabled=(threshold_profile in ["Overall", "Weighted"]) # Disable in Overall/Weighted modes ) # If in Overall/Weighted mode, add an explanation if threshold_profile in ["Overall", "Weighted"]: st.info(f"Using global {threshold_profile.lower()} threshold. Switch to Category-specific mode to adjust individual categories.") else: # In other modes, update the category threshold active_category_thresholds[category] = new_threshold threshold = new_threshold # Update the active_category_thresholds with the new value active_category_thresholds[category] = new_threshold # Very important: Update the threshold variable for this iteration of the loop # This ensures the new threshold is used in this category's display threshold = new_threshold # Determine which tags to display based on show_all_tags setting if show_all_tags: tags_to_display = all_tags_in_category else: # Refilter the tags based on the updated threshold # This ensures the display immediately reflects threshold changes tags_to_display = [(tag, prob) for tag, prob in all_tags_in_category if prob >= threshold] filtered_tags[category] = tags_to_display # Add per-category limit controls limit_col1, limit_col2 = st.columns([1, 2]) with limit_col1: # Generate a unique key for each category's limit checkbox limit_key = f"limit_{category}_tags" limit_tags_for_category = st.checkbox("Limit tags", value=False, key=limit_key) with limit_col2: # Generate a unique key for each category's tag count slider slider_key = f"top_n_{category}_tags" # Handle the slider values carefully to avoid Streamlit errors tag_count = len(tags_to_display) # Skip slider if there are no tags to display if tag_count > 0: # Set sensible min and max values that avoid errors min_value = 0 # Make sure max_value is at least 1 (Streamlit requires min < max) max_value = max(1, min(999, tag_count)) # Default value should be between min and max default_value = min(max_value, 5) # Default to 5 tags if available top_n_tags_for_category = st.slider( "Show top", min_value=min_value, max_value=max_value, value=default_value, step=1, disabled=not limit_tags_for_category, key=slider_key ) else: # If no tags, just set a default value without showing the slider top_n_tags_for_category = 5 st.write("No tags to display") st.markdown("---") # Add separator line after controls if not tags_to_display: st.info(f"No tags above {min_confidence:.2f} confidence threshold") continue # Apply the per-category limit if enabled original_count = len(tags_to_display) if limit_tags_for_category and tags_to_display: limited_tags_to_display = tags_to_display[:top_n_tags_for_category] display_count = len(limited_tags_to_display) else: limited_tags_to_display = tags_to_display display_count = original_count # Display tags based on view mode if compact_view: # Compact view - show tags in a comma-separated list with % tag_list = [] replace_underscores = st.session_state.settings.get('replace_underscores', False) for tag, prob in limited_tags_to_display: # Format: tag (percent%) percentage = int(prob * 100) # Apply underscore replacement for display if enabled display_tag = tag.replace('_', ' ') if replace_underscores else tag tag_list.append(f"{display_tag} ({percentage}%)") # Add tag to all_tags list if it passes threshold and category is selected # Note: always use original tag (with underscores) for all_tags collection if prob >= threshold and selected_categories.get(category, True): all_tags.append(tag) # Join with commas and display st.markdown(", ".join(tag_list)) else: # Expanded view with progress bars for tag, prob in limited_tags_to_display: # Get the replacement setting replace_underscores = st.session_state.settings.get('replace_underscores', False) # Apply underscore replacement for display if enabled display_tag = tag.replace('_', ' ') if replace_underscores else tag # Add tag to all_tags list if it passes threshold and category is selected if prob >= threshold and selected_categories.get(category, True): all_tags.append(tag) # Add original tag to all_tags tag_display = f"**{display_tag}**" # Bold for tags above threshold else: tag_display = display_tag # Regular for tags below threshold # Display tag with progress bar st.write(tag_display) st.markdown(display_progress_bar(prob), unsafe_allow_html=True) # If we're limiting tags and there are more available, show a message if limit_tags_for_category and original_count > display_count: st.caption(f"Showing top {display_count} of {original_count} qualifying tags.") # Show summary at bottom with the truly updated all_tags list st.markdown("---") st.subheader(f"All Tags ({len(all_tags)} total)") if all_tags: # Check if underscore replacement is enabled replace_underscores = st.session_state.settings.get('replace_underscores', False) if replace_underscores: # Create a new list with underscores replaced by spaces for display only display_tags = [tag.replace('_', ' ') for tag in all_tags] st.write(", ".join(display_tags)) else: # Display original tags st.write(", ".join(all_tags)) else: st.info("No tags detected above the threshold.") # Add Save Tags section st.markdown("---") st.subheader("Save Tags") # Create column for save options save_col = st.columns(1)[0] with save_col: # Option to save to custom location if 'custom_folders' not in st.session_state: # Initialize with default save locations st.session_state.custom_folders = get_default_save_locations() # Display folder selection dropdown 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 ) # Allow adding a new folder path new_folder = st.text_input("Or enter a new folder path:") if st.button("Add Folder", key="add_folder") and new_folder: if os.path.isdir(new_folder): if new_folder not in st.session_state.custom_folders: st.session_state.custom_folders.append(new_folder) st.success(f"Added folder: {new_folder}") st.rerun() else: st.info("This folder is already in the list.") else: try: # Try to create the folder if it doesn't exist os.makedirs(new_folder, exist_ok=True) st.session_state.custom_folders.append(new_folder) st.success(f"Created and added folder: {new_folder}") st.rerun() except Exception as e: st.error(f"Could not create folder: {str(e)}") # Save to selected location button if st.button("💾 Save to Selected Location"): try: # Get the original filename if it exists original_filename = st.session_state.original_filename if hasattr(st.session_state, 'original_filename') else None # Save tags to file in selected location 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 preview of saved file 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)}") st.code(traceback.format_exc()) if __name__ == "__main__": image_tagger_app()