[Feat] Volumetric (3D + channels) Affine warping

🚀 Feature Request

Add support for volumetric data (C, (D,) H, W) and affine transformations

Motivation

Volumetric transformations can be used for any kind of volumetric data (e.g. medical -> CT, MRI) and affine transformations are very common for augmentation transformations (e.g. in rising).

Pitch

Introduce a new base class GenericWarper from which DepthWarper , HomographyWarper and AffineWarper can be derived from.

The functional API of the warp functions shouldn’t be changed because it can automatically detect if the input tensor has an additional dimension or not.

I’m not quite sure how the volumetric support should be added for the future. One possibility would be to add a first version exists besides the 2D functions and has support for both 2D and 3D (could probably look similar to rising’s functionals) or add additional functions to specifically/only handle 3D data (this would probably lead to quite a lot of code duplication).

Additional context

@justusschock

[edgarriba]

my idea was to refactor kornia.geometry.warp with something like this. /cc @ducha-aiki @lferraz had some initial thoughts as well since he refactored few weeks ago to handle cases where you want to reuse the grid to speed up processes.

class BaseGridSampleWarper(nn.Module):
  """Encapsulates torch.grid_sample"""

  def compute_grid(self, transform: torch.Tensor) -> torch.Tensor:
    raise NotImplementedErrror("Overrride depending on your transform type)

  def forward(self, input: torch.Tensor, transform: torch.Tensor) -> torch.Tensor:
    grid: torch.Tensor = self.compute_grid(transform)
    assert grid.shape == (B, H, W, 2)
    return torch.nn.fuctional.grid_sample(input, grid, ....)

Then any type of transformation can be specialized from this by just adjusting the way it produces the grid depending on the desired behavior.

class AffineWarper(BaseGridSampleWarper):
  r"""Warp a tensor using a rigid 2d affine transformation."""
  def compute_grid(self, transform: torch.Tensor) -> torch.Tensor:
    # from https://github.com/kornia/kornia/blob/master/kornia/geometry/transform/imgwarp.py#L147
    B, C, H, W = src.size()
    dsize_src = (H, W)
    out_size = dsize
    # we generate a 3x3 transformation matrix from 2x3 affine
    M_3x3: torch.Tensor = convert_affinematrix_to_homography(M)
    dst_norm_trans_src_norm: torch.Tensor = normalize_homography(
        M_3x3, dsize_src, out_size)
    src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm)
    grid = F.affine_grid(src_norm_trans_dst_norm[:, :2, :],  # type: ignore
                         [B, C, out_size[0], out_size[1]],
                         align_corners=align_corners)
    return grid

class PerspectiveWarper(BaseGridSampleWarper):
  r"""Warp a tensor using a rigid 2d perspective transformation."""
  def compute_grid(self, transform: torch.Tensor) -> torch.Tensor:
    # we need to double check we can produce same with affine_grid
    return grid

class ProjectiveWarper(BaseGridSampleWarper):
  r"""Warp a tensor using a rigid 3d homogeneous transformation."""
  def compute_grid(self, transform: torch.Tensor) -> torch.Tensor:
    # use code from rising
    return grid

Then to follow OpenCV convention the high level interfaces would be:

out = K.warp_affine(input, mat2x3, (H, W))  # (BxCxHxW / Bx2x3) -> BxCxHxW
out = K.warp_perspective(input, mat3x3, (H, W))  # (BxCxHxW / Bx3x3) -> BxCxHxW
out = K.warp_projective(input, mat4x4, (D, H, W))  # (BxCxDxHxW / Bx4x4) -> BxCxDxHxW

[edgarriba]

@mibaumgartner I have internally implemented already the warp_projective functionality plus few utlils like get_projection_matrix which is an extension of get_rotation_matrix2d that basically does the logics to compute the transformation matrix given an angle vector in angle-axis format and the 3d coordinate to the origin from where you want to apply the transformation. Besides, I needed to implement some normalization tricks to compute the final transformation is we do for the 2d affine case in here but for the 3D case. Hope I can send a PR soon about it.