Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import cv2, torch | |
| from typing import Tuple, Dict, List | |
| 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 decode(simcc_x: np.ndarray, | |
| simcc_y: np.ndarray, | |
| simcc_split_ratio, | |
| use_torch=False) -> Tuple[np.ndarray, np.ndarray]: | |
| """Modulate simcc distribution with Gaussian. | |
| Args: | |
| simcc_x (np.ndarray[K, Wx]): model predicted simcc in x. | |
| simcc_y (np.ndarray[K, Wy]): model predicted simcc in y. | |
| simcc_split_ratio (int): The split ratio of simcc. | |
| Returns: | |
| tuple: A tuple containing center and scale. | |
| - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2) | |
| - np.ndarray[float32]: scores in shape (K,) or (n, K) | |
| """ | |
| keypoints, scores = get_simcc_maximum(simcc_x, simcc_y, use_torch=use_torch) | |
| keypoints /= simcc_split_ratio | |
| return keypoints, scores | |
| def _fix_aspect_ratio(bbox_scale: np.ndarray, | |
| aspect_ratio: float) -> np.ndarray: | |
| """Extend the scale to match the given aspect ratio. | |
| Args: | |
| scale (np.ndarray): The image scale (w, h) in shape (2, ) | |
| aspect_ratio (float): The ratio of ``w/h`` | |
| Returns: | |
| np.ndarray: The reshaped image scale in (2, ) | |
| """ | |
| w, h = np.hsplit(bbox_scale, [1]) | |
| bbox_scale = np.where(w > h * aspect_ratio, | |
| np.hstack([w, w / aspect_ratio]), | |
| np.hstack([h * aspect_ratio, h])) | |
| return bbox_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 | |
| bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / 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 | |
| def get_simcc_maximum(simcc_x: np.ndarray, | |
| simcc_y: np.ndarray, | |
| use_torch=False) -> Tuple[np.ndarray, np.ndarray]: | |
| """Get maximum response location and value from simcc representations. | |
| Note: | |
| instance number: N | |
| num_keypoints: K | |
| heatmap height: H | |
| heatmap width: W | |
| Args: | |
| simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx) | |
| simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy) | |
| Returns: | |
| tuple: | |
| - locs (np.ndarray): locations of maximum heatmap responses in shape | |
| (K, 2) or (N, K, 2) | |
| - vals (np.ndarray): values of maximum heatmap responses in shape | |
| (K,) or (N, K) | |
| """ | |
| N, K, Wx = simcc_x.shape | |
| simcc_x = simcc_x.reshape(N * K, -1) | |
| simcc_y = simcc_y.reshape(N * K, -1) | |
| # get maximum value locations | |
| x_locs = np.argmax(simcc_x, axis=1) | |
| y_locs = np.argmax(simcc_y, axis=1) | |
| locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) | |
| if use_torch: | |
| max_val_x = torch.max(simcc_x, dim=1)[0] | |
| max_val_y = torch.max(simcc_y, dim=1)[0] | |
| else: | |
| max_val_x = np.amax(simcc_x, axis=1) | |
| max_val_y = np.amax(simcc_y, axis=1) | |
| # get maximum value across x and y axis | |
| mask = max_val_x > max_val_y | |
| max_val_x[mask] = max_val_y[mask] | |
| vals = max_val_x | |
| locs[vals <= 0.] = -1 | |
| # reshape | |
| locs = locs.reshape(N, K, 2) | |
| vals = vals.reshape(N, K) | |
| return locs, vals |