Spaces:
Sleeping
Sleeping
File size: 4,157 Bytes
67fd17e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 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 |