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