Update drowsiness_detection.py

drowsiness_detection.py

@@ -7,18 +7,28 @@ import numpy as np
 import cv2 as cv
 import imutils
 import dlib
-import pygame
 import argparse
 import os
 
-# --- MODELS AND CONSTANTS ---
-# Use absolute paths relative to this script file for robustness
+# --- FIXED: Models and Constants with better error handling ---
 script_dir = os.path.dirname(os.path.abspath(__file__))
 haar_cascade_face_detector = os.path.join(script_dir, "haarcascade_frontalface_default.xml")
 dlib_facial_landmark_predictor = os.path.join(script_dir, "shape_predictor_68_face_landmarks.dat")
 
-face_detector = cv.CascadeClassifier(haar_cascade_face_detector)
-landmark_predictor = dlib.shape_predictor(dlib_facial_landmark_predictor)
+# Check if required files exist
+if not os.path.exists(haar_cascade_face_detector):
+    print(f"Warning: Face detector file not found at {haar_cascade_face_detector}")
+    # Try to use OpenCV's built-in cascade
+    face_detector = cv.CascadeClassifier(cv.data.haarcascades + 'haarcascade_frontalface_default.xml')
+else:
+    face_detector = cv.CascadeClassifier(haar_cascade_face_detector)
+
+if not os.path.exists(dlib_facial_landmark_predictor):
+    print(f"Error: Dlib predictor file not found at {dlib_facial_landmark_predictor}")
+    print("Please download shape_predictor_68_face_landmarks.dat from dlib's website")
+    landmark_predictor = None
+else:
+    landmark_predictor = dlib.shape_predictor(dlib_facial_landmark_predictor)
 
 font = cv.FONT_HERSHEY_SIMPLEX
 EYE_ASPECT_RATIO_THRESHOLD = 0.25
@@ -27,7 +37,7 @@ MOUTH_ASPECT_RATIO_THRESHOLD = 0.5
 MOUTH_OPEN_THRESHOLD = 15
 FACE_LOST_THRESHOLD = 25
 
-# --- GLOBAL STATE VARIABLES (managed by reset_counters) ---
+# --- GLOBAL STATE VARIABLES ---
 EYE_THRESH_COUNTER = 0
 DROWSY_COUNTER = 0
 drowsy_alert = False
@@ -38,39 +48,57 @@ FACE_LOST_COUNTER = 0
 HEAD_DOWN_COUNTER = 0
 head_down_alert = False
 
-# --- LAZY AUDIO INITIALIZATION ---
+# --- FIXED: Audio handling for cloud deployment ---
 _audio_initialized = False
-_drowsiness_sound = None
-_yawn_sound = None
+_audio_available = False
 
 def _initialize_audio():
-    """Initializes pygame mixer only when needed and handles errors."""
-    global _audio_initialized, _drowsiness_sound, _yawn_sound
+    """Initializes audio only if available (for local deployment)."""
+    global _audio_initialized, _audio_available
     if _audio_initialized:
         return
+    
     try:
-        pygame.mixer.init()
-        _drowsiness_sound = pygame.mixer.Sound(os.path.join(script_dir, "drowsiness-detected.mp3"))
-        _yawn_sound = pygame.mixer.Sound(os.path.join(script_dir, "yawning-detected.mp3"))
-        print("Audio initialized successfully.")
-    except pygame.error as e:
-        print(f"Warning: Could not initialize audio. Alert sounds will be disabled. Error: {e}")
+        # Check if we're in a cloud environment
+        if os.getenv("SPACE_ID") or os.getenv("HUGGINGFACE_SPACE"):
+            print("Cloud environment detected - audio disabled")
+            _audio_available = False
+        else:
+            import pygame
+            pygame.mixer.init()
+            _audio_available = True
+            print("Audio initialized successfully.")
+    except Exception as e:
+        print(f"Audio not available: {e}")
+        _audio_available = False
+    
     _audio_initialized = True
 
