How to use monai.networks.layers.grid_pull

[Constantin-Jehn]

Hello,

I tried moving from layer.Affine to layer.grid_pull to use higher order interpolation modes.
To do so I already installed monai from git and built it with cpp extension following https://docs.monai.io/en/latest/installation.html .

However, I have problems making proper use of layers.grid_pull to transform an image. I create a grid using https://pytorch.org/docs/stable/generated/torch.nn.functional.affine_grid.html and apply the transformation as follows.
I tried applying unitiy transform as sanity check first.

sample data:

10_3T_nody_002.nii.gz.zip

from torch.nn.functional import affine_grid
import monai
from monai.transforms import (
    AddChanneld,
    LoadImaged,
    ToTensord,
)
import torch as t
import numpy as np
import matplotlib.pyplot as plt

path = "10_3T_nody_002.nii.gz"
add_channel = AddChanneld(keys=["image"])
loader = LoadImaged(keys=["image"])
to_tensor = ToTensord(keys=["image"])
monai_dict = {"image": path}
monai_image_layer = loader(monai_dict)
monai_image_layer = add_channel(add_channel(monai_image_layer))
monai_image_layer = to_tensor(monai_image_layer)

image_tensor = monai_image_layer["image"]
input_2d = image_tensor[:,:,100:200,100:200,12]
input_2d.shape

plt.figure()
plt.imshow(input_2d[0,0,:,:])
plt.title('input')
plt.show()

theta = t.eye(3)[:2,:].unsqueeze(0)
theta[0,0,2] = 0.0
size = input_2d.shape
align_corners = False
grid = affine_grid(theta,size, align_corners=align_corners)

output_2d = monai.networks.layers.grid_pull(input_2d,grid,interpolation = 'cubic', bound = "zero", extrapolate = True)

plt.figure()
plt.imshow(output_2d[0,0,:,:])
plt.title('transformed')
plt.show()

Results looks unsatifactory:

I would appreciate any help possibly e.g. an example of correct usage of the transform.

Cheers.

[wyli]

The high-order interpolation in monai is not ready because it's missing the prefiltering step as tracked here #789.

(I put an example here in case it's somehow useful for understanding the API gaps)

import matplotlib.pyplot as plt
import numpy as np
import torch as t
from torch.nn.functional import affine_grid

import monai
from monai.networks.utils import normalize_transform
from monai.transforms import AddChanneld, LoadImaged, ToTensord

path = "10_3T_nody_002.nii.gz"
add_channel = AddChanneld(keys=["image"])
loader = LoadImaged(keys=["image"])
to_tensor = ToTensord(keys=["image"])
monai_dict = {"image": path}
monai_image_layer = loader(monai_dict)
monai_image_layer = add_channel(add_channel(monai_image_layer))
monai_image_layer = to_tensor(monai_image_layer)

image_tensor = monai_image_layer["image"]
input_2d = image_tensor[:, :, 100:200, 100:200, 12]
input_2d.shape

theta = t.eye(3)[:2, :].unsqueeze(0)
theta[0, 0, 2] = 0.0
size = input_2d.shape
align_corners = False
grid = affine_grid(theta, size, align_corners=align_corners)
input_shape = grid.shape
xform = t.inverse(normalize_transform(grid.shape[1:3], align_corners=align_corners, zero_centered=False, dtype=t.float32))
grid_t = t.movedim(grid, -1, 0)
grid = xform[0] @ t.concat([grid_t, t.ones(1, *grid_t.shape[1:])], dim=0).view(3, -1)
grid = grid[[1, 0]].reshape(2, *input_shape[:-1])
grid = t.movedim(grid, 0, -1)
output_2d = monai.networks.layers.grid_pull(input_2d, grid, interpolation="cubic", bound="zero", extrapolate=True)

plt.figure(figsize=(18, 6))
plt.subplot(1, 3, 1)
plt.imshow(input_2d[0, 0, :, :])
plt.title("original")
plt.subplot(1, 3, 2)
plt.imshow(output_2d[0, 0, :, :])
plt.title("transformed")
plt.subplot(1, 3, 3)
plt.imshow(np.abs(output_2d[0, 0, :, :] - input_2d[0, 0, :, :]))
plt.title("abs. diff")
plt.show()