199 lines
7.2 KiB
Python
199 lines
7.2 KiB
Python
from typing import Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def post_dark_udp(coords, batch_heatmaps, kernel=3):
|
|
"""DARK post-pocessing. Implemented by udp. Paper ref: Huang et al. The
|
|
Devil is in the Details: Delving into Unbiased Data Processing for Human
|
|
Pose Estimation (CVPR 2020). Zhang et al. Distribution-Aware Coordinate
|
|
Representation for Human Pose Estimation (CVPR 2020). Code borrowed from
|
|
easy_ViTPose: https://github.com/JunkyByte/easy_ViTPose.
|
|
|
|
Note:
|
|
- batch size: B
|
|
- num keypoints: K
|
|
- num persons: N
|
|
- height of heatmaps: H
|
|
- width of heatmaps: W
|
|
|
|
B=1 for bottom_up paradigm where all persons share the same heatmap.
|
|
B=N for top_down paradigm where each person has its own heatmaps.
|
|
|
|
Args:
|
|
coords (np.ndarray[N, K, 2]): Initial coordinates of human pose.
|
|
batch_heatmaps (np.ndarray[B, K, H, W]): batch_heatmaps
|
|
kernel (int): Gaussian kernel size (K) for modulation.
|
|
|
|
Returns:
|
|
np.ndarray([N, K, 2]): Refined coordinates.
|
|
"""
|
|
if not isinstance(batch_heatmaps, np.ndarray):
|
|
batch_heatmaps = batch_heatmaps.cpu().numpy()
|
|
B, K, H, W = batch_heatmaps.shape
|
|
N = coords.shape[0]
|
|
assert (B == 1 or B == N)
|
|
for heatmaps in batch_heatmaps:
|
|
for heatmap in heatmaps:
|
|
cv2.GaussianBlur(heatmap, (kernel, kernel), 0, heatmap)
|
|
np.clip(batch_heatmaps, 0.001, 50, batch_heatmaps)
|
|
np.log(batch_heatmaps, batch_heatmaps)
|
|
|
|
batch_heatmaps_pad = np.pad(batch_heatmaps,
|
|
((0, 0), (0, 0), (1, 1), (1, 1)),
|
|
mode='edge').flatten()
|
|
|
|
index = coords[..., 0] + 1 + (coords[..., 1] + 1) * (W + 2)
|
|
index += (W + 2) * (H + 2) * np.arange(0, B * K).reshape(-1, K)
|
|
index = index.astype(int).reshape(-1, 1)
|
|
i_ = batch_heatmaps_pad[index]
|
|
ix1 = batch_heatmaps_pad[index + 1]
|
|
iy1 = batch_heatmaps_pad[index + W + 2]
|
|
ix1y1 = batch_heatmaps_pad[index + W + 3]
|
|
ix1_y1_ = batch_heatmaps_pad[index - W - 3]
|
|
ix1_ = batch_heatmaps_pad[index - 1]
|
|
iy1_ = batch_heatmaps_pad[index - 2 - W]
|
|
|
|
dx = 0.5 * (ix1 - ix1_)
|
|
dy = 0.5 * (iy1 - iy1_)
|
|
derivative = np.concatenate([dx, dy], axis=1)
|
|
derivative = derivative.reshape(N, K, 2, 1)
|
|
dxx = ix1 - 2 * i_ + ix1_
|
|
dyy = iy1 - 2 * i_ + iy1_
|
|
dxy = 0.5 * (ix1y1 - ix1 - iy1 + i_ + i_ - ix1_ - iy1_ + ix1_y1_)
|
|
hessian = np.concatenate([dxx, dxy, dxy, dyy], axis=1)
|
|
hessian = hessian.reshape(N, K, 2, 2)
|
|
hessian = np.linalg.inv(hessian + np.finfo(np.float32).eps * np.eye(2))
|
|
coords -= np.einsum('ijmn,ijnk->ijmk', hessian, derivative).squeeze()
|
|
return coords
|
|
|
|
|
|
def get_simcc_maximum(simcc_x: np.ndarray,
|
|
simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Get maximum response location and value from simcc representations.
|
|
|
|
Note:
|
|
instance number: N
|
|
num_keypoints: K
|
|
heatmap height: H
|
|
heatmap width: W
|
|
|
|
Args:
|
|
simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
|
|
simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
|
|
|
|
Returns:
|
|
tuple:
|
|
- locs (np.ndarray): locations of maximum heatmap responses in shape
|
|
(K, 2) or (N, K, 2)
|
|
- vals (np.ndarray): values of maximum heatmap responses in shape
|
|
(K,) or (N, K)
|
|
"""
|
|
N, K, Wx = simcc_x.shape
|
|
simcc_x = simcc_x.reshape(N * K, -1)
|
|
simcc_y = simcc_y.reshape(N * K, -1)
|
|
|
|
# get maximum value locations
|
|
x_locs = np.argmax(simcc_x, axis=1)
|
|
y_locs = np.argmax(simcc_y, axis=1)
|
|
locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
|
|
max_val_x = np.amax(simcc_x, axis=1)
|
|
max_val_y = np.amax(simcc_y, axis=1)
|
|
|
|
# get maximum value across x and y axis
|
|
# mask = max_val_x > max_val_y
|
|
# max_val_x[mask] = max_val_y[mask]
|
|
vals = 0.5 * (max_val_x + max_val_y)
|
|
locs[vals <= 0.] = -1
|
|
|
|
# reshape
|
|
locs = locs.reshape(N, K, 2)
|
|
vals = vals.reshape(N, K)
|
|
|
|
return locs, vals
|
|
|
|
|
|
def get_simcc_maximum3d(simcc_x: np.ndarray, simcc_y: np.ndarray,
|
|
simcc_z: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Get maximum response location and value from simcc representations.
|
|
|
|
Note:
|
|
instance number: N
|
|
num_keypoints: K
|
|
heatmap height: H
|
|
heatmap width: W
|
|
|
|
Args:
|
|
simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
|
|
simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
|
|
simcc_z (np.ndarray): z-axis SimCC in shape (K, Wz) or (N, K, Wz)
|
|
|
|
Returns:
|
|
tuple:
|
|
- locs (np.ndarray): locations of maximum heatmap responses in shape
|
|
(K, 3) or (N, K, 3)
|
|
- vals (np.ndarray): values of maximum heatmap responses in shape
|
|
(K,) or (N, K)
|
|
"""
|
|
|
|
assert isinstance(simcc_x, np.ndarray), 'simcc_x should be numpy.ndarray'
|
|
assert isinstance(simcc_y, np.ndarray), 'simcc_y should be numpy.ndarray'
|
|
assert isinstance(simcc_z, np.ndarray), 'simcc_z should be numpy.ndarray'
|
|
assert simcc_x.ndim == 2 or simcc_x.ndim == 3, (
|
|
f'Invalid shape {simcc_x.shape}')
|
|
assert simcc_y.ndim == 2 or simcc_y.ndim == 3, (
|
|
f'Invalid shape {simcc_y.shape}')
|
|
assert simcc_z.ndim == 2 or simcc_z.ndim == 3, (
|
|
f'Invalid shape {simcc_z.shape}')
|
|
assert simcc_x.ndim == simcc_y.ndim == simcc_z.ndim, (
|
|
f'{simcc_x.shape} != {simcc_y.shape} or {simcc_z.shape}')
|
|
|
|
if simcc_x.ndim == 3:
|
|
n, k, _ = simcc_x.shape
|
|
simcc_x = simcc_x.reshape(n * k, -1)
|
|
simcc_y = simcc_y.reshape(n * k, -1)
|
|
simcc_z = simcc_z.reshape(n * k, -1)
|
|
else:
|
|
n = None
|
|
|
|
x_locs = np.argmax(simcc_x, axis=1)
|
|
y_locs = np.argmax(simcc_y, axis=1)
|
|
z_locs = np.argmax(simcc_z, axis=1)
|
|
locs = np.stack((x_locs, y_locs, z_locs), axis=-1).astype(np.float32)
|
|
max_val_x = np.amax(simcc_x, axis=1)
|
|
max_val_y = np.amax(simcc_y, axis=1)
|
|
|
|
mask = max_val_x > max_val_y
|
|
max_val_x[mask] = max_val_y[mask]
|
|
vals = max_val_x
|
|
locs[vals <= 0.] = -1
|
|
|
|
if n is not None:
|
|
locs = locs.reshape(n, k, 3)
|
|
vals = vals.reshape(n, k)
|
|
|
|
return locs, vals
|
|
|
|
|
|
def convert_coco_to_openpose(keypoints, scores):
|
|
keypoints_info = np.concatenate((keypoints, scores[..., None]), axis=-1)
|
|
|
|
# compute neck joint
|
|
neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
|
|
|
|
# neck score when visualizing pred
|
|
neck[:,
|
|
2:3] = np.where(keypoints_info[:, 5, 2:3] > keypoints_info[:, 6, 2:3],
|
|
keypoints_info[:, 6, 2:3], keypoints_info[:, 5, 2:3])
|
|
new_keypoints_info = np.insert(keypoints_info, 17, neck, axis=1)
|
|
|
|
mmpose_idx = [17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3]
|
|
openpose_idx = [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17]
|
|
new_keypoints_info[:, openpose_idx] = \
|
|
new_keypoints_info[:, mmpose_idx]
|
|
keypoints_info = new_keypoints_info
|
|
|
|
keypoints, scores = keypoints_info[..., :2], keypoints_info[..., 2]
|
|
return keypoints, scores
|