utils/preprocessing_ui.py

"""
Image Preprocessing UI Component for EmotionMirror application.

This module implements the UI components for showing image preprocessing options,
including the comparison between original and processed images.
Part of Step 4: Implementation of preprocessing techniques.
"""
import streamlit as st
import logging
import numpy as np
import cv2
from typing import Dict, Any, Optional
import io
from PIL import Image
import os
import time
import traceback

logger = logging.getLogger(__name__)

def show_preprocessing_ui(image_service, img: np.ndarray) -> Dict[str, Any]:
    """
    Complete UI handler for image preprocessing - this is the main entry point
    that should be called from app.py
    
    Args:
        image_service: The image service instance
        img: The image to preprocess
        
    Returns:
        Dict with information about the preprocessing state and user choices
    """
    try:
        # Save the original image in session state
        if hasattr(st.session_state, "original_image") == False:
            st.session_state.original_image = img.copy()
        
        # Initialize result
        result = {"success": True, "message": ""}
        
        # Set default selection to improved image if not already set
        if "selected_image_mode" not in st.session_state:
            st.session_state["selected_image_mode"] = "improved"
            st.session_state["use_improved_image"] = True
        
        # Create an expander for the suggested improvements
        with st.expander("Suggested improvements available", expanded=True):
            # Create a two-column layout for original and improved images
            col1, col2 = st.columns(2)
            
            # Apply basic enhancements and prepare the improved version
            try:
                # Use the new comprehensive enhancement method that combines
                # multiple advanced techniques for optimal facial detection
                improved_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                improved_bgr = image_service.enhance_image_for_facial_detection(improved_bgr)
                improved_rgb = cv2.cvtColor(improved_bgr, cv2.COLOR_BGR2RGB)
                
                # Get the specific parameters used for this image
                enhancement_params = image_service.get_last_enhancement_params()
                
                # Display images
                with col1:
                    st.markdown("**Original Image**")
                    st.image(img, use_column_width=True)
                
                with col2:
                    st.markdown("**Enhanced Image**")
                    st.image(improved_rgb, use_column_width=True)
                
                # Use a single row for buttons with current selection visually indicated
                col1, col2 = st.columns(2)
                
                with col1:
                    # Determine button style based on current selection
                    original_type = "primary" if st.session_state["selected_image_mode"] == "original" else "secondary"
                    if st.button("Continue with Original", key="continue_original_btn", type=original_type, use_container_width=True):
                        st.session_state["selected_image_mode"] = "original"
                        st.session_state["use_improved_image"] = False
                        st.session_state.current_image = img
                        st.experimental_rerun()
                
                with col2:
                    # Determine button style based on current selection
                    improved_type = "primary" if st.session_state["selected_image_mode"] == "improved" else "secondary"
                    if st.button("Use Improved Image", key="use_improved_btn", type=improved_type, use_container_width=True):
                        st.session_state["selected_image_mode"] = "improved"
                        st.session_state["use_improved_image"] = True
                        st.session_state.current_image = improved_rgb
                        st.experimental_rerun()
                
                # Single status message showing the current selection
                message = f"Using {'improved' if st.session_state['selected_image_mode'] == 'improved' else 'original'} image for analysis."
                st.info(message)
                
                # List the specific improvements made to this image
                st.markdown("**Dynamic Enhancements Applied:**")
                enhancement_params = image_service.get_last_enhancement_params()
                
                # Only show if we have parameters
                if enhancement_params:
                    # Format parameters for display
                    brightness_adj = enhancement_params.get("brightness_factor", 1.0)
                    contrast_adj = enhancement_params.get("contrast_factor", 1.0)
                    blur_kernel = enhancement_params.get("blur_kernel_size", 0)
                    gamma_val = enhancement_params.get("gamma", 1.0)
                    hist_eq = enhancement_params.get("needs_histogram_eq", False)
                    
                    # Display the specific adjustments
                    st.markdown(f"* **Brightness Adjustment**: {'Increased' if brightness_adj > 1.0 else 'Decreased'} to {brightness_adj:.2f}x")
                    st.markdown(f"* **Contrast Adjustment**: {'Increased' if contrast_adj > 1.0 else 'Decreased'} to {contrast_adj:.2f}x")
                    
                    if blur_kernel > 1:
                        st.markdown(f"* **Noise Reduction**: Applied with kernel size {blur_kernel}")
                    
                    if gamma_val != 1.0:
                        st.markdown(f"* **Gamma Correction**: {'Enhanced shadows' if gamma_val > 1.0 else 'Enhanced midtones'} (gamma={gamma_val:.2f})")
                    
                    if hist_eq:
                        st.markdown("* **Histogram Equalization**: Applied to improve contrast distribution")
                else:
                    # Fallback if no parameters available
                    st.markdown("* **Adaptive Image Processing**: Optimized for facial detection")
                
                # Dynamic explanation based on selected mode
                if st.session_state["selected_image_mode"] == "improved":
                    # Show explanation about improvements benefits
                    st.markdown("## Why these improvements help facial analysis")
                    
                    # Technical explanation
                    st.markdown("**Technical Benefits:**")
                    st.markdown("* **Balanced contrast:** Enhances the visibility of facial features while reducing shadows and highlights")
                    st.markdown("* **Optimal brightness:** Ensures facial features are clearly distinguishable without over-exposure")
                    st.markdown("* **Proper sizing:** Maintains ideal dimensions for detection algorithms to recognize facial landmarks")
                    
                    # Impact on emotion detection
                    st.markdown("**Impact on Emotion Detection:**")
                    st.markdown("* **More accurate emotion classification:** Cleaner input images lead to more reliable emotion detection")
                    st.markdown("* **Better feature extraction:** Facial features like eyes, mouth, and eyebrows are more clearly defined")
                    st.markdown("* **Reduced noise and artifacts:** Minimizes false detections and improves confidence scores")
                else:
                    # Show explanation about potential challenges with original images
                    st.markdown("## Potential challenges with original images")
                    
                    # Technical explanation
                    st.markdown("**Common issues with unprocessed images:**")
                    st.markdown("* **Variable lighting conditions:** Original images may have shadows or highlights that obscure facial features")
                    st.markdown("* **Inconsistent contrast:** Low contrast can make facial features harder to detect accurately")
                    st.markdown("* **Background noise:** Unprocessed images often contain visual elements that can distract detection algorithms")
                    
                    # Impact on detection
                    st.markdown("**Potential impact on detection quality:**")
                    st.markdown("* **Lower detection confidence:** Original images may result in less confident emotion classifications")
                    st.markdown("* **Feature detection challenges:** Some facial features might be missed or misidentified")
                    st.markdown("* **Increased false readings:** Environmental factors in the original image could lead to misinterpretations")
                    
                    # When to use original images
                    st.markdown("**When to use original images:**")
                    st.markdown("* When the original lighting and contrast are already optimal")
                    st.markdown("* When you want to analyze the image exactly as captured")
                    st.markdown("* For comparison with processed results")
                    
            except Exception as e:
                # Log detailed error information
                error_trace = traceback.format_exc()
                logger.error(f"Error in preprocessing UI: {str(e)}\n{error_trace}")
                
                # Show error to user - simple message
                st.error(f"Error processing image: {str(e)}")
                result["success"] = False
                result["message"] = f"Error: {str(e)}"
        
        return result
        
    except Exception as e:
        # Log detailed error information
        error_trace = traceback.format_exc()
        logger.error(f"Error in preprocessing UI: {str(e)}\n{error_trace}")
        
        # Show error to user - simple message
        st.error(f"Error processing image")
        
        # Return error information
        return {
            "success": False,
            "message": f"Error: {str(e)}",
            "error": str(e)
        }

