utils/ui_components.py

"""
UI components for the Image Tagger application.
"""

import os
import streamlit as st
from PIL import Image


def display_progress_bar(prob):
    """
    Create an HTML progress bar for displaying probability.
    
    Args:
        prob: Probability value between 0 and 1
        
    Returns:
        HTML string for the progress bar
    """
    # Convert probability to percentage
    percentage = int(prob * 100)
    
    # Choose color based on confidence level
    if prob >= 0.8:
        color = "green"
    elif prob >= 0.5:
        color = "orange"
    else:
        color = "red"
    
    # Return HTML for a styled progress bar
    return f"""
    <div style="margin-bottom: 5px; display: flex; align-items: center;">
        <div style="flex-grow: 1; background-color: #f0f0f0; border-radius: 3px; height: 8px; position: relative;">
            <div style="position: absolute; width: {percentage}%; background-color: {color}; height: 8px; border-radius: 3px;"></div>
        </div>
        <div style="margin-left: 8px; min-width: 40px; text-align: right; font-size: 0.9em;">{percentage}%</div>
    </div>
    """


def show_example_images(examples_dir):
    """
    Display example images from a directory.
    
    Args:
        examples_dir: Directory containing example images
        
    Returns:
        Selected image path or None
    """
    selected_image = None
    
    if os.path.exists(examples_dir):
        example_files = [f for f in os.listdir(examples_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        if example_files:
            st.write("Select an example image:")
            
            # Create a 2-column layout for examples
            example_cols = st.columns(2)
            
            for i, example_file in enumerate(example_files):
                col_idx = i % 2
                with example_cols[col_idx]:
                    example_path = os.path.join(examples_dir, example_file)
                    
                    # Display thumbnail
                    try:
                        img = Image.open(example_path)
                        st.image(img, width=150, caption=example_file)
                        
                        # Button to select this example
                        if st.button(f"Use", key=f"example_{i}"):
                            selected_image = example_path
                            st.session_state.original_filename = example_file
                            
                            # Display full image
                            st.image(img, use_container_width=True)
                            st.success(f"Example '{example_file}' selected!")
                    except Exception as e:
                        st.error(f"Error loading {example_file}: {str(e)}")
        else:
            st.info("No example images found.")
            st.write("Add some JPG or PNG images to the 'examples' directory.")
    else:
        st.info("Examples directory not found.")
        st.write("Create an 'examples' directory and add some JPG or PNG images.")
        
    return selected_image


def display_batch_results(batch_results):
    """
    Display batch processing results.
    
    Args:
        batch_results: Dictionary with batch processing results
    """
    if batch_results['success']:
        st.success(f"✅ Processed {batch_results['processed']} of {batch_results['total']} images")
        
        # Show details in an expander
        with st.expander("Batch Processing Results", expanded=True):
            # Count successes and failures
            successes = sum(1 for r in batch_results['results'].values() if r['success'])
            failures = batch_results['total'] - successes
            
            st.write(f"- Successfully tagged: {successes}")
            st.write(f"- Failed to process: {failures}")
            
            if failures > 0:
                # Show errors
                st.write("### Processing Errors")
                for img_path, result in batch_results['results'].items():
                    if not result['success']:
                        st.write(f"- **{os.path.basename(img_path)}**: {result.get('error', 'Unknown error')}")
            
            # Show the location of the output files
            if successes > 0:
                st.write("### Output Files")
                st.write(f"Tag files have been saved to the 'saved_tags' folder.")
                
                # Show the first few as examples
                st.write("Example outputs:")
                sample_results = [(path, res) for path, res in batch_results['results'].items() if res['success']][:3]
                for img_path, result in sample_results:
                    output_path = result.get('output_path', '')
                    if output_path and os.path.exists(output_path):
                        st.write(f"- **{os.path.basename(output_path)}**")
                        
                        # Show file contents in a collapsible code block
                        with open(output_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                        st.code(content, language='text')
    else:
        st.error(f"Batch processing failed: {batch_results.get('error', 'Unknown error')}")