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

# --- MediaPipe Initialization ---
mp_face_mesh = mp.solutions.face_mesh
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

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

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

# --- Analysis Functions (Keep exactly as before) ---
def extract_face_landmarks(image, face_mesh_instance):
    if image is None or face_mesh_instance is None: return None
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_rgb.flags.writeable = False
    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):
    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):
    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):
    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
    return max(0.0, min(1.0, normalized))

def estimate_head_pose(landmarks):
    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)); horizontal_tilt = max(-1.0, min(1.0, horizontal_tilt * 10))
    return vertical_tilt, horizontal_tilt

def calculate_metrics(landmarks):
    if not landmarks: return {metric: 0.5 for metric in metrics}
    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 {'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 ---
def update_metrics_visualization(metrics_values):
    if not metrics_values:
         fig, ax = plt.subplots(figsize=(10, 8)); ax.text(0.5, 0.5, "Waiting for analysis...", ha='center', va='center')
         ax.axis('off'); fig.patch.set_facecolor('#FFFFFF'); ax.set_facecolor('#FFFFFF'); return fig
    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'); axs = axs.flatten()
    colors = [(0.1, 0.1, 0.9), (0.9, 0.9, 0.1), (0.9, 0.1, 0.1)]; 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')
        r = 0.4; theta = np.linspace(np.pi, 0, 100); 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)
        value_angle = np.pi * (1 - value); 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)
        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)); lc = LineCollection(segments, cmap=cmap, norm=norm)
            lc.set_array(segment_values); lc.set_linewidth(5); ax.add_collection(lc)
        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
    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()

def process_frame(
    frame, analysis_freq, analyze_flag,
    metrics_data_state, last_analysis_time_state, latest_metrics_state, latest_landmarks_state
    ):
    # (This function remains the same as the previous working version)
    if frame is None:
        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

    if analyze_flag and face_mesh and (current_time - last_analysis_time_state >= analysis_freq):
        perform_analysis = True; last_analysis_time_state = current_time

    if perform_analysis:
        current_landmarks = extract_face_landmarks(frame, face_mesh)
        calculated_metrics = calculate_metrics(current_landmarks)
        latest_landmarks_state = current_landmarks; latest_metrics_state = calculated_metrics
        if current_landmarks:
            elapsed_time = current_time - app_start_time; new_row = {'timestamp': elapsed_time, **calculated_metrics}
            new_row_df = pd.DataFrame([new_row])
            if not isinstance(metrics_data_state, pd.DataFrame): metrics_data_state = initial_metrics_df.copy()
            metrics_data_state = pd.concat([metrics_data_state, new_row_df], ignore_index=True)

    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())

    metrics_plot = update_metrics_visualization(latest_metrics_state)
    return annotated_frame, metrics_plot, metrics_data_state, \
           metrics_data_state, last_analysis_time_state, \
           latest_metrics_state, latest_landmarks_state

# --- Function to Export DataFrame to CSV ---
def export_csv(data_df):
    """Saves the DataFrame to a CSV file and returns the file path."""
    if data_df is None or data_df.empty:
        print("No data to export.")
        # Return None or raise an error/warning for the UI? Gradio File handles None.
        return None
    # Define filename (consider making it unique if needed, e.g., with timestamp)
    csv_filename = "facial_analysis_log.csv"
    try:
        data_df.to_csv(csv_filename, index=False)
        print(f"Data exported successfully to {csv_filename}")
        return csv_filename # Return the path for the File component
    except Exception as e:
        print(f"Error exporting data to CSV: {e}")
        return None


# --- 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 ---
    metrics_data = gr.State(value=initial_metrics_df.copy())
    last_analysis_time = gr.State(value=time.time())
    latest_metrics = gr.State(value=None)
    latest_landmarks = gr.State(value=None)

    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") # This won't update live easily
        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) # Removed height

    # --- Add Export Button and File Output ---
    with gr.Row():
        with gr.Column(scale=1):
             export_button = gr.Button("Export Data to CSV")
        with gr.Column(scale=2):
             download_file_output = gr.File(label="Download CSV Log")

    # --- Define Stream Processing ---
    webcam_input.stream(
        fn=process_frame,
        inputs=[
            webcam_input, analysis_freq_slider, analyze_checkbox,
            metrics_data, last_analysis_time, latest_metrics, latest_landmarks
        ],
        outputs=[
            processed_output, metrics_plot_output, dataframe_output,
            metrics_data, last_analysis_time, latest_metrics, latest_landmarks
        ],
        # api_name="stream_frames" # Optional: Add API endpoint name
    )

    # --- Define Button Click Action ---
    export_button.click(
        fn=export_csv,
        inputs=[metrics_data], # Pass the DataFrame state to the export function
        outputs=[download_file_output], # Output the file path to the File component
        # api_name="export_data" # Optional: Add API endpoint name
    )

# --- Launch the App ---
if __name__ == "__main__":
    if face_mesh is None:
        print("WARNING: MediaPipe Face Mesh could not be initialized. Facial analysis will not work.")
    iface.launch(debug=True)