def __call__(self, spatial_size=None, grid=None): """ Args: spatial_size (list or tuple of int): output grid size. grid (ndarray): grid to be transformed. Shape must be (3, H, W) for 2D or (4, H, W, D) for 3D. """ if grid is None: if spatial_size is not None: grid = create_grid(spatial_size) else: raise ValueError("Either specify a grid or a spatial size to create a grid from.") spatial_dims = len(grid.shape) - 1 affine = np.eye(spatial_dims + 1) if self.rotate_params: affine = affine @ create_rotate(spatial_dims, self.rotate_params) if self.shear_params: affine = affine @ create_shear(spatial_dims, self.shear_params) if self.translate_params: affine = affine @ create_translate(spatial_dims, self.translate_params) if self.scale_params: affine = affine @ create_scale(spatial_dims, self.scale_params) affine = torch.as_tensor(np.ascontiguousarray(affine), device=self.device) grid = torch.tensor(grid) if not torch.is_tensor(grid) else grid.detach().clone() if self.device: grid = grid.to(self.device) grid = (affine.float() @ grid.reshape((grid.shape[0], -1)).float()).reshape([-1] + list(grid.shape[1:])) if self.as_tensor_output: return grid return grid.cpu().numpy()
def zoom_boxes(boxes: NdarrayOrTensor, zoom: Union[Sequence[float], float]): """ Zoom boxes Args: boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be StandardMode zoom: The zoom factor along the spatial axes. If a float, zoom is the same for each spatial axis. If a sequence, zoom should contain one value for each spatial axis. Returns: zoomed boxes, with same data type as ``boxes``, does not share memory with ``boxes`` Example: .. code-block:: python boxes = torch.ones(1,4) zoom_boxes(boxes, zoom=[0.5,2.2]) # will return tensor([[0.5, 2.2, 0.5, 2.2]]) """ spatial_dims = get_spatial_dims(boxes=boxes) # generate affine transform corresponding to ``zoom`` affine = create_scale(spatial_dims=spatial_dims, scaling_factor=zoom) return apply_affine_to_boxes(boxes=boxes, affine=affine)