#!/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(""" 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-v2 """) 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()