552 lines
18 KiB
Python
552 lines
18 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
this code is based on https://github.com/open-mmlab/mmpose
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import paddle.nn.functional as F
|
|
|
|
|
|
def get_affine_mat_kernel(h, w, s, inv=False):
|
|
if w < h:
|
|
w_ = s
|
|
h_ = int(np.ceil((s / w * h) / 64.) * 64)
|
|
scale_w = w
|
|
scale_h = h_ / w_ * w
|
|
|
|
else:
|
|
h_ = s
|
|
w_ = int(np.ceil((s / h * w) / 64.) * 64)
|
|
scale_h = h
|
|
scale_w = w_ / h_ * h
|
|
|
|
center = np.array([np.round(w / 2.), np.round(h / 2.)])
|
|
|
|
size_resized = (w_, h_)
|
|
trans = get_affine_transform(
|
|
center, np.array([scale_w, scale_h]), 0, size_resized, inv=inv)
|
|
|
|
return trans, size_resized
|
|
|
|
|
|
def get_affine_transform(center,
|
|
input_size,
|
|
rot,
|
|
output_size,
|
|
shift=(0., 0.),
|
|
inv=False):
|
|
"""Get the affine transform matrix, given the center/scale/rot/output_size.
|
|
|
|
Args:
|
|
center (np.ndarray[2, ]): Center of the bounding box (x, y).
|
|
input_size (np.ndarray[2, ]): Size of input feature (width, height).
|
|
rot (float): Rotation angle (degree).
|
|
output_size (np.ndarray[2, ]): Size of the destination heatmaps.
|
|
shift (0-100%): Shift translation ratio wrt the width/height.
|
|
Default (0., 0.).
|
|
inv (bool): Option to inverse the affine transform direction.
|
|
(inv=False: src->dst or inv=True: dst->src)
|
|
|
|
Returns:
|
|
np.ndarray: The transform matrix.
|
|
"""
|
|
assert len(center) == 2
|
|
assert len(output_size) == 2
|
|
assert len(shift) == 2
|
|
|
|
if not isinstance(input_size, (np.ndarray, list)):
|
|
input_size = np.array([input_size, input_size], dtype=np.float32)
|
|
scale_tmp = input_size
|
|
|
|
shift = np.array(shift)
|
|
src_w = scale_tmp[0]
|
|
dst_w = output_size[0]
|
|
dst_h = output_size[1]
|
|
|
|
rot_rad = np.pi * rot / 180
|
|
src_dir = rotate_point([0., src_w * -0.5], rot_rad)
|
|
dst_dir = np.array([0., dst_w * -0.5])
|
|
|
|
src = np.zeros((3, 2), dtype=np.float32)
|
|
|
|
src[0, :] = center + scale_tmp * shift
|
|
src[1, :] = center + src_dir + scale_tmp * shift
|
|
src[2, :] = _get_3rd_point(src[0, :], src[1, :])
|
|
|
|
dst = np.zeros((3, 2), dtype=np.float32)
|
|
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
|
|
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
|
|
dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
|
|
|
|
if inv:
|
|
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
|
else:
|
|
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
|
|
|
return trans
|
|
|
|
|
|
def get_warp_matrix(theta, size_input, size_dst, size_target):
|
|
"""This code is based on
|
|
https://github.com/open-mmlab/mmpose/blob/master/mmpose/core/post_processing/post_transforms.py
|
|
|
|
Calculate the transformation matrix under the constraint of unbiased.
|
|
Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased
|
|
Data Processing for Human Pose Estimation (CVPR 2020).
|
|
|
|
Args:
|
|
theta (float): Rotation angle in degrees.
|
|
size_input (np.ndarray): Size of input image [w, h].
|
|
size_dst (np.ndarray): Size of output image [w, h].
|
|
size_target (np.ndarray): Size of ROI in input plane [w, h].
|
|
|
|
Returns:
|
|
matrix (np.ndarray): A matrix for transformation.
|
|
"""
|
|
theta = np.deg2rad(theta)
|
|
matrix = np.zeros((2, 3), dtype=np.float32)
|
|
scale_x = size_dst[0] / size_target[0]
|
|
scale_y = size_dst[1] / size_target[1]
|
|
matrix[0, 0] = np.cos(theta) * scale_x
|
|
matrix[0, 1] = -np.sin(theta) * scale_x
|
|
matrix[0, 2] = scale_x * (
|
|
-0.5 * size_input[0] * np.cos(theta) + 0.5 * size_input[1] *
|
|
np.sin(theta) + 0.5 * size_target[0])
|
|
matrix[1, 0] = np.sin(theta) * scale_y
|
|
matrix[1, 1] = np.cos(theta) * scale_y
|
|
matrix[1, 2] = scale_y * (
|
|
-0.5 * size_input[0] * np.sin(theta) - 0.5 * size_input[1] *
|
|
np.cos(theta) + 0.5 * size_target[1])
|
|
return matrix
|
|
|
|
|
|
def _get_3rd_point(a, b):
|
|
"""To calculate the affine matrix, three pairs of points are required. This
|
|
function is used to get the 3rd point, given 2D points a & b.
|
|
|
|
The 3rd point is defined by rotating vector `a - b` by 90 degrees
|
|
anticlockwise, using b as the rotation center.
|
|
|
|
Args:
|
|
a (np.ndarray): point(x,y)
|
|
b (np.ndarray): point(x,y)
|
|
|
|
Returns:
|
|
np.ndarray: The 3rd point.
|
|
"""
|
|
assert len(
|
|
a) == 2, 'input of _get_3rd_point should be point with length of 2'
|
|
assert len(
|
|
b) == 2, 'input of _get_3rd_point should be point with length of 2'
|
|
direction = a - b
|
|
third_pt = b + np.array([-direction[1], direction[0]], dtype=np.float32)
|
|
|
|
return third_pt
|
|
|
|
|
|
def rotate_point(pt, angle_rad):
|
|
"""Rotate a point by an angle.
|
|
|
|
Args:
|
|
pt (list[float]): 2 dimensional point to be rotated
|
|
angle_rad (float): rotation angle by radian
|
|
|
|
Returns:
|
|
list[float]: Rotated point.
|
|
"""
|
|
assert len(pt) == 2
|
|
sn, cs = np.sin(angle_rad), np.cos(angle_rad)
|
|
new_x = pt[0] * cs - pt[1] * sn
|
|
new_y = pt[0] * sn + pt[1] * cs
|
|
rotated_pt = [new_x, new_y]
|
|
|
|
return rotated_pt
|
|
|
|
|
|
def transpred(kpts, h, w, s):
|
|
trans, _ = get_affine_mat_kernel(h, w, s, inv=True)
|
|
|
|
return warp_affine_joints(kpts[..., :2].copy(), trans)
|
|
|
|
|
|
def warp_affine_joints(joints, mat):
|
|
"""Apply affine transformation defined by the transform matrix on the
|
|
joints.
|
|
|
|
Args:
|
|
joints (np.ndarray[..., 2]): Origin coordinate of joints.
|
|
mat (np.ndarray[3, 2]): The affine matrix.
|
|
|
|
Returns:
|
|
matrix (np.ndarray[..., 2]): Result coordinate of joints.
|
|
"""
|
|
joints = np.array(joints)
|
|
shape = joints.shape
|
|
joints = joints.reshape(-1, 2)
|
|
return np.dot(np.concatenate(
|
|
(joints, joints[:, 0:1] * 0 + 1), axis=1),
|
|
mat.T).reshape(shape)
|
|
|
|
|
|
def affine_transform(pt, t):
|
|
new_pt = np.array([pt[0], pt[1], 1.]).T
|
|
new_pt = np.dot(t, new_pt)
|
|
return new_pt[:2]
|
|
|
|
|
|
def transform_preds(coords, center, scale, output_size):
|
|
target_coords = np.zeros(coords.shape)
|
|
trans = get_affine_transform(center, scale * 200, 0, output_size, inv=1)
|
|
for p in range(coords.shape[0]):
|
|
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
|
return target_coords
|
|
|
|
|
|
def oks_iou(g, d, a_g, a_d, sigmas=None, in_vis_thre=None):
|
|
if not isinstance(sigmas, np.ndarray):
|
|
sigmas = np.array([
|
|
.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07,
|
|
.87, .87, .89, .89
|
|
]) / 10.0
|
|
vars = (sigmas * 2)**2
|
|
xg = g[0::3]
|
|
yg = g[1::3]
|
|
vg = g[2::3]
|
|
ious = np.zeros((d.shape[0]))
|
|
for n_d in range(0, d.shape[0]):
|
|
xd = d[n_d, 0::3]
|
|
yd = d[n_d, 1::3]
|
|
vd = d[n_d, 2::3]
|
|
dx = xd - xg
|
|
dy = yd - yg
|
|
e = (dx**2 + dy**2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
|
|
if in_vis_thre is not None:
|
|
ind = list(vg > in_vis_thre) and list(vd > in_vis_thre)
|
|
e = e[ind]
|
|
ious[n_d] = np.sum(np.exp(-e)) / e.shape[0] if e.shape[0] != 0 else 0.0
|
|
return ious
|
|
|
|
|
|
def oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
|
"""greedily select boxes with high confidence and overlap with current maximum <= thresh
|
|
rule out overlap >= thresh
|
|
|
|
Args:
|
|
kpts_db (list): The predicted keypoints within the image
|
|
thresh (float): The threshold to select the boxes
|
|
sigmas (np.array): The variance to calculate the oks iou
|
|
Default: None
|
|
in_vis_thre (float): The threshold to select the high confidence boxes
|
|
Default: None
|
|
|
|
Return:
|
|
keep (list): indexes to keep
|
|
"""
|
|
|
|
if len(kpts_db) == 0:
|
|
return []
|
|
|
|
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
|
kpts = np.array(
|
|
[kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
|
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
|
|
|
order = scores.argsort()[::-1]
|
|
|
|
keep = []
|
|
while order.size > 0:
|
|
i = order[0]
|
|
keep.append(i)
|
|
|
|
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]],
|
|
sigmas, in_vis_thre)
|
|
|
|
inds = np.where(oks_ovr <= thresh)[0]
|
|
order = order[inds + 1]
|
|
|
|
return keep
|
|
|
|
|
|
def rescore(overlap, scores, thresh, type='gaussian'):
|
|
assert overlap.shape[0] == scores.shape[0]
|
|
if type == 'linear':
|
|
inds = np.where(overlap >= thresh)[0]
|
|
scores[inds] = scores[inds] * (1 - overlap[inds])
|
|
else:
|
|
scores = scores * np.exp(-overlap**2 / thresh)
|
|
|
|
return scores
|
|
|
|
|
|
def soft_oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
|
"""greedily select boxes with high confidence and overlap with current maximum <= thresh
|
|
rule out overlap >= thresh
|
|
|
|
Args:
|
|
kpts_db (list): The predicted keypoints within the image
|
|
thresh (float): The threshold to select the boxes
|
|
sigmas (np.array): The variance to calculate the oks iou
|
|
Default: None
|
|
in_vis_thre (float): The threshold to select the high confidence boxes
|
|
Default: None
|
|
|
|
Return:
|
|
keep (list): indexes to keep
|
|
"""
|
|
|
|
if len(kpts_db) == 0:
|
|
return []
|
|
|
|
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
|
kpts = np.array(
|
|
[kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
|
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
|
|
|
order = scores.argsort()[::-1]
|
|
scores = scores[order]
|
|
|
|
# max_dets = order.size
|
|
max_dets = 20
|
|
keep = np.zeros(max_dets, dtype=np.intp)
|
|
keep_cnt = 0
|
|
while order.size > 0 and keep_cnt < max_dets:
|
|
i = order[0]
|
|
|
|
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]],
|
|
sigmas, in_vis_thre)
|
|
|
|
order = order[1:]
|
|
scores = rescore(oks_ovr, scores[1:], thresh)
|
|
|
|
tmp = scores.argsort()[::-1]
|
|
order = order[tmp]
|
|
scores = scores[tmp]
|
|
|
|
keep[keep_cnt] = i
|
|
keep_cnt += 1
|
|
|
|
keep = keep[:keep_cnt]
|
|
|
|
return keep
|
|
|
|
|
|
def resize(input,
|
|
size=None,
|
|
scale_factor=None,
|
|
mode='nearest',
|
|
align_corners=None,
|
|
warning=True):
|
|
if warning:
|
|
if size is not None and align_corners:
|
|
input_h, input_w = tuple(int(x) for x in input.shape[2:])
|
|
output_h, output_w = tuple(int(x) for x in size)
|
|
if output_h > input_h or output_w > output_h:
|
|
if ((output_h > 1 and output_w > 1 and input_h > 1 and
|
|
input_w > 1) and (output_h - 1) % (input_h - 1) and
|
|
(output_w - 1) % (input_w - 1)):
|
|
warnings.warn(
|
|
f'When align_corners={align_corners}, '
|
|
'the output would more aligned if '
|
|
f'input size {(input_h, input_w)} is `x+1` and '
|
|
f'out size {(output_h, output_w)} is `nx+1`')
|
|
|
|
return F.interpolate(input, size, scale_factor, mode, align_corners)
|
|
|
|
|
|
def flip_back(output_flipped, flip_pairs, target_type='GaussianHeatmap'):
|
|
"""Flip the flipped heatmaps back to the original form.
|
|
Note:
|
|
- batch_size: N
|
|
- num_keypoints: K
|
|
- heatmap height: H
|
|
- heatmap width: W
|
|
Args:
|
|
output_flipped (np.ndarray[N, K, H, W]): The output heatmaps obtained
|
|
from the flipped images.
|
|
flip_pairs (list[tuple()): Pairs of keypoints which are mirrored
|
|
(for example, left ear -- right ear).
|
|
target_type (str): GaussianHeatmap or CombinedTarget
|
|
Returns:
|
|
np.ndarray: heatmaps that flipped back to the original image
|
|
"""
|
|
assert len(output_flipped.shape) == 4, \
|
|
'output_flipped should be [batch_size, num_keypoints, height, width]'
|
|
shape_ori = output_flipped.shape
|
|
channels = 1
|
|
if target_type.lower() == 'CombinedTarget'.lower():
|
|
channels = 3
|
|
output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
|
|
output_flipped = output_flipped.reshape((shape_ori[0], -1, channels,
|
|
shape_ori[2], shape_ori[3]))
|
|
output_flipped_back = output_flipped.clone()
|
|
|
|
# Swap left-right parts
|
|
for left, right in flip_pairs:
|
|
output_flipped_back[:, left, ...] = output_flipped[:, right, ...]
|
|
output_flipped_back[:, right, ...] = output_flipped[:, left, ...]
|
|
output_flipped_back = output_flipped_back.reshape(shape_ori)
|
|
# Flip horizontally
|
|
output_flipped_back = output_flipped_back[..., ::-1]
|
|
return output_flipped_back
|
|
|
|
|
|
def _calc_distances(preds, targets, mask, normalize):
|
|
"""Calculate the normalized distances between preds and target.
|
|
|
|
Note:
|
|
batch_size: N
|
|
num_keypoints: K
|
|
dimension of keypoints: D (normally, D=2 or D=3)
|
|
|
|
Args:
|
|
preds (np.ndarray[N, K, D]): Predicted keypoint location.
|
|
targets (np.ndarray[N, K, D]): Groundtruth keypoint location.
|
|
mask (np.ndarray[N, K]): Visibility of the target. False for invisible
|
|
joints, and True for visible. Invisible joints will be ignored for
|
|
accuracy calculation.
|
|
normalize (np.ndarray[N, D]): Typical value is heatmap_size
|
|
|
|
Returns:
|
|
np.ndarray[K, N]: The normalized distances. \
|
|
If target keypoints are missing, the distance is -1.
|
|
"""
|
|
N, K, _ = preds.shape
|
|
# set mask=0 when normalize==0
|
|
_mask = mask.copy()
|
|
_mask[np.where((normalize == 0).sum(1))[0], :] = False
|
|
distances = np.full((N, K), -1, dtype=np.float32)
|
|
# handle invalid values
|
|
normalize[np.where(normalize <= 0)] = 1e6
|
|
distances[_mask] = np.linalg.norm(
|
|
((preds - targets) / normalize[:, None, :])[_mask], axis=-1)
|
|
return distances.T
|
|
|
|
|
|
def _distance_acc(distances, thr=0.5):
|
|
"""Return the percentage below the distance threshold, while ignoring
|
|
distances values with -1.
|
|
|
|
Note:
|
|
batch_size: N
|
|
Args:
|
|
distances (np.ndarray[N, ]): The normalized distances.
|
|
thr (float): Threshold of the distances.
|
|
|
|
Returns:
|
|
float: Percentage of distances below the threshold. \
|
|
If all target keypoints are missing, return -1.
|
|
"""
|
|
distance_valid = distances != -1
|
|
num_distance_valid = distance_valid.sum()
|
|
if num_distance_valid > 0:
|
|
return (distances[distance_valid] < thr).sum() / num_distance_valid
|
|
return -1
|
|
|
|
|
|
def keypoint_pck_accuracy(pred, gt, mask, thr, normalize):
|
|
"""Calculate the pose accuracy of PCK for each individual keypoint and the
|
|
averaged accuracy across all keypoints for coordinates.
|
|
|
|
Note:
|
|
PCK metric measures accuracy of the localization of the body joints.
|
|
The distances between predicted positions and the ground-truth ones
|
|
are typically normalized by the bounding box size.
|
|
The threshold (thr) of the normalized distance is commonly set
|
|
as 0.05, 0.1 or 0.2 etc.
|
|
|
|
- batch_size: N
|
|
- num_keypoints: K
|
|
|
|
Args:
|
|
pred (np.ndarray[N, K, 2]): Predicted keypoint location.
|
|
gt (np.ndarray[N, K, 2]): Groundtruth keypoint location.
|
|
mask (np.ndarray[N, K]): Visibility of the target. False for invisible
|
|
joints, and True for visible. Invisible joints will be ignored for
|
|
accuracy calculation.
|
|
thr (float): Threshold of PCK calculation.
|
|
normalize (np.ndarray[N, 2]): Normalization factor for H&W.
|
|
|
|
Returns:
|
|
tuple: A tuple containing keypoint accuracy.
|
|
|
|
- acc (np.ndarray[K]): Accuracy of each keypoint.
|
|
- avg_acc (float): Averaged accuracy across all keypoints.
|
|
- cnt (int): Number of valid keypoints.
|
|
"""
|
|
distances = _calc_distances(pred, gt, mask, normalize)
|
|
|
|
acc = np.array([_distance_acc(d, thr) for d in distances])
|
|
valid_acc = acc[acc >= 0]
|
|
cnt = len(valid_acc)
|
|
avg_acc = valid_acc.mean() if cnt > 0 else 0
|
|
return acc, avg_acc, cnt
|
|
|
|
|
|
def keypoint_auc(pred, gt, mask, normalize, num_step=20):
|
|
"""Calculate the pose accuracy of PCK for each individual keypoint and the
|
|
averaged accuracy across all keypoints for coordinates.
|
|
|
|
Note:
|
|
- batch_size: N
|
|
- num_keypoints: K
|
|
|
|
Args:
|
|
pred (np.ndarray[N, K, 2]): Predicted keypoint location.
|
|
gt (np.ndarray[N, K, 2]): Groundtruth keypoint location.
|
|
mask (np.ndarray[N, K]): Visibility of the target. False for invisible
|
|
joints, and True for visible. Invisible joints will be ignored for
|
|
accuracy calculation.
|
|
normalize (float): Normalization factor.
|
|
|
|
Returns:
|
|
float: Area under curve.
|
|
"""
|
|
nor = np.tile(np.array([[normalize, normalize]]), (pred.shape[0], 1))
|
|
x = [1.0 * i / num_step for i in range(num_step)]
|
|
y = []
|
|
for thr in x:
|
|
_, avg_acc, _ = keypoint_pck_accuracy(pred, gt, mask, thr, nor)
|
|
y.append(avg_acc)
|
|
|
|
auc = 0
|
|
for i in range(num_step):
|
|
auc += 1.0 / num_step * y[i]
|
|
return auc
|
|
|
|
|
|
def keypoint_epe(pred, gt, mask):
|
|
"""Calculate the end-point error.
|
|
|
|
Note:
|
|
- batch_size: N
|
|
- num_keypoints: K
|
|
|
|
Args:
|
|
pred (np.ndarray[N, K, 2]): Predicted keypoint location.
|
|
gt (np.ndarray[N, K, 2]): Groundtruth keypoint location.
|
|
mask (np.ndarray[N, K]): Visibility of the target. False for invisible
|
|
joints, and True for visible. Invisible joints will be ignored for
|
|
accuracy calculation.
|
|
|
|
Returns:
|
|
float: Average end-point error.
|
|
"""
|
|
|
|
normalize = np.ones((pred.shape[0], pred.shape[2]), dtype=np.float32)
|
|
distances = _calc_distances(pred, gt, mask, normalize)
|
|
distance_valid = distances[distances != -1]
|
|
return distance_valid.sum() / max(1, len(distance_valid))
|