utils/image_visualization.py

"""
Image Visualization Module for EmotionMirror application.

This module contains functions for visualizing images with interactive controls
such as zoom, pan, and reset functionality.
"""
import streamlit as st
import numpy as np
from PIL import Image
import io
import logging
import base64
from typing import Dict, Any, Tuple, Optional

logger = logging.getLogger(__name__)

def get_image_download_link(img, filename="emotion_mirror_image.png", text="Download Image"):
    """
    Generate a download link for an image.
    
    Args:
        img: PIL Image or numpy array
        filename: Name of the file to download
        text: Text to display for the download link
        
    Returns:
        HTML string with download link
    """
    try:
        # Convert numpy array to PIL Image if necessary
        if isinstance(img, np.ndarray):
            img_pil = Image.fromarray(img)
        else:
            img_pil = img
            
        # Create a byte buffer
        buffered = io.BytesIO()
        img_pil.save(buffered, format="PNG")
        
        # Encode the bytes to base64
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        # Create download link HTML
        href = f'<a href="data:file/png;base64,{img_str}" download="{filename}">{text}</a>'
        return href
    except Exception as e:
        logger.error(f"Error creating image download link: {e}")
        return None

def display_image_with_controls(image, 
                              title: str = "Image Viewer", 
                              allow_zoom: bool = True,
                              allow_download: bool = True) -> Dict[str, Any]:
    """
    Display an image with zoom, pan and reset controls.
    
    Args:
        image: PIL Image or numpy array to display
        title: Title to display above the image
        allow_zoom: Whether to show zoom controls
        allow_download: Whether to show download option
        
    Returns:
        Dict containing the status and any relevant info
    """
    try:
        # Ensure image is not None
        if image is None:
            return {"success": False, "message": "No image provided"}
            
        # Convert numpy array to PIL Image if necessary
        if isinstance(image, np.ndarray):
            pil_image = Image.fromarray(image)
        else:
            pil_image = image
        
        # Store original image dimensions
        original_width, original_height = pil_image.size
        
        # Create container for image display
        st.subheader(title)
        
        # Image controls in columns for better layout
        col1, col2, col3 = st.columns([1, 2, 1])
        
        # Zoom functionality
        zoom_factor = 1.0
        if allow_zoom:
            with col1:
                zoom_factor = st.slider(
                    "Zoom", 
                    min_value=0.5, 
                    max_value=2.0, 
                    value=1.0, 
                    step=0.1,
                    help="Adjust the zoom level of the image"
                )
        
        # Reset button in the third column
        with col3:
            reset_pressed = st.button("Reset View", help="Reset zoom and view settings")
            if reset_pressed:
                zoom_factor = 1.0
                st.session_state.zoom_factor = 1.0  # Store in session state
                
        # Apply zoom if needed
        if zoom_factor != 1.0:
            # Calculate new dimensions
            new_width = int(original_width * zoom_factor)
            new_height = int(original_height * zoom_factor)
            
            # Resize the image
            resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
        else:
            resized_image = pil_image
        
        # Display the image
        st.image(resized_image, use_column_width=True)
        
        # Add download option
        if allow_download:
            download_link = get_image_download_link(pil_image)
            if download_link:
                st.markdown(download_link, unsafe_allow_html=True)
                
        return {
            "success": True, 
            "zoom_factor": zoom_factor,
            "image_dimensions": (original_width, original_height),
            "displayed_dimensions": resized_image.size
        }
        
    except Exception as e:
        logger.error(f"Error displaying image with controls: {e}")
        return {"success": False, "message": str(e)}

def handle_image_viewer(image_data, 
                      title: str = "Image Viewer", 
                      description: str = None,
                      allow_zoom: bool = True, 
                      allow_reset: bool = True,
                      allow_download: bool = True) -> Dict[str, Any]:
    """
    Complete image viewer with all controls and options.
    
    Args:
        image_data: PIL Image or numpy array
        title: Title for the image viewer
        description: Optional description text
        allow_zoom: Whether to include zoom controls
        allow_reset: Whether to include reset button
        allow_download: Whether to include download option
        
    Returns:
        Dict containing status and control information
    """
    st.subheader(title)
    
    if description:
        st.markdown(description)
    
    # Create container for the viewer
    viewer_container = st.container()
    
    with viewer_container:
        # Create a two-column layout for controls and image
        control_col, image_col = st.columns([1, 3])
        
        with control_col:
            st.markdown("### Controls")
            
            # Zoom controls
            zoom_factor = 1.0
            if allow_zoom:
                zoom_factor = st.slider(
                    "Zoom", 
                    min_value=0.5, 
                    max_value=3.0, 
                    value=1.0, 
                    step=0.1,
                    help="Adjust the zoom level"
                )
            
            # Reset button
            if allow_reset:
                if st.button("Reset View", key="reset_view"):
                    zoom_factor = 1.0
                    # Reset any other relevant state
            
            # Add some spacing
            st.markdown("<br>", unsafe_allow_html=True)
            
            # Download option
            if allow_download and image_data is not None:
                download_link = get_image_download_link(image_data)
                if download_link:
                    st.markdown("### Download")
                    st.markdown(download_link, unsafe_allow_html=True)
        
        with image_col:
            if image_data is not None:
                # Process the image based on controls
                try:
                    # Convert to PIL if numpy array
                    if isinstance(image_data, np.ndarray):
                        pil_image = Image.fromarray(image_data)
                    else:
                        pil_image = image_data
                    
                    # Apply zoom if needed
                    if zoom_factor != 1.0:
                        original_width, original_height = pil_image.size
                        new_width = int(original_width * zoom_factor)
                        new_height = int(original_height * zoom_factor)
                        displayed_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
                    else:
                        displayed_image = pil_image
                        
                    # Display the image
                    st.image(displayed_image, use_column_width=True)
                    
                    return {
                        "success": True,
                        "zoom_factor": zoom_factor,
                        "image_displayed": True
                    }
                    
                except Exception as e:
                    st.error(f"Error displaying image: {str(e)}")
                    logger.error(f"Error in image viewer: {e}")
                    return {"success": False, "message": str(e)}
            else:
                st.info("No image to display. Please upload an image first.")
                return {"success": False, "message": "No image data provided"}

def create_image_tabs(original_image, processed_image=None):
    """
    Create tabs to display original and processed images side by side.
    
    Args:
        original_image: The original image (PIL or numpy array)
        processed_image: The processed image (PIL or numpy array), optional
        
    Returns:
        Dict with status information
    """
    # Create tabs
    if processed_image is not None:
        tab1, tab2 = st.tabs(["Original Image", "Processed Image"])
        
        with tab1:
            result1 = display_image_with_controls(
                original_image, 
                title="Original Image", 
                allow_zoom=True,
                allow_download=True
            )
            
        with tab2:
            result2 = display_image_with_controls(
                processed_image, 
                title="Processed Image", 
                allow_zoom=True,
                allow_download=True
            )
            
        return {
            "success": result1.get("success", False) and result2.get("success", False),
            "tabs_created": True
        }
    else:
        # Only original image available
        result = display_image_with_controls(
            original_image, 
            title="Uploaded Image", 
            allow_zoom=True,
            allow_download=True
        )
        
        return {
            "success": result.get("success", False),
            "tabs_created": False
        }