Spaces:
Runtime error
Runtime error
File size: 4,881 Bytes
41b615f |
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 |
import cv2
import mediapipe as mp
import numpy as np
def draw_rounded_rect(img, rect_start, rect_end, corner_width, box_color):
x1, y1 = rect_start
x2, y2 = rect_end
w = corner_width
# draw filled rectangles
cv2.rectangle(img, (x1 + w, y1), (x2 - w, y1 + w), box_color, -1)
cv2.rectangle(img, (x1 + w, y2 - w), (x2 - w, y2), box_color, -1)
cv2.rectangle(img, (x1, y1 + w), (x1 + w, y2 - w), box_color, -1)
cv2.rectangle(img, (x2 - w, y1 + w), (x2, y2 - w), box_color, -1)
cv2.rectangle(img, (x1 + w, y1 + w), (x2 - w, y2 - w), box_color, -1)
# draw filled ellipses
cv2.ellipse(img, (x1 + w, y1 + w), (w, w),
angle = 0, startAngle = -90, endAngle = -180, color = box_color, thickness = -1)
cv2.ellipse(img, (x2 - w, y1 + w), (w, w),
angle = 0, startAngle = 0, endAngle = -90, color = box_color, thickness = -1)
cv2.ellipse(img, (x1 + w, y2 - w), (w, w),
angle = 0, startAngle = 90, endAngle = 180, color = box_color, thickness = -1)
cv2.ellipse(img, (x2 - w, y2 - w), (w, w),
angle = 0, startAngle = 0, endAngle = 90, color = box_color, thickness = -1)
return img
def draw_dotted_line(frame, lm_coord, start, end, line_color):
pix_step = 0
for i in range(start, end+1, 8):
cv2.circle(frame, (lm_coord[0], i+pix_step), 2, line_color, -1, lineType=cv2.LINE_AA)
return frame
def draw_text(
img,
msg,
width = 8,
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)))
img = draw_rounded_rect(img, rec_start, rec_end, width, text_color_bg)
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 |