utils/simple_face_labeling.py

"""
Simplified face labeling module.

Provides a streamlined UI for face detection and labeling with minimal complexity.
"""
import streamlit as st
import numpy as np
import cv2
import logging
import time
from typing import List, Dict, Tuple, Any, Set

# Configurar logging
logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)

def draw_numbered_faces(image: np.ndarray, faces: List[Tuple[int, int, int, int]], 
                      max_faces: int = 5) -> np.ndarray:
    """
    Draw numbered rectangles on detected faces.
    
    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 labeled 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
    
    # Get removed faces set (if exists)
    removed_faces = st.session_state.get("removed_faces", set())
    
    # Draw each face
    for i, (x, y, w, h) in enumerate(faces_to_draw):
        face_key = f"face_{i}"
        
        # Skip if this face was marked as removed
        if face_key in removed_faces:
            continue
            
        # Draw green rectangle
        cv2.rectangle(labeled_img, (x, y), (x+w, y+h), (0, 255, 0), 2)
        
        # Add numbered label
        label = f"Face {i+1}"
        cv2.putText(labeled_img, label, (x, y-10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
    return labeled_img

def extract_face_thumbnails(image: np.ndarray, faces: List[Tuple[int, int, int, int]],
                           max_faces: int = 5) -> Dict[int, np.ndarray]:
    """
    Extracts thumbnails of detected faces.
    
    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 process
        
    Returns:
        Dictionary with face index and its cropped image
    """
    thumbnails = {}
    
    # Limit to max_faces faces
    faces_to_extract = faces[:max_faces] if len(faces) > max_faces else faces
    
    # Extract each thumbnail
    for i, (x, y, w, h) in enumerate(faces_to_extract):
        # Apply a small margin around the face if possible
        margin = int(min(w, h) * 0.1)  # 10% margin
        
        # Ensure we don't go out of the image bounds
        img_h, img_w = image.shape[:2]
        x_start = max(0, x - margin)
        y_start = max(0, y - margin)
        x_end = min(img_w, x + w + margin)
        y_end = min(img_h, y + h + margin)
        
        # Extract the thumbnail with margin
        face_thumbnail = image[y_start:y_end, x_start:x_end]
        thumbnails[i] = face_thumbnail
    
    return thumbnails

def simple_face_labeling_ui(image: np.ndarray, faces: List[Tuple[int, int, int, int]], 
                           max_faces: int = 5) -> Dict[str, Any]:
    """
    Displays a simplified interface for labeling faces.
    
    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 process
        
    Returns:
        Dictionary with information about labeled faces
    """
    # Iniciar timestamp de sesión si no existe (para crear claves únicas)
    if "session_timestamp" not in st.session_state:
        st.session_state.session_timestamp = int(time.time())
    
    # Initialize session state for face labels and removed faces
    if "face_labels" not in st.session_state:
        st.session_state.face_labels = {}
    
    if "removed_faces" not in st.session_state:
        st.session_state.removed_faces = set()
    
    # Timestamp para generar IDs únicos para los botones en esta sesión
    timestamp = st.session_state.session_timestamp
    
    # Limit to max_faces faces
    faces_to_show = faces[:max_faces] if len(faces) > max_faces else faces
    num_faces = len(faces_to_show)
    
    # Display labeled image
    labeled_image = draw_numbered_faces(image, faces_to_show)
    st.image(labeled_image, caption="Detected Faces", use_column_width=True)
    
    # Only proceed if faces were detected
    if num_faces > 0:
        st.success(f"{num_faces} face(s) detected in the image")
        
        # Extract thumbnails
        thumbnails = extract_face_thumbnails(image, faces_to_show)
        
        # Create a form for labeling
        st.subheader("Enter names for detected faces")
        
        # Lista para seguir qué caras se muestran (para preparar el resultado)
        displayed_faces = []
        
        # Create a simple list of faces with names and remove buttons
        for i, (x, y, w, h) in enumerate(faces_to_show):
            face_key = f"face_{i}"
            
            # Skip if this face was marked as removed
            if face_key in st.session_state.removed_faces:
                continue
            
            displayed_faces.append((i, face_key, (x, y, w, h)))
                
            # Create a row with thumbnail, name field and remove button
            cols = st.columns([1, 3, 1])
            
            with cols[0]:
                # Display thumbnail
                if i in thumbnails:
                    st.image(thumbnails[i], caption=f"Face {i+1}", width=80)
            
            with cols[1]:
                # Input field for name
                label = st.text_input(
                    f"Name for Face {i+1}:", 
                    key=f"label_{face_key}_{timestamp}",
                    value=st.session_state.face_labels.get(face_key, "")
                )
                # Save to session state
                st.session_state.face_labels[face_key] = label
            
            with cols[2]:
                # Cada botón tiene una clave única para esta sesión
                remove_button_key = f"btn_remove_{face_key}_{timestamp}"
                if st.button("Remove", key=remove_button_key):
                    # Marcar la cara como eliminada
                    st.session_state.removed_faces.add(face_key)
                    # Eliminar la etiqueta si existe
                    if face_key in st.session_state.face_labels:
                        del st.session_state.face_labels[face_key]
                    # Forzar rerun para actualizar la interfaz
                    st.experimental_rerun()
        
        # Prepare result data
        result = {
            "success": True,
            "num_faces": num_faces,
            "labeled_image": labeled_image
        }
        
        # Add face data
        labeled_faces = {}
        
        # Usar solo las caras que se mostraron
        for i, face_key, coords in displayed_faces:
            # Check if this face has a label
            label = st.session_state.face_labels.get(face_key, "")
            if label:
                labeled_faces[face_key] = {
                    "index": i,
                    "label": label,
                    "coordinates": coords
                }
        
        # Add to result
        result["labeled_faces"] = labeled_faces
        result["can_proceed"] = len(labeled_faces) > 0
        
        # Show proceed button if at least one face is labeled
        if result["can_proceed"]:
            if st.button("Continue to Analysis", key=f"continue_to_analysis_{timestamp}"):
                result["proceed_to_analysis"] = True
            else:
                result["proceed_to_analysis"] = False
        else:
            st.warning("Please provide at least one name to continue to analysis.")
            result["proceed_to_analysis"] = False
        
        return result
    else:
        st.warning("No faces detected in the image.")
        return {
            "success": False,
            "num_faces": 0,
            "message": "No faces detected in the image."
        }

def simple_face_detection_and_labeling_ui(image: np.ndarray, face_service: Any) -> Dict[str, Any]:
    """
    Main function for simplified face detection and labeling.
    
    Args:
        image: Image in numpy array format (RGB)
        face_service: Face detection service
        
    Returns:
        Dictionary with processed results
    """
    # Ensure we have an image
    if image is None:
        st.warning("No image available for processing.")
        return {
            "success": False,
            "message": "No image available for processing."
        }
    
    # Set maximum faces
    max_faces = 5
    
    # Convert to BGR for detection if needed
    img_bgr = None
    if len(image.shape) == 3 and image.shape[2] == 3:
        img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    else:
        img_bgr = image.copy()
    
    # Perform face detection
    with st.spinner("Detecting faces..."):
        faces = face_service.detect_faces(img_bgr)
    
    # Check if any faces were detected
    if faces is None or len(faces) == 0:
        st.warning("No faces detected in the image.")
        st.image(image, caption="Uploaded image (no faces detected)", use_column_width=True)
        return {
            "success": False,
            "message": "No faces detected in the image."
        }
    
    # Save detected faces in session state
    st.session_state["detected_faces"] = faces
    
    # Show the simple labeling UI
    labeling_result = simple_face_labeling_ui(image, faces, max_faces)
    
    # Handle result
    if labeling_result.get("proceed_to_analysis", False):
        # Prepare data for analysis
        faces_to_analyze = []
        labeled_faces = labeling_result.get("labeled_faces", {})
        
        # Process each labeled face
        for face_key, face_info in labeled_faces.items():
            index = face_info.get("index", 0)
            label = face_info.get("label", "")
            coords = face_info.get("coordinates", (0, 0, 0, 0))
            
            # Extract thumbnail
            x, y, w, h = coords
            margin = int(min(w, h) * 0.1)
            img_h, img_w = image.shape[:2]
            x_start = max(0, x - margin)
            y_start = max(0, y - margin)
            x_end = min(img_w, x + w + margin)
            y_end = min(img_h, y + h + margin)
            thumbnail = image[y_start:y_end, x_start:x_end]
            
            # Add to faces to analyze
            faces_to_analyze.append({
                "key": face_key,
                "label": label,
                "coordinates": coords,
                "thumbnail": thumbnail
            })
        
        # Return analysis data
        return {
            "success": True,
            "proceed_to_analysis": True,
            "faces_to_analyze": faces_to_analyze
        }
    
    # Return result without proceeding
    return {
        "success": labeling_result.get("success", False),
        "proceed_to_analysis": False,
        "message": labeling_result.get("message", "")
    }