DDF output for data augmentation transforms

[mariakawula]

Hi all,

is there a way to output the dense displacement field that was used to deform images for data augmentation? I defined a Compose transform that contains, among others, RandAffine and BSpline based deformation field (RandomElasticDeformation from torchio).

I can extract the affine matrix that in principle could be used to get the ddf but I don't have anything for B-splines.

[wyli]

not sure if I understand your question correctly, but you may need some modification of the __call__ to store the ddf while doing the transform, an example is here #1793, more specifically

MONAI/monai/transforms/spatial/dictionary.py

Lines 844 to 856 in 4211162

	if self._do_transform:
	cpg = self.rand_2d_elastic.deform_grid(spatial_size=sp_size)
	cpg_w_affine = self.rand_2d_elastic.rand_affine_grid(grid=cpg)
	affine = self.rand_2d_elastic.rand_affine_grid.get_transformation_matrix()
	grid = self.cpg_to_dvf(cpg_w_affine, self.rand_2d_elastic.deform_grid.spacing, sp_size)
	extra_info: Optional[Dict] = {"cpg": deepcopy(cpg), "affine": deepcopy(affine)}
	else:
	grid = create_grid(spatial_size=sp_size)
	extra_info = None
	for key, mode, padding_mode in self.key_iterator(d, self.mode, self.padding_mode):
	self.push_transform(d, key, extra_info=extra_info)
	d[key] = self.rand_2d_elastic.resampler(d[key], grid, mode=mode, padding_mode=padding_mode)

cc @rijobro

[ebrahimebrahim]

For creating a DDF out of affine transforms specifically, the implementation of AffineHead, which is part of GlobalNet, could be helpful. See here:

MONAI/monai/networks/nets/regunet.py

Lines 267 to 278 in f15ff93

	def affine_transform(self, theta: torch.Tensor):
	# (spatial_dims, ...) -> (spatial_dims + 1, ...)
	grid_padded = torch.cat([self.grid, torch.ones_like(self.grid[:1])])
	# grid_warped[b,p,...] = sum_over_q(grid_padded[q,...] * theta[b,p,q]
	if self.spatial_dims == 2:
	grid_warped = torch.einsum("qij,bpq->bpij", grid_padded, theta.reshape(-1, 2, 3))
	elif self.spatial_dims == 3:
	grid_warped = torch.einsum("qijk,bpq->bpijk", grid_padded, theta.reshape(-1, 3, 4))
	else:
	raise ValueError(f"do not support spatial_dims={self.spatial_dims}")
	return grid_warped

[mariakawula]

Thank you all for your suggestions. I will try to implement them and let you know if it worked out

[mariakawula]

That would be the transform

from monai.transforms import RandAffined, AddChanneld, CenterSpatialCropd
from torchio.transforms import RandomElasticDeformation, Compose

train_transform = Compose(
        [
            LoadImaged(keys=["image", "label"], reader="itkreader"),
            AddChanneld(keys=["image", "label"]),

            RandAffined(keys=["image", "label"], prob=args.aug_prob, rotate_range=np.deg2rad(args.rotate_range),
                        padding_mode="zeros", translate_range=args.translate_range, scale_range=args.scale_factor),
            RandomElasticDeformation(include=["image", "label"], num_control_points=args.n_control_points,
                                     max_displacement=args.max_displacement, p=args.aug_prob),
            CenterSpatialCropd(keys=["image", "label"], roi_size=args.vol_dims),
            ToTensord(keys=["image", "label"])
        ]
    )