Update drowsiness_detection.py
Browse files- drowsiness_detection.py +93 -145
drowsiness_detection.py
CHANGED
@@ -1,4 +1,5 @@
|
|
1 |
-
#
|
|
|
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
|
10 |
import argparse
|
11 |
import os
|
12 |
|
13 |
-
# ---
|
|
|
|
|
|
|
|
|
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
|
44 |
head_down_alert = False
|
45 |
|
46 |
-
# --- AUDIO
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
67 |
-
|
68 |
-
|
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 |
-
|
79 |
-
|
80 |
-
|
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 |
-
|
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 |
-
|
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 |
-
|
103 |
-
|
104 |
-
left_ear
|
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 |
-
|
114 |
-
return mar
|
115 |
|
116 |
def final_mouth_aspect_ratio(shape):
|
117 |
(mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
|
118 |
-
|
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
|
152 |
-
final_mar
|
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 |
-
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
242 |
elif args.mode == 'video':
|
243 |
-
if not args.input:
|
244 |
-
print("Error:
|
245 |
-
elif not os.path.exists(args.input):
|
246 |
-
print(f"Error: Input file not found at {args.input}")
|
247 |
else:
|
248 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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.")
|