kaiko-eva 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

eva/vision/data/datasets/wsi.py

@@ -11,12 +11,12 @@ from torchvision import tv_tensors
 from torchvision.transforms.v2 import functional
 from typing_extensions import override
 
+from eva.core.data.datasets import base
 from eva.vision.data import wsi
-from eva.vision.data.datasets import vision
 from eva.vision.data.wsi.patching import samplers
 
 
-class WsiDataset(vision.VisionDataset):
+class WsiDataset(base.MapDataset):
     """Dataset class for reading patches from whole-slide images."""
 
     def __init__(
@@ -57,8 +57,8 @@ class WsiDataset(vision.VisionDataset):
     def __len__(self):
         return len(self._coords.x_y)
 
-    @override
     def filename(self, index: int) -> str:
+        """Returns the filename of the patch at the specified index."""
         return f"{self._file_path}_{index}"
 
     @property
@@ -103,7 +103,7 @@ class WsiDataset(vision.VisionDataset):
         return image
 
 
-class MultiWsiDataset(vision.VisionDataset):
+class MultiWsiDataset(base.MapDataset):
     """Dataset class for reading patches from multiple whole-slide images."""
 
     def __init__(
@@ -171,8 +171,8 @@ class MultiWsiDataset(vision.VisionDataset):
     def __getitem__(self, index: int) -> tv_tensors.Image:
         return self._concat_dataset[index]
 
-    @override
     def filename(self, index: int) -> str:
+        """Returns the filename of the patch at the specified index."""
         return os.path.basename(self._file_paths[self._get_dataset_idx(index)])
 
     def load_metadata(self, index: int) -> Dict[str, Any]:

eva/vision/data/transforms/__init__.py

@@ -1,6 +1,25 @@
 """Vision data transforms."""
 
-from eva.vision.data.transforms.common import ResizeAndClamp, ResizeAndCrop
-from eva.vision.data.transforms.normalization import Clamp, RescaleIntensity
+from eva.vision.data.transforms.common import ResizeAndCrop
+from eva.vision.data.transforms.croppad import CropForeground, RandCropByPosNegLabel, SpatialPad
+from eva.vision.data.transforms.intensity import (
+    RandScaleIntensity,
+    RandShiftIntensity,
+    ScaleIntensityRange,
+)
+from eva.vision.data.transforms.spatial import RandFlip, RandRotate90, Spacing
+from eva.vision.data.transforms.utility import EnsureChannelFirst
 
-__all__ = ["ResizeAndCrop", "ResizeAndClamp", "Clamp", "RescaleIntensity"]
+__all__ = [
+    "ResizeAndCrop",
+    "CropForeground",
+    "RandCropByPosNegLabel",
+    "SpatialPad",
+    "RandScaleIntensity",
+    "RandShiftIntensity",
+    "ScaleIntensityRange",
+    "RandFlip",
+    "RandRotate90",
+    "Spacing",
+    "EnsureChannelFirst",
+]

eva/vision/data/transforms/common/__init__.py

@@ -1,6 +1,5 @@
 """Common vision transforms."""
 
-from eva.vision.data.transforms.common.resize_and_clamp import ResizeAndClamp
 from eva.vision.data.transforms.common.resize_and_crop import ResizeAndCrop
 
-__all__ = ["ResizeAndClamp", "ResizeAndCrop"]
+__all__ = ["ResizeAndCrop"]

eva/vision/data/transforms/croppad/__init__.py

@@ -0,0 +1,11 @@
+"""Transforms for crop and pad operations."""
+
+from eva.vision.data.transforms.croppad.crop_foreground import CropForeground
+from eva.vision.data.transforms.croppad.rand_crop_by_pos_neg_label import RandCropByPosNegLabel
+from eva.vision.data.transforms.croppad.spatial_pad import SpatialPad
+
+__all__ = [
+    "CropForeground",
+    "RandCropByPosNegLabel",
+    "SpatialPad",
+]

eva/vision/data/transforms/croppad/crop_foreground.py

@@ -0,0 +1,110 @@
+"""Crop foreground transform."""
+
+import functools
+from typing import Any, Dict, List, Sequence
+
+import torch
+from monai.config import type_definitions
+from monai.transforms.croppad import array as monai_croppad_transforms
+from monai.utils.enums import PytorchPadMode
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class CropForeground(v2.Transform):
+    """Crop an image using a bounding box.
+
+    The bounding box is generated by selecting foreground using select_fn
+    at channels channel_indices. margin is added in each spatial dimension
+    of the bounding box. The typical usage is to help training and evaluation
+    if the valid part is small in the whole medical image.
+    """
+
+    def __init__(
+        self,
+        threshold: float = 0.0,
+        channel_indices: type_definitions.IndexSelection | None = None,
+        margin: Sequence[int] | int = 0,
+        allow_smaller: bool = True,
+        return_coords: bool = False,
+        k_divisible: Sequence[int] | int = 1,
+        mode: str = PytorchPadMode.CONSTANT,
+        **pad_kwargs,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            threshold: function to select expected foreground.
+            channel_indices: if defined, select foreground only on the specified channels
+                of image. if None, select foreground on the whole image.
+            margin: add margin value to spatial dims of the bounding box, if only 1 value provided,
+                use it for all dims.
+            allow_smaller: when computing box size with `margin`, whether to allow the image edges
+                to be smaller than the final box edges. If `False`, part of a padded output box
+                might be outside of the original image, if `True`, the image edges will be used as
+                the box edges. Default to `True`.
+            return_coords: whether return the coordinates of spatial bounding box for foreground.
+            k_divisible: make each spatial dimension to be divisible by k, default to 1.
+                if `k_divisible` is an int, the same `k` be applied to all the input spatial
+                dimensions.
+            mode: available modes for numpy array:{``"constant"``, ``"edge"``, ``"linear_ramp"``,
+                ``"maximum"``, ``"mean"``, ``"median"``, ``"minimum"``, ``"reflect"``,
+                ``"symmetric"``, ``"wrap"``, ``"empty"``} available modes for PyTorch Tensor:
+                {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}. One of the listed
+                string values or a user supplied function. Defaults to ``"constant"``.
+                See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
+                https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
+            pad_kwargs: other arguments for the `np.pad` or `torch.pad` function.
+                note that `np.pad` treats channel dimension as the first dimension.
+        """
+        super().__init__()
+
+        self._foreground_crop = monai_croppad_transforms.CropForeground(
+            select_fn=functools.partial(_threshold_fn, threshold=threshold),
+            channel_indices=channel_indices,
+            margin=margin,
+            allow_smaller=allow_smaller,
+            return_coords=return_coords,
+            k_divisible=k_divisible,
+            mode=mode,
+            lazy=False,
+            **pad_kwargs,
+        )
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        volume = next(inpt for inpt in flat_inputs if isinstance(inpt, eva_tv_tensors.Volume))
+        box_start, box_end = self._foreground_crop.compute_bounding_box(volume)
+        return {"box_start": box_start, "box_end": box_end}
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_foreground_cropped = self._foreground_crop.crop_pad(
+            inpt, params["box_start"], params["box_end"]
+        )
+        return tv_tensors.wrap(inpt_foreground_cropped, like=inpt)
+
+
+def _threshold_fn(image: torch.Tensor, threshold: int | float = 0.3) -> torch.Tensor:
+    """Applies a thresholding operation to an image tensor.
+
+    Pixels greater than the threshold are set to True, while others are False.
+
+    Args:
+        image: Input image tensor with pixel values.
+        threshold: Threshold value.
+
+    Returns:
+        A binary mask tensor of the same shape as `image`,
+        where True represents pixels above the threshold.
+    """
+    return image > threshold

eva/vision/data/transforms/croppad/rand_crop_by_pos_neg_label.py

@@ -0,0 +1,109 @@
+"""Crop foreground transform."""
+
+import functools
+from typing import Any, Dict, List, Sequence
+
+import torch
+from monai.config.type_definitions import NdarrayOrTensor
+from monai.transforms.croppad import array as monai_croppad_transforms
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class RandCropByPosNegLabel(v2.Transform):
+    """Crop random fixed sized regions with the center being a foreground or background voxel.
+
+    Its based on the Pos Neg Ratio and will return a list of arrays for all the cropped images.
+    For example, crop two (3 x 3) arrays from (5 x 5) array with pos/neg=1::
+
+        [[[0, 0, 0, 0, 0],
+          [0, 1, 2, 1, 0],            [[0, 1, 2],     [[2, 1, 0],
+          [0, 1, 3, 0, 0],     -->     [0, 1, 3],      [3, 0, 0],
+          [0, 0, 0, 0, 0],             [0, 0, 0]]      [0, 0, 0]]
+          [0, 0, 0, 0, 0]]]
+
+    If a dimension of the expected spatial size is larger than the input image size,
+    will not crop that dimension. So the cropped result may be smaller than expected
+    size, and the cropped results of several images may not have exactly same shape.
+    """
+
+    def __init__(
+        self,
+        spatial_size: Sequence[int] | int,
+        label: torch.Tensor | None = None,
+        pos: float = 1.0,
+        neg: float = 1.0,
+        num_samples: int = 1,
+        image: torch.Tensor | None = None,
+        image_threshold: float = 0.0,
+        fg_indices: NdarrayOrTensor | None = None,
+        bg_indices: NdarrayOrTensor | None = None,
+        allow_smaller: bool = False,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            spatial_size: the spatial size of the crop region e.g. [224, 224, 128].
+                if a dimension of ROI size is larger than image size, will not crop that dimension.
+                if components have non-positive values, corresponding size of `label` will be used.
+                for example: if the spatial size of input data is [40, 40, 40] and
+                `spatial_size=[32, 64, -1]`, the spatial size of output data will be [32, 40, 40].
+            label: the label image that is used for finding foreground/background, if None, must
+                set at `self.__call__`. Non-zero indicates foreground, zero indicates background.
+            pos: used with `neg` together to calculate the ratio ``pos / (pos + neg)`` for
+                the probability to pick a foreground voxel as center rather than background voxel.
+            neg: used with `pos` together to calculate the ratio ``pos / (pos + neg)`` for
+                the probability to pick a foreground voxel as center rather than background voxel.
+            num_samples: number of samples (crop regions) to take in each list.
+            image: optional image data to help select valid area, can be same as `img` or another.
+                if not None, use ``label == 0 & image > image_threshold`` to select the negative
+                sample (background) center. Crop center will only come from valid image areas.
+            image_threshold: if enabled `image`, use ``image > image_threshold`` to determine
+                the valid image content areas.
+            fg_indices: if provided pre-computed foreground indices of `label`, will ignore `image`
+                and `image_threshold`, randomly select crop centers based on them, need to provide
+                `fg_indices` and `bg_indices` together, expect to be 1 dim array of spatial indices.
+                a typical usage is to call `FgBgToIndices` transform first and cache the results.
+            bg_indices: if provided pre-computed background indices of `label`, will ignore `image`
+                and `image_threshold`, randomly select crop centers based on them, need to provide
+                `fg_indices` and `bg_indices` together, expect to be 1 dim array of spatial indices.
+                a typical usage is to call `FgBgToIndices` transform first and cache the results.
+            allow_smaller: if `False`, an exception will be raised if the image is smaller than
+                the requested ROI in any dimension. If `True`, any smaller dimensions will be set to
+                match the cropped size (i.e., no cropping in that dimension).
+        """
+        super().__init__()
+
+        self._rand_crop = monai_croppad_transforms.RandCropByPosNegLabel(
+            spatial_size=spatial_size,
+            label=label,
+            pos=pos,
+            neg=neg,
+            num_samples=num_samples,
+            image=image,
+            image_threshold=image_threshold,
+            fg_indices=fg_indices,
+            bg_indices=bg_indices,
+            allow_smaller=allow_smaller,
+            lazy=False,
+        )
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        mask = next(inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.Mask))
+        self._rand_crop.randomize(label=mask)
+        return {}
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_foreground_crops = self._rand_crop(img=inpt, randomize=False)
+        return [tv_tensors.wrap(crop, like=inpt) for crop in inpt_foreground_crops]

eva/vision/data/transforms/croppad/spatial_pad.py

@@ -0,0 +1,67 @@
+"""General purpose cropper to produce sub-volume region of interest (ROI)."""
+
+import functools
+from typing import Any, Dict, Sequence
+
+from monai.transforms.croppad import array as monai_croppad_transforms
+from monai.utils.enums import Method, PytorchPadMode
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class SpatialPad(v2.Transform):
+    """Performs padding to the data.
+
+    Padding is applied symmetric for all sides or all on one side for each dimension.
+    """
+
+    def __init__(
+        self,
+        spatial_size: Sequence[int] | int | tuple[tuple[int, ...] | int, ...],
+        method: str = Method.SYMMETRIC,
+        mode: str = PytorchPadMode.CONSTANT,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            spatial_size: The spatial size of output data after padding.
+                If a dimension of the input data size is larger than the
+                pad size, will not pad that dimension. If its components
+                have non-positive values, the corresponding size of input
+                image will be used (no padding). for example: if the spatial
+                size of input data is [30, 30, 30] and `spatial_size=[32, 25, -1]`,
+                the spatial size of output data will be [32, 30, 30].
+            method: {``"symmetric"``, ``"end"``}
+                Pad image symmetrically on every side or only pad at the
+                end sides. Defaults to ``"symmetric"``.
+            mode: available modes for numpy array:{``"constant"``, ``"edge"``,
+                ``"linear_ramp"``, ``"maximum"``, ``"mean"``, ``"median"``, ``"minimum"``,
+                ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
+                available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``,
+                ``"circular"``}. One of the listed string values or a user supplied function.
+                Defaults to ``"constant"``.
+                See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
+                https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
+        """
+        super().__init__()
+
+        self._spatial_pad = monai_croppad_transforms.SpatialPad(
+            spatial_size=spatial_size,
+            method=method,
+            mode=mode,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_padded = self._spatial_pad(inpt)
+        return tv_tensors.wrap(inpt_padded, like=inpt)

eva/vision/data/transforms/intensity/__init__.py

@@ -0,0 +1,11 @@
+"""Transforms for intensity adjustment."""
+
+from eva.vision.data.transforms.intensity.rand_scale_intensity import RandScaleIntensity
+from eva.vision.data.transforms.intensity.rand_shift_intensity import RandShiftIntensity
+from eva.vision.data.transforms.intensity.scale_intensity_ranged import ScaleIntensityRange
+
+__all__ = [
+    "RandScaleIntensity",
+    "RandShiftIntensity",
+    "ScaleIntensityRange",
+]

eva/vision/data/transforms/intensity/rand_scale_intensity.py

@@ -0,0 +1,59 @@
+"""Intensity scaling transform."""
+
+import functools
+from typing import Any, Dict
+
+import numpy as np
+from monai.config.type_definitions import DtypeLike
+from monai.transforms.intensity import array as monai_intensity_transforms
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class RandScaleIntensity(v2.Transform):
+    """Randomly scale the intensity of input image.
+
+    The factor is by ``v = v * (1 + factor)``, where
+    the `factor` is randomly picked.
+    """
+
+    def __init__(
+        self,
+        factors: tuple[float, float] | float,
+        prob: float = 0.1,
+        channel_wise: bool = False,
+        dtype: DtypeLike = np.float32,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            factors: factor range to randomly scale by ``v = v * (1 + factor)``.
+                if single number, factor value is picked from (-factors, factors).
+            prob: probability of scale.
+            channel_wise: if True, shift intensity on each channel separately.
+                For each channel, a random offset will be chosen. Please ensure
+                that the first dimension represents the channel of the image if True.
+            dtype: output data type, if None, same as input image. defaults to float32.
+        """
+        super().__init__()
+
+        self._rand_scale_intensity = monai_intensity_transforms.RandScaleIntensity(
+            factors=factors,
+            prob=prob,
+            channel_wise=channel_wise,
+            dtype=dtype,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
+        inpt_scaled = self._rand_scale_intensity(inpt)
+        return tv_tensors.wrap(inpt_scaled, like=inpt)

eva/vision/data/transforms/intensity/rand_shift_intensity.py

@@ -0,0 +1,55 @@
+"""Intensity shifting transform."""
+
+import functools
+from typing import Any, Dict
+
+from monai.transforms.intensity import array as monai_intensity_transforms
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class RandShiftIntensity(v2.Transform):
+    """Randomly shift intensity with randomly picked offset."""
+
+    def __init__(
+        self,
+        offsets: tuple[float, float] | float,
+        safe: bool = False,
+        prob: float = 0.1,
+        channel_wise: bool = False,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            offsets: Offset range to randomly shift.
+                if single number, offset value is picked from (-offsets, offsets).
+            safe: If `True`, then do safe dtype convert when intensity overflow.
+                E.g., `[256, -12]` -> `[array(0), array(244)]`. If `True`, then
+                `[256, -12]` -> `[array(255), array(0)]`.
+            prob: Probability of shift.
+            channel_wise: If True, shift intensity on each channel separately.
+                For each channel, a random offset will be chosen. Please ensure
+                that the first dimension represents the channel of the image if True.
+        """
+        super().__init__()
+
+        self._rand_swift_intensity = monai_intensity_transforms.RandShiftIntensity(
+            offsets=offsets,
+            safe=safe,
+            prob=prob,
+            channel_wise=channel_wise,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
+        inpt_scaled = self._rand_swift_intensity(inpt)
+        return tv_tensors.wrap(inpt_scaled, like=inpt)

eva/vision/data/transforms/intensity/scale_intensity_ranged.py

@@ -0,0 +1,56 @@
+"""Intensity scaling transform."""
+
+import functools
+from typing import Any, Dict, Tuple
+
+from monai.transforms.intensity import array as monai_intensity_transforms
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class ScaleIntensityRange(v2.Transform):
+    """Intensity scaling transform.
+
+    Scaling from [a_min, a_max] to [b_min, b_max] with clip option.
+
+    When `b_min` or `b_max` are `None`, `scaled_array * (b_max - b_min) + b_min`
+    will be skipped. If `clip=True`, when `b_min`/`b_max` is None, the clipping
+    is not performed on the corresponding edge.
+    """
+
+    def __init__(
+        self,
+        input_range: Tuple[float, float],
+        output_range: Tuple[float, float] | None = None,
+        clip: bool = True,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            input_range: Intensity original range min and max.
+            output_range: Intensity target range min and max.
+            clip: Whether to perform clip after scaling.
+        """
+        super().__init__()
+
+        self._scale_intensity_range = monai_intensity_transforms.ScaleIntensityRange(
+            a_min=input_range[0],
+            a_max=input_range[1],
+            b_min=output_range[0] if output_range else None,
+            b_max=output_range[1] if output_range else None,
+            clip=clip,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
+        inpt_scaled = self._scale_intensity_range(inpt)
+        return tv_tensors.wrap(inpt_scaled, like=inpt)

eva/vision/data/transforms/spatial/__init__.py

@@ -0,0 +1,7 @@
+"""Transforms for spatial operations."""
+
+from eva.vision.data.transforms.spatial.flip import RandFlip
+from eva.vision.data.transforms.spatial.rotate import RandRotate90
+from eva.vision.data.transforms.spatial.spacing import Spacing
+
+__all__ = ["Spacing", "RandFlip", "RandRotate90"]

eva/vision/data/transforms/spatial/flip.py

@@ -0,0 +1,72 @@
+"""Flip transforms."""
+
+import functools
+from typing import Any, Dict, List, Sequence
+
+import torch
+from monai.transforms.spatial import array as monai_spatial_transforms
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class RandFlip(v2.Transform):
+    """Randomly flips the image along axes."""
+
+    def __init__(
+        self,
+        prob: float = 0.1,
+        spatial_axes: Sequence[int] | int | None = None,
+        apply_per_axis: bool = True,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            prob: Probability of flipping.
+            spatial_axes: Spatial axes along which to flip over. Default is None.
+            apply_per_axis: If True, will apply a random flip transform to each
+                axis individually (if spatial_axes is a sequence of multiple axis).
+                If False, will apply a single random flip transform applied to all axes.
+        """
+        super().__init__()
+
+        if apply_per_axis:
+            if not isinstance(spatial_axes, (list, tuple)):
+                raise ValueError(
+                    "`spatial_axis` is expected to be sequence `apply_per_axis` "
+                    f"is enabled, got {type(spatial_axes)}"
+                )
+            self._flips = [
+                monai_spatial_transforms.RandFlip(prob=prob, spatial_axis=axis)
+                for axis in spatial_axes
+            ]
+        else:
+            self._flips = [monai_spatial_transforms.RandFlip(prob=prob, spatial_axis=spatial_axes)]
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        for flip in self._flips:
+            flip.randomize(None)
+        return {}
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_flipped = self._apply_flips(inpt)
+        return tv_tensors.wrap(inpt_flipped, like=inpt)
+
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_flipped = torch.tensor(self._apply_flips(inpt), dtype=torch.long)
+        return tv_tensors.wrap(inpt_flipped, like=inpt)
+
+    def _apply_flips(self, inpt: Any) -> Any:
+        for flip in self._flips:
+            inpt = flip(img=inpt, randomize=False)
+        return inpt