Spaces:
Runtime error
Runtime error
File size: 5,934 Bytes
19f420a edbbe80 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a a3c1bbf f37553c 19f420a f37553c |
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 |
# drive_paddy/detection/strategies/geometric.py
import cv2
import mediapipe as mp
import numpy as np
import math
from src.detection.base_processor import BaseProcessor
# --- Helper Functions (No changes here) ---
def calculate_ear(eye_landmarks, frame_shape):
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):
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]); v2 = np.linalg.norm(coords[2] - coords[6])
v3 = np.linalg.norm(coords[3] - coords[5]); h1 = np.linalg.norm(coords[0] - coords[4])
return (v1 + v2 + v3) / (2.0 * h1) if h1 > 0 else 0.0
class GeometricProcessor(BaseProcessor):
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)
self.counters = { "eye_closure": 0, "yawning": 0, "head_nod": 0, "looking_away": 0 }
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):
h, w, _ = frame.shape
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
brightness = np.mean(gray)
is_low_light = brightness < self.settings['low_light_thresh']
drowsiness_indicators = {
"drowsiness_level": "Awake", "lighting": "Good", "details": {}
}
face_landmarks = None
if is_low_light:
drowsiness_indicators["lighting"] = "Low"
else:
img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = self.face_mesh.process(img_rgb)
face_landmarks = results.multi_face_landmarks
if face_landmarks:
landmarks = face_landmarks[0].landmark
score = 0
weights = self.settings['indicator_weights']
# Eye Closure
ear = (calculate_ear([landmarks[i] for i in self.L_EYE],(h,w)) + calculate_ear([landmarks[i] for i in self.R_EYE],(h,w)))/2.0
if ear < self.settings['eye_ar_thresh']: self.counters['eye_closure']+=1
else: self.counters['eye_closure']=0
if self.counters['eye_closure'] >= self.settings['eye_ar_consec_frames']: score += weights['eye_closure']
# Yawning
mar = calculate_mar([landmarks[i] for i in self.MOUTH], (h, w))
if mar > self.settings['yawn_mar_thresh']: self.counters['yawning']+=1
else: self.counters['yawning']=0
if self.counters['yawning'] >= self.settings['yawn_consec_frames']: score += weights['yawning']
# Head Pose
face_3d = np.array([[0.0,0.0,0.0], [0.0,-330.0,-65.0], [-225.0,170.0,-135.0], [225.0,170.0,-135.0], [-150.0,-150.0,-125.0], [150.0,-150.0,-125.0]], dtype=np.float64)
face_2d = np.array([(landmarks[1].x*w, landmarks[1].y*h), (landmarks[152].x*w, landmarks[152].y*h), (landmarks[263].x*w, landmarks[263].y*h), (landmarks[33].x*w, landmarks[33].y*h), (landmarks[287].x*w, landmarks[287].y*h), (landmarks[57].x*w, landmarks[57].y*h)], 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]
if pitch > self.settings['head_nod_thresh']: self.counters['head_nod']+=1
else: self.counters['head_nod']=0
if self.counters['head_nod'] >= self.settings['head_pose_consec_frames']: score += weights['head_nod']
if abs(yaw) > self.settings['head_look_away_thresh']: self.counters['looking_away']+=1
else: self.counters['looking_away']=0
if self.counters['looking_away'] >= self.settings['head_pose_consec_frames']: score += weights['looking_away']
# Determine Drowsiness Level
levels = self.settings['drowsiness_levels']
if score >= levels['very_drowsy_threshold']: drowsiness_indicators['drowsiness_level'] = "Very Drowsy"
elif score >= levels['slightly_drowsy_threshold']: drowsiness_indicators['drowsiness_level'] = "Slightly Drowsy"
drowsiness_indicators['details']['Score'] = score
# --- Visualization on Video Frame ---
level = drowsiness_indicators['drowsiness_level']
score_val = drowsiness_indicators.get("details", {}).get("Score", 0)
color = (0, 255, 0) # Green for Awake
if drowsiness_indicators['lighting'] == "Low":
color = (0, 165, 255) # Orange for low light
cv2.putText(frame, "LOW LIGHT", (w // 2 - 120, h // 2), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3, cv2.LINE_AA)
elif level == "Slightly Drowsy":
color = (0, 255, 255) # Yellow
elif level == "Very Drowsy":
color = (0, 0, 255) # Red
# Draw a colored border around the frame
cv2.rectangle(frame, (0, 0), (w, h), color, 10)
# Display status text
status_text = f"Status: {level} (Score: {score_val:.2f})"
cv2.putText(frame, status_text, (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
return frame, drowsiness_indicators, face_landmarks |