simran0608 commited on
Commit
3e098b3
·
verified ·
1 Parent(s): 4ff02b4

Update drowsiness_detection.py

Browse files
Files changed (1) hide show
  1. drowsiness_detection.py +93 -145
drowsiness_detection.py CHANGED
@@ -1,4 +1,5 @@
1
- # PREP DEPENDENCIES
 
2
  from scipy.spatial import distance as dist
3
  from imutils import face_utils
4
  from threading import Thread
@@ -6,79 +7,80 @@ import numpy as np
6
  import cv2 as cv
7
  import imutils
8
  import dlib
9
- import pygame # Used for playing alarm sounds cross-platform
10
  import argparse
11
  import os
12
 
13
- # --- INITIALIZE MODELS AND CONSTANTS ---
 
 
 
 
14
 
15
- # Haar cascade classifier for face detection
16
- haar_cascade_face_detector = "haarcascade_frontalface_default.xml"
17
  face_detector = cv.CascadeClassifier(haar_cascade_face_detector)
18
-
19
- # Dlib facial landmark detector
20
- dlib_facial_landmark_predictor = "shape_predictor_68_face_landmarks.dat"
21
  landmark_predictor = dlib.shape_predictor(dlib_facial_landmark_predictor)
22
 
23
- # Important Variables
24
  font = cv.FONT_HERSHEY_SIMPLEX
25
- # --- INITIALIZE MODELS AND CONSTANTS ---
26
- # Eye Drowsiness Detection
27
  EYE_ASPECT_RATIO_THRESHOLD = 0.25
28
  EYE_CLOSED_THRESHOLD = 20
 
 
 
 
 
29
  EYE_THRESH_COUNTER = 0
30
  DROWSY_COUNTER = 0
31
  drowsy_alert = False
32
-
33
- # Mouth Yawn Detection
34
- MOUTH_ASPECT_RATIO_THRESHOLD = 0.5
35
- MOUTH_OPEN_THRESHOLD = 15
36
  YAWN_THRESH_COUNTER = 0
37
  YAWN_COUNTER = 0
38
  yawn_alert = False
39
-
40
- # NEW: Head Not Visible Detection
41
- FACE_LOST_THRESHOLD = 25 # Conseq. frames face must be lost to trigger alert
42
  FACE_LOST_COUNTER = 0
43
- HEAD_DOWN_COUNTER = 0 # Renaming for clarity
44
  head_down_alert = False
45
 
46
- # --- AUDIO SETUP (using Pygame) ---
47
- # pygame.mixer.init()
48
- # drowsiness_sound = pygame.mixer.Sound("drowsiness-detected.mp3")
49
- # yawn_sound = pygame.mixer.Sound("yawning-detected.mp3")
50
- # head_down_sound = pygame.mixer.Sound("dependencies/audio/head-down-detected.mp3")
51
-
52
- # --- CORE FUNCTIONS ---
53
- # def play_alarm(sound_to_play):
54
- # if not pygame.mixer.get_busy():
55
- # sound_to_play.play()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
 
57
  def generate_alert(final_eye_ratio, final_mouth_ratio):
58
- global EYE_THRESH_COUNTER, YAWN_THRESH_COUNTER
59
- global drowsy_alert, yawn_alert
60
- global DROWSY_COUNTER, YAWN_COUNTER
61
-
62
- # Drowsiness check
63
  if final_eye_ratio < EYE_ASPECT_RATIO_THRESHOLD:
64
  EYE_THRESH_COUNTER += 1
65
- if EYE_THRESH_COUNTER >= EYE_CLOSED_THRESHOLD:
66
- if not drowsy_alert:
67
- DROWSY_COUNTER += 1
68
- drowsy_alert = True
69
- # Thread(target=play_alarm, args=(drowsiness_sound,)).start()
70
  else:
71
  EYE_THRESH_COUNTER = 0
72
  drowsy_alert = False
73
 
74
- # Yawn check
75
  if final_mouth_ratio > MOUTH_ASPECT_RATIO_THRESHOLD:
