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

eva/vision/callbacks/loggers/batch/base.py

@@ -0,0 +1,130 @@
+"""Base batch callback logger."""
+
+import abc
+
+from lightning import pytorch as pl
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from typing_extensions import override
+
+from eva.core.models.modules.typings import INPUT_TENSOR_BATCH
+
+
+class BatchLogger(pl.Callback, abc.ABC):
+    """Logs training and validation batch assets."""
+
+    _batch_idx_to_log: int = 0
+    """The batch index log."""
+
+    def __init__(
+        self,
+        log_every_n_epochs: int | None = None,
+        log_every_n_steps: int | None = None,
+    ) -> None:
+        """Initializes the callback object.
+
+        Args:
+            log_every_n_epochs: Epoch-wise logging frequency.
+            log_every_n_steps: Step-wise logging frequency.
+        """
+        super().__init__()
+
+        if log_every_n_epochs is None and log_every_n_steps is None:
+            raise ValueError(
+                "Please configure the logging frequency though "
+                "`log_every_n_epochs` or `log_every_n_steps`."
+            )
+        if None not in [log_every_n_epochs, log_every_n_steps]:
+            raise ValueError(
+                "Arguments `log_every_n_epochs` and `log_every_n_steps` "
+                "are mutually exclusive. Please configure one of them."
+            )
+
+        self._log_every_n_epochs = log_every_n_epochs
+        self._log_every_n_steps = log_every_n_steps
+
+    @override
+    def on_train_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_TENSOR_BATCH,
+        batch_idx: int,
+    ) -> None:
+        if self._skip_logging(trainer, batch_idx if self._log_every_n_epochs else None):
+            return
+
+        self._log_batch(
+            trainer=trainer,
+            batch=batch,
+            outputs=outputs,
+            tag="BatchTrain",
+        )
+
+    @override
+    def on_validation_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_TENSOR_BATCH,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        if self._skip_logging(trainer, batch_idx):
+            return
+
+        self._log_batch(
+            trainer=trainer,
+            batch=batch,
+            outputs=outputs,
+            tag="BatchValidation",
+        )
+
+    @abc.abstractmethod
+    def _log_batch(
+        self,
+        trainer: pl.Trainer,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_TENSOR_BATCH,
+        tag: str,
+    ) -> None:
+        """Logs the batch data.
+
+        Args:
+            trainer: The trainer.
+            outputs: The output of the train / val step.
+            batch: The data batch.
+            tag: The log tag.
+        """
+
+    def _skip_logging(
+        self,
+        trainer: pl.Trainer,
+        batch_idx: int | None = None,
+    ) -> bool:
+        """Determines whether skip the logging step or not.
+
+        Args:
+            trainer: The trainer.
+            batch_idx: The batch index.
+
+        Returns:
+            A boolean indicating whether to skip the step execution.
+        """
+        if trainer.global_step in [0, 1]:
+            return False
+
+        skip_due_frequency = any(
+            [
+                (trainer.current_epoch + 1) % (self._log_every_n_epochs or 1) != 0,
+                (trainer.global_step + 1) % (self._log_every_n_steps or 1) != 0,
+            ]
+        )
+
+        conditions = [
+            skip_due_frequency,
+            not trainer.is_global_zero,
+            batch_idx != self._batch_idx_to_log if batch_idx else False,
+        ]
+        return any(conditions)

eva/vision/callbacks/loggers/batch/segmentation.py