def show_preprocessing_expandable(image_service, preprocessing_result: Dict[str, Any]) -> None:
    """
    Display preprocessing UI with an expandable section.
    This is the standalone implementation that doesn't depend on other functions.
    
    Args:
        image_service: The image service instance
        preprocessing_result: Dict with preprocessing results
    """
    # Only show if there are improvements
    if not preprocessing_result or "improvements" not in preprocessing_result or not preprocessing_result["improvements"]:
        logger.info("No improvements to display")
        return
        
    try:
        # Direct expandable implementation 
        with st.expander("Suggested improvements available", expanded=True):
            # Create columns for side-by-side comparison
            original_col, improved_col = st.columns(2)
            
            with original_col:
                st.markdown("**Original Image**")
                st.image(preprocessing_result["original_image"], use_column_width=True)
            
            with improved_col:
                st.markdown("**Improved Image**")
                st.image(preprocessing_result["processed_image"], use_column_width=True)
            
            # Display improvements applied
            st.markdown("**Improvements applied:**")
            for improvement in preprocessing_result["improvements"]:
                st.markdown(f"- {improvement}")
            
            # Add explanation about benefits
            st.markdown("### Why these improvements help facial analysis")
            st.markdown("""
            **Technical Benefits:**
            - **Balanced contrast:** Enhances the visibility of facial features while reducing shadows and highlights
            - **Optimal brightness:** Ensures facial features are clearly distinguishable without over-exposure
            - **Proper sizing:** Maintains ideal dimensions for detection algorithms to recognize facial landmarks
            
            **Impact on Emotion Detection:**
            - **More accurate emotion classification:** Cleaner input images lead to more reliable emotion detection
            - **Better feature extraction:** Facial features like eyes, mouth, and eyebrows are more clearly defined
            - **Reduced noise and artifacts:** Minimizes false detections and improves confidence scores
            """)
            
            # Add buttons to select image
            setup_image_selection_buttons(image_service, preprocessing_result)
    
    except Exception as e:
        logger.error(f"Error in expandable UI: {str(e)}")
        st.warning(f"Could not display image comparison: {str(e)}")
        
        # Fallback to old display method if needed
        try:
            display_preprocessing_comparison(preprocessing_result)
            setup_preprocessing_controls(image_service, preprocessing_result)
        except Exception as fallback_error:
            logger.error(f"Fallback display also failed: {str(fallback_error)}")

