Repository: Project-MONAI/MONAI

Issue:
add warp layer
**Is your feature request related to a problem? Please describe.**
Current repository does not support warping base on ddf, which is indispensable for registration.

**Describe the solution you'd like**
Add a Warp layer.

**Describe alternatives you've considered**
N/A

**Additional context**
[DeepReg implementation of Warp layer](https://github.com/DeepRegNet/DeepReg/blob/456731e51b69548b7f529c24b4de6d96c406ed60/deepreg/model/layer.py#L473-L505)



Discussion:
Note that this is related to #789, #189 and #95

Pytorch's [`grid-sample`](https://pytorch.org/docs/stable/nn.functional.html#grid-sample) function provides a GPU implementation of linear interpolation for image resampling. 

There is already an initial implementation of higher order resampling here but it needs more testing and use case examples:
https://github.com/Project-MONAI/MONAI/tree/master/monai/csrc/resample
> Note that this is related to #789, #189 and #95
> 
> Pytorch's [`grid-sample`](https://pytorch.org/docs/stable/nn.functional.html#grid-sample) function provides a GPU implementation of linear interpolation for image resampling.
> 
> There is already an initial implementation of higher order resampling here but it needs more testing and use case examples:
> https://github.com/Project-MONAI/MONAI/tree/master/monai/csrc/resample

Thank you for the useful information. I will try to figure out the best way to integrate both pytorch and MONAI implementations into the layer.
It would be super helpful if anyone could answer the following questions about the MONAI implementation of resampling:
- Is it true that in `grid_pull(input, grid)`, `grid[n, d, h, w]` specifies the `delta_z, delta_y, delta_x` displacements from the sample location to the target location ?
- What is the expected range of `grid`? Is it defined in the pixel_space, like defined in [`create_grid`](https://github.com/Project-MONAI/MONAI/blob/ec70d9637d7b3be4b3266c0060df25d08e0b6ecd/monai/transforms/utils.py#L394-L414)?

Thanks in advance. @tvercaut 
@brudfors is probably th ebest person to advice on this implementation as it stemmed from https://github.com/balbasty/nitorch
Hi @kate-sann5100 

I am happy to try to answer your questions.

The grid is defined in the target/output space and should specify the coordinates to sample in the input image. The grid coordinates should be defined in the voxel space of the input image. Here is some sample code that may, or may not, shed some more light on how its used:

https://github.com/balbasty/nitorch/blob/master/demo/demo_spatial.ipynb
@brudfors Thank you for the information, that completely answers my questions.
By the way, @brudfors is there a computational time comparison somewhere for pytorch `grid_sample` vs MONAI/nitorch `grid_pull` in say 3D linear interpolation mode?
Hi @tvercaut, 

The below Colab notebook shows the speed-up doing GPU resampling compared to CPU:

https://colab.research.google.com/drive/1qICvEDn-p8RnmaG-9v0sZu__6OQafM4A?usp=sharing

it is quite substantial.
 Thanks. I was more interested in GPU `grid_sample` (vanilla pytorch) vs GPU `grid_pull` (nitorch based).
Oh yes, sorry, I added that comparison to the notebook as well: vanilla pytorch and nitorch seem to perform similarly, but there is probably a more thorough validation that could be done.
I am reopening this issue following the comment above and the note from @wyli here:
https://github.com/Project-MONAI/MONAI/pull/1463#pullrequestreview-571046758

I would suggest only using `grid_pull` for interpolation orders strictly higher than linear as these are not supported by the vanilla `grid_sample`. For linear and nearest, unless I missed it, there is no benefit in using `grid_pull` over `grid_sample`, especially in view of the requirement for compilation.

@kate-sann5100: Can you adapt you warp layer accordingly?
nitorch does differentiable splatting, not sure if pytorch does?
@tvercaut That sounds promising. I will implement that. Thank you for the advice.
@brudfors I think there is no splatting function in pytorch. But in the Warp layer, the warp output should have the same shape as the input image, so splatting is not required.

Resolving patch:
diff --git a/monai/networks/blocks/warp.py b/monai/networks/blocks/warp.py
--- a/monai/networks/blocks/warp.py
+++ b/monai/networks/blocks/warp.py
@@ -4,9 +4,7 @@
 from torch import nn
 from torch.nn import functional as F
 
-from monai.config import USE_COMPILED
-from monai.networks.layers import grid_pull
-from monai.utils import GridSampleMode, GridSamplePadMode
+from monai.utils import GridSamplePadMode
 
 
 class Warp(nn.Module):
@@ -17,30 +15,38 @@ class Warp(nn.Module):
     def __init__(
         self,
         spatial_dims: int,
-        mode: Optional[Union[GridSampleMode, str]] = GridSampleMode.BILINEAR,
+        mode: int = 1,
         padding_mode: Optional[Union[GridSamplePadMode, str]] = GridSamplePadMode.ZEROS,
     ):
         """
         Args:
             spatial_dims: {2, 3}. number of spatial dimensions
-            mode: {``"bilinear"``, ``"nearest"``}
-                Interpolation mode to calculate output values. Defaults to ``"bilinear"``.
-                See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
+            mode: interpolation mode to calculate output values, defaults to 1.
+                Possible values are::
+
+                    - 0 or 'nearest'    or InterpolationType.nearest
+                    - 1 or 'linear'     or InterpolationType.linear
+                    - 2 or 'quadratic'  or InterpolationType.quadratic
+                    - 3 or 'cubic'      or InterpolationType.cubic
+                    - 4 or 'fourth'     or InterpolationType.fourth
+                    - etc.
             padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
                 Padding mode for outside grid values. Defaults to ``"border"``.
                 See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
         """
         super(Warp, self).__init__()
         if spatial_dims not in [2, 3]:
-            raise ValueError(f"got unsupported spatial_dims = {spatial_dims}, only support 2-d and 3-d input")
+            raise ValueError(f"got unsupported spatial_dims={spatial_dims}, only support 2-d and 3-d input")
         self.spatial_dims = spatial_dims
-        self.mode: GridSampleMode = GridSampleMode(mode)
+        if mode < 0:
+            raise ValueError(f"do not support negative mode, got mode={mode}")
+        self.mode = mode
         self.padding_mode: GridSamplePadMode = GridSamplePadMode(padding_mode)
 
     @staticmethod
     def get_reference_grid(ddf: torch.Tensor) -> torch.Tensor:
         mesh_points = [torch.arange(0, dim) for dim in ddf.shape[2:]]
-        grid = torch.stack(torch.meshgrid(*mesh_points[::-1]), dim=0)  # (spatial_dims, ...)
+        grid = torch.stack(torch.meshgrid(*mesh_points), dim=0)  # (spatial_dims, ...)
         grid = torch.stack([grid] * ddf.shape[0], dim=0)  # (batch, spatial_dims, ...)
         grid = grid.to(ddf)
         return grid
@@ -77,27 +83,31 @@ def forward(self, image: torch.Tensor, ddf: torch.Tensor) -> torch.Tensor:
         grid = self.get_reference_grid(ddf) + ddf
         grid = grid.permute([0] + list(range(2, 2 + self.spatial_dims)) + [1])  # (batch, ..., self.spatial_dims)
 
-        if USE_COMPILED:
-            _padding_mode = self.padding_mode.value
-            if _padding_mode == "zeros":
-                bound = 7
-            elif _padding_mode == "border":
-                bound = 0
-            else:
-                bound = 1
-            _interp_mode = self.mode.value
-            warped_image: torch.Tensor = grid_pull(
-                image,
-                grid,
-                bound=bound,
-                extrapolate=True,
-                interpolation=1 if _interp_mode == "bilinear" else _interp_mode,
-            )
+        if self.mode > 1:
+            raise ValueError(f"{self.mode}-order interpolation not yet implemented.")
+            # if not USE_COMPILED:
+            #     raise ValueError(f"cannot perform {self.mode}-order interpolation without C compile.")
+            # _padding_mode = self.padding_mode.value
+            # if _padding_mode == "zeros":
+            #     bound = 7
+            # elif _padding_mode == "border":
+            #     bound = 0
+            # else:
+            #     bound = 1
+            # warped_image: torch.Tensor = grid_pull(
+            #     image,
+            #     grid,
+            #     bound=bound,
+            #     extrapolate=True,
+            #     interpolation=self.mode,
+            # )
         else:
             grid = self.normalize_grid(grid)
             index_ordering: List[int] = list(range(self.spatial_dims - 1, -1, -1))
             grid = grid[..., index_ordering]  # z, y, x -> x, y, z
+            _interp_mode = "bilinear" if self.mode == 1 else "nearest"
             warped_image = F.grid_sample(
-                image, grid, mode=self.mode.value, padding_mode=self.padding_mode.value, align_corners=True
+                image, grid, mode=_interp_mode, padding_mode=self.padding_mode.value, align_corners=True
             )
+
         return warped_image


Test patch:
diff --git a/tests/test_warp.py b/tests/test_warp.py
--- a/tests/test_warp.py
+++ b/tests/test_warp.py
@@ -6,53 +6,22 @@
 
 from monai.networks.blocks.warp import Warp
 
-TEST_CASE = [
+LOW_POWER_TEST_CASES = [
     [
-        {"spatial_dims": 2, "mode": "bilinear", "padding_mode": "zeros"},
+        {"spatial_dims": 2, "mode": 0, "padding_mode": "zeros"},
         {"image": torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), "ddf": torch.zeros(1, 2, 2, 2)},
         torch.arange(4).reshape((1, 1, 2, 2)),
     ],
     [
-        {"spatial_dims": 2, "mode": "nearest", "padding_mode": "zeros"},
-        {"image": torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), "ddf": torch.zeros(1, 2, 4, 4)},
-        torch.tensor([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 2.0, 3.0, 0.0], [0.0, 0.0, 0.0, 0.0]])
-        .unsqueeze(0)
-        .unsqueeze(0),
-    ],
-    [
-        {"spatial_dims": 2, "mode": "bilinear", "padding_mode": "border"},
-        {"image": torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), "ddf": torch.zeros(1, 2, 4, 4)},
-        torch.tensor([[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0], [2.0, 2.0, 3, 3.0], [2.0, 2.0, 3.0, 3.0]])
-        .unsqueeze(0)
-        .unsqueeze(0),
-    ],
-    [
-        {"spatial_dims": 2, "mode": "nearest", "padding_mode": "reflection"},
-        {"image": torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), "ddf": torch.zeros(1, 2, 4, 4)},
-        torch.tensor([[3.0, 2.0, 3.0, 2.0], [1.0, 0.0, 1.0, 0.0], [3.0, 2.0, 3.0, 2.0], [1.0, 0.0, 1.0, 0.0]])
-        .unsqueeze(0)
-        .unsqueeze(0),
-    ],
-    [
-        {"spatial_dims": 3, "mode": "bilinear", "padding_mode": "zeros"},
-        {"image": torch.arange(8).reshape((1, 1, 2, 2, 2)).to(dtype=torch.float), "ddf": torch.zeros(1, 3, 2, 2, 2)},
-        torch.arange(8).reshape((1, 1, 2, 2, 2)).to(dtype=torch.float),
-    ],
-]
-
-TEST_CASES = [
-    [
-        {"spatial_dims": 2, "mode": "bilinear", "padding_mode": "zeros"},
-        {"image": torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), "ddf": torch.zeros(1, 2, 2, 2)},
-        torch.arange(4).reshape((1, 1, 2, 2)),
-    ],
-    [
-        {"spatial_dims": 2, "mode": "bilinear", "padding_mode": "zeros"},
+        {"spatial_dims": 2, "mode": 1, "padding_mode": "zeros"},
         {"image": torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), "ddf": torch.ones(1, 2, 2, 2)},
         torch.tensor([[[[3, 0], [0, 0]]]]),
     ],
