tabs/heart_rate_variability.py

import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import butter, filtfilt, find_peaks
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_stress_level(rmssd, hr_mean, hr_std):
    """
    Calculate stress level based on HRV parameters.
    Returns both numerical value (0-100) and category.
    """
    # RMSSD factor (lower RMSSD = higher stress)
    rmssd_normalized = max(0, min(100, (150 - rmssd) / 1.5))
    
    # Heart rate factor (higher HR = higher stress)
    hr_factor = max(0, min(100, (hr_mean - 60) * 2))
    
    # Heart rate variability factor (lower variability = higher stress)
    hr_variability_factor = max(0, min(100, hr_std * 5))
    
    # Combine factors with weights
    stress_value = (0.4 * rmssd_normalized + 
                   0.4 * hr_factor + 
                   0.2 * hr_variability_factor)
    
    # Determine category
    if stress_value < 30:
        category = "Low"
    elif stress_value < 60:
        category = "Moderate"
    else:
        category = "High"
        
    return stress_value, category

def get_anxiety_level(value):
    """Get anxiety level category based on value."""
    if value < 30:
        return "Low"
    elif value < 70:
        return "Moderate"
    else:
        return "High"

def calculate_anxiety_index(heart_rate, hrv):
    """Calculate anxiety index based on heart rate and HRV."""
    if len(heart_rate) < 2:
        return 0
    
    hr_mean = np.mean(heart_rate)
    hr_std = np.std(heart_rate)
    
    # Combine factors indicating anxiety
    hr_factor = min(100, max(0, (hr_mean - 60) / 0.4))
    variability_factor = min(100, (hr_std / 20) * 100)
    hrv_factor = min(100, max(0, (100 - hrv) / 1))
    
    anxiety_index = (hr_factor + variability_factor + hrv_factor) / 3
    return anxiety_index

def process_video_for_hrv(video_path):
    """Process video and extract HRV metrics focusing on stress and anxiety."""
    if not video_path:
        return None, None
    
    try:
        cap = cv2.VideoCapture(video_path)
        ppg_signal = []
        fps = cap.get(cv2.CAP_PROP_FPS)
        last_frame = None
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Extract green channel for PPG
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            green_channel = frame_rgb[:, :, 1]
            ppg_value = np.mean(green_channel)
            ppg_signal.append(ppg_value)
            
            # Store last frame for display
            last_frame = cv2.resize(frame_rgb, (320, 240))
        
        cap.release()
        
        if not ppg_signal or last_frame is None:
            return None, None
            
        # Process PPG signal
        ppg_signal = np.array(ppg_signal)
        filtered_signal = filtfilt(*butter(2, [0.5, 5], fs=fps, btype='band'), ppg_signal)
        
        # Find peaks for heart rate calculation
        peaks, _ = find_peaks(filtered_signal, distance=int(0.5 * fps))
        if len(peaks) < 2:
            return None, None
            
        # Calculate basic metrics
        rr_intervals = np.diff(peaks) / fps * 1000
        heart_rate = 60 * fps / np.diff(peaks)
        hrv_rmssd = np.sqrt(np.mean(np.diff(rr_intervals) ** 2))
        
        # Calculate stress and anxiety indices
        hr_mean = np.mean(heart_rate)
        hr_std = np.std(heart_rate)
        stress_value, stress_category = get_stress_level(hrv_rmssd, hr_mean, hr_std)
        anxiety_idx = calculate_anxiety_index(heart_rate, hrv_rmssd)
        
        # Create visualization
        fig = plt.figure(figsize=(12, 10))
        
        # Plot 1: Stress and Anxiety Levels (top)
        ax1 = plt.subplot(211)
        metrics = ['Stress Level', 'Anxiety Level']
        values = [stress_value, anxiety_idx]
        colors = ['#FF6B6B', '#4D96FF']  # Warm red for stress, cool blue for anxiety
        
        bars = ax1.bar(metrics, values, color=colors)
        ax1.set_ylim(0, 100)
        ax1.set_title('Stress and Anxiety Analysis', pad=20)
        ax1.set_ylabel('Level (%)')
        
        # Add value labels and status
        for bar, val, metric in zip(bars, values, metrics):
            height = val
            status = stress_category if metric == 'Stress Level' else get_anxiety_level(val)
            ax1.text(bar.get_x() + bar.get_width()/2., height + 1,
                    f'{val:.1f}%\n{status}',
                    ha='center', va='bottom')
        
        # Plot 2: Heart Rate and HRV Trends (bottom)
        ax2 = plt.subplot(212)
        time = np.linspace(0, len(heart_rate), len(heart_rate))
        ax2.plot(time, heart_rate, color='#2ECC71', label='Heart Rate', linewidth=2)
        ax2.set_title('Heart Rate Variation')
        ax2.set_xlabel('Beat Number')
        ax2.set_ylabel('Heart Rate (BPM)')
        ax2.grid(True, alpha=0.3)
        
        # Add metrics information with color-coded status
        def get_status_color(category):
            return {
                'Low': '#2ECC71',      # Green
                'Moderate': '#F1C40F',  # Yellow
                'High': '#E74C3C'       # Red
            }.get(category, 'black')
        
        info_text = (
            f'HRV (RMSSD): {hrv_rmssd:.1f} ms\n'
            f'Average HR: {hr_mean:.1f} BPM\n'
            f'Recording: {len(ppg_signal)/fps:.1f} s\n\n'
            f'Stress Status: {stress_category}\n'
            f'Anxiety Status: {get_anxiety_level(anxiety_idx)}'
        )
        
        # Add metrics box with gradient background
        bbox_props = dict(
            boxstyle='round,pad=0.5',
            facecolor='white',
            alpha=0.8,
            edgecolor='gray'
        )
        
        ax2.text(0.02, 0.98, info_text,
                transform=ax2.transAxes,
                verticalalignment='top',
                bbox=bbox_props,
                fontsize=10)
        
        plt.tight_layout()
        
        return last_frame, fig
        
    except Exception as e:
        logger.error(f"Error processing video: {str(e)}")
        return None, None

def create_heart_rate_variability_tab():
    with gr.Row():
        with gr.Column(scale=1):
            input_video = gr.Video()
            gr.Markdown("""
            ### Stress and Anxiety Analysis
            
            **Measurements:**
            - Stress Level (0-100%)
            - Anxiety Level (0-100%)
            - Heart Rate Variability (HRV)
            
            **Status Levels:**
            🟢 Low: Normal state
            🟡 Moderate: Elevated levels
            🔴 High: Significant elevation
            
            **For best results:**
            1. Ensure good lighting
            2. Minimize movement
            3. Face the camera directly
            """)
            gr.Examples(["./assets/videos/fitness.mp4"], inputs=[input_video])
            
        with gr.Column(scale=2):
            output_frame = gr.Image(label="Face Detection", height=240)
            hrv_plot = gr.Plot(label="Stress and Anxiety Analysis")
    
    # Automatically trigger analysis on video upload
    input_video.change(
        fn=process_video_for_hrv,
        inputs=[input_video],
        outputs=[output_frame, hrv_plot]
    )

    return input_video, output_frame, hrv_plot