Spaces:
Runtime error
Runtime error
File size: 6,372 Bytes
19f420a 929f736 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a f37553c 19f420a 9ae3b11 d78d683 9ae3b11 19f420a f37553c 19f420a f37553c 19f420a f37553c 929f736 f37553c 19f420a f37553c 19f420a 9ae3b11 929f736 f37553c 19f420a f37553c 19f420a 929f736 f37553c 929f736 f37553c 26440c6 f37553c 929f736 f37553c 19f420a f37553c 19f420a 929f736 |
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 |
# 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 (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):
# Create a writable copy to prevent read-only errors from Gradio/OpenCV
frame = frame.copy()
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 not is_low_light:
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']
# --- Draw Facial Landmarks (Logic Added Back) ---
# This will draw the green dots for eyes and mouth to show what is being tracked.
eye_mouth_landmarks_indices = self.L_EYE + self.R_EYE + self.MOUTH
for idx in eye_mouth_landmarks_indices:
lm = landmarks[idx]
x, y = int(lm.x * w), int(lm.y * h)
cv2.circle(frame, (x, y), 1, (0, 255, 0), -1)
# --- Drowsiness Calculations ---
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']
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']
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']
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
else: # is_low_light is True
drowsiness_indicators["lighting"] = "Low"
# --- 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
cv2.rectangle(frame, (0, 0), (w, h), color, 10)
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
|