app.py

import gradio as gr
import cv2
import numpy as np
import pandas as pd
import time
import mediapipe as mp
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.collections import LineCollection
import os # Potentially needed if saving plots temporarily

# --- MediaPipe Initialization (Keep as is) ---
mp_face_mesh = mp.solutions.face_mesh
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

# Create Face Mesh instance globally (or manage creation/closing if resource intensive)
# Using try-except block for safer initialization if needed in complex setups
try:
    face_mesh = mp_face_mesh.FaceMesh(
        max_num_faces=1,
        refine_landmarks=True,
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5)
except Exception as e:
    print(f"Error initializing MediaPipe Face Mesh: {e}")
    face_mesh = None # Handle potential initialization errors

# --- Metrics Definition (Keep as is) ---
metrics = [
    "valence", "arousal", "dominance", "cognitive_load",
    "emotional_stability", "openness", "agreeableness",
    "neuroticism", "conscientiousness", "extraversion",
    "stress_index", "engagement_level"
]
# Initial DataFrame structure for the state
initial_metrics_df = pd.DataFrame(columns=['timestamp'] + metrics)


# --- Analysis Functions (Keep exactly as you provided) ---
# Ensure these functions handle None input for landmarks gracefully
def extract_face_landmarks(image, face_mesh_instance):
    if image is None or face_mesh_instance is None:
        return None
    # Process the image
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_rgb.flags.writeable = False # Optimize
    results = face_mesh_instance.process(image_rgb)
    image_rgb.flags.writeable = True

    if results.multi_face_landmarks:
        return results.multi_face_landmarks[0]
    return None

def calculate_ear(landmarks): # Keep as is
    if not landmarks: return 0.0
    LEFT_EYE = [33, 160, 158, 133, 153, 144]
    RIGHT_EYE = [362, 385, 387, 263, 373, 380]
    def get_landmark_coords(landmark_indices):
        return np.array([(landmarks.landmark[idx].x, landmarks.landmark[idx].y) for idx in landmark_indices])
    left_eye_points = get_landmark_coords(LEFT_EYE)
    right_eye_points = get_landmark_coords(RIGHT_EYE)
    def eye_aspect_ratio(eye_points):
        v1 = np.linalg.norm(eye_points[1] - eye_points[5])
        v2 = np.linalg.norm(eye_points[2] - eye_points[4])
        h = np.linalg.norm(eye_points[0] - eye_points[3])
        return (v1 + v2) / (2.0 * h) if h > 0 else 0.0
    left_ear = eye_aspect_ratio(left_eye_points)
    right_ear = eye_aspect_ratio(right_eye_points)
    return (left_ear + right_ear) / 2.0

def calculate_mar(landmarks): # Keep as is
    if not landmarks: return 0.0
    MOUTH_OUTLINE = [61, 291, 39, 181, 0, 17, 269, 405]
    mouth_points = np.array([(landmarks.landmark[idx].x, landmarks.landmark[idx].y) for idx in MOUTH_OUTLINE])
    height = np.mean([
        np.linalg.norm(mouth_points[1] - mouth_points[5]),
        np.linalg.norm(mouth_points[2] - mouth_points[6]),
        np.linalg.norm(mouth_points[3] - mouth_points[7])
    ])
    width = np.linalg.norm(mouth_points[0] - mouth_points[4])
    return height / width if width > 0 else 0.0

def calculate_eyebrow_position(landmarks): # Keep as is
    if not landmarks: return 0.0
    LEFT_EYEBROW = 107; RIGHT_EYEBROW = 336
    LEFT_EYE = 159; RIGHT_EYE = 386
    left_eyebrow_y = landmarks.landmark[LEFT_EYEBROW].y
    right_eyebrow_y = landmarks.landmark[RIGHT_EYEBROW].y
    left_eye_y = landmarks.landmark[LEFT_EYE].y
    right_eye_y = landmarks.landmark[RIGHT_EYE].y
    left_distance = left_eye_y - left_eyebrow_y
    right_distance = right_eye_y - right_eyebrow_y
    avg_distance = (left_distance + right_distance) / 2.0
    normalized = (avg_distance - 0.02) / 0.06 # Approximate normalization
    return max(0.0, min(1.0, normalized))

