|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
colors = ( |
|
|
track_to_colors(track_ids) |
|
|
if track_ids is not None |
|
|
else torch.ones(B, 3, device) * 0.5 |
|
|
) |
|
|
|
|
|
|
|
|
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) |
|
|
""" |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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]] |
|
|
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) |
|
|
|
|
|
zs = mean[:, 2] |
|
|
counts = vis_mask[:, :T].sum(dim=-1) |
|
|
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 |
|
|
|
|
|
|
|
|
def get_colors(): |
|
|
|
|
|
color_file = os.path.abspath(os.path.join(__file__, "../colors.txt")) |
|
|
RGB_tuples = np.vstack( |
|
|
[ |
|
|
np.loadtxt(color_file, skiprows=0), |
|
|
|
|
|
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) |
|
|
cur_verts[:, 0] += c1 |
|
|
cur_verts[:, 2] += c2 |
|
|
else: |
|
|
cur_verts = np.stack([us, vs, zs], axis=-1) |
|
|
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) |
|
|
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", |
|
|
): |
|
|
|
|
|
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": |
|
|
|
|
|
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": |
|
|
|
|
|
skeleton = [] |
|
|
|
|
|
pose_link_color = palette[[]] |
|
|
pose_kpt_color = palette[[19] * 68] |
|
|
kpt_score_thr = 0 |
|
|
|
|
|
elif dataset == "FaceAFLWDataset": |
|
|
|
|
|
skeleton = [] |
|
|
|
|
|
pose_link_color = palette[[]] |
|
|
pose_kpt_color = palette[[19] * 19] |
|
|
kpt_score_thr = 0 |
|
|
|
|
|
elif dataset == "FaceCOFWDataset": |
|
|
|
|
|
skeleton = [] |
|
|
|
|
|
pose_link_color = palette[[]] |
|
|
pose_kpt_color = palette[[19] * 29] |
|
|
kpt_score_thr = 0 |
|
|
|
|
|
elif dataset == "FaceWFLWDataset": |
|
|
|
|
|
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] |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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 |