@@ -0,0 +1,188 @@
+"""Segmentation datasets related data loggers."""
+
+from typing import List, Tuple
+
+import torch
+import torchvision
+from lightning import pytorch as pl
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from torch.nn import functional
+from typing_extensions import override
+
+from eva.core.loggers import log
+from eva.core.models.modules.typings import INPUT_TENSOR_BATCH
+from eva.core.utils import to_cpu
+from eva.vision.callbacks.loggers.batch import base
+from eva.vision.utils import colormap, convert
+
+
+class SemanticSegmentationLogger(base.BatchLogger):
+    """Log the segmentation batch."""
+
+    def __init__(
+        self,
+        max_samples: int = 10,
+        number_of_images_per_subgrid_row: int = 1,
+        log_images: bool = True,
+        mean: Tuple[float, ...] = (0.0, 0.0, 0.0),
+        std: Tuple[float, ...] = (1.0, 1.0, 1.0),
+        log_every_n_epochs: int | None = None,
+        log_every_n_steps: int | None = None,
+    ) -> None:
+        """Initializes the callback object.
+
+        Args:
+            max_samples: The maximum number of images displayed in the grid.
+            number_of_images_per_subgrid_row: Number of images displayed in each
+                row of each sub-grid (that is images, targets and predictions).
+            log_images: Whether to log the input batch images.
+            mean: The mean of the input images to de-normalize from.
+            std: The std of the input images to de-normalize from.
+            log_every_n_epochs: Epoch-wise logging frequency.
+            log_every_n_steps: Step-wise logging frequency.
+        """
+        super().__init__(
+            log_every_n_epochs=log_every_n_epochs,
+            log_every_n_steps=log_every_n_steps,
+        )
+
+        self._max_samples = max_samples
+        self._number_of_images_per_subgrid_row = number_of_images_per_subgrid_row
+        self._log_images = log_images
+        self._mean = mean
+        self._std = std
+
+    @override
+    def _log_batch(
+        self,
+        trainer: pl.Trainer,
+        outputs: STEP_OUTPUT,
+        batch: INPUT_TENSOR_BATCH,
+        tag: str,
+    ) -> None:
+        predictions = outputs.get("predictions") if isinstance(outputs, dict) else None
+        if predictions is None:
+            raise ValueError("Key `predictions` is missing from the output.")
+
+        data_batch, target_batch = batch[0], batch[1]
+        data, targets, predictions = _subsample_tensors(
+            tensors_stack=[data_batch, target_batch, predictions],
+            max_samples=self._max_samples,
+        )
+        data, targets, predictions = to_cpu([data, targets, predictions])
+        predictions = torch.argmax(predictions, dim=1)
+
+        target_images = list(map(_draw_semantic_mask, targets))
+        prediction_images = list(map(_draw_semantic_mask, predictions))
+        image_groups = [target_images, prediction_images]
+
+        if self._log_images:
+            images = list(map(self._format_image, data))
+            overlay_targets = [
+                _overlay_mask(image, mask) for image, mask in zip(images, targets, strict=False)
+            ]
+            overlay_predictions = [
+                _overlay_mask(image, mask) for image, mask in zip(images, predictions, strict=False)
+            ]
+            image_groups = [images, overlay_targets, overlay_predictions] + image_groups
+
+        image_grid = _make_grid_from_image_groups(
+            image_groups, self._number_of_images_per_subgrid_row
+        )
+
+        log.log_image(
+            trainer.loggers,
+            image=image_grid,
+            tag=tag,
+            step=trainer.global_step,
+        )
+
+    def _format_image(self, image: torch.Tensor) -> torch.Tensor:
+        """Descaled an image tensor to (0, 255) uint8 tensor."""
+        return convert.descale_and_denorm_image(image, mean=self._mean, std=self._std)
+
+
+def _subsample_tensors(
+    tensors_stack: List[torch.Tensor],
+    max_samples: int,
+) -> List[torch.Tensor]:
+    """Sub-samples tensors from a list of tensors in-place.
+
+    Args:
+        tensors_stack: A list of tensors.
+        max_samples: The maximum number of images
+            displayed in the grid.
+
+    Returns:
+        A sub-sample of the input tensors stack.
+    """
+    for i, tensor in enumerate(tensors_stack):
+        tensors_stack[i] = tensor[:max_samples]
+    return tensors_stack
+
+
+def _draw_semantic_mask(tensor: torch.Tensor) -> torch.Tensor:
+    """Draws a semantic mask to an image RGB tensor.
+
+    The input semantic mask is a (H x W) shaped tensor with
+    integer values which represent the pixel class id.
+
+    Args:
+        tensor: An image tensor of range [0., 1.].
+
+    Returns:
+        The image as a tensor of range [0., 255.].
+    """
+    tensor = torch.squeeze(tensor)
+    height, width = tensor.shape[-2], tensor.shape[-1]
+    red, green, blue = torch.zeros((3, height, width), dtype=torch.uint8)
+    for class_id, color in colormap.COLORMAP.items():
+        indices = tensor == class_id
+        red[indices], green[indices], blue[indices] = color
+    return torch.stack([red, green, blue])
+
+
+def _overlay_mask(image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Overlays a segmentation mask onto an image.
+
+    Args:
+        image: A 3D tensor of shape (C, H, W) representing the image.
+        mask: A 2D tensor of shape (H, W) representing the segmentation mask.
+            Each pixel in the mask corresponds to a class label.
+
+    Returns:
+        A tensor of the same shape as the input image (C, H, W) with the
+        segmentation mask overlaid on top. The output image retains the
+        original color channels but with the mask applied, using the colors
+        from the predefined colormap.
+    """
+    binary_masks = functional.one_hot(mask).permute(2, 0, 1).to(dtype=torch.bool)
+    return torchvision.utils.draw_segmentation_masks(
+        image, binary_masks[1:], alpha=0.65, colors=colormap.COLORS[1:]  # type: ignore
+    )
+
+
+def _make_grid_from_image_groups(
+    image_groups: List[List[torch.Tensor]],
+    number_of_images_per_subgrid_row: int = 2,
+) -> torch.Tensor:
+    """Creates a single image grid from image groups.
+
+    For example, it can combine the input images, targets predictions into a single image.
+
+    Args:
+        image_groups: A list of lists of image tensors of shape (C x H x W)
+            all of the same size.
+        number_of_images_per_subgrid_row: Number of images displayed in each
+            row of the sub-grid.
+
+    Returns:
+        An image grid as a `torch.Tensor`.
+    """
+    return torchvision.utils.make_grid(
+        [
+            torchvision.utils.make_grid(image_group, nrow=number_of_images_per_subgrid_row)
+            for image_group in image_groups
+        ],
+        nrow=len(image_groups),
+    )

eva/vision/data/datasets/__init__.py

@@ -1,15 +1,42 @@
 """Vision Datasets API."""
 
-from eva.vision.data.datasets.classification import BACH, CRC, MHIST, PatchCamelyon
-from eva.vision.data.datasets.segmentation import ImageSegmentation, TotalSegmentator2D
+from eva.vision.data.datasets.classification import (
+    BACH,
+    CRC,
+    MHIST,
+    PANDA,
+    Camelyon16,
+    PatchCamelyon,
+    WsiClassificationDataset,
+)
+from eva.vision.data.datasets.segmentation import (
+    BCSS,
+    CoNSeP,
+    EmbeddingsSegmentationDataset,
+    ImageSegmentation,
+    LiTS,
+    MoNuSAC,
+    TotalSegmentator2D,
+)
 from eva.vision.data.datasets.vision import VisionDataset
+from eva.vision.data.datasets.wsi import MultiWsiDataset, WsiDataset
 
 __all__ = [
     "BACH",
+    "BCSS",
     "CRC",
     "MHIST",
-    "ImageSegmentation",
+    "PANDA",
+    "Camelyon16",
     "PatchCamelyon",
+    "WsiClassificationDataset",
+    "CoNSeP",
+    "EmbeddingsSegmentationDataset",
+    "ImageSegmentation",
+    "LiTS",
+    "MoNuSAC",
     "TotalSegmentator2D",
     "VisionDataset",
+    "MultiWsiDataset",
+    "WsiDataset",
 ]

eva/vision/data/datasets/_validators.py

@@ -13,7 +13,7 @@ _SUFFIX_ERROR_MESSAGE = "Please verify that the data are properly downloaded and
 def check_dataset_integrity(
     dataset: vision.VisionDataset,
     *,
-    length: int,
+    length: int | None,
     n_classes: int,
     first_and_last_labels: Tuple[str, str],
 ) -> None:
@@ -23,7 +23,7 @@ def check_dataset_integrity(
         ValueError: If the input dataset's values do not
             match the expected ones.
     """
-    if len(dataset) != length:
+    if length and len(dataset) != length:
         raise ValueError(
             f"Dataset's '{dataset.__class__.__qualname__}' length "
             f"({len(dataset)}) does not match the expected one ({length}). "
@@ -57,3 +57,16 @@ def check_dataset_exists(dataset_dir: str, download_available: bool) -> None:
         if download_available:
             error_message += " You can set `download=True` to download the dataset automatically."
         raise FileNotFoundError(error_message)
+
+
+def check_number_of_files(file_paths: List[str], expected_length: int, split: str | None) -> None:
+    """Verifies the number of files in the dataset.
+
+    Raise:
+        ValueError: If the number of files in the dataset does not match the expected one.
+    """
+    if len(file_paths) != expected_length:
+        raise ValueError(
+            f"Expected {expected_length} files, for split '{split}' found {len(file_paths)}. "
+            f"{_SUFFIX_ERROR_MESSAGE}"
+        )

eva/vision/data/datasets/classification/__init__.py

@@ -1,8 +1,19 @@
 """Image classification datasets API."""
 
 from eva.vision.data.datasets.classification.bach import BACH
+from eva.vision.data.datasets.classification.camelyon16 import Camelyon16
 from eva.vision.data.datasets.classification.crc import CRC
 from eva.vision.data.datasets.classification.mhist import MHIST
+from eva.vision.data.datasets.classification.panda import PANDA
 from eva.vision.data.datasets.classification.patch_camelyon import PatchCamelyon
+from eva.vision.data.datasets.classification.wsi import WsiClassificationDataset
 
-__all__ = ["BACH", "CRC", "MHIST", "PatchCamelyon"]
+__all__ = [
+    "BACH",
+    "CRC",
+    "MHIST",
+    "PatchCamelyon",
+    "WsiClassificationDataset",
+    "PANDA",
+    "Camelyon16",
+]

eva/vision/data/datasets/classification/bach.py

@@ -3,7 +3,8 @@
 import os
 from typing import Callable, Dict, List, Literal, Tuple
 
-import numpy as np
+import torch
+from torchvision import tv_tensors
 from torchvision.datasets import folder, utils
 from typing_extensions import override
 
@@ -52,8 +53,7 @@ class BACH(base.ImageClassification):
         root: str,
         split: Literal["train", "val"] | None = None,
         download: bool = False,
-        image_transforms: Callable | None = None,
-        target_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initialize the dataset.
 
@@ -68,15 +68,10 @@ class BACH(base.ImageClassification):
                 Note that the download will be executed only by additionally
                 calling the :meth:`prepare_data` method and if the data does
                 not yet exist on disk.
-            image_transforms: A function/transform that takes in an image
-                and returns a transformed version.
-            target_transforms: A function/transform that takes in the target
-                and transforms it.
+            transforms: A function/transform which returns a transformed
+                version of the raw data samples.
         """
-        super().__init__(
-            image_transforms=image_transforms,
-            target_transforms=target_transforms,
-        )
+        super().__init__(transforms=transforms)
 
         self._root = root
         self._split = split
@@ -130,14 +125,14 @@ class BACH(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> np.ndarray:
+    def load_image(self, index: int) -> tv_tensors.Image:
         image_path, _ = self._samples[self._indices[index]]
-        return io.read_image(image_path)
+        return io.read_image_as_tensor(image_path)
 
     @override
-    def load_target(self, index: int) -> np.ndarray:
+    def load_target(self, index: int) -> torch.Tensor:
         _, target = self._samples[self._indices[index]]
-        return np.asarray(target, dtype=np.int64)
+        return torch.tensor(target, dtype=torch.long)
 
     @override
     def __len__(self) -> int:

eva/vision/data/datasets/classification/base.py

@@ -3,32 +3,29 @@
 import abc
 from typing import Any, Callable, Dict, List, Tuple
 
-import numpy as np
+import torch
+from torchvision import tv_tensors
 from typing_extensions import override
 
 from eva.vision.data.datasets import vision
 
 
-class ImageClassification(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc.ABC):
+class ImageClassification(vision.VisionDataset[Tuple[tv_tensors.Image, torch.Tensor]], abc.ABC):
     """Image classification abstract dataset."""
 
     def __init__(
         self,
-        image_transforms: Callable | None = None,
-        target_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initializes the image classification dataset.
 
         Args:
-            image_transforms: A function/transform that takes in an image
-                and returns a transformed version.
-            target_transforms: A function/transform that takes in the target
-                and transforms it.
+            transforms: A function/transform which returns a transformed
+                version of the raw data samples.
         """
         super().__init__()
 
-        self._image_transforms = image_transforms
-        self._target_transforms = target_transforms
+        self._transforms = transforms
 
     @property
     def classes(self) -> List[str] | None:
@@ -38,19 +35,18 @@ class ImageClassification(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], a
     def class_to_idx(self) -> Dict[str, int] | None:
         """Returns a mapping of the class name to its target index."""
 
-    def load_metadata(self, index: int | None) -> Dict[str, Any] | List[Dict[str, Any]] | None:
+    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.
-                If `None`, it will return the metadata of the current dataset.
 
         Returns:
             The sample metadata.
         """
 
     @abc.abstractmethod
-    def load_image(self, index: int) -> np.ndarray:
+    def load_image(self, index: int) -> tv_tensors.Image:
         """Returns the `index`'th image sample.
 
         Args:
@@ -61,7 +57,7 @@ class ImageClassification(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], a
         """
 
     @abc.abstractmethod
-    def load_target(self, index: int) -> np.ndarray:
+    def load_target(self, index: int) -> torch.Tensor:
         """Returns the `index`'th target sample.
 
         Args:
@@ -77,14 +73,15 @@ class ImageClassification(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], a
         raise NotImplementedError
 
     @override
-    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
+    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, torch.Tensor, Dict[str, Any]]:
         image = self.load_image(index)
         target = self.load_target(index)
-        return self._apply_transforms(image, target)
+        image, target = self._apply_transforms(image, target)
+        return image, target, self.load_metadata(index) or {}
 
     def _apply_transforms(
-        self, image: np.ndarray, target: np.ndarray
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, image: tv_tensors.Image, target: torch.Tensor
+    ) -> Tuple[tv_tensors.Image, torch.Tensor]:
         """Applies the transforms to the provided data and returns them.
 
         Args:
@@ -94,10 +91,6 @@ class ImageClassification(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], a
         Returns:
             A tuple with the image and the target transformed.
         """
-        if self._image_transforms is not None:
-            image = self._image_transforms(image)
-
-        if self._target_transforms is not None:
-            target = self._target_transforms(target)
-
+        if self._transforms is not None:
+            image, target = self._transforms(image, target)
         return image, target