Spaces:
Runtime error
Runtime error
File size: 4,844 Bytes
a7eca0b | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | #data/utils.py
import numpy as np
from mediapipe.python.solutions import pose
from visualization import draw_text_on_image
class Arm:
def __init__(
self,
side: str,
visibility: float = 0.5,
) -> None:
if side == "left":
self.shoulde_idx = pose.PoseLandmark.LEFT_SHOULDER.value
self.elbow_idx = pose.PoseLandmark.LEFT_ELBOW.value
self.wrist_idx = pose.PoseLandmark.LEFT_WRIST.value
elif side == "right":
self.shoulde_idx = pose.PoseLandmark.RIGHT_SHOULDER.value
self.elbow_idx = pose.PoseLandmark.RIGHT_ELBOW.value
self.wrist_idx = pose.PoseLandmark.RIGHT_WRIST.value
else:
raise ValueError("Side must be either 'left' or 'right'")
self.visibility = visibility
self.is_up = False
self.num_up_frames = 0
self.num_down_frames = 0
self.start_time = 0
self.end_time = 0
self.shoulder = None
self.elbow = None
self.wrist = None
self.angle = 0
def reset_state(self) -> None:
self.is_up = False
self.num_up_frames = 0
self.num_down_frames = 0
self.start_time = 0
self.end_time = 0
self.shoulder = None
self.elbow = None
self.wrist = None
self.angle = 0
def set_pose(self, landmarks) -> bool:
if landmarks[self.shoulde_idx].visibility < self.visibility:
return False
self.shoulder = (
landmarks[self.shoulde_idx].x,
landmarks[self.shoulde_idx].y,
)
if landmarks[self.elbow_idx].visibility < self.visibility:
return False
self.elbow = (
landmarks[self.elbow_idx].x,
landmarks[self.elbow_idx].y,
)
if landmarks[self.wrist_idx].visibility < self.visibility:
return False
self.wrist = (
landmarks[self.wrist_idx].x,
landmarks[self.wrist_idx].y,
)
self.angle = calculate_angle(self.shoulder, self.elbow, self.wrist)
return True
def visualize(
self,
frame: np.ndarray,
position: tuple = (20, 50),
prefix: str = "Angle",
color: tuple = (0, 0, 255),
) -> np.ndarray:
text = prefix + ": " + str(round(self.angle, 2))
return draw_text_on_image(
image=frame,
text=text,
position=position,
color=color,
font_size=20,
)
def get_sample_timestamp(left_arm: Arm, right_arm: Arm) -> tuple:
start_time, end_time = 0, 0
left_arm_available = left_arm.start_time > 0 and left_arm.end_time > 0
right_arm_available = right_arm.start_time > 0 and right_arm.end_time > 0
if left_arm_available and right_arm.start_time == 0:
start_time = left_arm.start_time
end_time = left_arm.end_time
if right_arm_available and left_arm.start_time == 0:
start_time = right_arm.start_time
end_time = right_arm.end_time
if all((
left_arm_available, not left_arm.is_up,
right_arm_available, not right_arm.is_up,
)):
start_time = min(left_arm.start_time, right_arm.start_time)
end_time = max(left_arm.end_time, right_arm.end_time)
# Convert seconds to milliseconds
start_time /= 1000
end_time /= 1000
return start_time, end_time
def calculate_angle(a: tuple, b: tuple, c: tuple) -> float:
a = np.array(a) # First
b = np.array(b) # Mid
c = np.array(c) # End
radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
angle = np.abs(radians * 180.0 / np.pi)
return 360 - angle if angle > 180 else angle
def ok_to_get_frame(
arm: Arm,
angle_threshold: int,
min_num_up_frames: int,
min_num_down_frames: int,
current_time: int,
delay: int,
) -> bool:
if 0 < arm.angle < angle_threshold:
if arm.is_up:
arm.num_down_frames = 0
arm.end_time = 0
else:
if arm.num_up_frames == min_num_up_frames:
arm.is_up = True
arm.num_up_frames = 0
else:
if arm.num_up_frames == 0:
arm.start_time = current_time - delay
arm.num_up_frames += 1
return False
else:
if arm.is_up:
if arm.num_down_frames == min_num_down_frames:
arm.is_up = False
arm.num_down_frames = 0
else:
if arm.num_down_frames == 0:
arm.end_time = current_time + delay
arm.num_down_frames += 1
return True
else:
arm.num_up_frames = 0
arm.start_time = 0
return arm.is_up
|