transforms to have multi-sample trait (#6003)

Fixes #6002



### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

Author: wyli

monai/transforms/__init__.py

@@ -453,18 +453,8 @@
     ZoomD,
     ZoomDict,
 )
-from .transform import (
-    LazyTrait,
-    LazyTransform,
-    MapTransform,
-    MultiSampleTrait,
-    Randomizable,
-    RandomizableTrait,
-    RandomizableTransform,
-    ThreadUnsafe,
-    Transform,
-    apply_transform,
-)
+from .traits import LazyTrait, MultiSampleTrait, RandomizableTrait, ThreadUnsafe
+from .transform import LazyTransform, MapTransform, Randomizable, RandomizableTransform, Transform, apply_transform
 from .utility.array import (
     AddChannel,
     AddCoordinateChannels,

monai/transforms/croppad/array.py

@@ -30,6 +30,7 @@
 from monai.data.meta_tensor import MetaTensor
 from monai.data.utils import get_random_patch, get_valid_patch_size
 from monai.transforms.inverse import InvertibleTransform, TraceableTransform
+from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import Randomizable, Transform
 from monai.transforms.utils import (
     compute_divisible_spatial_size,
@@ -683,7 +684,7 @@ def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor:
         return super().__call__(img=img, randomize=randomize)
 
 
-class RandSpatialCropSamples(Randomizable, TraceableTransform):
+class RandSpatialCropSamples(Randomizable, TraceableTransform, MultiSampleTrait):
     """
     Crop image with random size or specific size ROI to generate a list of N samples.
     It can crop at a random position as center or at the image center. And allows to set
@@ -893,7 +894,7 @@ def inverse(self, img: MetaTensor) -> MetaTensor:
         return super().inverse(inv)
 
 
-class RandWeightedCrop(Randomizable, TraceableTransform):
+class RandWeightedCrop(Randomizable, TraceableTransform, MultiSampleTrait):
     """
     Samples a list of `num_samples` image patches according to the provided `weight_map`.
 
@@ -958,7 +959,7 @@ def __call__(
         return results
 
 
-class RandCropByPosNegLabel(Randomizable, TraceableTransform):
+class RandCropByPosNegLabel(Randomizable, TraceableTransform, MultiSampleTrait):
     """
     Crop random fixed sized regions with the center being a foreground or background voxel
     based on the Pos Neg Ratio.
@@ -1118,7 +1119,7 @@ def __call__(
         return results
 
 
-class RandCropByLabelClasses(Randomizable, TraceableTransform):
+class RandCropByLabelClasses(Randomizable, TraceableTransform, MultiSampleTrait):
     """
     Crop random fixed sized regions with the center being a class based on the specified ratios of every class.
     The label data can be One-Hot format array or Argmax data. And will return a list of arrays for all the

monai/transforms/croppad/dictionary.py

@@ -47,6 +47,7 @@
     SpatialPad,
 )
 from monai.transforms.inverse import InvertibleTransform
+from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import MapTransform, Randomizable
 from monai.transforms.utils import is_positive
 from monai.utils import MAX_SEED, Method, PytorchPadMode, deprecated_arg_default, ensure_tuple_rep
@@ -528,7 +529,7 @@ def __init__(
         super().__init__(keys, cropper=cropper, allow_missing_keys=allow_missing_keys)
 
 
-class RandSpatialCropSamplesd(Randomizable, MapTransform):
+class RandSpatialCropSamplesd(Randomizable, MapTransform, MultiSampleTrait):
     """
     Dictionary-based version :py:class:`monai.transforms.RandSpatialCropSamples`.
     Crop image with random size or specific size ROI to generate a list of N samples.
@@ -682,7 +683,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torc
         return d
 
 
-class RandWeightedCropd(Randomizable, MapTransform):
+class RandWeightedCropd(Randomizable, MapTransform, MultiSampleTrait):
     """
     Samples a list of `num_samples` image patches according to the provided `weight_map`.
 
@@ -739,7 +740,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> list[dict[Hashable,
         return ret
 
 
-class RandCropByPosNegLabeld(Randomizable, MapTransform):
+class RandCropByPosNegLabeld(Randomizable, MapTransform, MultiSampleTrait):
     """
     Dictionary-based version :py:class:`monai.transforms.RandCropByPosNegLabel`.
     Crop random fixed sized regions with the center being a foreground or background voxel
@@ -860,7 +861,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> list[dict[Hashable,
         return ret
 
 
-class RandCropByLabelClassesd(Randomizable, MapTransform):
+class RandCropByLabelClassesd(Randomizable, MapTransform, MultiSampleTrait):
     """
     Dictionary-based version :py:class:`monai.transforms.RandCropByLabelClasses`.
     Crop random fixed sized regions with the center being a class based on the specified ratios of every class.

monai/transforms/nvtx.py

@@ -14,7 +14,8 @@
 
 from __future__ import annotations
 
-from monai.transforms.transform import RandomizableTrait, Transform
+from monai.transforms.traits import RandomizableTrait
+from monai.transforms.transform import Transform
 from monai.utils import optional_import
 
 _nvtx, _ = optional_import("torch._C._nvtx", descriptor="NVTX is not installed. Are you sure you have a CUDA build?")

monai/transforms/spatial/array.py

@@ -35,6 +35,7 @@
 from monai.transforms.croppad.array import CenterSpatialCrop, ResizeWithPadOrCrop
 from monai.transforms.intensity.array import GaussianSmooth
 from monai.transforms.inverse import InvertibleTransform
+from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import Randomizable, RandomizableTransform, Transform
 from monai.transforms.utils import (
     convert_pad_mode,
@@ -3045,7 +3046,7 @@ def __call__(
         return self.grid_distortion(img, distort_steps=self.distort_steps, mode=mode, padding_mode=padding_mode)
 
 
-class GridSplit(Transform):
+class GridSplit(Transform, MultiSampleTrait):
     """
     Split the image into patches based on the provided grid in 2D.
 
@@ -3130,7 +3131,7 @@ def _get_params(self, image_size: Sequence[int] | np.ndarray, size: Sequence[int
         return size, steps
 
 
-class GridPatch(Transform):
+class GridPatch(Transform, MultiSampleTrait):
     """
     Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps.
     It can sort the patches and return all or a subset of them.
@@ -3257,7 +3258,7 @@ def __call__(self, array: NdarrayOrTensor):
         return output
 
 
-class RandGridPatch(GridPatch, RandomizableTransform):
+class RandGridPatch(GridPatch, RandomizableTransform, MultiSampleTrait):
     """
     Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps,
     and with random offset for the minimal corner of the image, (0,0) for 2D and (0,0,0) for 3D.

monai/transforms/spatial/dictionary.py

@@ -53,6 +53,7 @@
     SpatialResample,
     Zoom,
 )
+from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import MapTransform, RandomizableTransform
 from monai.transforms.utils import create_grid
 from monai.utils import (
@@ -1779,7 +1780,7 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torc
         return d
 
 
-class GridSplitd(MapTransform):
+class GridSplitd(MapTransform, MultiSampleTrait):
     """
     Split the image into patches based on the provided grid in 2D.
 
@@ -1820,7 +1821,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> list[dict[Hashab
         return output
 
 
-class GridPatchd(MapTransform):
+class GridPatchd(MapTransform, MultiSampleTrait):
     """
     Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps.
     It can sort the patches and return all or a subset of them.
@@ -1884,7 +1885,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
         return d
 
 
-class RandGridPatchd(RandomizableTransform, MapTransform):
+class RandGridPatchd(RandomizableTransform, MapTransform, MultiSampleTrait):
     """
     Extract all the patches sweeping the entire image in a row-major sliding-window manner with possible overlaps,
     and with random offset for the minimal corner of the image, (0,0) for 2D and (0,0,0) for 3D.

monai/transforms/traits.py

@@ -0,0 +1,80 @@
+# Copyright (c) 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.
+"""
+A collection of generic traits for MONAI transforms.
+"""
+
+from __future__ import annotations
+
+__all__ = ["LazyTrait", "RandomizableTrait", "MultiSampleTrait", "ThreadUnsafe"]
+
+
+class LazyTrait:
+    """
+    An interface to indicate that the transform has the capability to execute using
+    MONAI's lazy resampling feature. In order to do this, the implementing class needs
+    to be able to describe its operation as an affine matrix or grid with accompanying metadata.
+    This interface can be extended from by people adapting transforms to the MONAI framework as
+    well as by implementors of MONAI transforms.
+    """
+
+    @property
+    def lazy_evaluation(self):
+        """
+        Get whether lazy_evaluation is enabled for this transform instance.
+        Returns:
+            True if the transform is operating in a lazy fashion, False if not.
+        """
+        raise NotImplementedError()
+
+    @lazy_evaluation.setter
+    def lazy_evaluation(self, enabled: bool):
+        """
+        Set whether lazy_evaluation is enabled for this transform instance.
+        Args:
+            enabled: True if the transform should operate in a lazy fashion, False if not.
+        """
+        raise NotImplementedError()
+
+
+class RandomizableTrait:
+    """
+    An interface to indicate that the transform has the capability to perform
+    randomized transforms to the data that it is called upon. This interface
+    can be extended from by people adapting transforms to the MONAI framework as well as by
+    implementors of MONAI transforms.
+    """
+
+    pass
+
+
+class MultiSampleTrait:
+    """
+    An interface to indicate that the transform has the capability to return multiple samples
+    given an input, such as when performing random crops of a sample. This interface can be
+    extended from by people adapting transforms to the MONAI framework as well as by implementors
+    of MONAI transforms.
+    """
+
+    pass
+
+
+class ThreadUnsafe:
+    """
+    A class to denote that the transform will mutate its member variables,
+    when being applied. Transforms inheriting this class should be used
+    cautiously in a multi-thread context.
+
+    This type is typically used by :py:class:`monai.data.CacheDataset` and
+    its extensions, where the transform cache is built with multiple threads.
+    """
+
+    pass

monai/transforms/transform.py

@@ -25,16 +25,14 @@
 from monai import config, transforms
 from monai.config import KeysCollection
 from monai.data.meta_tensor import MetaTensor
+from monai.transforms.traits import LazyTrait, RandomizableTrait, ThreadUnsafe
 from monai.utils import MAX_SEED, ensure_tuple, first
 from monai.utils.enums import TransformBackends
 from monai.utils.misc import MONAIEnvVars
 
 __all__ = [
     "ThreadUnsafe",
     "apply_transform",
-    "LazyTrait",
-    "RandomizableTrait",
-    "MultiSampleTrait",
     "Randomizable",
     "LazyTransform",
     "RandomizableTransform",
@@ -132,69 +130,6 @@ def _log_stats(data, prefix: str | None = "Data"):
         raise RuntimeError(f"applying transform {transform}") from e
 
 
-class LazyTrait:
-    """
-    An interface to indicate that the transform has the capability to execute using
-    MONAI's lazy resampling feature. In order to do this, the implementing class needs
-    to be able to describe its operation as an affine matrix or grid with accompanying metadata.
-    This interface can be extended from by people adapting transforms to the MONAI framework as
-    well as by implementors of MONAI transforms.
-    """
-
-    @property
-    def lazy_evaluation(self):
-        """
-        Get whether lazy_evaluation is enabled for this transform instance.
-        Returns:
-            True if the transform is operating in a lazy fashion, False if not.
-        """
-        raise NotImplementedError()
-
-    @lazy_evaluation.setter
-    def lazy_evaluation(self, enabled: bool):
-        """
-        Set whether lazy_evaluation is enabled for this transform instance.
-        Args:
-            enabled: True if the transform should operate in a lazy fashion, False if not.
-        """
-        raise NotImplementedError()
-
-
-class RandomizableTrait:
-    """
-    An interface to indicate that the transform has the capability to perform
-    randomized transforms to the data that it is called upon. This interface
-    can be extended from by people adapting transforms to the MONAI framework as well as by
-    implementors of MONAI transforms.
-    """
-
-    pass
-
-
-class MultiSampleTrait:
-    """
-    An interface to indicate that the transform has the capability to return multiple samples
-    given an input, such as when performing random crops of a sample. This interface can be
-    extended from by people adapting transforms to the MONAI framework as well as by implementors
-    of MONAI transforms.
-    """
-
-    pass
-
-
-class ThreadUnsafe:
-    """
-    A class to denote that the transform will mutate its member variables,
-    when being applied. Transforms inheriting this class should be used
-    cautiously in a multi-thread context.
-
-    This type is typically used by :py:class:`monai.data.CacheDataset` and
-    its extensions, where the transform cache is built with multiple threads.
-    """
-
-    pass
-
-
 class Randomizable(ThreadUnsafe, RandomizableTrait):
     """
     An interface for handling random state locally, currently based on a class

monai/transforms/utility/array.py

@@ -44,6 +44,7 @@
     median_filter,
 )
 from monai.transforms.inverse import InvertibleTransform
+from monai.transforms.traits import MultiSampleTrait
 from monai.transforms.transform import Randomizable, RandomizableTrait, RandomizableTransform, Transform
 from monai.transforms.utils import (
     extreme_points_to_image,
@@ -343,7 +344,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         return out
 
 
-class SplitDim(Transform):
+class SplitDim(Transform, MultiSampleTrait):
     """
     Given an image of size X along a certain dimension, return a list of length X containing
     images. Useful for converting 3D images into a stack of 2D images, splitting multichannel inputs into
@@ -976,7 +977,7 @@ def __call__(
         return data
 
 
-class FgBgToIndices(Transform):
+class FgBgToIndices(Transform, MultiSampleTrait):
     """
     Compute foreground and background of the input label data, return the indices.
     If no output_shape specified, output data will be 1 dim indices after flattening.
@@ -1017,7 +1018,7 @@ def __call__(
         return fg_indices, bg_indices
 
 
-class ClassesToIndices(Transform):
+class ClassesToIndices(Transform, MultiSampleTrait):
     backend = [TransformBackends.NUMPY, TransformBackends.TORCH]
 
     def __init__(