kaiko-eva 0.1.8__py3-none-any.whl → 0.2.1__py3-none-any.whl

eva/vision/data/datasets/segmentation/monusac.py

@@ -16,12 +16,11 @@ from torchvision.datasets import utils
 from typing_extensions import override
 
 from eva.core.utils.progress_bar import tqdm
-from eva.vision.data.datasets import _validators, structs
-from eva.vision.data.datasets.segmentation import base
+from eva.vision.data.datasets import _validators, structs, vision
 from eva.vision.utils import io
 
 
-class MoNuSAC(base.ImageSegmentation):
+class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
     """MoNuSAC2020: A Multi-organ Nuclei Segmentation and Classification Challenge.
 
     Webpage: https://monusac-2020.grand-challenge.org/
@@ -112,13 +111,13 @@ class MoNuSAC(base.ImageSegmentation):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_path = self._image_files[index]
         image_rgb_array = io.read_image(image_path)
         return tv_tensors.Image(image_rgb_array.transpose(2, 0, 1))
 
     @override
-    def load_mask(self, index: int) -> tv_tensors.Mask:
+    def load_target(self, index: int) -> tv_tensors.Mask:
         semantic_labels = (
             self._load_semantic_mask_file(index)
             if self._export_masks

eva/vision/data/datasets/segmentation/total_segmentator_2d.py

@@ -17,12 +17,12 @@ from typing_extensions import override
 
 from eva.core.utils import io as core_io
 from eva.core.utils import multiprocessing
-from eva.vision.data.datasets import _validators, structs
-from eva.vision.data.datasets.segmentation import _total_segmentator, base
+from eva.vision.data.datasets import _validators, structs, vision
+from eva.vision.data.datasets.segmentation import _total_segmentator
 from eva.vision.utils import io
 
 
-class TotalSegmentator2D(base.ImageSegmentation):
+class TotalSegmentator2D(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
     """TotalSegmentator 2D segmentation dataset."""
 
     _expected_dataset_lengths: Dict[str, int] = {
@@ -206,19 +206,20 @@ class TotalSegmentator2D(base.ImageSegmentation):
         return len(self._indices)
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         sample_index, slice_index = self._indices[index]
         image_path = self._get_image_path(sample_index)
-        image_array = io.read_nifti(image_path, slice_index)
+        image_nii = io.read_nifti(image_path, slice_index)
+        image_array = io.nifti_to_array(image_nii)
         image_array = self._fix_orientation(image_array)
         return tv_tensors.Image(image_array.copy().transpose(2, 0, 1))
 
     @override
-    def load_mask(self, index: int) -> tv_tensors.Mask:
+    def load_target(self, index: int) -> tv_tensors.Mask:
         if self._optimize_mask_loading:
             mask = self._load_semantic_label_mask(index)
         else:
-            mask = self._load_mask(index)
+            mask = self._load_target(index)
         mask = self._fix_orientation(mask)
         return tv_tensors.Mask(mask.copy().squeeze(), dtype=torch.int64)  # type: ignore
 
@@ -227,14 +228,15 @@ class TotalSegmentator2D(base.ImageSegmentation):
         _, slice_index = self._indices[index]
         return {"slice_index": slice_index}
 
-    def _load_mask(self, index: int) -> npt.NDArray[Any]:
+    def _load_target(self, index: int) -> npt.NDArray[Any]:
         sample_index, slice_index = self._indices[index]
         return self._load_masks_as_semantic_label(sample_index, slice_index)
 
     def _load_semantic_label_mask(self, index: int) -> npt.NDArray[Any]:
         """Loads the segmentation mask from a semantic label NifTi file."""
         sample_index, slice_index = self._indices[index]
-        return io.read_nifti(self._get_optimized_masks_file(sample_index), slice_index)
+        nii = io.read_nifti(self._get_optimized_masks_file(sample_index), slice_index)
+        return io.nifti_to_array(nii)
 
     def _load_masks_as_semantic_label(
         self, sample_index: int, slice_index: int | None = None
@@ -248,7 +250,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
         masks_dir = self._get_masks_dir(sample_index)
         classes = self._class_mappings.keys() if self._class_mappings else self.classes[1:]
         mask_paths = [os.path.join(masks_dir, f"{label}.nii.gz") for label in classes]
-        binary_masks = [io.read_nifti(path, slice_index) for path in mask_paths]
+        binary_masks = [io.nifti_to_array(io.read_nifti(path, slice_index)) for path in mask_paths]
 
         if self._class_mappings:
             mapped_binary_masks = [np.zeros_like(binary_masks[0], dtype=np.bool_)] * len(

eva/vision/data/datasets/vision.py

@@ -1,17 +1,92 @@
 """Vision Dataset base class."""
 
 import abc
-from typing import Generic, TypeVar
+from typing import Any, Callable, Dict, Generic, List, Tuple, TypeVar
 
 from eva.core.data.datasets import base
 
-DataSample = TypeVar("DataSample")
-"""The data sample type."""
+InputType = TypeVar("InputType")
+"""The input data type."""
 
+TargetType = TypeVar("TargetType")
+"""The target data type."""
 
-class VisionDataset(base.MapDataset, abc.ABC, Generic[DataSample]):
+
+class VisionDataset(
+    base.MapDataset[Tuple[InputType, TargetType, Dict[str, Any]]],
+    abc.ABC,
+    Generic[InputType, TargetType],
+):
     """Base dataset class for vision tasks."""
 
+    def __init__(
+        self,
+        transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            transforms: A function/transform which returns a transformed
+                version of the raw data samples.
+        """
+        super().__init__()
+
+        self._transforms = transforms
+
+    @property
+    def classes(self) -> List[str] | None:
+        """Returns the list with names of the dataset names."""
+
+    @property
+    def class_to_idx(self) -> Dict[str, int] | None:
+        """Returns a mapping of the class name to its target index."""
+
+    def __getitem__(self, index: int) -> Tuple[InputType, TargetType, Dict[str, Any]]:
+        """Returns the `index`'th data sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            A tuple with the image, the target and the metadata.
+        """
+        image = self.load_data(index)
+        target = self.load_target(index)
+        image, target = self._apply_transforms(image, target)
+        return image, target, self.load_metadata(index) or {}
+
+    def load_metadata(self, index: int) -> Dict[str, Any] | None:
+        """Returns the dataset metadata.
+
+        Args:
+            index: The index of the data sample to return the metadata of.
+
+        Returns:
+            The sample metadata.
+        """
+
+    @abc.abstractmethod
+    def load_data(self, index: int) -> InputType:
+        """Returns the `index`'th data sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The sample data.
+        """
+
+    @abc.abstractmethod
+    def load_target(self, index: int) -> TargetType:
+        """Returns the `index`'th target sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The sample target.
+        """
+
     @abc.abstractmethod
     def filename(self, index: int) -> str:
         """Returns the filename of the `index`'th data sample.
@@ -24,3 +99,19 @@ class VisionDataset(base.MapDataset, abc.ABC, Generic[DataSample]):
         Returns:
             The filename of the `index`'th data sample.
         """
+
+    def _apply_transforms(
+        self, image: InputType, target: TargetType
+    ) -> Tuple[InputType, TargetType]:
+        """Applies the transforms to the provided data and returns them.
+
+        Args:
+            image: The desired image.
+            target: The target of the image.
+
+        Returns:
+            A tuple with the image and the target transformed.
+        """
+        if self._transforms is not None:
+            image, target = self._transforms(image, target)
+        return image, target

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)