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:

Patch:
diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst
index 96513066..a9a1e9fa 100644
--- a/docs/source/transforms.rst
+++ b/docs/source/transforms.rst
@@ -297,6 +297,11 @@ Vanilla Transforms
     :members:
     :special-members: __call__
 
+`SqueezeDim`
+~~~~~~~~~~~~
+.. autoclass:: SqueezeDim
+    :members:
+    :special-members: __call__
 
 Dictionary-based Composables
 ----------------------------
@@ -538,6 +543,12 @@ Dictionary-based Composables
     :members:
     :special-members: __call__
 
+`SqueezeDimd`
+~~~~~~~~~~~~~
+.. autoclass:: SqueezeDimd
+    :members:
+    :special-members: __call__
+
 Transform Adaptors
 ------------------
 
diff --git a/monai/transforms/composables.py b/monai/transforms/composables.py
index b59e1fae..9e03135a 100644
--- a/monai/transforms/composables.py
+++ b/monai/transforms/composables.py
@@ -25,7 +25,8 @@ from monai.transforms.transforms import (AddChannel, AsChannelFirst, Flip, LoadN
                                          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 @@ class DeleteKeysd(MapTransform):
         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 @@ RandRotateD = RandRotateDict = RandRotated
 ZoomD = ZoomDict = Zoomd
 RandZoomD = RandZoomDict = RandZoomd
 DeleteKeysD = DeleteKeysDict = DeleteKeysd
+SqueezeDimD = SqueezeDimDict = SqueezeDimd
diff --git a/monai/transforms/transforms.py b/monai/transforms/transforms.py
index 8853130f..80926c75 100644
--- a/monai/transforms/transforms.py
+++ b/monai/transforms/transforms.py
@@ -1738,3 +1738,26 @@ class Rand3DElastic(Randomizable, Transform):
             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
index 00000000..e5dc852d
--- /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
index 00000000..d0df400a
--- /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()