Testys commited on
Commit
ea86682
·
verified ·
1 Parent(s): d918bfb

Trying to implement analysis only seperately

Browse files
Files changed (1) hide show
  1. src/detection/strategies/geometric.py +107 -0
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