File size: 5,727 Bytes
b9be998
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
"""
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"