+]
+
+HIGH_POWER_TEST_CASES = [
     [
-        {"spatial_dims": 3, "mode": "nearest", "padding_mode": "border"},
+        {"spatial_dims": 3, "mode": 2, "padding_mode": "border"},
         {
             "image": torch.arange(8).reshape((1, 1, 2, 2, 2)).to(dtype=torch.float),
             "ddf": torch.ones(1, 3, 2, 2, 2) * -1,
@@ -60,12 +29,16 @@
         torch.tensor([[[[[0, 0], [0, 0]], [[0, 0], [0, 0]]]]]),
     ],
     [
-        {"spatial_dims": 3, "mode": "nearest", "padding_mode": "reflection"},
+        {"spatial_dims": 3, "mode": 3, "padding_mode": "reflection"},
         {"image": torch.arange(8).reshape((1, 1, 2, 2, 2)).to(dtype=torch.float), "ddf": torch.ones(1, 3, 2, 2, 2)},
         torch.tensor([[[[[7, 6], [5, 4]], [[3, 2], [1, 0]]]]]),
     ],
 ]
 
+TEST_CASES = LOW_POWER_TEST_CASES
+# if USE_COMPILED:
+#     TEST_CASES += HIGH_POWER_TEST_CASES
+
 
 class TestWarp(unittest.TestCase):
     @parameterized.expand(TEST_CASES)