full-body drawing

app.py

@@ -3,7 +3,7 @@ import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 from ultralytics import YOLO
 import torch
-import infer
+import utils
 
 
 @st.cache_resource()
@@ -30,7 +30,7 @@ def draw_output(image_pil: Image.Image, keypoints: dict):
     
     # draw extended left and right eye lines
     if ear and eye:
-        left_new_point = infer.extend_line(ear, eye, 3)
+        left_new_point = utils.extend_line(ear, eye, 3)
         l1 = [ear, left_new_point]
         draw.line(l1, fill='red', width=line_width)
         # draw a horizontal line from ear forwards
@@ -44,11 +44,11 @@ def draw_output(image_pil: Image.Image, keypoints: dict):
         l = [ear, tuple(p2.tolist())]
         draw.line(l, fill='gray', width=line_width//2)
         # draw angle
-        angle = infer.calculate_angle_to_horizontal(l1_vector)
+        angle = utils.calculate_angle_to_horizontal(l1_vector)
         draw.text(ear, f'{angle:.2f}', fill='red', font=font)
 
     # draw elbow angles
-    left_elbow_angle, right_elbow_angle = infer.get_elbow_angles(keypoints)
+    left_elbow_angle, right_elbow_angle = utils.get_elbow_angles(keypoints)
     if left_elbow_angle:
         draw.text(keypoints['left_elbow'], f'{left_elbow_angle:.2f}', fill='red', font=font)
         # draw polyline for left arm
@@ -84,12 +84,12 @@ if img is not None:
         model = load_model()
         pred = model(img)[0]
     st.markdown('**Results:**')
-    keypoints = infer.get_keypoints(pred)
+    keypoints = utils.get_keypoints(pred)
     if keypoints is not None:
         img = draw_output(img, keypoints)
         st.image(img, caption='Predicted image', use_column_width=True)
-        lea, rea = infer.get_eye_angles(keypoints)
-        lba, rba = infer.get_elbow_angles(keypoints)
+        lea, rea = utils.get_eye_angles(keypoints)
+        lba, rba = utils.get_elbow_angles(keypoints)
         st.write('Angles:')
         st.json({'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba})
         st.write('Raw keypoints:')

drawing.py

@@ -0,0 +1,109 @@
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import utils
+
+
+skeleton = np.array([
+    [16, 14],
+    [14, 12],
+    [17, 15],
+    [15, 13],
+    [12, 13],
+    [6, 12],
+    [7, 13],
+    [6, 7],
+    [6, 8],
+    [7, 9],
+    [8, 10],
+    [9, 11],
+    [2, 3],
+    [1, 2],
+    [1, 3],
+    [2, 4],
+    [3, 5],
+    [4, 6],
+    [5, 7],
+]) - 1
+
+
+def draw_keypoints(img: Image.Image, keypoints: dict):
+    draw = ImageDraw.Draw(img)
+
+    font_size = max(round(sum(img.size) / 2 * 0.035), 12)
+    font = ImageFont.truetype("DejaVuSerif.ttf", font_size)
+
+    color = (0, 255, 0)
+    lw = max(round(sum(img.size) / 2 * 0.005), 2)
+    lw_poly = int(round(0.75*lw))
+    r = max(round(sum(img.size) / 2 * 0.0132), 8)
+    text_margin = r
+
+    # draw skeleton
+    keypoint_list = list(keypoints.values())
+    for i, sk in enumerate(skeleton):
+        if keypoint_list[sk[0]] and keypoint_list[sk[1]]:
+            draw.line([keypoint_list[sk[0]], keypoint_list[sk[1]]], fill=color, width=lw)
+            # draw rectangle polygons around keypoint_list
+            poly1 = [keypoint_list[sk[0]][0] - r, keypoint_list[sk[0]][1] - r, keypoint_list[sk[0]][0] + r, keypoint_list[sk[0]][1] + r]
+            poly1 = [(poly1[0], poly1[1]), (poly1[0], poly1[3]), (poly1[2], poly1[3]), (poly1[2], poly1[1])]
+            draw.polygon(poly1, outline=color, width=lw_poly)
+            poly2 = [keypoint_list[sk[1]][0] - r, keypoint_list[sk[1]][1] - r, keypoint_list[sk[1]][0] + r, keypoint_list[sk[1]][1] + r]
+            poly2 = [(poly2[0], poly2[1]), (poly2[0], poly2[3]), (poly2[2], poly2[3]), (poly2[2], poly2[1])]
+            draw.polygon(poly2, outline=color, width=lw_poly)
+
+    # draw angles
+    ear, eye = None, None
+    is_left = None
+    if keypoints["left_ear"] and keypoints["left_eye"]:
+        ear = keypoints["left_ear"]
+        eye = keypoints["left_eye"]
+        is_left = True
+    elif keypoints["right_ear"] and keypoints["right_eye"]:
+        ear = keypoints["right_ear"]
+        eye = keypoints["right_eye"]
+        is_left = False
+
+    # draw extended left and right eye lines
+    if ear and eye:
+        left_new_point = utils.extend_line(ear, eye, 3)
+        l1 = [ear, left_new_point]
+        draw.line(l1, fill='red', width=lw)
+        # draw a horizontal line from ear forwards
+        ear = np.array(ear)
+        l1 = np.array(l1)
+        l1_vector = l1[1] - l1[0]
+        x_s = np.sign(l1_vector)[0]
+        length_l1 = np.linalg.norm(l1_vector)
+        p2 = ear + np.array([length_l1*x_s, 0])
+        ear = tuple(ear.tolist())
+        l = [ear, tuple(p2.tolist())]
+        # draw.line(l, fill='gray', width=lw//2)
+        # draw angle
+        angle = utils.calculate_angle_to_horizontal(l1_vector)
+        txt = f'{angle:.0f}°'
+        txt_w, txt_h = font.getbbox(txt)[2:]
+        if is_left:
+            text_coords = (ear[0] - txt_w - text_margin, ear[1] - txt_h - text_margin)
+        else:
+            text_coords = (ear[0] + text_margin, ear[1] - txt_h - text_margin)
+        draw.text(text_coords, txt, fill='red', font=font)
+
+    # draw elbow angles
+    left_elbow_angle, right_elbow_angle = utils.get_elbow_angles(keypoints)
+    if left_elbow_angle and is_left is not False:
+        txt = f'{left_elbow_angle:.0f}°'
+        txt_w, txt_h = font.getbbox(txt)[2:]
+        # adjust upwards and leftwards
+        text_coords = (keypoints['left_elbow'][0] - txt_w - text_margin, keypoints['left_elbow'][1] - txt_h - text_margin)
+        draw.text(text_coords, txt, fill='red', font=font)
+        # draw polyline for left arm
+        # draw.line([keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']], fill='blue', width=lw)
+    if right_elbow_angle and is_left is not True:
+        txt = f'{right_elbow_angle:.0f}°'
+        txt_w, txt_h = font.getbbox(txt)[2:]
+        text_coords = (keypoints['right_elbow'][0] + text_margin, keypoints['right_elbow'][1] - txt_h - text_margin)
+        draw.text(text_coords, txt, fill='red', font=font)
+        # draw polyline for right arm
+        # draw.line([keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']], fill='blue', width=lw)
+
+    return img