feat(utils): calculate arm angle to detect sign better

utils.py

@@ -1,9 +1,10 @@
 import cv2
-import numpy as np
 import torch
+import numpy as np
 from mediapipe.python.solutions import (drawing_styles, drawing_utils,
                                         holistic, pose)
 from torchvision.transforms.v2 import Compose, UniformTemporalSubsample
+from transformers import VideoMAEForVideoClassification
 
 
 def draw_skeleton_on_image(
@@ -11,7 +12,7 @@ def draw_skeleton_on_image(
     detection_results,
     resize_to: tuple[int, int] = None,
 ) -> np.ndarray:
-    """
+    '''
     Draw skeleton on the image.
 
     Parameters
@@ -27,7 +28,7 @@ def draw_skeleton_on_image(
     -------
     np.ndarray
         Annotated image with skeleton.
-    """
+    '''
     annotated_image = np.copy(image)
 
     # Draw pose connections
@@ -63,24 +64,69 @@ def draw_skeleton_on_image(
     return annotated_image
 
 
-def are_hands_down(pose_landmarks: list) -> bool:
-    """
+def calculate_angle(
+    shoulder: list,
+    elbow: list,
+    wrist: list,
+) -> float:
+    '''
+    Calculate the angle between the shoulder, elbow, and wrist.
+
+    Parameters
+    ----------
+    shoulder : list
+        Shoulder coordinates.
+    elbow : list
+        Elbow coordinates.
+    wrist : list
+        Wrist coordinates.
+
+    Returns
+    -------
+    float
+        Angle in degree between the shoulder, elbow, and wrist.
+    '''
+    shoulder = np.array(shoulder)
+    elbow = np.array(elbow)
+    wrist = np.array(wrist)
+
+    radians = np.arctan2(wrist[1] - elbow[1], wrist[0] - elbow[0]) \
+        - np.arctan2(shoulder[1] - elbow[1], shoulder[0] - elbow[0])
+    angle = np.abs(radians * 180.0 / np.pi)
+
+    if angle > 180.0:
+        angle = 360 - angle
+    return angle
+
+
+def do_hands_relax(
+    pose_landmarks: list,
+    angle_threshold: float = 160.0,
+) -> bool:
+    '''
     Check if the hand is down.
 
     Parameters
     ----------
     hand_landmarks : list
         Hand landmarks.
+    angle_threshold : float, optional
+        Angle threshold, by default 160.0.
 
     Returns
     -------
     bool
         True if the hand is down, False otherwise.
-    """
+    '''
     if pose_landmarks is None:
         return True
 
     landmarks = pose_landmarks.landmark
+    left_shoulder = [
+        landmarks[pose.PoseLandmark.LEFT_SHOULDER.value].x,
+        landmarks[pose.PoseLandmark.LEFT_SHOULDER.value].y,
+        landmarks[pose.PoseLandmark.LEFT_SHOULDER.value].visibility,
+    ]
     left_elbow = [
         landmarks[pose.PoseLandmark.LEFT_ELBOW.value].x,
         landmarks[pose.PoseLandmark.LEFT_ELBOW.value].y,
@@ -91,6 +137,13 @@ def are_hands_down(pose_landmarks: list) -> bool:
         landmarks[pose.PoseLandmark.LEFT_WRIST.value].y,
         landmarks[pose.PoseLandmark.LEFT_SHOULDER.value].visibility,
     ]
+    left_angle = calculate_angle(left_shoulder, left_elbow, left_wrist)
+
+    right_shoulder = [
+        landmarks[pose.PoseLandmark.RIGHT_SHOULDER.value].x,
+        landmarks[pose.PoseLandmark.RIGHT_SHOULDER.value].y,
+        landmarks[pose.PoseLandmark.RIGHT_SHOULDER.value].visibility,
+    ]
     right_elbow = [
         landmarks[pose.PoseLandmark.RIGHT_ELBOW.value].x,
         landmarks[pose.PoseLandmark.RIGHT_ELBOW.value].y,
@@ -101,18 +154,50 @@ def are_hands_down(pose_landmarks: list) -> bool:
         landmarks[pose.PoseLandmark.RIGHT_WRIST.value].y,
         landmarks[pose.PoseLandmark.RIGHT_SHOULDER.value].visibility,
     ]
+    right_angle = calculate_angle(right_shoulder, right_elbow, right_wrist)
 
     is_visible = all(
-        [left_elbow[2] > 0, left_wrist[2] > 0, right_elbow[2] > 0, right_wrist[2] > 0]
+        [
+            left_shoulder[2] > 0,
+            left_elbow[2] > 0,
+            left_wrist[2] > 0,
+            right_shoulder[2] > 0,
+            right_elbow[2] > 0,
+            right_wrist[2] > 0,
+        ]
+    )
+
+    return all(
+        [
+            is_visible,
+            left_angle < angle_threshold,
+            right_angle < angle_threshold,
+        ]
     )
-    return is_visible and left_wrist[1] > left_elbow[1] and right_wrist[1] > right_elbow[1]
 
 
 def get_predictions(
     inputs: dict,
-    model,
+    model: VideoMAEForVideoClassification,
     k: int = 3,
 ) -> list:
+    '''
+    Get the top-k predictions.
+
+    Parameters
+    ----------
+    inputs : dict
+        Model inputs.
+    model : VideoMAEForVideoClassification
+        Model to get predictions from.
+    k : int, optional
+        Number of predictions to return, by default 3.
+
+    Returns
+    -------
+    list
+        Top-k predictions.
+    '''
     if inputs is None:
         return []
 
@@ -143,6 +228,31 @@ def preprocess(
     device: str,
     transform: Compose,
 ) -> dict:
+    '''
+    Preprocess the video.
+
+    Parameters
+    ----------
+    model_num_frames : int
+        Number of frames in the model.
+    keypoints_detector
+        Keypoints detector.
+    source : str
+        Video source.
+    model_input_height : int
+        Model input height.
+    model_input_width : int
+        Model input width.
+    device : str
+        Device to use.
+    transform : Compose
+        Transform to apply.
+
+    Returns
+    -------
+    dict
+        Model inputs.
+    '''
     skeleton_video = []
     did_sample_start = False
 
@@ -165,7 +275,7 @@ def preprocess(
         skeleton_frame = transform(torch.tensor(skeleton_frame).permute(2, 0, 1))
 
         # Extract sign video.
-        if not are_hands_down(detection_results.pose_landmarks):
+        if not do_hands_relax(detection_results.pose_landmarks):
             if not did_sample_start:
                 did_sample_start = True
         elif did_sample_start: