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)