Update utils.py

utils.py

@@ -2,6 +2,11 @@ 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
@@ -32,8 +37,6 @@ def draw_rounded_rect(img, rect_start, rect_end, corner_width, box_color):
     return img
 
 
-
-
 def draw_dotted_line(frame, lm_coord, start, end, line_color):
     pix_step = 0
 
@@ -46,7 +49,7 @@ def draw_dotted_line(frame, lm_coord, start, end, line_color):
 def draw_text(
     img,
     msg,
-    width = 8,
+    width = 7,
     font=cv2.FONT_HERSHEY_SIMPLEX,
     pos=(0, 0),
     font_scale=1,
@@ -54,35 +57,49 @@ def draw_text(
     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)))
-    
-    img = draw_rounded_rect(img, rec_start, rec_end, width, text_color_bg)
+
+    resize_height = 0
+
+    if overlay_image:
+        resize_height = rec_end[1] - rec_start[1]
+        
+        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] + 6), int(y + text_h + font_scale - 1)), 
+        (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
@@ -96,8 +113,6 @@ def find_angle(p1, p2, ref_pt = np.array([0,0])):
 
 
 
-
-
 def get_landmark_array(pose_landmark, key, frame_width, frame_height):
 
     denorm_x = int(pose_landmark[key].x * frame_width)
@@ -107,7 +122,6 @@ def get_landmark_array(pose_landmark, key, frame_width, frame_height):
 
 
 
-
 def get_landmark_features(kp_results, dict_features, feature, frame_width, frame_height):
 
     if feature == 'nose':
@@ -115,12 +129,12 @@ def get_landmark_features(kp_results, dict_features, feature, frame_width, frame
 
     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)
+        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)
+        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