Squeezing functionality in Monai transforms #299
Description
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