Spaces:
No application file
No application file
File size: 1,877 Bytes
030618c |
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 |
import cv2
import numpy as np
class PoseEstimator:
# Existing initialization and methods...
def rotation_matrix_to_angles(self, rotation_vector):
"""Convert rotation vector to pitch, yaw, and roll angles."""
rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
sy = np.sqrt(rotation_matrix[0, 0]**2 + rotation_matrix[1, 0]**2)
singular = sy < 1e-6
if not singular:
pitch = np.arctan2(rotation_matrix[2, 1], rotation_matrix[2, 2])
yaw = np.arctan2(-rotation_matrix[2, 0], sy)
roll = np.arctan2(rotation_matrix[1, 0], rotation_matrix[0, 0])
else:
pitch = np.arctan2(-rotation_matrix[1, 2], rotation_matrix[1, 1])
yaw = np.arctan2(-rotation_matrix[2, 0], sy)
roll = 0
return np.degrees(pitch), np.degrees(yaw), np.degrees(roll)
def is_distracted(self, rotation_vector):
"""Determine if the user is distracted based on head pose angles."""
pitch, yaw, roll = self.rotation_matrix_to_angles(rotation_vector)
# Define thresholds (you may need to adjust these based on testing)
pitch_threshold = 15 # Up/Down threshold
yaw_threshold = 20 # Left/Right threshold
roll_threshold = 10 # Tilt threshold
# Check if head is facing roughly forward
if abs(pitch) < pitch_threshold and abs(yaw) < yaw_threshold and abs(roll) < roll_threshold:
return False # Focused
else:
return True # Distracted
def detect_distraction(self, points):
"""Solve pose and detect distraction status based on pose."""
rotation_vector, translation_vector = self.solve(points)
distraction_status = self.is_distracted(rotation_vector)
return distraction_status, (rotation_vector, translation_vector)
|