Spaces:
Sleeping
Sleeping
Trying to implement analysis only seperately
Browse files
src/detection/strategies/geometric.py
CHANGED
@@ -193,3 +193,110 @@ class GeometricProcessor(BaseProcessor):
|
|
193 |
|
194 |
return frame
|
195 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
193 |
|
194 |
return frame
|
195 |
|
196 |
+
def analyse_frame(self, frame):
|
197 |
+
# self.frame_counter += 1
|
198 |
+
|
199 |
+
# --- FIX: More efficient frame skipping ---
|
200 |
+
# # Adaptive skipping: process more frequently if drowsiness is detected.
|
201 |
+
# last_level = self.last_indicators.get("drowsiness_level", "Awake")
|
202 |
+
# skip_n = 1 if last_level != "Awake" else self.default_skip
|
203 |
+
|
204 |
+
# if self.frame_counter % skip_n != 0:
|
205 |
+
# # If we have a cached frame, return it to avoid re-drawing.
|
206 |
+
# if self.last_drawn_frame is not None:
|
207 |
+
# return self.last_drawn_frame, self.last_indicators
|
208 |
+
# # Fallback if the first frame was skipped (unlikely but safe)
|
209 |
+
# else:
|
210 |
+
# return frame.copy(), self.last_indicators
|
211 |
+
|
212 |
+
# --- CORE FRAME PROCESSING ---
|
213 |
+
original_frame = frame.copy()
|
214 |
+
h_orig, w_orig, _ = original_frame.shape
|
215 |
+
|
216 |
+
# Optimization: Downscale frame for faster processing
|
217 |
+
small_frame = cv2.resize(original_frame, (0, 0), fx=self.downscale_factor, fy=self.downscale_factor, interpolation=cv2.INTER_AREA)
|
218 |
+
h, w, _ = small_frame.shape
|
219 |
+
|
220 |
+
# All processing is done on the `small_frame` for speed.
|
221 |
+
gray = cv2.cvtColor(small_frame, cv2.COLOR_BGR2GRAY)
|
222 |
+
brightness = np.mean(gray)
|
223 |
+
|
224 |
+
drowsiness_indicators = {"drowsiness_level": "Awake", "lighting": "Good", "details": {}}
|
225 |
+
face_landmarks_data = None
|
226 |
+
|
227 |
+
if brightness < self.settings['low_light_thresh']:
|
228 |
+
drowsiness_indicators["lighting"] = "Low"
|
229 |
+
else:
|
230 |
+
# Convert the SMALL frame to RGB for MediaPipe
|
231 |
+
img_rgb = cv2.cvtColor(small_frame, cv2.COLOR_BGR2RGB)
|
232 |
+
img_rgb.flags.writeable = False # Performance enhancement
|
233 |
+
results = self.face_mesh.process(img_rgb)
|
234 |
+
img_rgb.flags.writeable = True
|
235 |
+
|
236 |
+
if results.multi_face_landmarks:
|
237 |
+
face_landmarks_data = results.multi_face_landmarks[0]
|
238 |
+
landmarks = face_landmarks_data.landmark
|
239 |
+
score = 0
|
240 |
+
weights = self.settings['indicator_weights']
|
241 |
+
|
242 |
+
# --- Drowsiness Calculations (on small frame dimensions 'h', 'w') ---
|
243 |
+
ear_left = calculate_ear([landmarks[i] for i in self.L_EYE],(h,w))
|
244 |
+
ear_right = calculate_ear([landmarks[i] for i in self.R_EYE],(h,w))
|
245 |
+
ear = (ear_left + ear_right) / 2.0
|
246 |
+
|
247 |
+
if ear < self.settings['eye_ar_thresh']: self.counters['eye_closure']+=1
|
248 |
+
else: self.counters['eye_closure']=0
|
249 |
+
if self.counters['eye_closure'] >= self.settings['eye_ar_consec_frames']: score += weights['eye_closure']
|
250 |
+
|
251 |
+
mar = calculate_mar([landmarks[i] for i in self.MOUTH], (h, w))
|
252 |
+
if mar > self.settings['yawn_mar_thresh']: self.counters['yawning']+=1
|
253 |
+
else: self.counters['yawning']=0
|
254 |
+
if self.counters['yawning'] >= self.settings['yawn_consec_frames']: score += weights['yawning']
|
255 |
+
|
256 |
+
# --- Head Pose Estimation (on small frame dimensions 'h', 'w') ---
|
257 |
+
face_3d_model = np.array([
|
258 |
+
[0.0, 0.0, 0.0], # Nose tip
|
259 |
+
[0.0, -330.0, -65.0], # Chin
|
260 |
+
[-225.0, 170.0, -135.0], # Left eye left corner
|
261 |
+
[225.0, 170.0, -135.0], # Right eye right corner
|
262 |
+
[-150.0, -150.0, -125.0], # Left Mouth corner
|
263 |
+
[150.0, -150.0, -125.0] # Right mouth corner
|
264 |
+
], dtype=np.float32)
|
265 |
+
|
266 |
+
face_2d_points = np.array([(landmarks[i].x * w, landmarks[i].y * h) for i in self.HEAD_POSE_LANDMARKS], dtype=np.float32)
|
267 |
+
cam_matrix = np.array([[w, 0, w/2], [0, w, h/2], [0, 0, 1]], dtype=np.float32)
|
268 |
+
|
269 |
+
_, rvec, _ = cv2.solvePnP(face_3d_model, face_2d_points, cam_matrix, self.zeros_4x1, flags=cv2.SOLVEPNP_EPNP)
|
270 |
+
rmat, _ = cv2.Rodrigues(rvec)
|
271 |
+
angles, _, _, _, _, _ = cv2.RQDecomp3x3(rmat)
|
272 |
+
pitch, yaw = angles[0], angles[1]
|
273 |
+
|
274 |
+
if pitch > self.settings['head_nod_thresh']: self.counters['head_nod']+=1
|
275 |
+
else: self.counters['head_nod']=0
|
276 |
+
if self.counters['head_nod'] >= self.settings['head_pose_consec_frames']: score += weights['head_nod']
|
277 |
+
|
278 |
+
if abs(yaw) > self.settings['head_look_away_thresh']: self.counters['looking_away']+=1
|
279 |
+
else: self.counters['looking_away']=0
|
280 |
+
if self.counters['looking_away'] >= self.settings['head_pose_consec_frames']: score += weights['looking_away']
|
281 |
+
|
282 |
+
# Determine final drowsiness level based on score
|
283 |
+
levels = self.settings['drowsiness_levels']
|
284 |
+
if score >= levels['very_drowsy_threshold']:
|
285 |
+
drowsiness_indicators['drowsiness_level'] = "Very Drowsy"
|
286 |
+
elif score >= levels['slightly_drowsy_threshold']:
|
287 |
+
drowsiness_indicators['drowsiness_level'] = "Slightly Drowsy"
|
288 |
+
|
289 |
+
drowsiness_indicators['details']['Score'] = score
|
290 |
+
|
291 |
+
# --- Update state for next frame (skipped or processed) ---
|
292 |
+
self.last_indicators = drowsiness_indicators
|
293 |
+
self.last_landmarks = face_landmarks_data
|
294 |
+
|
295 |
+
# --- Draw visuals on the ORIGINAL frame for high-quality output ---
|
296 |
+
# processed_frame = self.draw_visuals(original_frame, drowsiness_indicators, face_landmarks_data)
|
297 |
+
|
298 |
+
# --- FIX: Cache the newly drawn frame ---
|
299 |
+
# self.last_drawn_frame = processed_frame
|
300 |
+
|
301 |
+
# --- FIX: Return only the two values expected by the Gradio app ---
|
302 |
+
return drowsiness_indicators
|