forked from PulseFocusPlatform/PulseFocusPlatform
303 lines
8.9 KiB
Python
303 lines
8.9 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.
|
|
|
|
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).
|
|
scale (np.ndarray[2, ]): Scale of the bounding box
|
|
wrt [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(input_size) == 2
|
|
assert len(output_size) == 2
|
|
assert len(shift) == 2
|
|
|
|
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_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
|
|
assert len(b) == 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
|