Upload utils.py

utils.py

@@ -0,0 +1,168 @@
+import cv2
+import mediapipe as mp
+import numpy as np
+
+correct = cv2.imread('right.png')
+correct = cv2.cvtColor(correct, cv2.COLOR_BGR2RGB)
+incorrect = cv2.imread('wrong.png')
+incorrect = cv2.cvtColor(incorrect, cv2.COLOR_BGR2RGB)
+
+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 = 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]
+        # print("Height: ", resize_height)
+        # print("Width: ", rec_end[0] - rec_start[0])
+        img = draw_rounded_rect(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 = draw_rounded_rect(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_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