utils/face_visualization.py

"""
Face visualization module for dynamic UI controls and visual feedback.

This module provides functionality for rendering face detection results
with various visual states (selected, unselected, removed) without requiring
page reloads.
"""
import streamlit as st
import cv2
import numpy as np
from typing import List, Dict, Tuple, Any, Optional, Set

# Color constants for face boxes
COLOR_SELECTED = (0, 255, 0)  # Green
COLOR_UNSELECTED = (100, 100, 100)  # Gray
COLOR_HIGHLIGHT = (255, 165, 0)  # Orange - for highlighting changes

def initialize_face_visualization_state():
    """
    Initialize session state variables used for face visualization.
    Should be called at the beginning of the app.
    """
    if "face_display_mode" not in st.session_state:
        st.session_state.face_display_mode = "show_all"  # Options: show_all, hide_unselected, hide_removed
    
    if "removed_faces" not in st.session_state:
        st.session_state.removed_faces = set()
    
    if "selected_faces" not in st.session_state:
        st.session_state.selected_faces = {}
    
    if "face_labels" not in st.session_state:
        st.session_state.face_labels = {}
    
    if "recent_action" not in st.session_state:
        st.session_state.recent_action = None  # For tracking UI changes

def draw_faces_with_state(
    image: np.ndarray, 
    faces: List[Tuple[int, int, int, int]], 
    max_faces: int = 5
) -> np.ndarray:
    """
    Draw faces on image with different visual states based on session state.
    
    Args:
        image: Image in numpy array format (RGB)
        faces: List of tuples (x, y, w, h) with face coordinates
        max_faces: Maximum number of faces to display
        
    Returns:
        Image with visualized faces according to current display mode
    """
    # Get current display mode
    display_mode = st.session_state.get("face_display_mode", "show_all")
    removed_faces = st.session_state.get("removed_faces", set())
    selected_faces = st.session_state.get("selected_faces", {})
    
    # Work with a copy to avoid modifying the original
    labeled_img = image.copy()
    
    # Limit to max_faces faces
    faces_to_draw = faces[:max_faces] if len(faces) > max_faces else faces
    
    # Draw each face based on its state
    for i, (x, y, w, h) in enumerate(faces_to_draw):
        face_key = f"face_{i}"
        
        # Skip faces marked as false positives if not showing all
        if face_key in removed_faces and display_mode != "show_all":
            continue
            
        # Determine if the face is selected
        is_selected = selected_faces.get(face_key, True)
        
        # Skip unselected faces if in hide_unselected mode
        if not is_selected and display_mode == "hide_unselected":
            continue
        
        # Determine color based on state
        if face_key in removed_faces:
            # Skip drawing removed faces entirely
            continue
        elif is_selected:
            color = COLOR_SELECTED
        else:
            color = COLOR_UNSELECTED
            
        # Draw rectangle with appropriate color
        cv2.rectangle(labeled_img, (x, y), (x+w, y+h), color, 2)
        
        # Add numbered label with same color
        label = f"Face {i+1}"
        cv2.putText(labeled_img, label, (x, y-10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
    
    return labeled_img

def face_control_panel():
    """
    Renders a control panel for face visualization options.
    Returns the current display mode.
    """
    st.markdown("### Display Options")
    
    # Create a row of buttons for display modes
    col1, col2, col3 = st.columns(3)
    
    with col1:
        if st.button("Show All Faces", key="btn_show_all"):
            st.session_state.face_display_mode = "show_all"
    
    with col2:
        if st.button("Hide Unselected", key="btn_hide_unselected"):
            st.session_state.face_display_mode = "hide_unselected"
    
    with col3:
        if st.button("Reset Removed Faces", key="btn_reset_removed"):
            # Clear any faces that were marked as removed
            st.session_state.removed_faces = set()
    
    # Display current statistics
    num_detected = len(st.session_state.get("detected_faces", []))
    num_selected = sum(1 for v in st.session_state.get("selected_faces", {}).values() if v)
    num_removed = len(st.session_state.get("removed_faces", set()))
    
    st.markdown(f"**Status:** {num_detected} detected, {num_selected} selected, {num_removed} removed")
    
    return st.session_state.face_display_mode

def on_face_select(face_idx: int, value: bool):
    """
    Callback when a face is selected/unselected.
    Updates session state without page reload.
    
    Args:
        face_idx: Index of the face
        value: New selection state (True/False)
    """
    face_key = f"face_{face_idx}"
    st.session_state.selected_faces[face_key] = value
    
    # Update recent action for UI feedback
    st.session_state.recent_action = "select" if value else "unselect"

def on_face_remove(face_idx: int):
    """
    Callback when a face is marked as a false positive.
    Updates session state without page reload.
    
    Args:
        face_idx: Index of the face to remove
    """
    face_key = f"face_{face_idx}"
    
    # Mark as not selected
    st.session_state.selected_faces[face_key] = False
    
    # Remove any existing label
    if face_key in st.session_state.face_labels:
        del st.session_state.face_labels[face_key]
    
    # Add to set of removed faces
    st.session_state.removed_faces.add(face_key)
    
    # Update recent action for UI feedback
    st.session_state.recent_action = "remove"