def display_preprocessing_comparison(preprocessing_result: Dict[str, Any]) -> None:
    """
    Display the comparison between original and processed images.
    
    Args:
        preprocessing_result: Dictionary containing preprocessing results
    """
    if not preprocessing_result or "improvements" not in preprocessing_result:
        return
        
    # Only show if there are improvements applied
    if preprocessing_result["improvements"]:
        st.subheader("Image Enhancement Options")
        
        # Display side-by-side comparison
        before_col, after_col = st.columns(2)
        
        with before_col:
            st.markdown("**Original Image**")
            st.image(preprocessing_result["original_image"], use_column_width=True)
        
        with after_col:
            st.markdown("**Improved Image**")
            st.image(preprocessing_result["processed_image"], use_column_width=True)
        
        # Display improvements applied
        st.markdown("**Improvements applied:**")
        for improvement in preprocessing_result["improvements"]:
            st.markdown(f"- {improvement}")
        
        # Add explanation about why these improvements are beneficial
        st.markdown("### Why these improvements help facial analysis")
        st.markdown("""
        **Technical Benefits:**
        - **Balanced contrast:** Enhances the visibility of facial features while reducing shadows and highlights
        - **Optimal brightness:** Ensures facial features are clearly distinguishable without over-exposure
        - **Proper sizing:** Maintains ideal dimensions for detection algorithms to recognize facial landmarks
        
        **Impact on Emotion Detection:**
        - **More accurate emotion classification:** Cleaner input images lead to more reliable emotion detection
        - **Better feature extraction:** Facial features like eyes, mouth, and eyebrows are more clearly defined
        - **Reduced noise and artifacts:** Minimizes false detections and improves confidence scores
        
        These improvements help our algorithms perform more consistently across different lighting conditions and image sources.
        """)

def setup_image_selection_buttons(image_service, preprocessing_result: Dict[str, Any]) -> None:
    """
    Set up buttons for selecting between original and improved images.
    
    Args:
        image_service: The image service instance
        preprocessing_result: Dictionary containing preprocessing results
    """
    # Add buttons to use original or improved image
    col1, col2 = st.columns(2)
    
    # Button for original image
    with col1:
        if st.button("Continue with Original"):
            # Set session state to use original
            st.session_state["using_preprocessed_image"] = False
            st.session_state["image_processing_status"] = "using_original"
            
            # Display confirmation message
            st.markdown("""
            <div style="background-color: #17a2b8; color: white; padding: 10px; border-radius: 5px; margin-bottom: 10px;">
                <strong>ℹ️ Using original image for analysis.</strong>
            </div>
            """, unsafe_allow_html=True)
    
    # Button for improved image
    with col2:
        if st.button("Use Improved Image"):
            try:
                # Save the processed image to a temporary file
                temp_path = image_service.save_processed_image(
                    preprocessing_result["processed_image"]
                )
                
                # Update session state
                st.session_state["preprocessed_image_path"] = temp_path
                st.session_state["using_preprocessed_image"] = True
                st.session_state["image_processing_status"] = "using_improved"
                
                # Display confirmation
                st.markdown("""
                <div style="background-color: #28a745; color: white; padding: 10px; border-radius: 5px; margin-bottom: 10px;">
                    <strong>✅ Using improved image for analysis!</strong>
                </div>
                """, unsafe_allow_html=True)
            except Exception as e:
                logger.error(f"Error saving processed image: {str(e)}")
                st.error(f"Could not save processed image: {str(e)}")

