Repository: Project-MONAI/MONAI

Issue:
lambda transform
**Is your feature request related to a problem? Please describe.**
this feature request is about providing a transform that supports use-defined simple lambda functions.

a use case came from the experiments is, for example,
only retaining the first four classes from a segmentation ground truth:

assuming the input
`input_data={'image': array, 'label': np.ones((10, 2, 2))}`
a lambda transform could generate:
`data={'image': array, 'label': input_data['label'][:4, :, :]}`

**Describe the solution you'd like**
for that example, the proposed transform would be (writing in-place):
```python
monai.transforms.Lambda(keys='label', func=lambda x: x[:4, :, :])
```

**Describe alternatives you've considered**
Since transform is just a callable, it could be done with the build-in lambda:
```python
lambda data_dict: data_dict.update({'label': data_dict['label'][:4,:,:]}) or data_dict
```
but the code is difficult to follow

**Additional context**
see also torchvision Lambda https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.Lambda


Discussion:

Patch:
diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst
index 673d4bae..b02db767 100644
--- a/docs/source/transforms.rst
+++ b/docs/source/transforms.rst
@@ -343,6 +343,12 @@ Vanilla Transforms
     :members:
     :special-members: __call__
 
+`Lambda`
+~~~~~~~~
+.. autoclass:: Lambda
+    :members:
+    :special-members: __call__
+
 `SplitChannel`
 ~~~~~~~~~~~~~~
 .. autoclass:: SplitChannel
@@ -695,6 +701,12 @@ Dictionary-based Transforms
     :members:
     :special-members: __call__
 
+`Lambdad`
+~~~~~~~~~
+.. autoclass:: Lambdad
+    :members:
+    :special-members: __call__
+
 
 Transform Adaptors
 ------------------
diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py
index 5b2eb426..a82c8f8b 100644
--- a/monai/transforms/utility/array.py
+++ b/monai/transforms/utility/array.py
@@ -291,3 +291,29 @@ class SimulateDelay(Transform):
     def __call__(self, img, delay_time=None):
         time.sleep(self.delay_time if delay_time is None else delay_time)
         return img
+
+
+class Lambda(Transform):
+    """
+    Apply a user-defined lambda as a transform.
+
+    For example:
+
+    .. code-block:: python
+        :emphasize-lines: 2
+
+        image = np.ones((10, 2, 2))
+        lambd = Lambda(func=lambda x: x[:4, :, :])
+        print(lambd(image).shape)
+        (4, 2, 2)
+
+    Args:
+        func: Lambda/function to be applied.
+    """
+
+    def __init__(self, func: Callable) -> None:
+        assert callable(func), "func must be callable"
+        self.func = func
+
+    def __call__(self, img):
+        return self.func(img)
diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py
index 8ad2d211..7b8d9f84 100644
--- a/monai/transforms/utility/dictionary.py
+++ b/monai/transforms/utility/dictionary.py
@@ -16,7 +16,7 @@ Class names are ended with 'd' to denote dictionary-based transforms.
 """
 
 from logging import Handler
-from typing import Optional
+from typing import Optional, Callable
 import copy
 import torch
 import numpy as np
@@ -36,6 +36,7 @@ from monai.transforms.utility.array import (
     DataStats,
     SimulateDelay,
     Identity,
+    Lambda,
 )
 
 
@@ -413,6 +414,37 @@ class ConcatItemsd(MapTransform):
         return d
 
 
+class Lambdad(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.Lambda`.
+
+    For example:
+
+    .. code-block:: python
+        :emphasize-lines: 2
+
+        input_data={'image': np.zeros((10, 2, 2)), 'label': np.ones((10, 2, 2))}
+        lambd = Lambdad(keys='label', func=lambda x: x[:4, :, :])
+        print(lambd(input_data)['label'].shape)
+        (4, 2, 2)
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: :py:class:`monai.transforms.compose.MapTransform`
+        func: Lambda/function to be applied.
+    """
+
+    def __init__(self, keys: KeysCollection, func: Callable) -> None:
+        super().__init__(keys)
+        self.lambd = Lambda(func)
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.keys:
+            d[key] = self.lambd(d[key])
+        return d
+
+
 IdentityD = IdentityDict = Identityd
 AsChannelFirstD = AsChannelFirstDict = AsChannelFirstd
 AsChannelLastD = AsChannelLastDict = AsChannelLastd
@@ -426,3 +458,4 @@ DataStatsD = DataStatsDict = DataStatsd
 SimulateDelayD = SimulateDelayDict = SimulateDelayd
 CopyItemsD = CopyItemsDict = CopyItemsd
 ConcatItemsD = ConcatItemsDict = ConcatItemsd
+LambdaD = LambdaDict = Lambdad


Test patch:
diff --git a/tests/test_lambda.py b/tests/test_lambda.py
new file mode 100644
index 00000000..96ebd887
--- /dev/null
+++ b/tests/test_lambda.py
@@ -0,0 +1,41 @@
+# 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 monai.transforms.utility.array import Lambda
+from tests.utils import NumpyImageTestCase2D
+
+
+class TestLambda(NumpyImageTestCase2D):
+    def test_lambda_identity(self):
+        img = self.imt
+
+        def identity_func(x):
+            return x
+
+        lambd = Lambda(func=identity_func)
+        self.assertTrue(np.allclose(identity_func(img), lambd(img)))
+
+    def test_lambda_slicing(self):
+        img = self.imt
+
+        def slice_func(x):
+            return x[:, :, :6, ::-2]
+
+        lambd = Lambda(func=slice_func)
+        self.assertTrue(np.allclose(slice_func(img), lambd(img)))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_lambdad.py b/tests/test_lambdad.py
new file mode 100644
index 00000000..fbebe081
--- /dev/null
+++ b/tests/test_lambdad.py
@@ -0,0 +1,49 @@
+# 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 monai.transforms.utility.dictionary import Lambdad
+from tests.utils import NumpyImageTestCase2D
+
+
+class TestLambdad(NumpyImageTestCase2D):
+    def test_lambdad_identity(self):
+        img = self.imt
+        data = dict()
+        data["img"] = img
+
+        def identity_func(x):
+            return x
+
+        lambd = Lambdad(keys=data.keys(), func=identity_func)
+        expected = data
+        expected["img"] = identity_func(data["img"])
+        self.assertTrue(np.allclose(expected["img"], lambd(data)["img"]))
+
+    def test_lambdad_slicing(self):
+        img = self.imt
+        data = dict()
+        data["img"] = img
+
+        def slice_func(x):
+            return x[:, :, :6, ::-2]
+
+        lambd = Lambdad(keys=data.keys(), func=slice_func)
+        expected = dict()
+        expected["img"] = slice_func(data["img"])
+        self.assertTrue(np.allclose(expected["img"], lambd(data)["img"]))
+
+
+if __name__ == "__main__":
+    unittest.main()