stanfordmimi/Synthpose-Markerless-MoCap-VitPose

1### Visualization for advanced user2import math3 4import cv25import numpy as np6 7 8def draw_points(9    image,10    keypoints,11    scores,12    pose_keypoint_color,13    keypoint_score_threshold,14    radius,15    show_keypoint_weight,16):17    if pose_keypoint_color is not None:18        assert len(pose_keypoint_color) == len(keypoints)19    for kid, (kpt, kpt_score) in enumerate(zip(keypoints, scores)):20        x_coord, y_coord = int(kpt[0]), int(kpt[1])21        if kpt_score > keypoint_score_threshold:22            color = tuple(int(c) for c in pose_keypoint_color[kid])23            if show_keypoint_weight:24                cv2.circle(image, (int(x_coord), int(y_coord)), radius, color, -1)25                transparency = max(0, min(1, kpt_score))26                cv2.addWeighted(27                    image, transparency, image, 1 - transparency, 0, dst=image28                )29            else:30                cv2.circle(image, (int(x_coord), int(y_coord)), radius, color, -1)31 32 33def draw_links(34    image,35    keypoints,36    scores,37    keypoint_edges,38    link_colors,39    keypoint_score_threshold,40    thickness,41    show_keypoint_weight,42    stick_width=2,43):44    height, width, _ = image.shape45    if keypoint_edges is not None and link_colors is not None:46        assert len(link_colors) == len(keypoint_edges)47        for sk_id, sk in enumerate(keypoint_edges):48            x1, y1, score1 = (49                int(keypoints[sk[0], 0]),50                int(keypoints[sk[0], 1]),51                scores[sk[0]],52            )53            x2, y2, score2 = (54                int(keypoints[sk[1], 0]),55                int(keypoints[sk[1], 1]),56                scores[sk[1]],57            )58            if (59                x1 > 060                and x1 < width61                and y1 > 062                and y1 < height63                and x2 > 064                and x2 < width65                and y2 > 066                and y2 < height67                and score1 > keypoint_score_threshold68                and score2 > keypoint_score_threshold69            ):70                color = tuple(int(c) for c in link_colors[sk_id])71                if show_keypoint_weight:72                    X = (x1, x2)73                    Y = (y1, y2)74                    mean_x = np.mean(X)75                    mean_y = np.mean(Y)76                    length = ((Y[0] - Y[1]) ** 2 + (X[0] - X[1]) ** 2) ** 0.577                    angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1]))78                    polygon = cv2.ellipse2Poly(79                        (int(mean_x), int(mean_y)),80                        (int(length / 2), int(stick_width)),81                        int(angle),82                        0,83                        360,84                        1,85                    )86                    cv2.fillConvexPoly(image, polygon, color)87                    transparency = max(88                        0, min(1, 0.5 * (keypoints[sk[0], 2] + keypoints[sk[1], 2]))89                    )90                    cv2.addWeighted(91                        image, transparency, image, 1 - transparency, 0, dst=image92                    )93                else:94                    cv2.line(image, (x1, y1), (x2, y2), color, thickness=thickness)95 96 97palette = np.array(98    [99        [255, 128, 0],100        [255, 153, 51],101        [255, 178, 102],102        [230, 230, 0],103        [255, 153, 255],104        [153, 204, 255],105        [255, 102, 255],106        [255, 51, 255],107        [102, 178, 255],108        [51, 153, 255],109        [255, 153, 153],110        [255, 102, 102],111        [255, 51, 51],112        [153, 255, 153],113        [102, 255, 102],114        [51, 255, 51],115        [0, 255, 0],116        [0, 0, 255],117        [255, 0, 0],118        [255, 255, 255],119    ]120)121 122link_colors = palette[[0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16]]123keypoint_colors = palette[124    [16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0] + [4] * (52 - 17)125]126