def setup_preprocessing_controls(image_service, preprocessing_result: Dict[str, Any]) -> None:
    """
    Set up the controls for selecting between original and processed images.
    
    Args:
        image_service: The image service instance
        preprocessing_result: Dictionary containing preprocessing results
    """
    if not preprocessing_result or "improvements" not in preprocessing_result:
        return
        
    # Only show if there are improvements applied
    if preprocessing_result["improvements"]:
        # Add buttons to use original or improved image
        col1, col2 = st.columns(2)
        with col1:
            use_original = st.button("Continue with Original")
        with col2:
            use_improved = st.button("Use Improved Image")
        
        # Handle the user's choice
        if use_improved:
            # Save the processed image to a temporary file
            temp_path = image_service.save_processed_image(
                preprocessing_result["processed_image"]
            )
            
            # Update session state to use the processed image
            st.session_state["preprocessed_image_path"] = temp_path
            st.session_state["using_preprocessed_image"] = True
            
            # Display a more prominent success message
            st.markdown("""
            <div style="background-color: #28a745; color: white; padding: 10px; border-radius: 5px; margin-bottom: 10px;">
                <strong>✅ Using improved image for analysis!</strong>
            </div>
            """, unsafe_allow_html=True)
            
            # Store confirmation message in session state for persistence
            st.session_state["image_processing_status"] = "using_improved"
            
            # Small delay to ensure UI updates
            time.sleep(0.5)
            
        elif use_original:
            # Set session state to use original
            st.session_state["using_preprocessed_image"] = False
            st.session_state["image_processing_status"] = "using_original"
            
            # Display a clear message
            st.markdown("""
            <div style="background-color: #17a2b8; color: white; padding: 10px; border-radius: 5px; margin-bottom: 10px;">
                <strong>ℹ️ Using original image for analysis.</strong>
            </div>
            """, unsafe_allow_html=True)

def display_processing_status() -> None:
    """
    Display the current image processing status (original or improved).
    """
    # Display persistent status indicator at the top of the interface
    if "image_processing_status" in st.session_state:
        if st.session_state["image_processing_status"] == "using_improved":
            st.markdown("""
            <div style="background-color: #28a745; color: white; padding: 5px; border-radius: 5px; margin-bottom: 10px;">
                <strong>✅ Currently using improved image for analysis</strong>
            </div>
            """, unsafe_allow_html=True)
        elif st.session_state["image_processing_status"] == "using_original":
            st.markdown("""
            <div style="background-color: #17a2b8; color: white; padding: 5px; border-radius: 5px; margin-bottom: 10px;">
                <strong>ℹ️ Using original image for analysis</strong>
            </div>
            """, unsafe_allow_html=True)

def get_processing_image(image_service, original_image: np.ndarray) -> np.ndarray:
    """
    Get the appropriate image for processing (preprocessed or original).
    
    Args:
        image_service: The image service instance
        original_image: The original image as backup
        
    Returns:
        The appropriate image to use for processing
    """
    if "using_preprocessed_image" in st.session_state and st.session_state["using_preprocessed_image"] and "preprocessed_image_path" in st.session_state:
        try:
            # Load the preprocessed image from the temporary file
            preprocessed_path = st.session_state["preprocessed_image_path"]
            img = image_service.load_image_from_path(preprocessed_path)
            
            # Check if image was loaded successfully
            if img is None or img.size == 0:
                logger.warning("Could not load preprocessed image. Using original instead.")
                return original_image
            
            return img
        except Exception as e:
            # Log error and fallback to original
            logger.error(f"Error loading preprocessed image: {e}")
            return original_image
    else:
        # Return the original image
        return original_image