kaiko-eva 0.0.2__py3-none-any.whl → 0.1.0__py3-none-any.whl

eva/vision/data/datasets/segmentation/{total_segmentator.py → total_segmentator_2d.py}

@@ -3,24 +3,30 @@
 import functools
 import os
 from glob import glob
-from typing import Callable, Dict, List, Literal, Tuple
+from typing import Any, Callable, Dict, List, Literal, Tuple
 
 import numpy as np
+import numpy.typing as npt
+import torch
+import tqdm
 from torchvision import tv_tensors
 from torchvision.datasets import utils
 from typing_extensions import override
 
-from eva.vision.data.datasets import _utils, _validators, structs
+from eva.vision.data.datasets import _validators, structs
 from eva.vision.data.datasets.segmentation import base
-from eva.vision.utils import convert, io
+from eva.vision.utils import io
 
 
 class TotalSegmentator2D(base.ImageSegmentation):
     """TotalSegmentator 2D segmentation dataset."""
 
     _expected_dataset_lengths: Dict[str, int] = {
-        "train_small": 29892,
-        "val_small": 6480,
+        "train_small": 35089,
+        "val_small": 1283,
+        "train_full": 278190,
+        "val_full": 14095,
+        "test_full": 25578,
     }
     """Dataset version and split to the expected size."""
 
@@ -45,13 +51,20 @@ class TotalSegmentator2D(base.ImageSegmentation):
     ]
     """Resources for the small dataset version."""
 
+    _license: str = (
+        "Creative Commons Attribution 4.0 International "
+        "(https://creativecommons.org/licenses/by/4.0/deed.en)"
+    )
+    """Dataset license."""
+
     def __init__(
         self,
         root: str,
-        split: Literal["train", "val"] | None,
-        version: Literal["small", "full"] | None = "small",
+        split: Literal["train", "val", "test"] | None,
+        version: Literal["small", "full"] | None = "full",
         download: bool = False,
-        as_uint8: bool = True,
+        classes: List[str] | None = None,
+        optimize_mask_loading: bool = True,
         transforms: Callable | None = None,
     ) -> None:
         """Initialize dataset.
@@ -66,7 +79,12 @@ class TotalSegmentator2D(base.ImageSegmentation):
                 Note that the download will be executed only by additionally
                 calling the :meth:`prepare_data` method and if the data does not
                 exist yet on disk.
-            as_uint8: Whether to convert and return the images as a 8-bit.
+            classes: Whether to configure the dataset with a subset of classes.
+                If `None`, it will use all of them.
+            optimize_mask_loading: Whether to pre-process the segmentation masks
+                in order to optimize the loading time. In the `setup` method, it
+                will reformat the binary one-hot masks to a semantic mask and store
+                it on disk.
             transforms: A function/transforms that takes in an image and a target
                 mask and returns the transformed versions of both.
         """
@@ -76,7 +94,13 @@ class TotalSegmentator2D(base.ImageSegmentation):
         self._split = split
         self._version = version
         self._download = download
-        self._as_uint8 = as_uint8
+        self._classes = classes
+        self._optimize_mask_loading = optimize_mask_loading
+
+        if self._optimize_mask_loading and self._classes is not None:
+            raise ValueError(
+                "To use customize classes please set the optimize_mask_loading to `False`."
+            )
 
         self._samples_dirs: List[str] = []
         self._indices: List[Tuple[int, int]] = []
@@ -91,7 +115,13 @@ class TotalSegmentator2D(base.ImageSegmentation):
         first_sample_labels = os.path.join(
             self._root, self._samples_dirs[0], "segmentations", "*.nii.gz"
         )
-        return sorted(map(get_filename, glob(first_sample_labels)))
+        all_classes = sorted(map(get_filename, glob(first_sample_labels)))
+        if self._classes:
+            is_subset = all(name in all_classes for name in self._classes)
+            if not is_subset:
+                raise ValueError("Provided class names are not subset of the dataset onces.")
+
+        return all_classes if self._classes is None else self._classes
 
     @property
     @override
@@ -99,7 +129,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
         return {label: index for index, label in enumerate(self.classes)}
 
     @override
