Spaces:
Running
Running
import cv2 | |
import mediapipe as mp | |
import numpy as np | |
correct = cv2.imread('images/correct.png') | |
correct = cv2.cvtColor(correct, cv2.COLOR_BGR2RGB) | |
incorrect = cv2.imread('images/wrong.png') | |
incorrect = cv2.cvtColor(incorrect, cv2.COLOR_BGR2RGB) | |
def rectangle_box(img, rect_start, rect_end, corner_width, box_color): | |
x1, y1 = rect_start | |
x2, y2 = rect_end | |
w = corner_width | |
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) | |
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 dashed_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 = 7, | |
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), | |
overlay_image = False, | |
overlay_type = None | |
): | |
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))) | |
resize_height = 0 | |
if overlay_image: | |
resize_height = rec_end[1] - rec_start[1] | |
img = rectangle_box(img, rec_start, (rec_end[0]+resize_height, rec_end[1]), width, text_color_bg) | |
if overlay_type == "correct": | |
overlay_res = cv2.resize(correct, (resize_height, resize_height), interpolation = cv2.INTER_AREA) | |
elif overlay_type == "incorrect": | |
overlay_res = cv2.resize(incorrect, (resize_height, resize_height), interpolation = cv2.INTER_AREA) | |
img[rec_start[1]:rec_start[1]+resize_height, rec_start[0]+width:rec_start[0]+width+resize_height] = overlay_res | |
else: | |
img = rectangle_box(img, rec_start, rec_end, width, text_color_bg) | |
cv2.putText( | |
img, | |
msg, | |
(int(rec_start[0]+resize_height + 8), 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_imp_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 |