Spaces:
Running
Running
# Useful visualization functions from SLAHMR and WHAM | |
import os | |
import cv2 | |
import numpy as np | |
import torch | |
from PIL import Image | |
def read_image(path, scale=1): | |
im = Image.open(path) | |
if scale == 1: | |
return np.array(im) | |
W, H = im.size | |
w, h = int(scale * W), int(scale * H) | |
return np.array(im.resize((w, h), Image.ANTIALIAS)) | |
def transform_torch3d(T_c2w): | |
""" | |
:param T_c2w (*, 4, 4) | |
returns (*, 3, 3), (*, 3) | |
""" | |
R1 = torch.tensor( | |
[[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, 1.0],], device=T_c2w.device, | |
) | |
R2 = torch.tensor( | |
[[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0],], device=T_c2w.device, | |
) | |
cam_R, cam_t = T_c2w[..., :3, :3], T_c2w[..., :3, 3] | |
cam_R = torch.einsum("...ij,jk->...ik", cam_R, R1) | |
cam_t = torch.einsum("ij,...j->...i", R2, cam_t) | |
return cam_R, cam_t | |
def transform_pyrender(T_c2w): | |
""" | |
:param T_c2w (*, 4, 4) | |
""" | |
T_vis = torch.tensor( | |
[ | |
[1.0, 0.0, 0.0, 0.0], | |
[0.0, -1.0, 0.0, 0.0], | |
[0.0, 0.0, -1.0, 0.0], | |
[0.0, 0.0, 0.0, 1.0], | |
], | |
device=T_c2w.device, | |
) | |
return torch.einsum( | |
"...ij,jk->...ik", torch.einsum("ij,...jk->...ik", T_vis, T_c2w), T_vis | |
) | |
def smpl_to_geometry(verts, faces, vis_mask=None, track_ids=None): | |
""" | |
:param verts (B, T, V, 3) | |
:param faces (F, 3) | |
:param vis_mask (optional) (B, T) visibility of each person | |
:param track_ids (optional) (B,) | |
returns list of T verts (B, V, 3), faces (F, 3), colors (B, 3) | |
where B is different depending on the visibility of the people | |
""" | |
B, T = verts.shape[:2] | |
device = verts.device | |
# (B, 3) | |
colors = ( | |
track_to_colors(track_ids) | |
if track_ids is not None | |
else torch.ones(B, 3, device) * 0.5 | |
) | |
# list T (B, V, 3), T (B, 3), T (F, 3) | |
return filter_visible_meshes(verts, colors, faces, vis_mask) | |
def filter_visible_meshes(verts, colors, faces, vis_mask=None, vis_opacity=False): | |
""" | |
:param verts (B, T, V, 3) | |
:param colors (B, 3) | |
:param faces (F, 3) | |
:param vis_mask (optional tensor, default None) (B, T) ternary mask | |
-1 if not in frame | |
0 if temporarily occluded | |
1 if visible | |
:param vis_opacity (optional bool, default False) | |
if True, make occluded people alpha=0.5, otherwise alpha=1 | |
returns a list of T lists verts (Bi, V, 3), colors (Bi, 4), faces (F, 3) | |
""" | |
# import ipdb; ipdb.set_trace() | |
B, T = verts.shape[:2] | |
faces = [faces for t in range(T)] | |
if vis_mask is None: | |
verts = [verts[:, t] for t in range(T)] | |
colors = [colors for t in range(T)] | |
return verts, colors, faces | |
# render occluded and visible, but not removed | |
vis_mask = vis_mask >= 0 | |
if vis_opacity: | |
alpha = 0.5 * (vis_mask[..., None] + 1) | |
else: | |
alpha = (vis_mask[..., None] >= 0).float() | |
vert_list = [verts[vis_mask[:, t], t] for t in range(T)] | |
colors = [ | |
torch.cat([colors[vis_mask[:, t]], alpha[vis_mask[:, t], t]], dim=-1) | |
for t in range(T) | |
] | |
bounds = get_bboxes(verts, vis_mask) | |
return vert_list, colors, faces, bounds | |
def get_bboxes(verts, vis_mask): | |
""" | |
return bb_min, bb_max, and mean for each track (B, 3) over entire trajectory | |
:param verts (B, T, V, 3) | |
:param vis_mask (B, T) | |
""" | |
B, T, *_ = verts.shape | |
bb_min, bb_max, mean = [], [], [] | |
for b in range(B): | |
v = verts[b, vis_mask[b, :T]] # (Tb, V, 3) | |
bb_min.append(v.amin(dim=(0, 1))) | |
bb_max.append(v.amax(dim=(0, 1))) | |
mean.append(v.mean(dim=(0, 1))) | |
bb_min = torch.stack(bb_min, dim=0) | |
bb_max = torch.stack(bb_max, dim=0) | |
mean = torch.stack(mean, dim=0) | |
# point to a track that's long and close to the camera | |
zs = mean[:, 2] | |
counts = vis_mask[:, :T].sum(dim=-1) # (B,) | |
mask = counts < 0.8 * T | |
zs[mask] = torch.inf | |
sel = torch.argmin(zs) | |
return bb_min.amin(dim=0), bb_max.amax(dim=0), mean[sel] | |
def track_to_colors(track_ids): | |
""" | |
:param track_ids (B) | |
""" | |
color_map = torch.from_numpy(get_colors()).to(track_ids) | |
return color_map[track_ids] / 255 # (B, 3) | |
def get_colors(): | |
# color_file = os.path.abspath(os.path.join(__file__, "../colors_phalp.txt")) | |
color_file = os.path.abspath(os.path.join(__file__, "../colors.txt")) | |
RGB_tuples = np.vstack( | |
[ | |
np.loadtxt(color_file, skiprows=0), | |
# np.loadtxt(color_file, skiprows=1), | |
np.random.uniform(0, 255, size=(10000, 3)), | |
[[0, 0, 0]], | |
] | |
) | |
b = np.where(RGB_tuples == 0) | |
RGB_tuples[b] = 1 | |
return RGB_tuples.astype(np.float32) | |
def checkerboard_geometry( | |
length=12.0, | |
color0=[0.8, 0.9, 0.9], | |
color1=[0.6, 0.7, 0.7], | |
tile_width=0.5, | |
alpha=1.0, | |
up="y", | |
c1=0.0, | |
c2=0.0, | |
): | |
assert up == "y" or up == "z" | |
color0 = np.array(color0 + [alpha]) | |
color1 = np.array(color1 + [alpha]) | |
radius = length / 2.0 | |
num_rows = num_cols = max(2, int(length / tile_width)) | |
vertices = [] | |
vert_colors = [] | |
faces = [] | |
face_colors = [] | |
for i in range(num_rows): | |
for j in range(num_cols): | |
u0, v0 = j * tile_width - radius, i * tile_width - radius | |
us = np.array([u0, u0, u0 + tile_width, u0 + tile_width]) | |
vs = np.array([v0, v0 + tile_width, v0 + tile_width, v0]) | |
zs = np.zeros(4) | |
if up == "y": | |
cur_verts = np.stack([us, zs, vs], axis=-1) # (4, 3) | |
cur_verts[:, 0] += c1 | |
cur_verts[:, 2] += c2 | |
else: | |
cur_verts = np.stack([us, vs, zs], axis=-1) # (4, 3) | |
cur_verts[:, 0] += c1 | |
cur_verts[:, 1] += c2 | |
cur_faces = np.array( | |
[[0, 1, 3], [1, 2, 3], [0, 3, 1], [1, 3, 2]], dtype=np.int64 | |
) | |
cur_faces += 4 * (i * num_cols + j) # the number of previously added verts | |
use_color0 = (i % 2 == 0 and j % 2 == 0) or (i % 2 == 1 and j % 2 == 1) | |
cur_color = color0 if use_color0 else color1 | |
cur_colors = np.array([cur_color, cur_color, cur_color, cur_color]) | |
vertices.append(cur_verts) | |
faces.append(cur_faces) | |
vert_colors.append(cur_colors) | |
face_colors.append(cur_colors) | |
vertices = np.concatenate(vertices, axis=0).astype(np.float32) | |
vert_colors = np.concatenate(vert_colors, axis=0).astype(np.float32) | |
faces = np.concatenate(faces, axis=0).astype(np.float32) | |
face_colors = np.concatenate(face_colors, axis=0).astype(np.float32) | |
return vertices, faces, vert_colors, face_colors | |
def camera_marker_geometry(radius, height, up): | |
assert up == "y" or up == "z" | |
if up == "y": | |
vertices = np.array( | |
[ | |
[-radius, -radius, 0], | |
[radius, -radius, 0], | |
[radius, radius, 0], | |
[-radius, radius, 0], | |
[0, 0, height], | |
] | |
) | |
else: | |
vertices = np.array( | |
[ | |
[-radius, 0, -radius], | |
[radius, 0, -radius], | |
[radius, 0, radius], | |
[-radius, 0, radius], | |
[0, -height, 0], | |
] | |
) | |
faces = np.array( | |
[[0, 3, 1], [1, 3, 2], [0, 1, 4], [1, 2, 4], [2, 3, 4], [3, 0, 4],] | |
) | |
face_colors = np.array( | |
[ | |
[1.0, 1.0, 1.0, 1.0], | |
[1.0, 1.0, 1.0, 1.0], | |
[0.0, 1.0, 0.0, 1.0], | |
[1.0, 0.0, 0.0, 1.0], | |
[0.0, 1.0, 0.0, 1.0], | |
[1.0, 0.0, 0.0, 1.0], | |
] | |
) | |
return vertices, faces, face_colors | |
def vis_keypoints( | |
keypts_list, | |
img_size, | |
radius=6, | |
thickness=3, | |
kpt_score_thr=0.3, | |
dataset="TopDownCocoDataset", | |
): | |
""" | |
Visualize keypoints | |
From ViTPose/mmpose/apis/inference.py | |
""" | |
palette = np.array( | |
[ | |
[255, 128, 0], | |
[255, 153, 51], | |
[255, 178, 102], | |
[230, 230, 0], | |
[255, 153, 255], | |
[153, 204, 255], | |
[255, 102, 255], | |
[255, 51, 255], | |
[102, 178, 255], | |
[51, 153, 255], | |
[255, 153, 153], | |
[255, 102, 102], | |
[255, 51, 51], | |
[153, 255, 153], | |
[102, 255, 102], | |
[51, 255, 51], | |
[0, 255, 0], | |
[0, 0, 255], | |
[255, 0, 0], | |
[255, 255, 255], | |
] | |
) | |
if dataset in ( | |
"TopDownCocoDataset", | |
"BottomUpCocoDataset", | |
"TopDownOCHumanDataset", | |
"AnimalMacaqueDataset", | |
): | |
# show the results | |
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_link_color = palette[ | |
[0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16] | |
] | |
pose_kpt_color = palette[ | |
[16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0] | |
] | |
elif dataset == "TopDownCocoWholeBodyDataset": | |
# show the results | |
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], | |
[15, 17], | |
[15, 18], | |
[15, 19], | |
[16, 20], | |
[16, 21], | |
[16, 22], | |
[91, 92], | |
[92, 93], | |
[93, 94], | |
[94, 95], | |
[91, 96], | |
[96, 97], | |
[97, 98], | |
[98, 99], | |
[91, 100], | |
[100, 101], | |
[101, 102], | |
[102, 103], | |
[91, 104], | |
[104, 105], | |
[105, 106], | |
[106, 107], | |
[91, 108], | |
[108, 109], | |
[109, 110], | |
[110, 111], | |
[112, 113], | |
[113, 114], | |
[114, 115], | |
[115, 116], | |
[112, 117], | |
[117, 118], | |
[118, 119], | |
[119, 120], | |
[112, 121], | |
[121, 122], | |
[122, 123], | |
[123, 124], | |
[112, 125], | |
[125, 126], | |
[126, 127], | |
[127, 128], | |
[112, 129], | |
[129, 130], | |
[130, 131], | |
[131, 132], | |
] | |
pose_link_color = palette[ | |
[0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16] | |
+ [16, 16, 16, 16, 16, 16] | |
+ [0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16] | |
+ [0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16] | |
] | |
pose_kpt_color = palette[ | |
[16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0] | |
+ [0, 0, 0, 0, 0, 0] | |
+ [19] * (68 + 42) | |
] | |
elif dataset == "TopDownAicDataset": | |
skeleton = [ | |
[2, 1], | |
[1, 0], | |
[0, 13], | |
[13, 3], | |
[3, 4], | |
[4, 5], | |
[8, 7], | |
[7, 6], | |
[6, 9], | |
[9, 10], | |
[10, 11], | |
[12, 13], | |
[0, 6], | |
[3, 9], | |
] | |
pose_link_color = palette[[9, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 0, 7, 7]] | |
pose_kpt_color = palette[[9, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 0, 0]] | |
elif dataset == "TopDownMpiiDataset": | |
skeleton = [ | |
[0, 1], | |
[1, 2], | |
[2, 6], | |
[6, 3], | |
[3, 4], | |
[4, 5], | |
[6, 7], | |
[7, 8], | |
[8, 9], | |
[8, 12], | |
[12, 11], | |
[11, 10], | |
[8, 13], | |
[13, 14], | |
[14, 15], | |
] | |
pose_link_color = palette[[16, 16, 16, 16, 16, 16, 7, 7, 0, 9, 9, 9, 9, 9, 9]] | |
pose_kpt_color = palette[[16, 16, 16, 16, 16, 16, 7, 7, 0, 0, 9, 9, 9, 9, 9, 9]] | |
elif dataset == "TopDownMpiiTrbDataset": | |
skeleton = [ | |
[12, 13], | |
[13, 0], | |
[13, 1], | |
[0, 2], | |
[1, 3], | |
[2, 4], | |
[3, 5], | |
[0, 6], | |
[1, 7], | |
[6, 7], | |
[6, 8], | |
[7, 9], | |
[8, 10], | |
[9, 11], | |
[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], | |
] | |
pose_link_color = palette[[16] * 14 + [19] * 13] | |
pose_kpt_color = palette[[16] * 14 + [0] * 26] | |
elif dataset in ("OneHand10KDataset", "FreiHandDataset", "PanopticDataset"): | |
skeleton = [ | |
[0, 1], | |
[1, 2], | |
[2, 3], | |
[3, 4], | |
[0, 5], | |
[5, 6], | |
[6, 7], | |
[7, 8], | |
[0, 9], | |
[9, 10], | |
[10, 11], | |
[11, 12], | |
[0, 13], | |
[13, 14], | |
[14, 15], | |
[15, 16], | |
[0, 17], | |
[17, 18], | |
[18, 19], | |
[19, 20], | |
] | |
pose_link_color = palette[ | |
[0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16] | |
] | |
pose_kpt_color = palette[ | |
[0, 0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16] | |
] | |
elif dataset == "InterHand2DDataset": | |
skeleton = [ | |
[0, 1], | |
[1, 2], | |
[2, 3], | |
[4, 5], | |
[5, 6], | |
[6, 7], | |
[8, 9], | |
[9, 10], | |
[10, 11], | |
[12, 13], | |
[13, 14], | |
[14, 15], | |
[16, 17], | |
[17, 18], | |
[18, 19], | |
[3, 20], | |
[7, 20], | |
[11, 20], | |
[15, 20], | |
[19, 20], | |
] | |
pose_link_color = palette[ | |
[0, 0, 0, 4, 4, 4, 8, 8, 8, 12, 12, 12, 16, 16, 16, 0, 4, 8, 12, 16] | |
] | |
pose_kpt_color = palette[ | |
[0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16, 0] | |
] | |
elif dataset == "Face300WDataset": | |
# show the results | |
skeleton = [] | |
pose_link_color = palette[[]] | |
pose_kpt_color = palette[[19] * 68] | |
kpt_score_thr = 0 | |
elif dataset == "FaceAFLWDataset": | |
# show the results | |
skeleton = [] | |
pose_link_color = palette[[]] | |
pose_kpt_color = palette[[19] * 19] | |
kpt_score_thr = 0 | |
elif dataset == "FaceCOFWDataset": | |
# show the results | |
skeleton = [] | |
pose_link_color = palette[[]] | |
pose_kpt_color = palette[[19] * 29] | |
kpt_score_thr = 0 | |
elif dataset == "FaceWFLWDataset": | |
# show the results | |
skeleton = [] | |
pose_link_color = palette[[]] | |
pose_kpt_color = palette[[19] * 98] | |
kpt_score_thr = 0 | |
elif dataset == "AnimalHorse10Dataset": | |
skeleton = [ | |
[0, 1], | |
[1, 12], | |
[12, 16], | |
[16, 21], | |
[21, 17], | |
[17, 11], | |
[11, 10], | |
[10, 8], | |
[8, 9], | |
[9, 12], | |
[2, 3], | |
[3, 4], | |
[5, 6], | |
[6, 7], | |
[13, 14], | |
[14, 15], | |
[18, 19], | |
[19, 20], | |
] | |
pose_link_color = palette[[4] * 10 + [6] * 2 + [6] * 2 + [7] * 2 + [7] * 2] | |
pose_kpt_color = palette[ | |
[4, 4, 6, 6, 6, 6, 6, 6, 4, 4, 4, 4, 4, 7, 7, 7, 4, 4, 7, 7, 7, 4] | |
] | |
elif dataset == "AnimalFlyDataset": | |
skeleton = [ | |
[1, 0], | |
[2, 0], | |
[3, 0], | |
[4, 3], | |
[5, 4], | |
[7, 6], | |
[8, 7], | |
[9, 8], | |
[11, 10], | |
[12, 11], | |
[13, 12], | |
[15, 14], | |
[16, 15], | |
[17, 16], | |
[19, 18], | |
[20, 19], | |
[21, 20], | |
[23, 22], | |
[24, 23], | |
[25, 24], | |
[27, 26], | |
[28, 27], | |
[29, 28], | |
[30, 3], | |
[31, 3], | |
] | |
pose_link_color = palette[[0] * 25] | |
pose_kpt_color = palette[[0] * 32] | |
elif dataset == "AnimalLocustDataset": | |
skeleton = [ | |
[1, 0], | |
[2, 1], | |
[3, 2], | |
[4, 3], | |
[6, 5], | |
[7, 6], | |
[9, 8], | |
[10, 9], | |
[11, 10], | |
[13, 12], | |
[14, 13], | |
[15, 14], | |
[17, 16], | |
[18, 17], | |
[19, 18], | |
[21, 20], | |
[22, 21], | |
[24, 23], | |
[25, 24], | |
[26, 25], | |
[28, 27], | |
[29, 28], | |
[30, 29], | |
[32, 31], | |
[33, 32], | |
[34, 33], | |
] | |
pose_link_color = palette[[0] * 26] | |
pose_kpt_color = palette[[0] * 35] | |
elif dataset == "AnimalZebraDataset": | |
skeleton = [[1, 0], [2, 1], [3, 2], [4, 2], [5, 7], [6, 7], [7, 2], [8, 7]] | |
pose_link_color = palette[[0] * 8] | |
pose_kpt_color = palette[[0] * 9] | |
elif dataset in "AnimalPoseDataset": | |
skeleton = [ | |
[0, 1], | |
[0, 2], | |
[1, 3], | |
[0, 4], | |
[1, 4], | |
[4, 5], | |
[5, 7], | |
[6, 7], | |
[5, 8], | |
[8, 12], | |
[12, 16], | |
[5, 9], | |
[9, 13], | |
[13, 17], | |
[6, 10], | |
[10, 14], | |
[14, 18], | |
[6, 11], | |
[11, 15], | |
[15, 19], | |
] | |
pose_link_color = palette[[0] * 20] | |
pose_kpt_color = palette[[0] * 20] | |
else: | |
NotImplementedError() | |
img_w, img_h = img_size | |
img = 255 * np.ones((img_h, img_w, 3), dtype=np.uint8) | |
img = imshow_keypoints( | |
img, | |
keypts_list, | |
skeleton, | |
kpt_score_thr, | |
pose_kpt_color, | |
pose_link_color, | |
radius, | |
thickness, | |
) | |
alpha = 255 * (img != 255).any(axis=-1, keepdims=True).astype(np.uint8) | |
return np.concatenate([img, alpha], axis=-1) | |
def imshow_keypoints( | |
img, | |
pose_result, | |
skeleton=None, | |
kpt_score_thr=0.3, | |
pose_kpt_color=None, | |
pose_link_color=None, | |
radius=4, | |
thickness=1, | |
show_keypoint_weight=False, | |
): | |
"""Draw keypoints and links on an image. | |
From ViTPose/mmpose/core/visualization/image.py | |
Args: | |
img (H, W, 3) array | |
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. | |
pose_kpt_color (np.array[Nx3]`): Color of N keypoints. If None, | |
the keypoint will not be drawn. | |
pose_link_color (np.array[Mx3]): Color of M links. If None, the | |
links will not be drawn. | |
thickness (int): Thickness of lines. | |
show_keypoint_weight (bool): If True, opacity indicates keypoint score | |
""" | |
import math | |
img_h, img_w, _ = img.shape | |
idcs = [0, 16, 15, 18, 17, 5, 2, 6, 3, 7, 4, 12, 9, 13, 10, 14, 11] | |
for kpts in pose_result: | |
kpts = np.array(kpts, copy=False)[idcs] | |
# draw each point on image | |
if pose_kpt_color is not None: | |
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: | |
color = tuple(int(c) for c in pose_kpt_color[kid]) | |
if show_keypoint_weight: | |
img_copy = img.copy() | |
cv2.circle( | |
img_copy, (int(x_coord), int(y_coord)), radius, color, -1 | |
) | |
transparency = max(0, min(1, kpt_score)) | |
cv2.addWeighted( | |
img_copy, transparency, img, 1 - transparency, 0, dst=img | |
) | |
else: | |
cv2.circle(img, (int(x_coord), int(y_coord)), radius, color, -1) | |
# draw links | |
if skeleton is not None and pose_link_color is not None: | |
assert len(pose_link_color) == len(skeleton) | |
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 | |
and pos1[0] < img_w | |
and pos1[1] > 0 | |
and pos1[1] < img_h | |
and pos2[0] > 0 | |
and pos2[0] < img_w | |
and pos2[1] > 0 | |
and pos2[1] < img_h | |
and kpts[sk[0], 2] > kpt_score_thr | |
and kpts[sk[1], 2] > kpt_score_thr | |
): | |
color = tuple(int(c) for c in pose_link_color[sk_id]) | |
if show_keypoint_weight: | |
img_copy = img.copy() | |
X = (pos1[0], pos2[0]) | |
Y = (pos1[1], pos2[1]) | |
mX = np.mean(X) | |
mY = np.mean(Y) | |
length = ((Y[0] - Y[1]) ** 2 + (X[0] - X[1]) ** 2) ** 0.5 | |
angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1])) | |
stickwidth = 2 | |
polygon = cv2.ellipse2Poly( | |
(int(mX), int(mY)), | |
(int(length / 2), int(stickwidth)), | |
int(angle), | |
0, | |
360, | |
1, | |
) | |
cv2.fillConvexPoly(img_copy, polygon, color) | |
transparency = max( | |
0, min(1, 0.5 * (kpts[sk[0], 2] + kpts[sk[1], 2])) | |
) | |
cv2.addWeighted( | |
img_copy, transparency, img, 1 - transparency, 0, dst=img | |
) | |
else: | |
cv2.line(img, pos1, pos2, color, thickness=thickness) | |
return img |