def estimate_head_pose(landmarks): # Keep as is
    if not landmarks: return 0.0, 0.0
    NOSE_TIP = 4; LEFT_EYE = 159; RIGHT_EYE = 386
    nose = np.array([landmarks.landmark[NOSE_TIP].x, landmarks.landmark[NOSE_TIP].y, landmarks.landmark[NOSE_TIP].z])
    left_eye = np.array([landmarks.landmark[LEFT_EYE].x, landmarks.landmark[LEFT_EYE].y, landmarks.landmark[LEFT_EYE].z])
    right_eye = np.array([landmarks.landmark[RIGHT_EYE].x, landmarks.landmark[RIGHT_EYE].y, landmarks.landmark[RIGHT_EYE].z])
    eye_level = (left_eye[1] + right_eye[1]) / 2.0
    vertical_tilt = nose[1] - eye_level
    horizontal_mid = (left_eye[0] + right_eye[0]) / 2.0
    horizontal_tilt = nose[0] - horizontal_mid
    vertical_tilt = max(-1.0, min(1.0, vertical_tilt * 10)) # Normalize approx
    horizontal_tilt = max(-1.0, min(1.0, horizontal_tilt * 10)) # Normalize approx
    return vertical_tilt, horizontal_tilt

def calculate_metrics(landmarks): # Keep as is
    if not landmarks:
        # Return default/neutral values when no face is detected
        return {metric: 0.5 for metric in metrics}
    # --- Calculations --- (Same as before)
    ear = calculate_ear(landmarks)
    mar = calculate_mar(landmarks)
    eyebrow_position = calculate_eyebrow_position(landmarks)
    vertical_tilt, horizontal_tilt = estimate_head_pose(landmarks)
    cognitive_load = max(0, min(1, 1.0 - ear * 2.5))
    valence = max(0, min(1, mar * 2.0 * (1.0 - eyebrow_position)))
    arousal = max(0, min(1, (mar + (1.0 - ear) + eyebrow_position) / 3.0))
    dominance = max(0, min(1, 0.5 + vertical_tilt))
    neuroticism = max(0, min(1, (cognitive_load * 0.6) + ((1.0 - valence) * 0.4)))
    emotional_stability = 1.0 - neuroticism
    extraversion = max(0, min(1, (arousal * 0.5) + (valence * 0.5)))
    openness = max(0, min(1, 0.5 + ((mar - 0.5) * 0.5)))
    agreeableness = max(0, min(1, (valence * 0.7) + ((1.0 - arousal) * 0.3)))
    conscientiousness = max(0, min(1, (1.0 - abs(arousal - 0.5)) * 0.7 + (emotional_stability * 0.3)))
    stress_index = max(0, min(1, (cognitive_load * 0.5) + (eyebrow_position * 0.3) + ((1.0 - valence) * 0.2)))
    engagement_level = max(0, min(1, (arousal * 0.7) + ((1.0 - abs(horizontal_tilt)) * 0.3)))
    # --- Return dictionary ---
    return {
        'valence': valence, 'arousal': arousal, 'dominance': dominance,
        'cognitive_load': cognitive_load, 'emotional_stability': emotional_stability,
        'openness': openness, 'agreeableness': agreeableness, 'neuroticism': neuroticism,
        'conscientiousness': conscientiousness, 'extraversion': extraversion,
        'stress_index': stress_index, 'engagement_level': engagement_level
    }


