Spaces:
Runtime error
Runtime error
File size: 5,892 Bytes
19f420a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
# drive_paddy/detection/strategies/geometric.py
import cv2
import mediapipe as mp
import numpy as np
import math
from ..base_processor import BaseProcessor
# --- Helper Functions ---
def calculate_ear(eye_landmarks, frame_shape):
"""Calculates the Eye Aspect Ratio (EAR)."""
# ... (implementation remains the same)
coords = np.array([(lm.x * frame_shape[1], lm.y * frame_shape[0]) for lm in eye_landmarks])
v1 = np.linalg.norm(coords[1] - coords[5]); v2 = np.linalg.norm(coords[2] - coords[4])
h1 = np.linalg.norm(coords[0] - coords[3])
return (v1 + v2) / (2.0 * h1) if h1 > 0 else 0.0
def calculate_mar(mouth_landmarks, frame_shape):
"""Calculates the Mouth Aspect Ratio (MAR) for yawn detection."""
coords = np.array([(lm.x * frame_shape[1], lm.y * frame_shape[0]) for lm in mouth_landmarks])
v1 = np.linalg.norm(coords[1] - coords[7]) # Vertical distances
v2 = np.linalg.norm(coords[2] - coords[6])
v3 = np.linalg.norm(coords[3] - coords[5])
h1 = np.linalg.norm(coords[0] - coords[4]) # Horizontal distance
return (v1 + v2 + v3) / (2.0 * h1) if h1 > 0 else 0.0
class GeometricProcessor(BaseProcessor):
"""
Drowsiness detection using a combination of facial landmarks:
- Eye Aspect Ratio (EAR) for eye closure.
- Mouth Aspect Ratio (MAR) for yawning.
- Head Pose Estimation for nodding off or looking away.
"""
def __init__(self, config):
self.settings = config['geometric_settings']
self.face_mesh = mp.solutions.face_mesh.FaceMesh(
max_num_faces=1, refine_landmarks=True,
min_detection_confidence=0.5, min_tracking_confidence=0.5)
# State counters
self.counters = {
"eye_closure": 0, "yawning": 0,
"head_nod": 0, "looking_away": 0
}
# Landmark indices
self.L_EYE = [362, 385, 387, 263, 373, 380]
self.R_EYE = [33, 160, 158, 133, 153, 144]
self.MOUTH = [61, 291, 39, 181, 0, 17, 84, 178]
def process_frame(self, frame):
img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
h, w, _ = frame.shape
results = self.face_mesh.process(img_rgb)
drowsiness_indicators = {
"eye_closure": False, "yawning": False,
"head_nod": False, "looking_away": False, "details": {}
}
if results.multi_face_landmarks:
landmarks = results.multi_face_landmarks[0].landmark
# --- Eye Closure Detection (EAR) ---
left_ear = calculate_ear([landmarks[i] for i in self.L_EYE], (h, w))
right_ear = calculate_ear([landmarks[i] for i in self.R_EYE], (h, w))
ear = (left_ear + right_ear) / 2.0
if ear < self.settings['eye_ar_thresh']:
self.counters['eye_closure'] += 1
if self.counters['eye_closure'] >= self.settings['eye_ar_consec_frames']:
drowsiness_indicators['eye_closure'] = True
else:
self.counters['eye_closure'] = 0
drowsiness_indicators['details']['EAR'] = ear
# --- Yawn Detection (MAR) ---
mar = calculate_mar([landmarks[i] for i in self.MOUTH], (h, w))
if mar > self.settings['yawn_mar_thresh']:
self.counters['yawning'] += 1
if self.counters['yawning'] >= self.settings['yawn_consec_frames']:
drowsiness_indicators['yawning'] = True
else:
self.counters['yawning'] = 0
drowsiness_indicators['details']['MAR'] = mar
# --- Head Pose Estimation ---
face_3d = np.array([
[0.0, 0.0, 0.0], # Nose tip
[0.0, -330.0, -65.0], # Chin
[-225.0, 170.0, -135.0], # Left eye left corner
[225.0, 170.0, -135.0], # Right eye right corner
[-150.0, -150.0, -125.0], # Left Mouth corner
[150.0, -150.0, -125.0] # Right mouth corner
], dtype=np.float64)
face_2d = np.array([
(landmarks[1].x * w, landmarks[1].y * h), # Nose tip
(landmarks[152].x * w, landmarks[152].y * h), # Chin
(landmarks[263].x * w, landmarks[263].y * h), # Left eye corner
(landmarks[33].x * w, landmarks[33].y * h), # Right eye corner
(landmarks[287].x * w, landmarks[287].y * h), # Left mouth corner
(landmarks[57].x * w, landmarks[57].y * h) # Right mouth corner
], dtype=np.float64)
cam_matrix = np.array([[w, 0, w / 2], [0, w, h / 2], [0, 0, 1]], dtype=np.float64)
_, rot_vec, _ = cv2.solvePnP(face_3d, face_2d, cam_matrix, np.zeros((4, 1), dtype=np.float64))
rmat, _ = cv2.Rodrigues(rot_vec)
angles, _, _, _, _, _ = cv2.RQDecomp3x3(rmat)
pitch, yaw = angles[0], angles[1]
drowsiness_indicators['details']['Pitch'] = pitch
drowsiness_indicators['details']['Yaw'] = yaw
if pitch > self.settings['head_nod_thresh']:
self.counters['head_nod'] += 1
if self.counters['head_nod'] >= self.settings['head_pose_consec_frames']:
drowsiness_indicators['head_nod'] = True
else:
self.counters['head_nod'] = 0
if abs(yaw) > self.settings['head_look_away_thresh']:
self.counters['looking_away'] += 1
if self.counters['looking_away'] >= self.settings['head_pose_consec_frames']:
drowsiness_indicators['looking_away'] = True
else:
self.counters['looking_away'] = 0
# This processor now returns the frame and a dictionary of indicators
return frame, drowsiness_indicators
|