scripts/gaze_tracker.py

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import cv2
import time
import numpy as np
from scripts.inference import GazePredictor
from utils.ear_utils import BlinkDetector
from pygame import mixer

def smooth_values(history, current_value, window_size=5):
    if current_value is not None:
        history.append(current_value)
    if len(history) > window_size:
        history.pop(0)
    return np.mean(history, axis=0) if isinstance(current_value, np.ndarray) and history else current_value if current_value is not None else 0

def track_gaze(model_path):
    gaze_predictor = GazePredictor(model_path)
    blink_detector = BlinkDetector()
    cap = cv2.VideoCapture(0)
    
    if not cap.isOpened():
        print("Error: Could not open webcam.")
        return
    
    GAZE_STABILITY_THRESHOLD = 0.5
    TIME_THRESHOLD = 15
    BLINK_RATE_THRESHOLD = 1
    EYE_CLOSURE_THRESHOLD = 10
    HEAD_STABILITY_THRESHOLD = 0.05
    
    gaze_history = []
    head_history = []
    ear_history = []
    stable_gaze_time = 0
    stable_head_time = 0
    eye_closed_time = 0
    blink_count = 0
    start_time = time.time()
    is_unconscious = False
    
    # Initialize pygame mixer
    mixer.init()
    ALARM_PATH = os.path.normpath(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "sounds", "alarm.wav")))
    if not os.path.exists(ALARM_PATH):
        print(f"Warning: Alarm sound file not found at {ALARM_PATH}. No sound will play.")
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("Failed to capture frame")
            break
        
        head_pose_gaze, gaze_h, gaze_v = gaze_predictor.predict_gaze(frame)
        current_gaze = np.array([gaze_h, gaze_v])
        smoothed_gaze = smooth_values(gaze_history, current_gaze)
        
        ear, left_eye, right_eye, head_pose, left_iris, right_iris = blink_detector.detect_blinks(frame)
        if ear is None:
            cv2.putText(frame, "No face detected", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            print("No face detected")
            smoothed_head = smooth_values(head_history, None)
            smoothed_ear = smooth_values(ear_history, None)
        else:
            print(f"EAR: {ear:.2f}, Head Pose: {head_pose:.2f}, Gaze: [{smoothed_gaze[0]:.2f}, {smoothed_gaze[1]:.2f}]")
            smoothed_head = smooth_values(head_history, head_pose)
            smoothed_ear = smooth_values(ear_history, ear)
            if smoothed_ear >= blink_detector.EAR_THRESHOLD:
                cv2.drawMarker(frame, left_iris, (0, 255, 0), markerType=cv2.MARKER_CROSS, markerSize=10, thickness=2)
                cv2.drawMarker(frame, right_iris, (0, 255, 0), markerType=cv2.MARKER_CROSS, markerSize=10, thickness=2)
        
        cv2.putText(frame, f"Gaze H: {smoothed_gaze[0]:.2f}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(frame, f"Gaze V: {smoothed_gaze[1]:.2f}", (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(frame, f"Head Pose: {smoothed_head:.2f}", (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(frame, f"EAR: {smoothed_ear:.2f}", (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        if len(gaze_history) > 1:
            gaze_diff = np.sqrt(np.sum((smoothed_gaze - gaze_history[-2])**2))
            print(f"Gaze Diff: {gaze_diff:.2f}")
            if gaze_diff < GAZE_STABILITY_THRESHOLD:
                if stable_gaze_time == 0:
                    stable_gaze_time = time.time()
                elif time.time() - stable_gaze_time > TIME_THRESHOLD:
                    cv2.putText(frame, "Gaze Fixed", (10, 180), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            else:
                stable_gaze_time = 0
        
        if len(head_history) > 1 and head_pose is not None:
            head_diff = abs(smoothed_head - head_history[-2])
            print(f"Head Diff: {head_diff:.2f}")
            if head_diff < HEAD_STABILITY_THRESHOLD:
                if stable_head_time == 0:
                    stable_head_time = time.time()
            else:
                stable_head_time = 0
        
        if ear is not None and smoothed_ear < blink_detector.EAR_THRESHOLD:
            if eye_closed_time == 0:
                eye_closed_time = time.time()
            elif time.time() - eye_closed_time > EYE_CLOSURE_THRESHOLD:
                cv2.putText(frame, "Eyes Closed", (10, 210), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        else:
            if eye_closed_time > 0 and time.time() - eye_closed_time < 0.5:
                blink_count += 1
            eye_closed_time = 0
        
        elapsed_minutes = (time.time() - start_time) / 60
        blink_rate = blink_count / elapsed_minutes if elapsed_minutes > 0 else 0
        cv2.putText(frame, f"Blink Rate: {blink_rate:.1f}/min", (10, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        unconscious_conditions = [
            stable_gaze_time > 0 and time.time() - stable_gaze_time > TIME_THRESHOLD,
            blink_rate < BLINK_RATE_THRESHOLD and elapsed_minutes > 1,
            eye_closed_time > 0 and time.time() - eye_closed_time > EYE_CLOSURE_THRESHOLD,
            stable_head_time > 0 and time.time() - stable_head_time > TIME_THRESHOLD
        ]
        print(f"Conditions: {unconscious_conditions}")
        if sum(unconscious_conditions) >= 2:
            if not is_unconscious and os.path.exists(ALARM_PATH):
                print(f"Attempting to play alarm at {ALARM_PATH}")
                try:
                    mixer.music.load(ALARM_PATH)
                    mixer.music.play()
                except Exception as e:
                    print(f"Error playing alarm sound: {e}")
            print("Unconscious detected!")
            cv2.putText(frame, "Unconscious Detected", (10, 270), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            is_unconscious = True
        else:
            is_unconscious = False
        
        cv2.imshow("Gaze Tracking", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()
    mixer.quit()

if __name__ == "__main__":
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
    MODEL_PATH = os.path.join(SCRIPT_DIR, "..", "models", "gaze_estimation_model.pth")
    if not os.path.exists(MODEL_PATH):
        print(f"Error: Missing model file at {MODEL_PATH}")
        sys.exit(1)
    track_gaze(MODEL_PATH)