# --- Visualization Function (Keep as is, ensure it handles None input) ---
def update_metrics_visualization(metrics_values):
    # Create a blank figure if no metrics are available
    if not metrics_values:
         fig, ax = plt.subplots(figsize=(10, 8)) # Match approx size
         ax.text(0.5, 0.5, "Waiting for analysis...", ha='center', va='center')
         ax.axis('off')
         # Ensure background matches Gradio theme potentially
         fig.patch.set_facecolor('#FFFFFF') # Set background if needed
         ax.set_facecolor('#FFFFFF')
         return fig

    # Calculate grid size
    num_metrics = len([k for k in metrics_values if k != 'timestamp'])
    nrows = (num_metrics + 2) // 3
    fig, axs = plt.subplots(nrows, 3, figsize=(10, nrows * 2.5), facecolor='#FFFFFF') # Match background
    axs = axs.flatten()

    # Colormap and normalization
    colors = [(0.1, 0.1, 0.9), (0.9, 0.9, 0.1), (0.9, 0.1, 0.1)] # Blue to Yellow to Red
    cmap = LinearSegmentedColormap.from_list("custom_cmap", colors, N=100)
    norm = plt.Normalize(0, 1)

    metric_idx = 0
    for key, value in metrics_values.items():
        if key == 'timestamp': continue

        ax = axs[metric_idx]
        ax.set_title(key.replace('_', ' ').title(), fontsize=10)
        ax.set_xlim(0, 1); ax.set_ylim(0, 0.5); ax.set_aspect('equal'); ax.axis('off')
        ax.set_facecolor('#FFFFFF') # Match background

        r = 0.4 # radius
        theta = np.linspace(np.pi, 0, 100) # Flipped for gauge direction
        x_bg = 0.5 + r * np.cos(theta); y_bg = 0.1 + r * np.sin(theta)
        ax.plot(x_bg, y_bg, 'k-', linewidth=3, alpha=0.2) # Background arc

        # Value arc calculation
        value_angle = np.pi * (1 - value) # Map value [0,1] to angle [pi, 0]
        # Ensure there are at least 2 points for the line segment, even for value=0
        num_points = max(2, int(100 * value))
        value_theta = np.linspace(np.pi, value_angle, num_points)
        x_val = 0.5 + r * np.cos(value_theta); y_val = 0.1 + r * np.sin(value_theta)

        # Create line segments for coloring if there are points to draw
        if len(x_val) > 1:
            points = np.array([x_val, y_val]).T.reshape(-1, 1, 2)
            segments = np.concatenate([points[:-1], points[1:]], axis=1)
            segment_values = np.linspace(0, value, len(segments)) # Color based on value
            lc = LineCollection(segments, cmap=cmap, norm=norm)
            lc.set_array(segment_values); lc.set_linewidth(5)
            ax.add_collection(lc)

        # Add value text
        ax.text(0.5, 0.15, f"{value:.2f}", ha='center', va='center', fontsize=11,
                fontweight='bold', bbox=dict(facecolor='white', alpha=0.7, boxstyle='round,pad=0.2'))
        metric_idx += 1

    # Hide unused subplots
    for i in range(metric_idx, len(axs)):
        axs[i].axis('off')

    plt.tight_layout(pad=0.5)
    return fig


# --- Gradio Processing Function ---
app_start_time = time.time() # Use a fixed start time for the app session

