Spaces:
Running
Running
| # ------------------------------------------------------------------------------ | |
| # Adapted from https://github.com/leoxiaobin/deep-high-resolution-net.pytorch | |
| # Original licence: Copyright (c) Microsoft, under the MIT License. | |
| # ------------------------------------------------------------------------------ | |
| import cv2 | |
| import numpy as np | |
| def get_affine_transform(center, | |
| scale, | |
| 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(scale) == 2 | |
| assert len(output_size) == 2 | |
| assert len(shift) == 2 | |
| # pixel_std is 200. | |
| scale_tmp = scale * 200.0 | |
| 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 affine_transform(pt, trans_mat): | |
| """Apply an affine transformation to the points. | |
| Args: | |
| pt (np.ndarray): a 2 dimensional point to be transformed | |
| trans_mat (np.ndarray): 2x3 matrix of an affine transform | |
| Returns: | |
| np.ndarray: Transformed points. | |
| """ | |
| assert len(pt) == 2 | |
| new_pt = np.array(trans_mat) @ np.array([pt[0], pt[1], 1.]) | |
| return new_pt | |
| 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 | |