Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import math | |
| from typing import Tuple | |
| import cv2 | |
| import numpy as np | |
| def bbox_xyxy2xywh(bbox_xyxy: np.ndarray) -> np.ndarray: | |
| """Transform the bbox format from x1y1x2y2 to xywh. | |
| Args: | |
| bbox_xyxy (np.ndarray): Bounding boxes (with scores), shaped (n, 4) or | |
| (n, 5). (left, top, right, bottom, [score]) | |
| Returns: | |
| np.ndarray: Bounding boxes (with scores), | |
| shaped (n, 4) or (n, 5). (left, top, width, height, [score]) | |
| """ | |
| bbox_xywh = bbox_xyxy.copy() | |
| bbox_xywh[:, 2] = bbox_xywh[:, 2] - bbox_xywh[:, 0] | |
| bbox_xywh[:, 3] = bbox_xywh[:, 3] - bbox_xywh[:, 1] | |
| return bbox_xywh | |
| def bbox_xywh2xyxy(bbox_xywh: np.ndarray) -> np.ndarray: | |
| """Transform the bbox format from xywh to x1y1x2y2. | |
| Args: | |
| bbox_xywh (ndarray): Bounding boxes (with scores), | |
| shaped (n, 4) or (n, 5). (left, top, width, height, [score]) | |
| Returns: | |
| np.ndarray: Bounding boxes (with scores), shaped (n, 4) or | |
| (n, 5). (left, top, right, bottom, [score]) | |
| """ | |
| bbox_xyxy = bbox_xywh.copy() | |
| bbox_xyxy[:, 2] = bbox_xyxy[:, 2] + bbox_xyxy[:, 0] | |
| bbox_xyxy[:, 3] = bbox_xyxy[:, 3] + bbox_xyxy[:, 1] | |
| return bbox_xyxy | |
| 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, :] | |
| scale = (bbox[..., 2:] - bbox[..., :2]) * padding | |
| center = (bbox[..., 2:] + bbox[..., :2]) * 0.5 | |
| if dim == 1: | |
| center = center[0] | |
| scale = scale[0] | |
| return center, scale | |
| def bbox_xywh2cs(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 (x, y, h, w) | |
| 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, :] | |
| x, y, w, h = np.hsplit(bbox, [1, 2, 3]) | |
| center = np.hstack([x + w * 0.5, y + h * 0.5]) | |
| scale = np.hstack([w, h]) * padding | |
| if dim == 1: | |
| center = center[0] | |
| scale = scale[0] | |
| return center, scale | |
| def bbox_cs2xyxy(center: np.ndarray, | |
| scale: np.ndarray, | |
| padding: float = 1.) -> np.ndarray: | |
| """Transform the bbox format from (center, scale) to (x1,y1,x2,y2). | |
| Args: | |
| center (ndarray): BBox center (x, y) in shape (2,) or (n, 2) | |
| scale (ndarray): BBox scale (w, h) in shape (2,) or (n, 2) | |
| padding (float): BBox padding factor that will be multilied to scale. | |
| Default: 1.0 | |
| Returns: | |
| ndarray[float32]: BBox (x1, y1, x2, y2) in shape (4, ) or (n, 4) | |
| """ | |
| dim = center.ndim | |
| assert scale.ndim == dim | |
| if dim == 1: | |
| center = center[None, :] | |
| scale = scale[None, :] | |
| wh = scale / padding | |
| xy = center - 0.5 * wh | |
| bbox = np.hstack((xy, xy + wh)) | |
| if dim == 1: | |
| bbox = bbox[0] | |
| return bbox | |
| def bbox_cs2xywh(center: np.ndarray, | |
| scale: np.ndarray, | |
| padding: float = 1.) -> np.ndarray: | |
| """Transform the bbox format from (center, scale) to (x,y,w,h). | |
| Args: | |
| center (ndarray): BBox center (x, y) in shape (2,) or (n, 2) | |
| scale (ndarray): BBox scale (w, h) in shape (2,) or (n, 2) | |
| padding (float): BBox padding factor that will be multilied to scale. | |
| Default: 1.0 | |
| Returns: | |
| ndarray[float32]: BBox (x, y, w, h) in shape (4, ) or (n, 4) | |
| """ | |
| dim = center.ndim | |
| assert scale.ndim == dim | |
| if dim == 1: | |
| center = center[None, :] | |
| scale = scale[None, :] | |
| wh = scale / padding | |
| xy = center - 0.5 * wh | |
| bbox = np.hstack((xy, wh)) | |
| if dim == 1: | |
| bbox = bbox[0] | |
| return bbox | |
| def bbox_xyxy2corner(bbox: np.ndarray): | |
| """Convert bounding boxes from xyxy format to corner format. | |
| Given a numpy array containing bounding boxes in the format | |
| (xmin, ymin, xmax, ymax), this function converts the bounding | |
| boxes to the corner format, where each box is represented by four | |
| corner points (top-left, top-right, bottom-right, bottom-left). | |
| Args: | |
| bbox (numpy.ndarray): Input array of shape (N, 4) representing | |
| N bounding boxes. | |
| Returns: | |
| numpy.ndarray: An array of shape (N, 4, 2) containing the corner | |
| points of the bounding boxes. | |
| Example: | |
| bbox = np.array([[0, 0, 100, 50], [10, 20, 200, 150]]) | |
| corners = bbox_xyxy2corner(bbox) | |
| """ | |
| dim = bbox.ndim | |
| if dim == 1: | |
| bbox = bbox[None] | |
| bbox = np.tile(bbox, 2).reshape(-1, 4, 2) | |
| bbox[:, 1:3, 0] = bbox[:, 0:2, 0] | |
| if dim == 1: | |
| bbox = bbox[0] | |
| return bbox | |
| def bbox_corner2xyxy(bbox: np.ndarray): | |
| """Convert bounding boxes from corner format to xyxy format. | |
| Given a numpy array containing bounding boxes in the corner | |
| format (four corner points for each box), this function converts | |
| the bounding boxes to the (xmin, ymin, xmax, ymax) format. | |
| Args: | |
| bbox (numpy.ndarray): Input array of shape (N, 4, 2) representing | |
| N bounding boxes. | |
| Returns: | |
| numpy.ndarray: An array of shape (N, 4) containing the bounding | |
| boxes in xyxy format. | |
| Example: | |
| corners = np.array([[[0, 0], [100, 0], [100, 50], [0, 50]], | |
| [[10, 20], [200, 20], [200, 150], [10, 150]]]) | |
| bbox = bbox_corner2xyxy(corners) | |
| """ | |
| if bbox.shape[-1] == 8: | |
| bbox = bbox.reshape(*bbox.shape[:-1], 4, 2) | |
| dim = bbox.ndim | |
| if dim == 2: | |
| bbox = bbox[None] | |
| bbox = np.concatenate((bbox.min(axis=1), bbox.max(axis=1)), axis=1) | |
| if dim == 2: | |
| bbox = bbox[0] | |
| return bbox | |
| def bbox_clip_border(bbox: np.ndarray, shape: Tuple[int, int]) -> np.ndarray: | |
| """Clip bounding box coordinates to fit within a specified shape. | |
| Args: | |
| bbox (np.ndarray): Bounding box coordinates of shape (..., 4) | |
| or (..., 2). | |
| shape (Tuple[int, int]): Shape of the image to which bounding | |
| boxes are being clipped in the format of (w, h) | |
| Returns: | |
| np.ndarray: Clipped bounding box coordinates. | |
| Example: | |
| >>> bbox = np.array([[10, 20, 30, 40], [40, 50, 80, 90]]) | |
| >>> shape = (50, 50) # Example image shape | |
| >>> clipped_bbox = bbox_clip_border(bbox, shape) | |
| """ | |
| width, height = shape[:2] | |
| if bbox.shape[-1] == 2: | |
| bbox[..., 0] = np.clip(bbox[..., 0], a_min=0, a_max=width) | |
| bbox[..., 1] = np.clip(bbox[..., 1], a_min=0, a_max=height) | |
| else: | |
| bbox[..., ::2] = np.clip(bbox[..., ::2], a_min=0, a_max=width) | |
| bbox[..., 1::2] = np.clip(bbox[..., 1::2], a_min=0, a_max=height) | |
| return bbox | |
| def flip_bbox(bbox: np.ndarray, | |
| image_size: Tuple[int, int], | |
| bbox_format: str = 'xywh', | |
| direction: str = 'horizontal') -> np.ndarray: | |
| """Flip the bbox in the given direction. | |
| Args: | |
| bbox (np.ndarray): The bounding boxes. The shape should be (..., 4) | |
| if ``bbox_format`` is ``'xyxy'`` or ``'xywh'``, and (..., 2) if | |
| ``bbox_format`` is ``'center'`` | |
| image_size (tuple): The image shape in [w, h] | |
| bbox_format (str): The bbox format. Options are ``'xywh'``, ``'xyxy'`` | |
| and ``'center'``. | |
| direction (str): The flip direction. Options are ``'horizontal'``, | |
| ``'vertical'`` and ``'diagonal'``. Defaults to ``'horizontal'`` | |
| Returns: | |
| np.ndarray: The flipped bounding boxes. | |
| """ | |
| direction_options = {'horizontal', 'vertical', 'diagonal'} | |
| assert direction in direction_options, ( | |
| f'Invalid flipping direction "{direction}". ' | |
| f'Options are {direction_options}') | |
| format_options = {'xywh', 'xyxy', 'center'} | |
| assert bbox_format in format_options, ( | |
| f'Invalid bbox format "{bbox_format}". ' | |
| f'Options are {format_options}') | |
| bbox_flipped = bbox.copy() | |
| w, h = image_size | |
| # TODO: consider using "integer corner" coordinate system | |
| if direction == 'horizontal': | |
| if bbox_format == 'xywh' or bbox_format == 'center': | |
| bbox_flipped[..., 0] = w - bbox[..., 0] - 1 | |
| elif bbox_format == 'xyxy': | |
| bbox_flipped[..., ::2] = w - bbox[..., -2::-2] - 1 | |
| elif direction == 'vertical': | |
| if bbox_format == 'xywh' or bbox_format == 'center': | |
| bbox_flipped[..., 1] = h - bbox[..., 1] - 1 | |
| elif bbox_format == 'xyxy': | |
| bbox_flipped[..., 1::2] = h - bbox[..., ::-2] - 1 | |
| elif direction == 'diagonal': | |
| if bbox_format == 'xywh' or bbox_format == 'center': | |
| bbox_flipped[..., :2] = [w, h] - bbox[..., :2] - 1 | |
| elif bbox_format == 'xyxy': | |
| bbox_flipped[...] = [w, h, w, h] - bbox - 1 | |
| bbox_flipped = np.concatenate( | |
| (bbox_flipped[..., 2:], bbox_flipped[..., :2]), axis=-1) | |
| return bbox_flipped | |
| def get_udp_warp_matrix( | |
| center: np.ndarray, | |
| scale: np.ndarray, | |
| rot: float, | |
| output_size: Tuple[int, int], | |
| ) -> np.ndarray: | |
| """Calculate the affine transformation matrix under the unbiased | |
| constraint. See `UDP (CVPR 2020)`_ for details. | |
| Note: | |
| - The bbox number: N | |
| 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 (tuple): Size ([w, h]) of the output image | |
| Returns: | |
| np.ndarray: A 2x3 transformation matrix | |
| .. _`UDP (CVPR 2020)`: https://arxiv.org/abs/1911.07524 | |
| """ | |
| assert len(center) == 2 | |
| assert len(scale) == 2 | |
| assert len(output_size) == 2 | |
| input_size = center * 2 | |
| rot_rad = np.deg2rad(rot) | |
| warp_mat = np.zeros((2, 3), dtype=np.float32) | |
| scale_x = (output_size[0] - 1) / scale[0] | |
| scale_y = (output_size[1] - 1) / scale[1] | |
| warp_mat[0, 0] = math.cos(rot_rad) * scale_x | |
| warp_mat[0, 1] = -math.sin(rot_rad) * scale_x | |
| warp_mat[0, 2] = scale_x * (-0.5 * input_size[0] * math.cos(rot_rad) + | |
| 0.5 * input_size[1] * math.sin(rot_rad) + | |
| 0.5 * scale[0]) | |
| warp_mat[1, 0] = math.sin(rot_rad) * scale_y | |
| warp_mat[1, 1] = math.cos(rot_rad) * scale_y | |
| warp_mat[1, 2] = scale_y * (-0.5 * input_size[0] * math.sin(rot_rad) - | |
| 0.5 * input_size[1] * math.cos(rot_rad) + | |
| 0.5 * scale[1]) | |
| return warp_mat | |
| 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, | |
| fix_aspect_ratio: bool = True, | |
| ) -> 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) | |
| fix_aspect_ratio (bool): Whether to fix aspect ratio during transform. | |
| Defaults to True. | |
| Returns: | |
| np.ndarray: A 2x3 transformation matrix | |
| """ | |
| assert len(center) == 2 | |
| assert len(scale) == 2 | |
| assert len(output_size) == 2 | |
| assert len(shift) == 2 | |
| shift = np.array(shift) | |
| src_w, src_h = scale[:2] | |
| dst_w, dst_h = output_size[:2] | |
| rot_rad = np.deg2rad(rot) | |
| src_dir = _rotate_point(np.array([src_w * -0.5, 0.]), rot_rad) | |
| dst_dir = np.array([dst_w * -0.5, 0.]) | |
| src = np.zeros((3, 2), dtype=np.float32) | |
| src[0, :] = center + scale * shift | |
| src[1, :] = center + src_dir + scale * shift | |
| 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 | |
| if fix_aspect_ratio: | |
| src[2, :] = _get_3rd_point(src[0, :], src[1, :]) | |
| dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :]) | |
| else: | |
| src_dir_2 = _rotate_point(np.array([0., src_h * -0.5]), rot_rad) | |
| dst_dir_2 = np.array([0., dst_h * -0.5]) | |
| src[2, :] = center + src_dir_2 + scale * shift | |
| dst[2, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir_2 | |
| 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 get_pers_warp_matrix(center: np.ndarray, translate: np.ndarray, | |
| scale: float, rot: float, | |
| shear: np.ndarray) -> np.ndarray: | |
| """Compute a perspective warp matrix based on specified transformations. | |
| Args: | |
| center (np.ndarray): Center of the transformation. | |
| translate (np.ndarray): Translation vector. | |
| scale (float): Scaling factor. | |
| rot (float): Rotation angle in degrees. | |
| shear (np.ndarray): Shearing angles in degrees along x and y axes. | |
| Returns: | |
| np.ndarray: Perspective warp matrix. | |
| Example: | |
| >>> center = np.array([0, 0]) | |
| >>> translate = np.array([10, 20]) | |
| >>> scale = 1.2 | |
| >>> rot = 30.0 | |
| >>> shear = np.array([15.0, 0.0]) | |
| >>> warp_matrix = get_pers_warp_matrix(center, translate, | |
| scale, rot, shear) | |
| """ | |
| translate_mat = np.array([[1, 0, translate[0] + center[0]], | |
| [0, 1, translate[1] + center[1]], [0, 0, 1]], | |
| dtype=np.float32) | |
| shear_x = math.radians(shear[0]) | |
| shear_y = math.radians(shear[1]) | |
| shear_mat = np.array([[1, np.tan(shear_x), 0], [np.tan(shear_y), 1, 0], | |
| [0, 0, 1]], | |
| dtype=np.float32) | |
| rotate_angle = math.radians(rot) | |
| rotate_mat = np.array([[np.cos(rotate_angle), -np.sin(rotate_angle), 0], | |
| [np.sin(rotate_angle), | |
| np.cos(rotate_angle), 0], [0, 0, 1]], | |
| dtype=np.float32) | |
| scale_mat = np.array([[scale, 0, 0], [0, scale, 0], [0, 0, 1]], | |
| dtype=np.float32) | |
| recover_center_mat = np.array([[1, 0, -center[0]], [0, 1, -center[1]], | |
| [0, 0, 1]], | |
| dtype=np.float32) | |
| warp_matrix = np.dot( | |
| np.dot( | |
| np.dot(np.dot(translate_mat, shear_mat), rotate_mat), scale_mat), | |
| recover_center_mat) | |
| return warp_matrix | |
| 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): | |
| """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 | |