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

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

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

@@ -0,0 +1,53 @@
+"""Intensity level scaling 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
+
+from eva.vision.data.transforms.normalization import functional
+
+
+class RescaleIntensity(torch_transforms.Transform):
+    """Stretches or shrinks the image intensity levels."""
+
+    def __init__(
+        self,
+        in_range: Tuple[float, float] | None = None,
+        out_range: Tuple[float, float] = (0.0, 1.0),
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            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.
+        """
+        super().__init__()
+
+        self._in_range = in_range
+        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 functional.rescale_intensity(
+            inpt, in_range=self._in_range, out_range=self._out_range
+        )
+
+    @_transform.register(tv_tensors.Image)
+    def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
+        scaled_inpt = functional.rescale_intensity(inpt, out_range=self._out_range)
+        return tv_tensors.wrap(scaled_inpt, 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/wsi/__init__.py

@@ -0,0 +1,16 @@
+"""WSI API."""
+
+from eva.vision.data.wsi.backends import Wsi, get_cached_wsi, wsi_backend
+from eva.vision.data.wsi.patching.coordinates import PatchCoordinates, get_cached_coords
+from eva.vision.data.wsi.patching.mask import Mask, get_mask, get_mask_level
+
+__all__ = [
+    "Wsi",
+    "PatchCoordinates",
+    "Mask",
+    "get_cached_coords",
+    "wsi_backend",
+    "get_cached_wsi",
+    "get_mask",
+    "get_mask_level",
+]

eva/vision/data/wsi/backends/__init__.py

@@ -0,0 +1,69 @@
+"""WSI Backends API."""
+
+import functools
+import importlib.util
+from typing import Callable
+
+from eva.vision.data.wsi.backends.base import Wsi
+
+LRU_CACHE_SIZE = 32
+
+
+def _is_openslide_available() -> bool:
+    """Whether the OpenSlide library is available."""
+    return importlib.util.find_spec("openslide") is not None
+
+
+def _is_tiffslide_available() -> bool:
+    """Whether the TiffSlide library is available."""
+    return importlib.util.find_spec("tiffslide") is not None
+
+
+def is_backend_available(backend: str) -> bool:
+    """Whether the specified backend is available."""
+    match backend:
+        case "openslide":
+            return _is_openslide_available()
+        case "tiffslide":
+            return _is_tiffslide_available()
+    return False
+
+
+def wsi_backend(backend: str = "openslide") -> Callable[..., Wsi]:
+    """Returns the backend to use for reading the whole-slide images."""
+    match backend:
+        case "openslide":
+            if _is_openslide_available():
+                from eva.vision.data.wsi.backends.openslide import WsiOpenslide
+
+                return WsiOpenslide
+            else:
+                raise ValueError(
+                    "Missing optional dependency: openslide.\n"
+                    "Please install using `pip install openslide-python`."
+                )
+        case "tiffslide":
+            if _is_tiffslide_available():
+                from eva.vision.data.wsi.backends.tiffslide import WsiTiffslide
+
+                return WsiTiffslide
+            else:
+                raise ValueError(
+                    "Missing optional dependency: tiffslide.\n"
+                    "Please install using `pip install tiffslide`."
+                )
+        case "pil":
+            from eva.vision.data.wsi.backends.pil import PILImage
+
+            return PILImage
+        case _:
+            raise ValueError(f"Unknown WSI backend selected: {backend}")
+
+
+@functools.lru_cache(LRU_CACHE_SIZE)
+def get_cached_wsi(file_path: str, backend: str, overwrite_mpp: float | None = None) -> Wsi:
+    """Returns a cached instance of the whole-slide image backend reader."""
+    return wsi_backend(backend)(file_path, overwrite_mpp)
+
+
+__all__ = ["Wsi", "wsi_backend", "get_cached_wsi", "_is_openslide_available"]

eva/vision/data/wsi/backends/base.py

@@ -0,0 +1,115 @@
+"""Base Module for loading data from WSI files."""
+
+import abc
+from typing import Any, Sequence, Tuple
+
+import numpy as np
+
+
+class Wsi(abc.ABC):
+    """Base class for loading data from Whole Slide Image (WSI) files."""
+
+    def __init__(self, file_path: str, overwrite_mpp: float | None = None):
+        """Initializes a Wsi object.
+
+        Args:
+            file_path: The path to the WSI file.
+            overwrite_mpp: The microns per pixel (mpp) value to use when missing in WSI metadata.
+        """
+        self._wsi = self.open_file(file_path)
+        self._overwrite_mpp = overwrite_mpp
+
+    @abc.abstractmethod
+    def open_file(self, file_path: str) -> Any:
+        """Opens the WSI file.
+
+        Args:
+            file_path: The path to the WSI file.
+        """
+
+    @property
+    @abc.abstractmethod
+    def level_dimensions(self) -> Sequence[Tuple[int, int]]:
+        """A list of (width, height) tuples for each level, from highest to lowest resolution."""
+
+    @property
+    @abc.abstractmethod
+    def level_downsamples(self) -> Sequence[float]:
+        """A list of downsampling factors for each level, relative to the highest resolution."""
+
+    @property
+    @abc.abstractmethod
+    def mpp(self) -> float:
+        """Microns per pixel at the highest resolution (level 0)."""
+
+    @abc.abstractmethod
+    def _read_region(
+        self, location: Tuple[int, int], level: int, size: Tuple[int, int]
+    ) -> np.ndarray:
+        """Abstract method to read a region at a specified zoom level."""
+
+    def read_region(
+        self, location: Tuple[int, int], level: int, size: Tuple[int, int]
+    ) -> np.ndarray:
+        """Reads and returns image data for a specified region and zoom level.
+
+        Args:
+            location: Top-left corner (x, y) to start reading at level 0.
+            level: WSI level to read from.
+            size: Region size as (width, height) in pixels at the selected read level.
+                Remember to scale the size correctly.
+        """
+        self._verify_location(location, size)
+        data = self._read_region(location, level, size)
+        return self._read_postprocess(data)
+
+    def get_closest_level(self, target_mpp: float) -> int:
+        """Calculate the slide level that is closest to the target mpp.
+
+        Args:
+            slide: The whole-slide image object.
+            target_mpp: The target microns per pixel (mpp) value.
+        """
+        # Calculate the mpp for each level
+        level_mpps = self.mpp * np.array(self.level_downsamples)
+
+        # Ignore levels with higher mpp
+        level_mpps_filtered = level_mpps.copy()
+        level_mpps_filtered[level_mpps_filtered > target_mpp] = 0
+
+        if level_mpps_filtered.max() == 0:
+            # When all levels have higher mpp than target_mpp return the level with lowest mpp
+            level_idx = np.argmin(level_mpps)
+        else:
+            level_idx = np.argmax(level_mpps_filtered)
+
+        return int(level_idx)
+
+    def _verify_location(self, location: Tuple[int, int], size: Tuple[int, int]) -> None:
+        """Verifies that the requested region is within the slide dimensions.
+
+        Args:
+            location: Top-left corner (x, y) to start reading at level 0.
+            size: Region size as (width, height) in pixels at the selected read level.
+        """
+        x_max, y_max = self.level_dimensions[0]
+        x_scale = x_max / self.level_dimensions[0][0]
+        y_scale = y_max / self.level_dimensions[0][1]
+
+        if (
+            int(location[0] + x_scale * size[0]) > x_max
+            or int(location[1] + y_scale * size[1]) > y_max
+        ):
+            raise ValueError(f"Out of bounds region: {location}, {size}")
+
+    def _read_postprocess(self, data: np.ndarray) -> np.ndarray:
+        """Post-processes the read region data.
+
+        Args:
+            data: The read region data as a numpy array of shape (height, width, channels).
+        """
+        # Change color to white where the alpha channel is 0
+        if data.shape[2] == 4:
+            data[data[:, :, 3] == 0] = 255
+
+        return data[:, :, :3]