squats-analysis / utils.py
Kunal7's picture
Uploaded scripts and sample video
9456869
raw
history blame
3.49 kB
import cv2
import mediapipe as mp
import numpy as np
def draw_text(
img,
# count_text,
msg,
font=cv2.FONT_HERSHEY_SIMPLEX,
pos=(0, 0),
font_scale=1,
font_thickness=2,
text_color=(0, 255, 0),
text_color_bg=(0, 0, 0),
box_offset=(20, 10),
):
offset = box_offset
x, y = pos
text_size, _ = cv2.getTextSize(msg, font, font_scale, font_thickness)
text_w, text_h = text_size
rec_start = tuple(p - o for p, o in zip(pos, offset))
rec_end = tuple(m + n - o for m, n, o in zip((x + text_w, y + text_h), offset, (25, 0)))
cv2.rectangle(img, rec_start, rec_end, text_color_bg, -1)
cv2.putText(
img,
msg,
(int(rec_start[0] + 6), int(y + text_h + font_scale - 1)),
font,
font_scale,
text_color,
font_thickness,
cv2.LINE_AA,
)
return text_size
def find_angle(p1, p2, ref_pt = np.array([0,0])):
p1_ref = p1 - ref_pt
p2_ref = p2 - ref_pt
cos_theta = (np.dot(p1_ref,p2_ref)) / (1.0 * np.linalg.norm(p1_ref) * np.linalg.norm(p2_ref))
theta = np.arccos(np.clip(cos_theta, -1.0, 1.0))
degree = int(180 / np.pi) * theta
return int(degree)
def get_landmark_array(pose_landmark, key, frame_width, frame_height):
denorm_x = int(pose_landmark[key].x * frame_width)
denorm_y = int(pose_landmark[key].y * frame_height)
return np.array([denorm_x, denorm_y])
def get_landmark_features(kp_results, dict_features, feature, frame_width, frame_height):
if feature == 'nose':
return get_landmark_array(kp_results, dict_features[feature], frame_width, frame_height)
elif feature == 'left' or 'right':
shldr_coord = get_landmark_array(kp_results, dict_features[feature]['shoulder'], frame_width, frame_height)
elbow_coord = get_landmark_array(kp_results, dict_features[feature]['elbow'], frame_width, frame_height)
wrist_coord = get_landmark_array(kp_results, dict_features[feature]['wrist'], frame_width, frame_height)
hip_coord = get_landmark_array(kp_results, dict_features[feature]['hip'], frame_width, frame_height)
knee_coord = get_landmark_array(kp_results, dict_features[feature]['knee'], frame_width, frame_height)
ankle_coord = get_landmark_array(kp_results, dict_features[feature]['ankle'], frame_width, frame_height)
foot_coord = get_landmark_array(kp_results, dict_features[feature]['foot'], frame_width, frame_height)
return shldr_coord, elbow_coord, wrist_coord, hip_coord, knee_coord, ankle_coord, foot_coord
else:
raise ValueError("feature needs to be either 'nose', 'left' or 'right")
def get_mediapipe_pose(
static_image_mode = False,
model_complexity = 1,
smooth_landmarks = True,
min_detection_confidence = 0.5,
min_tracking_confidence = 0.5
):
pose = mp.solutions.pose.Pose(
static_image_mode = static_image_mode,
model_complexity = model_complexity,
smooth_landmarks = smooth_landmarks,
min_detection_confidence = min_detection_confidence,
min_tracking_confidence = min_tracking_confidence
)
return pose