Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| import cv2 | |
| import numpy as np | |
| skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9], [8, 10], | |
| [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]] | |
| pose_kpt_color = [[51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0], | |
| [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0], | |
| [0, 255, 0], [255, 128, 0], [0, 255, 0], [255, 128, 0]] | |
| pose_link_color = [[0, 255, 0], [0, 255, 0], [255, 128, 0], [255, 128, 0], | |
| [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], [0, 255, 0], [255, 128, 0], | |
| [0, 255, 0], [255, 128, 0], [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255], | |
| [51, 153, 255], [51, 153, 255], [51, 153, 255]] | |
| def imshow_keypoints(img, | |
| pose_result, | |
| kpt_score_thr=0.1, | |
| radius=2, | |
| thickness=2): | |
| """Draw keypoints and links on an image. | |
| Args: | |
| img (ndarry): The image to draw poses on. | |
| pose_result (list[kpts]): The poses to draw. Each element kpts is | |
| a set of K keypoints as an Kx3 numpy.ndarray, where each | |
| keypoint is represented as x, y, score. | |
| kpt_score_thr (float, optional): Minimum score of keypoints | |
| to be shown. Default: 0.3. | |
| thickness (int): Thickness of lines. | |
| """ | |
| img_h, img_w, _ = img.shape | |
| img = np.zeros(img.shape) | |
| for idx, kpts in enumerate(pose_result): | |
| if idx > 1: | |
| continue | |
| kpts = kpts['keypoints'] | |
| # print(kpts) | |
| kpts = np.array(kpts, copy=False) | |
| # draw each point on image | |
| assert len(pose_kpt_color) == len(kpts) | |
| for kid, kpt in enumerate(kpts): | |
| x_coord, y_coord, kpt_score = int(kpt[0]), int(kpt[1]), kpt[2] | |
| if kpt_score < kpt_score_thr or pose_kpt_color[kid] is None: | |
| # skip the point that should not be drawn | |
| continue | |
| color = tuple(int(c) for c in pose_kpt_color[kid]) | |
| cv2.circle(img, (int(x_coord), int(y_coord)), radius, color, -1) | |
| # draw links | |
| for sk_id, sk in enumerate(skeleton): | |
| pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1])) | |
| pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1])) | |
| if (pos1[0] <= 0 or pos1[0] >= img_w or pos1[1] <= 0 or pos1[1] >= img_h or pos2[0] <= 0 | |
| or pos2[0] >= img_w or pos2[1] <= 0 or pos2[1] >= img_h or kpts[sk[0], 2] < kpt_score_thr | |
| or kpts[sk[1], 2] < kpt_score_thr or pose_link_color[sk_id] is None): | |
| # skip the link that should not be drawn | |
| continue | |
| color = tuple(int(c) for c in pose_link_color[sk_id]) | |
| cv2.line(img, pos1, pos2, color, thickness=thickness) | |
| return img | |