-def play_alarm(sound_to_play):
-    """Plays an alarm sound if the audio system is available."""
-    _initialize_audio() # Ensure audio is initialized
-    if sound_to_play and not pygame.mixer.get_busy():
-        sound_to_play.play()
+def play_alarm(sound_file=None):
+    """Plays an alarm sound if audio is available."""
+    _initialize_audio()
+    if not _audio_available:
+        return
+    
+    try:
+        import pygame
+        if sound_file and os.path.exists(sound_file) and not pygame.mixer.get_busy():
+            sound = pygame.mixer.Sound(sound_file)
+            sound.play()
+    except Exception as e:
+        print(f"Could not play sound: {e}")
 
 def generate_alert(final_eye_ratio, final_mouth_ratio):
     global EYE_THRESH_COUNTER, YAWN_THRESH_COUNTER, drowsy_alert, yawn_alert, DROWSY_COUNTER, YAWN_COUNTER
+    
     if final_eye_ratio < EYE_ASPECT_RATIO_THRESHOLD:
         EYE_THRESH_COUNTER += 1
         if EYE_THRESH_COUNTER >= EYE_CLOSED_THRESHOLD and not drowsy_alert:
             DROWSY_COUNTER += 1
             drowsy_alert = True
-            Thread(target=play_alarm, args=(_drowsiness_sound,)).start()
+            # Try to play sound if available
+            drowsiness_sound = os.path.join(script_dir, "drowsiness-detected.mp3")
+            Thread(target=play_alarm, args=(drowsiness_sound,)).start()
     else:
         EYE_THRESH_COUNTER = 0
         drowsy_alert = False
@@ -80,23 +108,31 @@ def generate_alert(final_eye_ratio, final_mouth_ratio):
         if YAWN_THRESH_COUNTER >= MOUTH_OPEN_THRESHOLD and not yawn_alert:
             YAWN_COUNTER += 1
             yawn_alert = True
-            Thread(target=play_alarm, args=(_yawn_sound,)).start()
+            # Try to play sound if available
+            yawn_sound = os.path.join(script_dir, "yawning-detected.mp3")
+            Thread(target=play_alarm, args=(yawn_sound,)).start()
     else:
         YAWN_THRESH_COUNTER = 0
         yawn_alert = False
 
 def detect_facial_landmarks(x, y, w, h, gray_frame):
+    """Detect facial landmarks using dlib predictor."""
+    if landmark_predictor is None:
+        return None
+    
     face = dlib.rectangle(int(x), int(y), int(x + w), int(y + h))
     face_landmarks = landmark_predictor(gray_frame, face)
     return face_utils.shape_to_np(face_landmarks)
 
 def eye_aspect_ratio(eye):
+    """Calculate eye aspect ratio."""
     A = dist.euclidean(eye[1], eye[5])
     B = dist.euclidean(eye[2], eye[4])
     C = dist.euclidean(eye[0], eye[3])
     return (A + B) / (2.0 * C)
 
 def final_eye_aspect_ratio(shape):
