Squeezing functionality in Monai transforms

[martaranzini]

Is your feature request related to a problem? Please describe.
I would like to feed 2D patches extracted from 3D images into a 2D U-Net, and I could not find a straightforward way to remove the redundant dimension. E.g. from size (B, C, X, Y, 1) to (B, C, X, Y), with B being batch size and C number of channels.

Describe the solution you'd like
A Monai transform that could perform similar to numpy.squeeze and composable for data preprocessing.

Additional context
My attempt to such transform:

class SqueezeDim(Transform):
    """
    Squeeze unnecessary unitary dimensions
    """

    def __init__(self, dim=None):
        """
        Args:
            dim (int): dimension to be squeezed.
                Default: None (all dimensions of size 1 will be removed)
        """
        if dim is not None:
            assert isinstance(dim, int) and dim >= -1, 'invalid channel dimension.'
        self.dim = dim

    def __call__(self, img):
        """
        Args:
            data (dict): dictionary of numpy arrays with dim removed,
        """
        return np.squeeze(img, self.dim)

    
class SqueezeDimd(MapTransform):
    """
    Dictionary-based wrapper of :py:class:SqueezeDim`.
    """

    def __init__(self, keys, dim=None):
        """
        Args:
            keys (hashable items): keys of the corresponding items to be transformed.
                See also: :py:class:`monai.transforms.compose.MapTransform`
            dim (int): dimension to be squeezed.
                Default: None (all dimensions of size 1 will be removed)
        """
        super().__init__(keys)
        self.converter = SqueezeDim(dim=dim)

    def __call__(self, data):
        d = dict(data)
        for key in self.keys:
            d[key] = self.converter(d[key])
        return d

[Nic-Ma]

Hi @martaranzini ,

Thanks for your feedback and detailed sample code.
I still don't quite understand how you "extract 2D patches from 3D images"? Is it a transform?
Why not squeeze directly when extracting 2D patches?
Thanks.

[martaranzini]

Hi @Nic-Ma,

Thank you for getting back to me. Basically I want to train on 2D slices from 3D medical images, and all the images have with different size. I created the following chain of transforms to create the dataset for PyTorch DataLoader:

train_transforms = Compose([
    LoadNiftid(keys=['img', 'seg']),
    AddChanneld(keys=['img', 'seg']),
    NormalizeIntensityd(keys=['img']),
    Resized(keys=['img'], spatial_size=[96, 96], order=1),
    Resized(keys=['seg'], spatial_size=[96, 96], order=0, anti_aliasing=False),
    RandSpatialCropd(keys=['img', 'seg'], roi_size=[96, 96, 1], random_size=False),
    ToTensord(keys=['img', 'seg'])
])
check_ds = monai.data.Dataset(data=check_train_files, transform=train_transforms)

The Resized transform resizes my inputs to (batch, channel, 96, 96, N) with N being the original number of slices along z, which is different for each image. I used RandSpatialCrop to extract the single slices (i.e. my 2D patches) along the z direction. However, if I set roi_size=[96, 96] instead of roi_size=[96, 96, 1], it preserves the size of z for each image, and thus it crashes as the images are not of the same size. I get this error message:
RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 35 and 45 in dimension 4 at /pytorch/aten/src/TH/generic/THTensor.cpp:612
Is there another way to extract the slices randomly or squeeze directly from RandSpatialCrop?

With the transform I posted previously the train_transform for my case becomes

train_transforms = Compose([
    LoadNiftid(keys=['img', 'seg']),
    AddChanneld(keys=['img', 'seg']),
    NormalizeIntensityd(keys=['img']),
    Resized(keys=['img'], spatial_size=[96, 96], order=1),
    Resized(keys=['seg'], spatial_size=[96, 96], order=0, anti_aliasing=False),
    RandSpatialCropd(keys=['img', 'seg'], roi_size=[96, 96, 1], random_size=False),
    SqueezeDimd(keys=['img', 'seg'], dim=-1),
    ToTensord(keys=['img', 'seg'])
])

and in my case it was effective in changing the patches obtained from RandSpatialCropd from [96, 96, 1] to [96, 96].
Is there a different way of achieving this in Monai which I haven't perhaps identified?

[Nic-Ma]

Hi @martaranzini ,

Thanks for your overall explanation, now it makes a lot of sense to me.
I think your SqueezeDim would be a nice transform in MONAI, could you please help submit a PR directly? Your sample code is almost good enough, just need to:

1. add 2 unit tests for them(you can refer to other test cases of transforms)
2. add SqueezeDimD = SqueezeDimDict = SqueezeDimd in the composables.py.
3. add SqueezeDim and SqueezeDimd to documentations
   (make sure the "~~~" match the string length of transform name)

Thanks.