Views
No views yet
1{
2 0: "Nose",
3 1: "L_Eye",
4 2: "R_Eye",
5 3: "L_Ear",
6 4: "R_Ear",
7 5: "L_Shoulder",
8 6: "R_Shoulder",
9 7: "L_Elbow",
10 8: "R_Elbow",
11 9: "L_Wrist",
12 10: "R_Wrist",
13 11: "L_Hip",
14 12: "R_Hip",
15 13: "L_Knee",
16 14: "R_Knee",
17 15: "L_Ankle",
18 16: "R_Ankle",
19 17: "sternum",
20 18: "rshoulder",
21 19: "lshoulder",
22 20: "r_lelbow",
23 21: "l_lelbow",
24 22: "r_melbow",
25 23: "l_melbow",
26 24: "r_lwrist",
27 25: "l_lwrist",
28 26: "r_mwrist",
29 27: "l_mwrist",
30 28: "r_ASIS",
31 29: "l_ASIS",
32 30: "r_PSIS",
33 31: "l_PSIS",
34 32: "r_knee",
35 33: "l_knee",
36 34: "r_mknee",
37 35: "l_mknee",
38 36: "r_ankle",
39 37: "l_ankle",
40 38: "r_mankle",
41 39: "l_mankle",
42 40: "r_5meta",
43 41: "l_5meta",
44 42: "r_toe",
45 43: "l_toe",
46 44: "r_big_toe",
47 45: "l_big_toe",
48 46: "l_calc",
49 47: "r_calc",
50 48: "C7",
51 49: "L2",
52 50: "T11",
53 51: "T6",
54}1import torch
2import requests
3import numpy as np
4
5from PIL import Image
6
7from transformers import (
8 AutoProcessor,
9 RTDetrForObjectDetection,
10 VitPoseForPoseEstimation,
11)
12
13device = "cuda" if torch.cuda.is_available() else "cpu"
14
15url = "http://farm4.staticflickr.com/3300/3416216247_f9c6dfc939_z.jpg"
16image = Image.open(requests.get(url, stream=True).raw)
17
18# ------------------------------------------------------------------------
19# Stage 1. Detect humans on the image
20# ------------------------------------------------------------------------
21
22# You can choose detector by your choice
23person_image_processor = AutoProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
24person_model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365", device_map=device)
25
26inputs = person_image_processor(images=image, return_tensors="pt").to(device)
27
28with torch.no_grad():
29 outputs = person_model(**inputs)
30
31results = person_image_processor.post_process_object_detection(
32 outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3
33)
34result = results[0] # take first image results
35
36# Human label refers 0 index in COCO dataset
37person_boxes = result["boxes"][result["labels"] == 0]
38person_boxes = person_boxes.cpu().numpy()
39
40# Convert boxes from VOC (x1, y1, x2, y2) to COCO (x1, y1, w, h) format
41person_boxes[:, 2] = person_boxes[:, 2] - person_boxes[:, 0]
42person_boxes[:, 3] = person_boxes[:, 3] - person_boxes[:, 1]
43
44# ------------------------------------------------------------------------
45# Stage 2. Detect keypoints for each person found
46# ------------------------------------------------------------------------
47
48image_processor = AutoProcessor.from_pretrained("yonigozlan/synthpose-vitpose-base-hf")
49model = VitPoseForPoseEstimation.from_pretrained("yonigozlan/synthpose-vitpose-base-hf", device_map=device)
50
51inputs = image_processor(image, boxes=[person_boxes], return_tensors="pt").to(device)
52
53with torch.no_grad():
54 outputs = model(**inputs)
55
56pose_results = image_processor.post_process_pose_estimation(outputs, boxes=[person_boxes])
57image_pose_result = pose_results[0] # results for first image1import supervision as sv
2
3xy = torch.stack([pose_result['keypoints'] for pose_result in image_pose_result]).cpu().numpy()
4scores = torch.stack([pose_result['scores'] for pose_result in image_pose_result]).cpu().numpy()
5
6key_points = sv.KeyPoints(
7 xy=xy, confidence=scores
8)
9
10vertex_annotator = sv.VertexAnnotator(
11 color=sv.Color.PINK,
12 radius=2
13)
14
15annotated_frame = vertex_annotator.annotate(
16 scene=image.copy(),
17 key_points=key_points
18)
19annotated_frame
1import math
2import cv2
3
4def draw_points(image, keypoints, scores, pose_keypoint_color, keypoint_score_threshold, radius, show_keypoint_weight):
5 if pose_keypoint_color is not None:
6 assert len(pose_keypoint_color) == len(keypoints)
7 for kid, (kpt, kpt_score) in enumerate(zip(keypoints, scores)):
8 x_coord, y_coord = int(kpt[0]), int(kpt[1])
9 if kpt_score > keypoint_score_threshold:
10 color = tuple(int(c) for c in pose_keypoint_color[kid])
11 if show_keypoint_weight:
12 cv2.circle(image, (int(x_coord), int(y_coord)), radius, color, -1)
13 transparency = max(0, min(1, kpt_score))
14 cv2.addWeighted(image, transparency, image, 1 - transparency, 0, dst=image)
15 else:
16 cv2.circle(image, (int(x_coord), int(y_coord)), radius, color, -1)
17
18def draw_links(image, keypoints, scores, keypoint_edges, link_colors, keypoint_score_threshold, thickness, show_keypoint_weight, stick_width = 2):
19 height, width, _ = image.shape
20 if keypoint_edges is not None and link_colors is not None:
21 assert len(link_colors) == len(keypoint_edges)
22 for sk_id, sk in enumerate(keypoint_edges):
23 x1, y1, score1 = (int(keypoints[sk[0], 0]), int(keypoints[sk[0], 1]), scores[sk[0]])
24 x2, y2, score2 = (int(keypoints[sk[1], 0]), int(keypoints[sk[1], 1]), scores[sk[1]])
25 if (
26 x1 > 0
27 and x1 < width
28 and y1 > 0
29 and y1 < height
30 and x2 > 0
31 and x2 < width
32 and y2 > 0
33 and y2 < height
34 and score1 > keypoint_score_threshold
35 and score2 > keypoint_score_threshold
36 ):
37 color = tuple(int(c) for c in link_colors[sk_id])
38 if show_keypoint_weight:
39 X = (x1, x2)
40 Y = (y1, y2)
41 mean_x = np.mean(X)
42 mean_y = np.mean(Y)
43 length = ((Y[0] - Y[1]) ** 2 + (X[0] - X[1]) ** 2) ** 0.5
44 angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1]))
45 polygon = cv2.ellipse2Poly(
46 (int(mean_x), int(mean_y)), (int(length / 2), int(stick_width)), int(angle), 0, 360, 1
47 )
48 cv2.fillConvexPoly(image, polygon, color)
49 transparency = max(0, min(1, 0.5 * (keypoints[sk[0], 2] + keypoints[sk[1], 2])))
50 cv2.addWeighted(image, transparency, image, 1 - transparency, 0, dst=image)
51 else:
52 cv2.line(image, (x1, y1), (x2, y2), color, thickness=thickness)
53
54
55# Note: keypoint_edges and color palette are dataset-specific
56keypoint_edges = model.config.edges
57
58palette = np.array(
59 [
60 [255, 128, 0],
61 [255, 153, 51],
62 [255, 178, 102],
63 [230, 230, 0],
64 [255, 153, 255],
65 [153, 204, 255],
66 [255, 102, 255],
67 [255, 51, 255],
68 [102, 178, 255],
69 [51, 153, 255],
70 [255, 153, 153],
71 [255, 102, 102],
72 [255, 51, 51],
73 [153, 255, 153],
74 [102, 255, 102],
75 [51, 255, 51],
76 [0, 255, 0],
77 [0, 0, 255],
78 [255, 0, 0],
79 [255, 255, 255],
80 ]
81)
82
83link_colors = palette[[0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16]]
84keypoint_colors = palette[[16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0]+[4]*(52-17)]
85
86numpy_image = np.array(image)
87
88for pose_result in image_pose_result:
89 scores = np.array(pose_result["scores"])
90 keypoints = np.array(pose_result["keypoints"])
91
92 # draw each point on image
93 draw_points(numpy_image, keypoints, scores, keypoint_colors, keypoint_score_threshold=0.3, radius=2, show_keypoint_weight=False)
94
95 # draw links
96 draw_links(numpy_image, keypoints, scores, keypoint_edges, link_colors, keypoint_score_threshold=0.3, thickness=1, show_keypoint_weight=False)
97
98pose_image = Image.fromarray(numpy_image)
99pose_image