+    """Calculate final eye aspect ratio from both eyes."""
     (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
     (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
     left_ear = eye_aspect_ratio(shape[lStart:lEnd])
@@ -104,12 +140,14 @@ def final_eye_aspect_ratio(shape):
     return (left_ear + right_ear) / 2.0
 
 def mouth_aspect_ratio(mouth):
+    """Calculate mouth aspect ratio."""
     A = dist.euclidean(mouth[2], mouth[10])
     B = dist.euclidean(mouth[4], mouth[8])
     C = dist.euclidean(mouth[0], mouth[6])
     return (A + B) / (2.0 * C)
 
 def final_mouth_aspect_ratio(shape):
+    """Calculate final mouth aspect ratio."""
     (mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
     return mouth_aspect_ratio(shape[mStart:mEnd])
 
@@ -129,31 +167,56 @@ def process_frame(frame):
     # The output frame will have a fixed width of 640px
     frame = imutils.resize(frame, width=640)
     gray_frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
-    faces = face_detector.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30), flags=cv.CASCADE_SCALE_IMAGE)
+    
+    # Detect faces
+    faces = face_detector.detectMultiScale(
+        gray_frame, 
+        scaleFactor=1.1, 
+        minNeighbors=5, 
+        minSize=(30, 30), 
+        flags=cv.CASCADE_SCALE_IMAGE
+    )
 
     if len(faces) > 0:
         FACE_LOST_COUNTER = 0
         head_down_alert = False
         (x, y, w, h) = faces[0]
+        
+        # Draw rectangle around face
+        cv.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
+        
+        # Detect landmarks if predictor is available
         face_landmarks = detect_facial_landmarks(x, y, w, h, gray_frame)
-        final_ear = final_eye_aspect_ratio(face_landmarks)
-        final_mar = final_mouth_aspect_ratio(face_landmarks)
-        generate_alert(final_ear, final_mar)
-        cv.putText(frame, f"EAR: {final_ear:.2f}", (10, 30), font, 0.7, (0, 0, 255), 2)
-        cv.putText(frame, f"MAR: {final_mar:.2f}", (10, 60), font, 0.7, (0, 0, 255), 2)
+        
+        if face_landmarks is not None:
+            final_ear = final_eye_aspect_ratio(face_landmarks)
+            final_mar = final_mouth_aspect_ratio(face_landmarks)
+            generate_alert(final_ear, final_mar)
+            
+            # Display ratios
+            cv.putText(frame, f"EAR: {final_ear:.2f}", (10, 30), font, 0.7, (0, 0, 255), 2)
+            cv.putText(frame, f"MAR: {final_mar:.2f}", (10, 60), font, 0.7, (0, 0, 255), 2)
+        else:
+            # If no landmarks detected, show warning
+            cv.putText(frame, "Landmarks not available", (10, 30), font, 0.7, (0, 0, 255), 2)
     else:
         FACE_LOST_COUNTER += 1
         if FACE_LOST_COUNTER >= FACE_LOST_THRESHOLD and not head_down_alert:
             HEAD_DOWN_COUNTER += 1
             head_down_alert = True
             
-    # Draw status text
+    # Draw status information
     cv.putText(frame, f"Drowsy: {DROWSY_COUNTER}", (480, 30), font, 0.7, (255, 255, 0), 2)
     cv.putText(frame, f"Yawn: {YAWN_COUNTER}", (480, 60), font, 0.7, (255, 255, 0), 2)
     cv.putText(frame, f"Head Down: {HEAD_DOWN_COUNTER}", (480, 90), font, 0.7, (255, 255, 0), 2)
-    if drowsy_alert: cv.putText(frame, "DROWSINESS ALERT!", (150, 30), font, 0.9, (0, 0, 255), 2)
-    if yawn_alert: cv.putText(frame, "YAWN ALERT!", (200, 60), font, 0.9, (0, 0, 255), 2)
-    if head_down_alert: cv.putText(frame, "HEAD NOT VISIBLE!", (180, 90), font, 0.9, (0, 0, 255), 2)
+    
+    # Draw alerts
+    if drowsy_alert:
+        cv.putText(frame, "DROWSINESS ALERT!", (150, 30), font, 0.9, (0, 0, 255), 2)
+    if yawn_alert:
+        cv.putText(frame, "YAWN ALERT!", (200, 60), font, 0.9, (0, 0, 255), 2)
+    if head_down_alert:
+        cv.putText(frame, "HEAD NOT VISIBLE!", (180, 90), font, 0.9, (0, 0, 255), 2)
     
     return frame
 
@@ -164,6 +227,12 @@ if __name__ == "__main__":
     parser.add_argument('--input', type=str, help='Input video file path for video mode')
     args = parser.parse_args()
     
+    # Check if landmark predictor is available
+    if landmark_predictor is None:
+        print("Error: Dlib facial landmark predictor not found!")
+        print("Please download shape_predictor_68_face_landmarks.dat")
+        exit(1)
+    
     if args.mode == 'webcam':
         print("Starting webcam detection... Press 'q' to quit.")
         cap = cv.VideoCapture(0)
@@ -173,10 +242,12 @@ if __name__ == "__main__":
             reset_counters()
             while True:
                 ret, frame = cap.read()
-                if not ret: break
+                if not ret:
+                    break
                 processed_frame = process_frame(frame)
                 cv.imshow("Live Drowsiness Detection", processed_frame)
-                if cv.waitKey(1) & 0xFF == ord('q'): break
+                if cv.waitKey(1) & 0xFF == ord('q'):
+                    break
             cap.release()
             cv.destroyAllWindows()