Spaces:
Runtime error
Runtime error
| """ | |
| crop | |
| for numpy array | |
| Given image, bbox(center, bboxsize) | |
| return: cropped image, tform(used for transform the keypoint accordingly) | |
| only support crop to squared images | |
| """ | |
| import numpy as np | |
| from skimage.transform import estimate_transform, rescale, resize, warp | |
| def points2bbox(points, points_scale=None): | |
| # recover range | |
| if points_scale: | |
| points[:, 0] = points[:, 0] * points_scale[1] / 2 + points_scale[1] / 2 | |
| points[:, 1] = points[:, 1] * points_scale[0] / 2 + points_scale[0] / 2 | |
| left = np.min(points[:, 0]) | |
| right = np.max(points[:, 0]) | |
| top = np.min(points[:, 1]) | |
| bottom = np.max(points[:, 1]) | |
| size = max(right - left, bottom - top) | |
| # + old_size*0.1]) | |
| center = np.array([right - (right - left) / 2.0, bottom - (bottom - top) / 2.0]) | |
| return center, size | |
| # translate center | |
| def augment_bbox(center, bbox_size, scale=[1.0, 1.0], trans_scale=0.0): | |
| trans_scale = (np.random.rand(2) * 2 - 1) * trans_scale | |
| center = center + trans_scale * bbox_size # 0.5 | |
| scale = np.random.rand() * (scale[1] - scale[0]) + scale[0] | |
| size = int(bbox_size * scale) | |
| return center, size | |
| def crop_array(image, center, bboxsize, crop_size): | |
| """for single image only | |
| Args: | |
| image (numpy.Array): the reference array of shape HxWXC. | |
| size (Tuple[int, int]): a tuple with the height and width that will be | |
| used to resize the extracted patches. | |
| Returns: | |
| cropped_image | |
| tform: 3x3 affine matrix | |
| """ | |
| # points: top-left, top-right, bottom-right | |
| src_pts = np.array([ | |
| [center[0] - bboxsize / 2, center[1] - bboxsize / 2], | |
| [center[0] + bboxsize / 2, center[1] - bboxsize / 2], | |
| [center[0] + bboxsize / 2, center[1] + bboxsize / 2], | |
| ]) | |
| DST_PTS = np.array([[0, 0], [crop_size - 1, 0], [crop_size - 1, crop_size - 1]]) | |
| # estimate transformation between points | |
| tform = estimate_transform("similarity", src_pts, DST_PTS) | |
| # warp images | |
| cropped_image = warp(image, tform.inverse, output_shape=(crop_size, crop_size)) | |
| return cropped_image, tform.params.T | |
| class Cropper(object): | |
| def __init__(self, crop_size, scale=[1, 1], trans_scale=0.0): | |
| self.crop_size = crop_size | |
| self.scale = scale | |
| self.trans_scale = trans_scale | |
| def crop(self, image, points, points_scale=None): | |
| # points to bbox | |
| center, bbox_size = points2bbox(points, points_scale) | |
| # argument bbox. | |
| center, bbox_size = augment_bbox( | |
| center, bbox_size, scale=self.scale, trans_scale=self.trans_scale | |
| ) | |
| # crop | |
| cropped_image, tform = crop_array(image, center, bbox_size, self.crop_size) | |
| return cropped_image, tform | |