def process_frame(
    frame,
    analysis_freq,
    analyze_flag,
    # --- State variables ---
    metrics_data_state,
    last_analysis_time_state,
    latest_metrics_state,
    latest_landmarks_state
    ):

    if frame is None:
        # Return default/empty outputs if no frame
        default_plot = update_metrics_visualization(latest_metrics_state)
        return frame, default_plot, metrics_data_state, \
               metrics_data_state, last_analysis_time_state, \
               latest_metrics_state, latest_landmarks_state

    annotated_frame = frame.copy()
    current_time = time.time()
    perform_analysis = False
    current_landmarks = None # Landmarks detected in *this* frame run

    # --- Decide whether to perform analysis ---
    if analyze_flag and face_mesh and (current_time - last_analysis_time_state >= analysis_freq):
        perform_analysis = True
        last_analysis_time_state = current_time # Update time immediately

    # --- Perform Analysis (if flag is set and frequency met) ---
    if perform_analysis:
        current_landmarks = extract_face_landmarks(frame, face_mesh)
        calculated_metrics = calculate_metrics(current_landmarks)

        # Update state variables
        latest_landmarks_state = current_landmarks # Store landmarks from this run
        latest_metrics_state = calculated_metrics

        # Log data only if a face was detected
        if current_landmarks:
            elapsed_time = current_time - app_start_time
            new_row = {'timestamp': elapsed_time, **calculated_metrics}
            new_row_df = pd.DataFrame([new_row])
            metrics_data_state = pd.concat([metrics_data_state, new_row_df], ignore_index=True)

    # --- Drawing ---
    # Always try to draw the latest known landmarks stored in state
    landmarks_to_draw = latest_landmarks_state
    if landmarks_to_draw:
        mp_drawing.draw_landmarks(
            image=annotated_frame,
            landmark_list=landmarks_to_draw,
            connections=mp_face_mesh.FACEMESH_TESSELATION,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_tesselation_style())
        mp_drawing.draw_landmarks(
            image=annotated_frame,
            landmark_list=landmarks_to_draw,
            connections=mp_face_mesh.FACEMESH_CONTOURS,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_contours_style())

    # --- Generate Metrics Plot ---
    metrics_plot = update_metrics_visualization(latest_metrics_state)

    # --- Return updated values for outputs AND state ---
    return annotated_frame, metrics_plot, metrics_data_state, \
           metrics_data_state, last_analysis_time_state, \
           latest_metrics_state, latest_landmarks_state


# --- Create Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(), title="Gradio Facial Analysis") as iface:
    gr.Markdown("# Basic Facial Analysis (Gradio Version)")
    gr.Markdown("Analyzes webcam feed for facial landmarks and estimates metrics. *Estimations are for demonstration only.*")

    # Define State Variables
    # Need to initialize them properly
    metrics_data = gr.State(value=initial_metrics_df.copy())
    last_analysis_time = gr.State(value=time.time())
    latest_metrics = gr.State(value=None) # Initially no metrics
    latest_landmarks = gr.State(value=None) # Initially no landmarks

    with gr.Row():
        with gr.Column(scale=1):
            webcam_input = gr.Image(sources="webcam", streaming=True, label="Webcam Input", type="numpy")
            analysis_freq_slider = gr.Slider(minimum=0.5, maximum=5.0, step=0.5, value=1.0, label="Analysis Frequency (s)")
            analyze_checkbox = gr.Checkbox(value=True, label="Enable Analysis Calculation")
            status_text = gr.Markdown("Status: Analysis Enabled" if analyze_checkbox.value else "Status: Analysis Paused") # Initial status text

            # Update status text dynamically (though Gradio handles this implicitly via reruns)
            # Might need a more complex setup with event listeners if precise text update is needed without full rerun
        with gr.Column(scale=1):
            processed_output = gr.Image(label="Processed Feed", type="numpy")
            metrics_plot_output = gr.Plot(label="Estimated Metrics")
            dataframe_output = gr.Dataframe(label="Data Log", headers=['timestamp'] + metrics, wrap=True, height=300)


    # Define the connections for the live interface
    webcam_input.stream(
        fn=process_frame,
        inputs=[
            webcam_input,
            analysis_freq_slider,
            analyze_checkbox,
            # Pass state variables as inputs
            metrics_data,
            last_analysis_time,
            latest_metrics,
            latest_landmarks
        ],
        outputs=[
            processed_output,
            metrics_plot_output,
            dataframe_output,
            # Return updated state variables
            metrics_data,
            last_analysis_time,
            latest_metrics,
            latest_landmarks
        ]
    )

# --- Launch the App ---
if __name__ == "__main__":
    if face_mesh is None:
        print("Face Mesh could not be initialized. Gradio app might not function correctly.")
    iface.launch(debug=True) # Enable debug for more detailed errors if needed

    # Optional: Add cleanup logic if needed, although launching blocks execution
    # try:
    #     iface.launch()
    # finally:
    #     if face_mesh:
    #         face_mesh.close() # Close mediapipe resources if app is stopped
    #         print("MediaPipe FaceMesh closed.")