-    def filename(self, index: int) -> str:
+    def filename(self, index: int, segmented: bool = True) -> str:
         sample_idx, _ = self._indices[index]
         sample_dir = self._samples_dirs[sample_idx]
         return os.path.join(sample_dir, "ct.nii.gz")
@@ -113,17 +143,23 @@ class TotalSegmentator2D(base.ImageSegmentation):
     def configure(self) -> None:
         self._samples_dirs = self._fetch_samples_dirs()
         self._indices = self._create_indices()
+        if self._optimize_mask_loading:
+            self._export_semantic_label_masks()
 
     @override
     def validate(self) -> None:
-        if self._version is None:
+        if self._version is None or self._sample_every_n_slices is not None:
             return
 
         _validators.check_dataset_integrity(
             self,
             length=self._expected_dataset_lengths.get(f"{self._split}_{self._version}", 0),
-            n_classes=117,
-            first_and_last_labels=("adrenal_gland_left", "vertebrae_T9"),
+            n_classes=len(self._classes) if self._classes else 117,
+            first_and_last_labels=(
+                (self._classes[0], self._classes[-1])
+                if self._classes
+                else ("adrenal_gland_left", "vertebrae_T9")
+            ),
         )
 
     @override
@@ -134,25 +170,68 @@ class TotalSegmentator2D(base.ImageSegmentation):
     def load_image(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_slice(image_path, slice_index)
-        if self._as_uint8:
-            image_array = convert.to_8bit(image_array)
+        image_array = io.read_nifti(image_path, slice_index)
         image_rgb_array = image_array.repeat(3, axis=2)
         return tv_tensors.Image(image_rgb_array.transpose(2, 0, 1))
 
     @override
     def load_mask(self, index: int) -> tv_tensors.Mask:
+        if self._optimize_mask_loading:
+            return self._load_semantic_label_mask(index)
+        return self._load_mask(index)
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        _, slice_index = self._indices[index]
+        return {"slice_index": slice_index}
+
+    def _load_mask(self, index: int) -> tv_tensors.Mask:
+        """Loads and builds the segmentation mask from NifTi files."""
+        sample_index, slice_index = self._indices[index]
+        semantic_labels = self._load_masks_as_semantic_label(sample_index, slice_index)
+        return tv_tensors.Mask(semantic_labels, dtype=torch.int64)  # type: ignore[reportCallIssue]
+
+    def _load_semantic_label_mask(self, index: int) -> tv_tensors.Mask:
+        """Loads the segmentation mask from a semantic label NifTi file."""
         sample_index, slice_index = self._indices[index]
         masks_dir = self._get_masks_dir(sample_index)
-        mask_paths = (os.path.join(masks_dir, label + ".nii.gz") for label in self.classes)
-        one_hot_encoded = np.concatenate(
-            [io.read_nifti_slice(path, slice_index) for path in mask_paths],
-            axis=2,
-        )
-        background_mask = one_hot_encoded.sum(axis=2, keepdims=True) == 0
-        one_hot_encoded_with_bg = np.concatenate([background_mask, one_hot_encoded], axis=2)
-        segmentation_label = np.argmax(one_hot_encoded_with_bg, axis=2)
-        return tv_tensors.Mask(segmentation_label)
+        filename = os.path.join(masks_dir, "semantic_labels", "masks.nii.gz")
+        semantic_labels = io.read_nifti(filename, slice_index)
+        return tv_tensors.Mask(semantic_labels.squeeze(), dtype=torch.int64)  # type: ignore[reportCallIssue]
+
+    def _load_masks_as_semantic_label(
+        self, sample_index: int, slice_index: int | None = None
+    ) -> npt.NDArray[Any]:
+        """Loads binary masks as a semantic label mask.
+
+        Args:
+            sample_index: The data sample index.
+            slice_index: Whether to return only a specific slice.
+        """
+        masks_dir = self._get_masks_dir(sample_index)
+        mask_paths = [os.path.join(masks_dir, label + ".nii.gz") for label in self.classes]
+        binary_masks = [io.read_nifti(path, slice_index) for path in mask_paths]
+        background_mask = np.zeros_like(binary_masks[0])
+        return np.argmax([background_mask] + binary_masks, axis=0)
+
+    def _export_semantic_label_masks(self) -> None:
+        """Exports the segmentation binary masks (one-hot) to semantic labels."""
+        total_samples = len(self._samples_dirs)
+        masks_dirs = map(self._get_masks_dir, range(total_samples))
+        semantic_labels = [
+            (index, os.path.join(directory, "semantic_labels", "masks.nii.gz"))
+            for index, directory in enumerate(masks_dirs)
+        ]
+        to_export = filter(lambda x: not os.path.isfile(x[1]), semantic_labels)
+
+        for sample_index, filename in tqdm.tqdm(
+            list(to_export),
+            desc=">> Exporting optimized semantic masks",
+            leave=False,
+        ):
+            semantic_labels = self._load_masks_as_semantic_label(sample_index)
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            io.save_array_as_nifti(semantic_labels, filename)
 
     def _get_image_path(self, sample_index: int) -> str:
         """Returns the corresponding image path."""
@@ -164,10 +243,16 @@ class TotalSegmentator2D(base.ImageSegmentation):
         sample_dir = self._samples_dirs[sample_index]
         return os.path.join(self._root, sample_dir, "segmentations")
 
+    def _get_semantic_labels_filename(self, sample_index: int) -> str:
+        """Returns the semantic label filename."""
+        masks_dir = self._get_masks_dir(sample_index)
+        return os.path.join(masks_dir, "semantic_labels", "masks.nii.gz")
+
     def _get_number_of_slices_per_sample(self, sample_index: int) -> int:
         """Returns the total amount of slices of a sample."""
         image_path = self._get_image_path(sample_index)
-        return io.fetch_total_nifti_slices(image_path)
+        image_shape = io.fetch_nifti_shape(image_path)
+        return image_shape[-1]
 
     def _fetch_samples_dirs(self) -> List[str]:
         """Returns the name of all the samples of all the splits of the dataset."""
@@ -180,16 +265,20 @@ class TotalSegmentator2D(base.ImageSegmentation):
 
     def _get_split_indices(self) -> List[int]:
         """Returns the samples indices that corresponding the dataset split and version."""
-        key = f"{self._split}_{self._version}"
-        match key:
-            case "train_small":
-                index_ranges = [(0, 83)]
-            case "val_small":
-                index_ranges = [(83, 102)]
+        metadata_file = os.path.join(self._root, "meta.csv")
+        metadata = io.read_csv(metadata_file, delimiter=";", encoding="utf-8-sig")
+
+        match self._split:
+            case "train":
+                image_ids = [item["image_id"] for item in metadata if item["split"] == "train"]
+            case "val":
+                image_ids = [item["image_id"] for item in metadata if item["split"] == "val"]
+            case "test":
+                image_ids = [item["image_id"] for item in metadata if item["split"] == "test"]
             case _:
-                index_ranges = [(0, len(self._samples_dirs))]
+                image_ids = self._samples_dirs
 
-        return _utils.ranges_to_indices(index_ranges)
+        return sorted(map(self._samples_dirs.index, image_ids))
 
     def _create_indices(self) -> List[Tuple[int, int]]:
         """Builds the dataset indices for the specified split.
@@ -219,6 +308,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
                 f"Can't download data version '{self._version}'. Use 'small' or 'full'."
             )
 
+        self._print_license()
         for resource in resources:
             if os.path.isdir(self._root):
                 continue
@@ -229,3 +319,7 @@ class TotalSegmentator2D(base.ImageSegmentation):
                 filename=resource.filename,
                 remove_finished=True,
             )
+
+    def _print_license(self) -> None:
+        """Prints the dataset license."""
+        print(f"Dataset license: {self._license}")

eva/vision/data/datasets/wsi.py

@@ -0,0 +1,187 @@
+"""Dataset classes for whole-slide images."""
+
+import bisect
+import os
+from typing import Callable, List
+
+from loguru import logger
+from torch.utils.data import dataset as torch_datasets
+from torchvision import tv_tensors
+from torchvision.transforms.v2 import functional
+from typing_extensions import override
+
+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):
+    """Dataset class for reading patches from whole-slide images."""
+
+    def __init__(
+        self,
+        file_path: str,
+        width: int,
+        height: int,
+        sampler: samplers.Sampler,
+        target_mpp: float,
+        overwrite_mpp: float | None = None,
+        backend: str = "openslide",
+        image_transforms: Callable | None = None,
+    ):
+        """Initializes a new dataset instance.
+
+        Args:
+            file_path: Path to the whole-slide image file.
+            width: Width of the patches to be extracted, in pixels.
+            height: Height of the patches to be extracted, in pixels.
+            sampler: The sampler to use for sampling patch coordinates.
+            target_mpp: Target microns per pixel (mpp) for the patches.
+            overwrite_mpp: The microns per pixel (mpp) value to use when missing in WSI metadata.
+            backend: The backend to use for reading the whole-slide images.
+            image_transforms: Transforms to apply to the extracted image patches.
+        """
+        super().__init__()
+
+        self._file_path = file_path
+        self._width = width
+        self._height = height
+        self._sampler = sampler
+        self._target_mpp = target_mpp
+        self._overwrite_mpp = overwrite_mpp
+        self._backend = backend
+        self._image_transforms = image_transforms
+
+    @override
+    def __len__(self):
+        return len(self._coords.x_y)
+
+    @override
+    def filename(self, index: int) -> str:
+        return f"{self._file_path}_{index}"
+
+    @property
+    def _wsi(self) -> wsi.Wsi:
+        return wsi.get_cached_wsi(self._file_path, self._backend, self._overwrite_mpp)
+
+    @property
+    def _coords(self) -> wsi.PatchCoordinates:
+        return wsi.get_cached_coords(
+            file_path=self._file_path,
+            width=self._width,
+            height=self._height,
+            target_mpp=self._target_mpp,
+            overwrite_mpp=self._overwrite_mpp,
+            sampler=self._sampler,
+            backend=self._backend,
+        )
+
+    @override
+    def __getitem__(self, index: int) -> tv_tensors.Image:
+        x, y = self._coords.x_y[index]
+        width, height, level_idx = self._coords.width, self._coords.height, self._coords.level_idx
+        patch = self._wsi.read_region((x, y), level_idx, (width, height))
+        patch = functional.to_image(patch)
+        patch = self._apply_transforms(patch)
+        return patch
+
+    def _apply_transforms(self, image: tv_tensors.Image) -> tv_tensors.Image:
+        if self._image_transforms is not None:
+            image = self._image_transforms(image)
+        return image
+
+
+class MultiWsiDataset(vision.VisionDataset):
+    """Dataset class for reading patches from multiple whole-slide images."""
+
+    def __init__(
+        self,
+        root: str,
+        file_paths: List[str],
+        width: int,
+        height: int,
+        sampler: samplers.Sampler,
+        target_mpp: float,
+        overwrite_mpp: float | None = None,
+        backend: str = "openslide",
+        image_transforms: Callable | None = None,
+    ):
+        """Initializes a new dataset instance.
+
+        Args:
+            root: Root directory of the dataset.
+            file_paths: List of paths to the whole-slide image files, relative to the root.
+            width: Width of the patches to be extracted, in pixels.
+            height: Height of the patches to be extracted, in pixels.
+            target_mpp: Target microns per pixel (mpp) for the patches.
+            overwrite_mpp: The microns per pixel (mpp) value to use when missing in WSI metadata.
+            sampler: The sampler to use for sampling patch coordinates.
+            backend: The backend to use for reading the whole-slide images.
+            image_transforms: Transforms to apply to the extracted image patches.
+        """
+        super().__init__()
+
+        self._root = root
+        self._file_paths = file_paths
+        self._width = width
+        self._height = height
+        self._target_mpp = target_mpp
+        self._overwrite_mpp = overwrite_mpp
+        self._sampler = sampler
+        self._backend = backend
+        self._image_transforms = image_transforms
+
+        self._concat_dataset: torch_datasets.ConcatDataset
+
+    @property
+    def datasets(self) -> List[WsiDataset]:
+        """Returns the list of WSI datasets."""
+        return self._concat_dataset.datasets  # type: ignore
+
+    @property
+    def cumulative_sizes(self) -> List[int]:
+        """Returns the cumulative sizes of the WSI datasets."""
+        return self._concat_dataset.cumulative_sizes
+
+    @override
+    def configure(self) -> None:
+        self._concat_dataset = torch_datasets.ConcatDataset(datasets=self._load_datasets())
+
+    @override
+    def __len__(self) -> int:
+        return len(self._concat_dataset)
+
+    @override
+    def __getitem__(self, index: int) -> tv_tensors.Image:
+        return self._concat_dataset[index]
+
+    @override
+    def filename(self, index: int) -> str:
+        return os.path.basename(self._file_paths[self._get_dataset_idx(index)])
+
+    def _load_datasets(self) -> list[WsiDataset]:
+        logger.info(f"Initializing dataset with {len(self._file_paths)} WSIs ...")
+        wsi_datasets = []
+        for file_path in self._file_paths:
+            file_path = (
+                os.path.join(self._root, file_path) if self._root not in file_path else file_path
+            )
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"File not found: {file_path}")
+
+            wsi_datasets.append(
+                WsiDataset(
+                    file_path=file_path,
+                    width=self._width,
+                    height=self._height,
+                    sampler=self._sampler,
+                    target_mpp=self._target_mpp,
+                    overwrite_mpp=self._overwrite_mpp,
+                    backend=self._backend,
+                    image_transforms=self._image_transforms,
+                )
+            )
+        return wsi_datasets
+
+    def _get_dataset_idx(self, index: int) -> int:
+        return bisect.bisect_right(self.cumulative_sizes, index)

eva/vision/data/transforms/__init__.py

@@ -1,5 +1,6 @@
 """Vision data transforms."""
 
-from eva.vision.data.transforms.common import ResizeAndCrop
+from eva.vision.data.transforms.common import ResizeAndClamp, ResizeAndCrop
+from eva.vision.data.transforms.normalization import Clamp, RescaleIntensity
 
-__all__ = ["ResizeAndCrop"]
+__all__ = ["ResizeAndCrop", "ResizeAndClamp", "Clamp", "RescaleIntensity"]

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

@@ -1,5 +1,6 @@
 """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__ = ["ResizeAndCrop"]
+__all__ = ["ResizeAndClamp", "ResizeAndCrop"]

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

@@ -0,0 +1,51 @@
+"""Specialized transforms for resizing, clamping and range normalizing."""
+
+from typing import Callable, Sequence, Tuple
+
+from torchvision.transforms import v2
+
+from eva.vision.data.transforms import normalization
+
+
+class ResizeAndClamp(v2.Compose):
+    """Resizes, crops, clamps and normalizes an input image."""
+
+    def __init__(
+        self,
+        size: int | Sequence[int] = 224,
+        clamp_range: Tuple[int, int] = (-1024, 1024),
+        mean: Sequence[float] = (0.0, 0.0, 0.0),
+        std: Sequence[float] = (1.0, 1.0, 1.0),
+    ) -> None:
+        """Initializes the transform object.
+
+        Args:
+            size: Desired output size of the crop. If size is an `int` instead
+                of sequence like (h, w), a square crop (size, size) is made.
+            clamp_range: The lower and upper bound to clamp the pixel values.
+            mean: Sequence of means for each image channel.
+            std: Sequence of standard deviations for each image channel.
+        """
+        self._size = size
+        self._clamp_range = clamp_range
+        self._mean = mean
+        self._std = std
+
+        super().__init__(transforms=self._build_transforms())
+
+    def _build_transforms(self) -> Sequence[Callable]:
+        """Builds and returns the list of transforms."""
+        transforms = [
+            v2.Resize(size=self._size),
+            v2.CenterCrop(size=self._size),
+            normalization.Clamp(out_range=self._clamp_range),
+            normalization.RescaleIntensity(
+                in_range=self._clamp_range,
+                out_range=(0.0, 1.0),
+            ),
+            v2.Normalize(
+                mean=self._mean,
+                std=self._std,
+            ),
+        ]
+        return transforms

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

@@ -3,10 +3,10 @@
 from typing import Callable, Sequence
 
 import torch
-import torchvision.transforms.v2 as torch_transforms
+from torchvision.transforms import v2
 
 
-class ResizeAndCrop(torch_transforms.Compose):
+class ResizeAndCrop(v2.Compose):
     """Resizes, crops and normalizes an input image while preserving its aspect ratio."""
 
     def __init__(
@@ -32,11 +32,10 @@ class ResizeAndCrop(torch_transforms.Compose):
     def _build_transforms(self) -> Sequence[Callable]:
         """Builds and returns the list of transforms."""
         transforms = [
-            torch_transforms.ToImage(),
-            torch_transforms.Resize(size=self._size),
-            torch_transforms.CenterCrop(size=self._size),
-            torch_transforms.ToDtype(torch.float32, scale=True),
-            torch_transforms.Normalize(
+            v2.Resize(size=self._size),
+            v2.CenterCrop(size=self._size),
+            v2.ToDtype(torch.float32, scale=True),
+            v2.Normalize(
                 mean=self._mean,
                 std=self._std,
             ),

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

@@ -0,0 +1,6 @@
+"""Normalization related transformations."""
+
+from eva.vision.data.transforms.normalization.clamp import Clamp
+from eva.vision.data.transforms.normalization.rescale_intensity import RescaleIntensity
+
+__all__ = ["Clamp", "RescaleIntensity"]

eva/vision/data/transforms/normalization/clamp.py

@@ -0,0 +1,43 @@
+"""Image clamp transform."""
+
+import functools
+from typing import Any, Dict, Tuple
+
+import torch
+import torchvision.transforms.v2 as torch_transforms
+from torchvision import tv_tensors
+from typing_extensions import override
+
+
+class Clamp(torch_transforms.Transform):
+    """Clamps all elements in input into a specific range."""
+
+    def __init__(self, out_range: Tuple[int, int]) -> None:
+        """Initializes the transform.
+
+        Args:
+            out_range: The lower and upper bound of the range to
+                be clamped to.
+        """
+        super().__init__()
+
+        self._out_range = out_range
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(torch.Tensor)
+    def _(self, inpt: torch.Tensor, params: Dict[str, Any]) -> Any:
+        return torch.clamp(inpt, min=self._out_range[0], max=self._out_range[1])
+
+    @_transform.register(tv_tensors.Image)
+    def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
+        inpt_clamp = torch.clamp(inpt, min=self._out_range[0], max=self._out_range[1])
+        return tv_tensors.wrap(inpt_clamp, like=inpt)
+
+    @_transform.register(tv_tensors.BoundingBoxes)
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: tv_tensors.BoundingBoxes | tv_tensors.Mask, params: Dict[str, Any]) -> Any:
+        return inpt

eva/vision/data/transforms/normalization/functional/__init__.py

@@ -0,0 +1,5 @@
+"""Functional normalization related transformations API."""
+
+from eva.vision.data.transforms.normalization.functional.rescale_intensity import rescale_intensity
+
+__all__ = ["rescale_intensity"]

eva/vision/data/transforms/normalization/functional/rescale_intensity.py

@@ -0,0 +1,28 @@
+"""Intensity level functions."""
+
+import sys
+from typing import Tuple
+
+import torch
+
+
+def rescale_intensity(
+    image: torch.Tensor,
+    in_range: Tuple[float, float] | None = None,
+    out_range: Tuple[float, float] = (0.0, 1.0),
+) -> torch.Tensor:
+    """Stretches or shrinks the image intensity levels.
+
+    Args:
+        image: The image tensor as float-type.
+        in_range: The input data range. If `None`, it will
+            fetch the min and max of the input image.
+        out_range: The desired intensity range of the output.
+
+    Returns:
+        The image tensor after stretching or shrinking its intensity levels.
+    """
+    imin, imax = in_range or (image.min(), image.max())
+    omin, omax = out_range
+    image_scaled = (image - imin) / (imax - imin + sys.float_info.epsilon)
+    return image_scaled * (omax - omin) + omin