95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
import numpy as np
|
|
|
|
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
|