Repository: Project-MONAI/MONAI

Issue:
Squeezing functionality in Monai transforms
**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:
```python
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
```




Discussion:
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.
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:
```python
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
```python
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?
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](https://github.com/Project-MONAI/MONAI/blob/master/docs/source/transforms.rst)
(make sure the "~~~" match the string length of transform name)

Thanks.

Resolving patch:
diff --git a/monai/transforms/composables.py b/monai/transforms/composables.py
--- a/monai/transforms/composables.py
+++ b/monai/transforms/composables.py
@@ -25,7 +25,8 @@
                                          Rand2DElastic, Rand3DElastic, RandAffine, Resize, Rotate, Rotate90,
                                          ScaleIntensityRange, Spacing, SpatialCrop, Zoom, ToTensor, LoadPNG,
                                          AsChannelLast, ThresholdIntensity, AdjustContrast, CenterSpatialCrop,
-                                         CastToType, SpatialPad, RepeatChannel, ShiftIntensity, ScaleIntensity)
+                                         CastToType, SpatialPad, RepeatChannel, ShiftIntensity, ScaleIntensity,
+                                         SqueezeDim)
 from monai.transforms.utils import (create_grid, generate_pos_neg_label_crop_centers, generate_spatial_bounding_box)
 from monai.utils.misc import ensure_tuple
 
@@ -1404,6 +1405,29 @@ def __call__(self, data):
         return {key: val for key, val in data.items() if key not in self.keys}
 
 
+class SqueezeDimd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.transforms.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
+
+
 SpacingD = SpacingDict = Spacingd
 OrientationD = OrientationDict = Orientationd
 LoadNiftiD = LoadNiftiDict = LoadNiftid
@@ -1443,3 +1467,4 @@ def __call__(self, data):
 ZoomD = ZoomDict = Zoomd
 RandZoomD = RandZoomDict = RandZoomd
 DeleteKeysD = DeleteKeysDict = DeleteKeysd
+SqueezeDimD = SqueezeDimDict = SqueezeDimd
diff --git a/monai/transforms/transforms.py b/monai/transforms/transforms.py
--- a/monai/transforms/transforms.py
+++ b/monai/transforms/transforms.py
@@ -1738,3 +1738,26 @@ def __call__(self, img, spatial_size=None, mode=None):
             grid[:3] += gaussian(self.rand_offset[None])[0] * self.magnitude
             grid = self.rand_affine_grid(grid=grid)
         return self.resampler(img, grid, mode)
+
+
+class SqueezeDim(Transform):
+    """
+    Squeeze undesired 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:
+            img (ndarray): numpy arrays with required dimension `dim` removed
+        """
+        return np.squeeze(img, self.dim)


Test patch:
diff --git a/tests/test_squeezedim.py b/tests/test_squeezedim.py
new file mode 100644
--- /dev/null
+++ b/tests/test_squeezedim.py
@@ -0,0 +1,76 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import numpy as np
+from parameterized import parameterized
+from monai.transforms import SqueezeDim
+
+TEST_CASE_1 = [
+    {
+        'dim': None
+    },
+    np.random.rand(1, 2, 1, 3),
+    (2, 3)
+]
+
+TEST_CASE_2 = [
+    {
+        'dim': 2
+    },
+    np.random.rand(1, 2, 1, 8, 16),
+    (1, 2, 8, 16)
+]
+
+TEST_CASE_3 = [
+    {
+        'dim': -1
+    },
+    np.random.rand(1, 1, 16, 8, 1),
+    (1, 1, 16, 8)
+]
+
+TEST_CASE_4 = [
+    {},
+    np.random.rand(1, 2, 1, 3),
+    (2, 3)
+]
+
+TEST_CASE_5 = [
+    {
+        'dim': -2
+    },
+    np.random.rand(1, 1, 16, 8, 1),
+]
+
+TEST_CASE_6 = [
+    {
+        'dim': 0.5
+    },
+    np.random.rand(1, 1, 16, 8, 1),
+]
+
+
+class TestSqueezeDim(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    def test_shape(self, input_param, test_data, expected_shape):
+        result = SqueezeDim(**input_param)(test_data)
+        self.assertTupleEqual(result.shape, expected_shape)
+
+    @parameterized.expand([TEST_CASE_5, TEST_CASE_6])
+    def test_invalid_inputs(self, input_param, test_data):
+        with self.assertRaises(AssertionError):
+            result = SqueezeDim(**input_param)(test_data)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_squeezedimd.py b/tests/test_squeezedimd.py
new file mode 100644
--- /dev/null
+++ b/tests/test_squeezedimd.py
@@ -0,0 +1,103 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import numpy as np
+from parameterized import parameterized
+from monai.transforms import SqueezeDimd
+
+
+TEST_CASE_1 = [
+    {
+        'keys': ['img', 'seg'],
+        'dim': None
+    },
+    {
+        'img': np.random.rand(1, 2, 1, 3),
+        'seg': np.random.randint(0, 2, size=[1, 2, 1, 3]),
+    },
+    (2, 3)
+]
+
+TEST_CASE_2 = [
+    {
+        'keys': ['img', 'seg'],
+        'dim': 2
+    },
+    {
+        'img': np.random.rand(1, 2, 1, 8, 16),
+        'seg': np.random.randint(0, 2, size=[1, 2, 1, 8, 16]),
+    },
+    (1, 2, 8, 16)
+]
+
+TEST_CASE_3 = [
+    {
+        'keys': ['img', 'seg'],
+        'dim': -1
+    },
+    {
+        'img': np.random.rand(1, 1, 16, 8, 1),
+        'seg': np.random.randint(0, 2, size=[1, 1, 16, 8, 1]),
+    },
+    (1, 1, 16, 8)
+]
+
+TEST_CASE_4 = [
+    {
+        'keys': ['img', 'seg']
+    },
+    {
+        'img': np.random.rand(1, 2, 1, 3),
+        'seg': np.random.randint(0, 2, size=[1, 2, 1, 3]),
+    },
+    (2, 3)
+]
+
+TEST_CASE_5 = [
+    {
+        'keys': ['img', 'seg'],
+        'dim': -2
+    },
+    {
+        'img': np.random.rand(1, 1, 16, 8, 1),
+        'seg': np.random.randint(0, 2, size=[1, 1, 16, 8, 1]),
+    }
+]
+
+TEST_CASE_6 = [
+    {
+        'keys': ['img', 'seg'],
+        'dim': 0.5
+    },
+    {
+        'img': np.random.rand(1, 1, 16, 8, 1),
+        'seg': np.random.randint(0, 2, size=[1, 1, 16, 8, 1]),
+    }
+]
+
+
+class TestSqueezeDim(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    def test_shape(self, input_param, test_data, expected_shape):
+        result = SqueezeDimd(**input_param)(test_data)
+        self.assertTupleEqual(result['img'].shape, expected_shape)
+        self.assertTupleEqual(result['seg'].shape, expected_shape)
+
+    @parameterized.expand([TEST_CASE_5, TEST_CASE_6])
+    def test_invalid_inputs(self, input_param, test_data):
+        with self.assertRaises(AssertionError):
+            result = SqueezeDimd(**input_param)(test_data)
+
+
+if __name__ == '__main__':
+    unittest.main()