Spaces:
Sleeping
Sleeping
| 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 |