76
  YAWN_THRESH_COUNTER += 1
77
- if YAWN_THRESH_COUNTER >= MOUTH_OPEN_THRESHOLD:
78
- if not yawn_alert:
79
- YAWN_COUNTER += 1
80
- yawn_alert = True
81
- # Thread(target=play_alarm, args=(yawn_sound,)).start()
82
  else:
83
  YAWN_THRESH_COUNTER = 0
84
  yawn_alert = False
@@ -86,51 +88,33 @@ def generate_alert(final_eye_ratio, final_mouth_ratio):
86
  def detect_facial_landmarks(x, y, w, h, gray_frame):
87
  face = dlib.rectangle(int(x), int(y), int(x + w), int(y + h))
88
  face_landmarks = landmark_predictor(gray_frame, face)
89
- face_landmarks = face_utils.shape_to_np(face_landmarks)
90
- return face_landmarks
91
 
92
  def eye_aspect_ratio(eye):
93
  A = dist.euclidean(eye[1], eye[5])
94
  B = dist.euclidean(eye[2], eye[4])
95
  C = dist.euclidean(eye[0], eye[3])
96
- ear = (A + B) / (2.0 * C)
97
- return ear
98
 
99
  def final_eye_aspect_ratio(shape):
100
  (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
101
  (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
102
- left_eye = shape[lStart:lEnd]
103
- right_eye = shape[rStart:rEnd]
104
- left_ear = eye_aspect_ratio(left_eye)
105
- right_ear = eye_aspect_ratio(right_eye)
106
- final_ear = (left_ear + right_ear) / 2.0
107
- return final_ear, left_eye, right_eye
108
 
109
  def mouth_aspect_ratio(mouth):
110
  A = dist.euclidean(mouth[2], mouth[10])
111
  B = dist.euclidean(mouth[4], mouth[8])
112
  C = dist.euclidean(mouth[0], mouth[6])
113
- mar = (A + B) / (2.0 * C)
114
- return mar
115
 
116
  def final_mouth_aspect_ratio(shape):
117
  (mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
118
- mouth = shape[mStart:mEnd]
119
- return mouth_aspect_ratio(mouth), mouth
120
-
121
- def head_pose_ratio(shape):
122
- nose_tip = shape[30]
123
- chin_tip = shape[8]
124
- left_face_corner = shape[0]
125
- right_face_corner = shape[16]
126
- nose_to_chin_dist = dist.euclidean(nose_tip, chin_tip)
127
- face_width = dist.euclidean(left_face_corner, right_face_corner)
128
- if face_width == 0:
129
- return 0.0
130
- hpr = nose_to_chin_dist / face_width
131
- return hpr
132
 
133
  def reset_counters():
 
134
  global EYE_THRESH_COUNTER, YAWN_THRESH_COUNTER, FACE_LOST_COUNTER
135
  global DROWSY_COUNTER, YAWN_COUNTER, HEAD_DOWN_COUNTER
136
  global drowsy_alert, yawn_alert, head_down_alert
@@ -139,21 +123,21 @@ def reset_counters():
139
  drowsy_alert, yawn_alert, head_down_alert = False, False, False
140
 
141
  def process_frame(frame):
 
142
  global FACE_LOST_COUNTER, head_down_alert, HEAD_DOWN_COUNTER
 
 
143
  frame = imutils.resize(frame, width=640)
144
  gray_frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
145
  faces = face_detector.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30), flags=cv.CASCADE_SCALE_IMAGE)
 
146
  if len(faces) > 0:
147
  FACE_LOST_COUNTER = 0
148
  head_down_alert = False
149
  (x, y, w, h) = faces[0]
150
  face_landmarks = detect_facial_landmarks(x, y, w, h, gray_frame)
151
- final_ear, left_eye, right_eye = final_eye_aspect_ratio(face_landmarks)
152
- final_mar, mouth = final_mouth_aspect_ratio(face_landmarks)
153
- # left_eye_hull, right_eye_hull, mouth_hull = cv.convexHull(left_eye), cv.convexHull(right_eye), cv.convexHull(mouth)
154
- # cv.drawContours(frame, [left_eye_hull], -1, (0, 255, 0), 1)
155
- # cv.drawContours(frame, [right_eye_hull], -1, (0, 255, 0), 1)
156
- # cv.drawContours(frame, [mouth_hull], -1, (0, 255, 0), 1)
157
  generate_alert(final_ear, final_mar)
158
  cv.putText(frame, f"EAR: {final_ear:.2f}", (10, 30), font, 0.7, (0, 0, 255), 2)
159
  cv.putText(frame, f"MAR: {final_mar:.2f}", (10, 60), font, 0.7, (0, 0, 255), 2)
@@ -162,87 +146,51 @@ def process_frame(frame):
162
  if FACE_LOST_COUNTER >= FACE_LOST_THRESHOLD and not head_down_alert:
163
  HEAD_DOWN_COUNTER += 1
164
  head_down_alert = True
 
 
165
  cv.putText(frame, f"Drowsy: {DROWSY_COUNTER}", (480, 30), font, 0.7, (255, 255, 0), 2)
166
  cv.putText(frame, f"Yawn: {YAWN_COUNTER}", (480, 60), font, 0.7, (255, 255, 0), 2)
167
  cv.putText(frame, f"Head Down: {HEAD_DOWN_COUNTER}", (480, 90), font, 0.7, (255, 255, 0), 2)
168
  if drowsy_alert: cv.putText(frame, "DROWSINESS ALERT!", (150, 30), font, 0.9, (0, 0, 255), 2)
169
  if yawn_alert: cv.putText(frame, "YAWN ALERT!", (200, 60), font, 0.9, (0, 0, 255), 2)
170
  if head_down_alert: cv.putText(frame, "HEAD NOT VISIBLE!", (180, 90), font, 0.9, (0, 0, 255), 2)
 
171
  return frame
172
 
173
- def process_video(input_path, output_path=None):
174
- reset_counters()
175
- video_stream = cv.VideoCapture(input_path)
176
- if not video_stream.isOpened():
177
- print(f"Error: Could not open video file {input_path}")
178
- return False
179
-
180
- fps = int(video_stream.get(cv.CAP_PROP_FPS))
181
- width = int(video_stream.get(cv.CAP_PROP_FRAME_WIDTH))
182
- height = int(video_stream.get(cv.CAP_PROP_FRAME_HEIGHT))
183
-
184
- print(f"Processing video: {input_path}")
185
- print(f"Original Res: {width}x{height}, FPS: {fps}")
186
-
187
- video_writer = None
188
- if output_path:
189
- fourcc = cv.VideoWriter_fourcc(*'mp4v')
190
- # --- FIX: Calculate correct output dimensions to prevent corruption ---
191
- # The process_frame function resizes frames to a fixed width of 640.
192
- output_width = 640
193
- # Maintain aspect ratio
194
- output_height = int(height * (output_width / float(width)))
195
- output_dims = (output_width, output_height)
196
- video_writer = cv.VideoWriter(output_path, fourcc, fps, output_dims)
197
- print(f"Outputting video with Res: {output_dims[0]}x{output_dims[1]}")
198
-
199
- while True:
200
- ret, frame = video_stream.read()
201
- if not ret: break
202
-
203
- processed_frame = process_frame(frame)
204
- if video_writer: video_writer.write(processed_frame)
205
-
206
- video_stream.release()
207
- if video_writer: video_writer.release()
208
-
209
- print("Video processing complete!")
210
- print(f"Final Stats - Drowsy: {DROWSY_COUNTER}, Yawn: {YAWN_COUNTER}, Head Down: {HEAD_DOWN_COUNTER}")
211
- return True
212
-
213
- def run_webcam():
214
- reset_counters()
215
- video_stream = cv.VideoCapture(0)
216
- if not video_stream.isOpened():
217
- print("Error: Could not open webcam")
218
- return False
219
- while True:
220
- ret, frame = video_stream.read()
221
- if not ret:
222
- print("Failed to grab frame")
223
- break
224
- processed_frame = process_frame(frame)
225
- cv.imshow("Live Drowsiness and Yawn Detection", processed_frame)
226
- if cv.waitKey(1) & 0xFF == ord('q'): break
227
- video_stream.release()
228
- cv.destroyAllWindows()
229
- return True
230
-
231
- # --- MAIN EXECUTION LOOP ---
232
  if __name__ == "__main__":
233
- parser = argparse.ArgumentParser(description='Drowsiness Detection System')
234
  parser.add_argument('--mode', choices=['webcam', 'video'], default='webcam', help='Mode of operation')
235
  parser.add_argument('--input', type=str, help='Input video file path for video mode')
236
- parser.add_argument('--output', type=str, help='Output video file path for video mode')
237
  args = parser.parse_args()
238
 
239
  if args.mode == 'webcam':
240
- print("Starting webcam detection...")
241
- run_webcam()
 
 
 
 
 
 
 
 
 
 
 
 
 
242
  elif args.mode == 'video':
243
- if not args.input:
244
- print("Error: --input argument is required for video mode.")
245
- elif not os.path.exists(args.input):
246
- print(f"Error: Input file not found at {args.input}")
247
  else:
248
- process_video(args.input, args.output)
 
 
 
 
 
 
 
 
 
 
1
+ # drowsiness_detection.py
2
+
3
  from scipy.spatial import distance as dist
4
  from imutils import face_utils
5
  from threading import Thread
 
7
  import cv2 as cv
8
  import imutils
9
  import dlib
10
+ import pygame
11
  import argparse
12
  import os
13
 
14
+ # --- MODELS AND CONSTANTS ---
15
+ # Use absolute paths relative to this script file for robustness
16
+ script_dir = os.path.dirname(os.path.abspath(__file__))
17
+ haar_cascade_face_detector = os.path.join(script_dir, "haarcascade_frontalface_default.xml")
18
+ dlib_facial_landmark_predictor = os.path.join(script_dir, "shape_predictor_68_face_landmarks.dat")
19
 
 
 
20
  face_detector = cv.CascadeClassifier(haar_cascade_face_detector)
 
 
 
21
  landmark_predictor = dlib.shape_predictor(dlib_facial_landmark_predictor)
22
 
 
23
  font = cv.FONT_HERSHEY_SIMPLEX
 
 
24
  EYE_ASPECT_RATIO_THRESHOLD = 0.25
25
  EYE_CLOSED_THRESHOLD = 20
26
+ MOUTH_ASPECT_RATIO_THRESHOLD = 0.5
27
+ MOUTH_OPEN_THRESHOLD = 15
28
+ FACE_LOST_THRESHOLD = 25
29
+
30
+ # --- GLOBAL STATE VARIABLES (managed by reset_counters) ---
31
  EYE_THRESH_COUNTER = 0
32
  DROWSY_COUNTER = 0
33
  drowsy_alert = False
 
 
 
 
34
  YAWN_THRESH_COUNTER = 0
35
  YAWN_COUNTER = 0
36
  yawn_alert = False
 
 
 
37
  FACE_LOST_COUNTER = 0
38
+ HEAD_DOWN_COUNTER = 0
39
  head_down_alert = False
40
 
41
+ # --- LAZY AUDIO INITIALIZATION ---
42
+ _audio_initialized = False
43
+ _drowsiness_sound = None
44
+ _yawn_sound = None
45
+
46
+ def _initialize_audio():
47
+ """Initializes pygame mixer only when needed and handles errors."""
48
+ global _audio_initialized, _drowsiness_sound, _yawn_sound
49
+ if _audio_initialized:
50
+ return
51
+ try:
52
+ pygame.mixer.init()
53
+ _drowsiness_sound = pygame.mixer.Sound(os.path.join(script_dir, "drowsiness-detected.mp3"))
54
+ _yawn_sound = pygame.mixer.Sound(os.path.join(script_dir, "yawning-detected.mp3"))
55
+ print("Audio initialized successfully.")
56
+ except pygame.error as e:
57
+ print(f"Warning: Could not initialize audio. Alert sounds will be disabled. Error: {e}")
58
+ _audio_initialized = True
59
+
60
+ def play_alarm(sound_to_play):
61
+ """Plays an alarm sound if the audio system is available."""
62
+ _initialize_audio() # Ensure audio is initialized
63
+ if sound_to_play and not pygame.mixer.get_busy():
64
+ sound_to_play.play()
65
 
66
  def generate_alert(final_eye_ratio, final_mouth_ratio):
67
+ global EYE_THRESH_COUNTER, YAWN_THRESH_COUNTER, drowsy_alert, yawn_alert, DROWSY_COUNTER, YAWN_COUNTER
 
 
 
 
68
  if final_eye_ratio < EYE_ASPECT_RATIO_THRESHOLD:
69
  EYE_THRESH_COUNTER += 1
70
+ if EYE_THRESH_COUNTER >= EYE_CLOSED_THRESHOLD and not drowsy_alert:
71
+ DROWSY_COUNTER += 1
72
+ drowsy_alert = True
73
+ Thread(target=play_alarm, args=(_drowsiness_sound,)).start()
 
74
  else:
75
  EYE_THRESH_COUNTER = 0
76
  drowsy_alert = False
77
 
 
78
  if final_mouth_ratio > MOUTH_ASPECT_RATIO_THRESHOLD:
79
  YAWN_THRESH_COUNTER += 1
80
+ if YAWN_THRESH_COUNTER >= MOUTH_OPEN_THRESHOLD and not yawn_alert:
81
+ YAWN_COUNTER += 1
82
+ yawn_alert = True
83
+ Thread(target=play_alarm, args=(_yawn_sound,)).start()
 
84
  else:
85
  YAWN_THRESH_COUNTER = 0
86
  yawn_alert = False
 
88
  def detect_facial_landmarks(x, y, w, h, gray_frame):
89
  face = dlib.rectangle(int(x), int(y), int(x + w), int(y + h))
90
  face_landmarks = landmark_predictor(gray_frame, face)
91
+ return face_utils.shape_to_np(face_landmarks)
 
92
 
93
  def eye_aspect_ratio(eye):
94
  A = dist.euclidean(eye[1], eye[5])
95
  B = dist.euclidean(eye[2], eye[4])
96
  C = dist.euclidean(eye[0], eye[3])
97
+ return (A + B) / (2.0 * C)
 
98
 
99
  def final_eye_aspect_ratio(shape):
100
  (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
101
  (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]
102
+ left_ear = eye_aspect_ratio(shape[lStart:lEnd])
103
+ right_ear = eye_aspect_ratio(shape[rStart:rEnd])
104
+ return (left_ear + right_ear) / 2.0
 
 
 
105
 
106
  def mouth_aspect_ratio(mouth):
107
  A = dist.euclidean(mouth[2], mouth[10])
108
  B = dist.euclidean(mouth[4], mouth[8])
109
  C = dist.euclidean(mouth[0], mouth[6])
110
+ return (A + B) / (2.0 * C)
 
111
 
112
  def final_mouth_aspect_ratio(shape):
113
  (mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
114
+ return mouth_aspect_ratio(shape[mStart:mEnd])
 
 
 
 
 
 
 
 
 
 
 
 
 
115
 
116
  def reset_counters():
117
+ """Resets all global counters and alerts for a new processing session."""
118
  global EYE_THRESH_COUNTER, YAWN_THRESH_COUNTER, FACE_LOST_COUNTER
119
  global DROWSY_COUNTER, YAWN_COUNTER, HEAD_DOWN_COUNTER
120
  global drowsy_alert, yawn_alert, head_down_alert
 
123
  drowsy_alert, yawn_alert, head_down_alert = False, False, False
124
 
125
  def process_frame(frame):
126
+ """Processes a single frame to detect drowsiness, yawns, and head position."""
127
  global FACE_LOST_COUNTER, head_down_alert, HEAD_DOWN_COUNTER
128
+
129
+ # The output frame will have a fixed width of 640px
130
  frame = imutils.resize(frame, width=640)
131
  gray_frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
132
  faces = face_detector.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30), flags=cv.CASCADE_SCALE_IMAGE)
133
+
134
  if len(faces) > 0:
135
  FACE_LOST_COUNTER = 0
136
  head_down_alert = False
137
  (x, y, w, h) = faces[0]
138
  face_landmarks = detect_facial_landmarks(x, y, w, h, gray_frame)
139
+ final_ear = final_eye_aspect_ratio(face_landmarks)
140
+ final_mar = final_mouth_aspect_ratio(face_landmarks)
 
 
 
 
141
  generate_alert(final_ear, final_mar)
142
  cv.putText(frame, f"EAR: {final_ear:.2f}", (10, 30), font, 0.7, (0, 0, 255), 2)
143
  cv.putText(frame, f"MAR: {final_mar:.2f}", (10, 60), font, 0.7, (0, 0, 255), 2)
 
146
  if FACE_LOST_COUNTER >= FACE_LOST_THRESHOLD and not head_down_alert:
147
  HEAD_DOWN_COUNTER += 1
148
  head_down_alert = True
149
+
150
+ # Draw status text
151
  cv.putText(frame, f"Drowsy: {DROWSY_COUNTER}", (480, 30), font, 0.7, (255, 255, 0), 2)
152
  cv.putText(frame, f"Yawn: {YAWN_COUNTER}", (480, 60), font, 0.7, (255, 255, 0), 2)
153
  cv.putText(frame, f"Head Down: {HEAD_DOWN_COUNTER}", (480, 90), font, 0.7, (255, 255, 0), 2)
154
  if drowsy_alert: cv.putText(frame, "DROWSINESS ALERT!", (150, 30), font, 0.9, (0, 0, 255), 2)
155
  if yawn_alert: cv.putText(frame, "YAWN ALERT!", (200, 60), font, 0.9, (0, 0, 255), 2)
156
  if head_down_alert: cv.putText(frame, "HEAD NOT VISIBLE!", (180, 90), font, 0.9, (0, 0, 255), 2)
157
+
158
  return frame
159
 
160
+ # --- Command-line execution for local testing ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  if __name__ == "__main__":
162
+ parser = argparse.ArgumentParser(description='Drowsiness Detection System (Local Runner)')
163
  parser.add_argument('--mode', choices=['webcam', 'video'], default='webcam', help='Mode of operation')
164
  parser.add_argument('--input', type=str, help='Input video file path for video mode')
 
165
  args = parser.parse_args()
166
 
167
  if args.mode == 'webcam':
168
+ print("Starting webcam detection... Press 'q' to quit.")
169
+ cap = cv.VideoCapture(0)
170
+ if not cap.isOpened():
171
+ print("Error: Could not open webcam.")
172
+ else:
173
+ reset_counters()
174
+ while True:
175
+ ret, frame = cap.read()
176
+ if not ret: break
177
+ processed_frame = process_frame(frame)
178
+ cv.imshow("Live Drowsiness Detection", processed_frame)
179
+ if cv.waitKey(1) & 0xFF == ord('q'): break
180
+ cap.release()
181
+ cv.destroyAllWindows()
182
+
183
  elif args.mode == 'video':
184
+ if not args.input or not os.path.exists(args.input):
185
+ print("Error: Please provide a valid --input video file path.")
 
 
186
  else:
187
+ from video_processor import process_video_with_progress
188
+ output_file = args.input.replace('.mp4', '_processed.mp4')
189
+ print(f"Processing video {args.input}, output will be {output_file}")
190
+
191
+ def cli_progress(current, total):
192
+ percent = int((current / total) * 100)
193
+ print(f"\rProcessing: {percent}%", end="")
194
+
195
+ process_video_with_progress(args.input, output_file, progress_callback=cli_progress)
196
+ print("\nDone.")