166 lines
5.5 KiB
Python
166 lines
5.5 KiB
Python
from typing import Tuple
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def bbox_xyxy2cs(bbox: np.ndarray,
|
|
padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Transform the bbox format from (x,y,w,h) into (center, scale)
|
|
|
|
Args:
|
|
bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
|
|
as (left, top, right, bottom)
|
|
padding (float): BBox padding factor that will be multilied to scale.
|
|
Default: 1.0
|
|
|
|
Returns:
|
|
tuple: A tuple containing center and scale.
|
|
- np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
|
|
(n, 2)
|
|
- np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
|
|
(n, 2)
|
|
"""
|
|
# convert single bbox from (4, ) to (1, 4)
|
|
dim = bbox.ndim
|
|
if dim == 1:
|
|
bbox = bbox[None, :]
|
|
|
|
# get bbox center and scale
|
|
x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
|
|
center = np.hstack([x1 + x2, y1 + y2]) * 0.5
|
|
scale = np.hstack([x2 - x1, y2 - y1]) * padding
|
|
|
|
if dim == 1:
|
|
center = center[0]
|
|
scale = scale[0]
|
|
|
|
return center, scale
|
|
|
|
|
|
def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
|
|
"""Rotate a point by an angle.
|
|
|
|
Args:
|
|
pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
|
|
angle_rad (float): rotation angle in radian
|
|
|
|
Returns:
|
|
np.ndarray: Rotated point in shape (2, )
|
|
"""
|
|
sn, cs = np.sin(angle_rad), np.cos(angle_rad)
|
|
rot_mat = np.array([[cs, -sn], [sn, cs]])
|
|
return rot_mat @ pt
|
|
|
|
|
|
def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
|
|
"""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): The 1st point (x,y) in shape (2, )
|
|
b (np.ndarray): The 2nd point (x,y) in shape (2, )
|
|
|
|
Returns:
|
|
np.ndarray: The 3rd point.
|
|
"""
|
|
direction = a - b
|
|
c = b + np.r_[-direction[1], direction[0]]
|
|
return c
|
|
|
|
|
|
def get_warp_matrix(center: np.ndarray,
|
|
scale: np.ndarray,
|
|
rot: float,
|
|
output_size: Tuple[int, int],
|
|
shift: Tuple[float, float] = (0., 0.),
|
|
inv: bool = False) -> np.ndarray:
|
|
"""Calculate the affine transformation matrix that can warp the bbox area
|
|
in the input image to the 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, ] | list(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: A 2x3 transformation matrix
|
|
"""
|
|
|
|
shift = np.array(shift)
|
|
src_w = scale[0]
|
|
dst_w = output_size[0]
|
|
dst_h = output_size[1]
|
|
|
|
# compute transformation matrix
|
|
rot_rad = np.deg2rad(rot)
|
|
src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
|
|
dst_dir = np.array([0., dst_w * -0.5])
|
|
|
|
# get four corners of the src rectangle in the original image
|
|
src = np.zeros((3, 2), dtype=np.float32)
|
|
src[0, :] = center + scale * shift
|
|
src[1, :] = center + src_dir + scale * shift
|
|
src[2, :] = _get_3rd_point(src[0, :], src[1, :])
|
|
|
|
# get four corners of the dst rectangle in the input image
|
|
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:
|
|
warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
|
else:
|
|
warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
|
|
|
return warp_mat
|
|
|
|
|
|
def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict,
|
|
img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Get the bbox image as the model input by affine transform.
|
|
|
|
Args:
|
|
input_size (dict): The input size of the model.
|
|
bbox_scale (dict): The bbox scale of the img.
|
|
bbox_center (dict): The bbox center of the img.
|
|
img (np.ndarray): The original image.
|
|
|
|
Returns:
|
|
tuple: A tuple containing center and scale.
|
|
- np.ndarray[float32]: img after affine transform.
|
|
- np.ndarray[float32]: bbox scale after affine transform.
|
|
"""
|
|
w, h = input_size
|
|
warp_size = (int(w), int(h))
|
|
|
|
# reshape bbox to fixed aspect ratio
|
|
aspect_ratio = w / h
|
|
b_w, b_h = np.hsplit(bbox_scale, [1])
|
|
bbox_scale = np.where(b_w > b_h * aspect_ratio,
|
|
np.hstack([b_w, b_w / aspect_ratio]),
|
|
np.hstack([b_h * aspect_ratio, b_h]))
|
|
|
|
# get the affine matrix
|
|
center = bbox_center
|
|
scale = bbox_scale
|
|
rot = 0
|
|
warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
|
|
|
|
# do affine transform
|
|
img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
|
|
|
|
return img, bbox_scale
|