You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
342 lines
10 KiB
342 lines
10 KiB
# Copyright (c) 2022 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 |
|
|
|
|
|
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
|
|
|