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

eva/vision/data/wsi/patching/samplers/_utils.py

@@ -0,0 +1,50 @@
+import random
+from typing import Tuple
+
+import numpy as np
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+def get_grid_coords_and_indices(
+    layer_shape: Tuple[int, int],
+    width: int,
+    height: int,
+    overlap: Tuple[int, int],
+    shuffle: bool = True,
+    seed: int = 42,
+):
+    """Get grid coordinates and indices.
+
+    Args:
+        layer_shape: The shape of the layer.
+        width: The width of the patches.
+        height: The height of the patches.
+        overlap: The overlap between patches in the grid.
+        shuffle: Whether to shuffle the indices.
+        seed: The random seed.
+    """
+    x_range = range(0, layer_shape[0] - width + 1, width - overlap[0])
+    y_range = range(0, layer_shape[1] - height + 1, height - overlap[1])
+    x_y = [(x, y) for x in x_range for y in y_range]
+
+    indices = list(range(len(x_y)))
+    if shuffle:
+        set_seed(seed)
+        np.random.shuffle(indices)
+    return x_y, indices
+
+
+def validate_dimensions(width: int, height: int, layer_shape: Tuple[int, int]) -> None:
+    """Checks if the width / height is bigger than the layer shape.
+
+    Args:
+        width: The width of the patches.
+        height: The height of the patches.
+        layer_shape: The shape of the layer.
+    """
+    if width > layer_shape[0] or height > layer_shape[1]:
+        raise ValueError("The width / height cannot be bigger than the layer shape.")

eva/vision/data/wsi/patching/samplers/base.py

@@ -0,0 +1,48 @@
+"""Base classes for samplers."""
+
+import abc
+from typing import Generator, Tuple
+
+from eva.vision.data.wsi.patching.mask import Mask
+
+
+class Sampler(abc.ABC):
+    """Base class for samplers."""
+
+    @abc.abstractmethod
+    def sample(
+        self,
+        width: int,
+        height: int,
+        layer_shape: Tuple[int, int],
+        mask: Mask | None = None,
+    ) -> Generator[Tuple[int, int], None, None]:
+        """Sample patche coordinates.
+
+        Args:
+            width: The width of the patches.
+            height: The height of the patches.
+            layer_shape: The shape of the layer.
+            mask: Tuple containing the mask array and the scaling factor with respect to the
+                provided layer_shape. Optional, only required for samplers with foreground
+                filtering.
+
+        Returns:
+            A generator producing sampled patch coordinates.
+        """
+
+
+class ForegroundSampler(Sampler):
+    """Base class for samplers with foreground filtering capabilities."""
+
+    @abc.abstractmethod
+    def is_foreground(
+        self,
+        mask: Mask,
+        x: int,
+        y: int,
+        width: int,
+        height: int,
+        min_foreground_ratio: float,
+    ) -> bool:
+        """Check if a patch contains sufficient foreground."""

eva/vision/data/wsi/patching/samplers/foreground_grid.py

@@ -0,0 +1,99 @@
+"""Foreground grid sampler."""
+
+from typing import Tuple
+
+from eva.vision.data.wsi.patching.mask import Mask
+from eva.vision.data.wsi.patching.samplers import _utils, base
+
+
+class ForegroundGridSampler(base.ForegroundSampler):
+    """Sample patches based on a grid, only returning patches containing foreground."""
+
+    def __init__(
+        self,
+        max_samples: int = 20,
+        overlap: Tuple[int, int] = (0, 0),
+        min_foreground_ratio: float = 0.35,
+        seed: int = 42,
+    ) -> None:
+        """Initializes the sampler.
+
+        Args:
+            max_samples: The maximum number of samples to return.
+            overlap: The overlap between patches in the grid.
+            min_foreground_ratio: The minimum amount of foreground
+                within a sampled patch.
+            seed: The random seed.
+        """
+        self.max_samples = max_samples
+        self.overlap = overlap
+        self.min_foreground_ratio = min_foreground_ratio
+        self.seed = seed
+
+    def sample(
+        self,
+        width: int,
+        height: int,
+        layer_shape: Tuple[int, int],
+        mask: Mask,
+    ):
+        """Sample patches from a grid containing foreground.
+
+        Args:
+            width: The width of the patches.
+            height: The height of the patches.
+            layer_shape: The shape of the layer.
+            mask: The mask of the image.
+        """
+        _utils.validate_dimensions(width, height, layer_shape)
+        x_y, indices = _utils.get_grid_coords_and_indices(
+            layer_shape, width, height, self.overlap, seed=self.seed
+        )
+
+        count = 0
+        for i in indices:
+            if count >= self.max_samples:
+                break
+
+            if self.is_foreground(
+                mask=mask,
+                x=x_y[i][0],
+                y=x_y[i][1],
+                width=width,
+                height=height,
+                min_foreground_ratio=self.min_foreground_ratio,
+            ):
+                count += 1
+                yield x_y[i]
+
+    def is_foreground(
+        self,
+        mask: Mask,
+        x: int,
+        y: int,
+        width: int,
+        height: int,
+        min_foreground_ratio: float,
+    ) -> bool:
+        """Check if a patch contains sufficient foreground.
+
+        Args:
+            mask: The mask of the image.
+            x: The x-coordinate of the patch.
+            y: The y-coordinate of the patch.
+            width: The width of the patch.
+            height: The height of the patch.
+            min_foreground_ratio: The minimum amount of foreground in the patch.
+        """
+        x_, y_ = self._scale_coords(x, y, mask.scale_factors)
+        width_, height_ = self._scale_coords(width, height, mask.scale_factors)
+        patch_mask = mask.mask_array[y_ : y_ + height_, x_ : x_ + width_]
+        return patch_mask.sum() / patch_mask.size >= min_foreground_ratio
+
+    def _scale_coords(
+        self,
+        x: int,
+        y: int,
+        scale_factors: Tuple[float, float],
+    ) -> Tuple[int, int]:
+        return int(x / scale_factors[0]), int(y / scale_factors[1])

eva/vision/data/wsi/patching/samplers/grid.py

@@ -0,0 +1,47 @@
+"""Grid sampler."""
+
+from typing import Generator, Tuple
+
+from eva.vision.data.wsi.patching.samplers import _utils, base
+
+
+class GridSampler(base.Sampler):
+    """Sample patches based on a grid.
+
+    Args:
+        max_samples: The maximum number of samples to return.
+        overlap: The overlap between patches in the grid.
+        seed: The random seed.
+    """
+
+    def __init__(
+        self,
+        max_samples: int | None = None,
+        overlap: Tuple[int, int] = (0, 0),
+        seed: int = 42,
+    ):
+        """Initializes the sampler."""
+        self.max_samples = max_samples
+        self.overlap = overlap
+        self.seed = seed
+
+    def sample(
+        self,
+        width: int,
+        height: int,
+        layer_shape: Tuple[int, int],
+    ) -> Generator[Tuple[int, int], None, None]:
+        """Sample patches from a grid.
+
+        Args:
+            width: The width of the patches.
+            height: The height of the patches.
+            layer_shape: The shape of the layer.
+        """
+        _utils.validate_dimensions(width, height, layer_shape)
+        x_y, indices = _utils.get_grid_coords_and_indices(
+            layer_shape, width, height, self.overlap, seed=self.seed
+        )
+        max_samples = len(indices) if self.max_samples is None else self.max_samples
+        for i in indices[:max_samples]:
+            yield x_y[i]

eva/vision/data/wsi/patching/samplers/random.py

@@ -0,0 +1,41 @@
+"""Random sampler."""
+
+import random
+from typing import Generator, Tuple
+
+from eva.vision.data.wsi.patching.samplers import _utils, base
+
+
+class RandomSampler(base.Sampler):
+    """Sample patch coordinates randomly.
+
+    Args:
+        n_samples: The number of samples to return.
+        seed: The random seed.
+    """
+
+    def __init__(self, n_samples: int = 1, seed: int = 42):
+        """Initializes the sampler."""
+        self.seed = seed
+        self.n_samples = n_samples
+
+    def sample(
+        self,
+        width: int,
+        height: int,
+        layer_shape: Tuple[int, int],
+    ) -> Generator[Tuple[int, int], None, None]:
+        """Sample random patches.
+
+        Args:
+            width: The width of the patches.
+            height: The height of the patches.
+            layer_shape: The shape of the layer.
+        """
+        _utils.validate_dimensions(width, height, layer_shape)
+        _utils.set_seed(self.seed)
+
+        x_max, y_max = layer_shape[0], layer_shape[1]
+        for _ in range(self.n_samples):
+            x, y = random.randint(0, x_max - width), random.randint(0, y_max - height)  # nosec
+            yield x, y

eva/vision/losses/__init__.py

@@ -0,0 +1,5 @@
+"""Loss functions API."""
+
+from eva.vision.losses.dice import DiceLoss
+
+__all__ = ["DiceLoss"]

eva/vision/losses/dice.py

@@ -0,0 +1,40 @@
+"""Dice loss."""
+
+import torch
+from monai import losses
+from monai.networks import one_hot  # type: ignore
+from typing_extensions import override
+
+
+class DiceLoss(losses.DiceLoss):  # type: ignore
+    """Computes the average Dice loss between two tensors.
+
+    Extends the implementation from MONAI
+        - to support semantic target labels (meaning targets of shape BHW)
+        - to support `ignore_index` functionality
+    """
+
+    def __init__(self, *args, ignore_index: int | None = None, **kwargs) -> None:
+        """Initialize the DiceLoss with support for ignore_index.
+
+        Args:
+            args: Positional arguments from the base class.
+            ignore_index: Specifies a target value that is ignored and
+                does not contribute to the input gradient.
+            kwargs: Key-word arguments from the base class.
+        """
+        super().__init__(*args, **kwargs)
+
+        self.ignore_index = ignore_index
+
+    @override
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:  # noqa
+        if self.ignore_index is not None:
+            mask = targets != self.ignore_index
+            targets = targets * mask
+            inputs = torch.mul(inputs, mask.unsqueeze(1) if mask.ndim == 3 else mask)
+
+        if targets.ndim == 3:
+            targets = one_hot(targets[:, None, ...], num_classes=inputs.shape[1])
+
+        return super().forward(inputs, targets)

eva/vision/models/__init__.py

@@ -1,5 +1,7 @@
 """Vision Models API."""
 
-from eva.vision.models import networks
+from eva.vision.models import networks, wrappers
+from eva.vision.models.networks import backbones
+from eva.vision.models.wrappers import ModelFromRegistry, TimmModel
 
-__all__ = ["networks"]
+__all__ = ["networks", "wrappers", "backbones", "ModelFromRegistry", "TimmModel"]

eva/vision/models/modules/__init__.py

@@ -0,0 +1,5 @@
+"""Vision modules API."""
+
+from eva.vision.models.modules.semantic_segmentation import SemanticSegmentationModule
+
+__all__ = ["SemanticSegmentationModule"]

eva/vision/models/modules/semantic_segmentation.py

@@ -0,0 +1,161 @@
+""""Neural Network Semantic Segmentation Module."""
+
+from typing import Any, Callable, Dict, Iterable, List, Tuple
+
+import torch
+from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from torch import nn, optim
+from torch.optim import lr_scheduler
+from typing_extensions import override
+
+from eva.core.metrics import structs as metrics_lib
+from eva.core.models.modules import module
+from eva.core.models.modules.typings import INPUT_BATCH, INPUT_TENSOR_BATCH
+from eva.core.models.modules.utils import batch_postprocess, grad
+from eva.core.utils import parser
+from eva.vision.models.networks import decoders
+
+
+class SemanticSegmentationModule(module.ModelModule):
+    """Neural network semantic segmentation module for training on patch embeddings."""
+
+    def __init__(
+        self,
+        decoder: decoders.Decoder,
+        criterion: Callable[..., torch.Tensor],
+        encoder: Dict[str, Any] | Callable[[torch.Tensor], List[torch.Tensor]] | None = None,
+        lr_multiplier_encoder: float = 0.0,
+        optimizer: OptimizerCallable = optim.AdamW,
+        lr_scheduler: LRSchedulerCallable = lr_scheduler.ConstantLR,
+        metrics: metrics_lib.MetricsSchema | None = None,
+        postprocess: batch_postprocess.BatchPostProcess | None = None,
+    ) -> None:
+        """Initializes the neural net head module.
+
+        Args:
+            decoder: The decoder model.
+            criterion: The loss function to use.
+            encoder: The encoder model. If `None`, it will be expected
+                that the input batch returns the features directly.
+                If pass as a dictionary, it will be parsed to an object
+                during the `configure_model` step.
+            lr_multiplier_encoder: The learning rate multiplier for the
+                encoder parameters. If `0`, it will freeze the encoder.
+            optimizer: The optimizer to use.
+            lr_scheduler: The learning rate scheduler to use.
+            metrics: The metric groups to track.
+            postprocess: A list of helper functions to apply after the
+                loss and before the metrics calculation to the model
+                predictions and targets.
+        """
+        super().__init__(metrics=metrics, postprocess=postprocess)
+
+        self.decoder = decoder
+        self.criterion = criterion
+        self.encoder = encoder  # type: ignore
+        self.lr_multiplier_encoder = lr_multiplier_encoder
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+
+    @override
+    def configure_model(self) -> None:
+        self._freeze_encoder()
+
+        if isinstance(self.encoder, dict):
+            self.encoder: Callable[[torch.Tensor], List[torch.Tensor]] = parser.parse_object(
+                self.encoder,
+                expected_type=nn.Module,
+            )
+
+    @override
+    def configure_optimizers(self) -> Any:
+        optimizer = self.optimizer(
+            [
+                {"params": self.decoder.parameters()},
+                {
+                    "params": self._encoder_trainable_parameters(),
+                    "lr": self._base_lr * self.lr_multiplier_encoder,
+                },
+            ]
+        )
+        lr_scheduler = self.lr_scheduler(optimizer)
+        return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}
+
+    @override
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        to_size: Tuple[int, int] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """Maps the input tensor (image tensor or embeddings) to masks.
+
+        If `inputs` is image tensor, then the `self.encoder`
+        should be implemented, otherwise it will be interpreted
+        as embeddings, where the `to_size` should be given.
+        """
+        if self.encoder is None and to_size is None:
+            raise ValueError(
+                "Please provide the expected `to_size` that the "
+                "decoder should map the embeddings (`inputs`) to."
+            )
+
+        patch_embeddings = self.encoder(inputs) if self.encoder else inputs
+        return self.decoder(patch_embeddings, to_size or inputs.shape[-2:])
+
+    @override
+    def training_step(self, batch: INPUT_TENSOR_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def validation_step(self, batch: INPUT_TENSOR_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def test_step(self, batch: INPUT_TENSOR_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        return self._batch_step(batch)
+
+    @override
+    def predict_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> torch.Tensor:
+        tensor = INPUT_BATCH(*batch).data
+        return self.encoder(tensor) if isinstance(self.encoder, nn.Module) else tensor
+
+    @property
+    def _base_lr(self) -> float:
+        """Returns the base learning rate."""
+        base_optimizer = self.optimizer(self.parameters())
+        return base_optimizer.param_groups[-1]["lr"]
+
+    def _encoder_trainable_parameters(self) -> Iterable[torch.Tensor]:
+        """Returns the trainable parameters of the encoder."""
+        return (
+            self.encoder.parameters()
+            if isinstance(self.encoder, nn.Module) and self.lr_multiplier_encoder > 0
+            else iter(())
+        )
+
+    def _freeze_encoder(self) -> None:
+        """If initialized, it freezes the encoder network."""
+        if isinstance(self.encoder, nn.Module) and self.lr_multiplier_encoder == 0:
+            grad.deactivate_requires_grad(self.encoder)
+
+    def _batch_step(self, batch: INPUT_TENSOR_BATCH) -> STEP_OUTPUT:
+        """Performs a model forward step and calculates the loss.
+
+        Args:
+            batch: The desired batch to process.
+
+        Returns:
+            The batch step output.
+        """
+        data, targets, metadata = INPUT_TENSOR_BATCH(*batch)
+        predictions = self(data, to_size=targets.shape[-2:])
+        loss = self.criterion(predictions, targets)
+        return {
+            "loss": loss,
+            "targets": targets,
+            "predictions": predictions,
+            "metadata": metadata,
+        }

eva/vision/models/networks/__init__.py

@@ -1,6 +1,5 @@
 """Vision Networks API."""
 
-from eva.vision.models.networks import postprocesses
 from eva.vision.models.networks.abmil import ABMIL
 
-__all__ = ["postprocesses", "ABMIL"]
+__all__ = ["ABMIL"]

eva/vision/models/networks/backbones/__init__.py

@@ -0,0 +1,6 @@
+"""Vision Model Backbones API."""
+
+from eva.vision.models.networks.backbones import pathology, timm, universal
+from eva.vision.models.networks.backbones.registry import BackboneModelRegistry, register_model
+
+__all__ = ["pathology", "timm", "universal", "BackboneModelRegistry", "register_model"]

eva/vision/models/networks/backbones/_utils.py

@@ -0,0 +1,39 @@
+"""Utilis for backbone networks."""
+
+from typing import Any, Dict, Tuple
+
+from torch import nn
+
+from eva import models
+from eva.core.models import transforms
+
+
+def load_hugingface_model(
+    model_name: str,
+    out_indices: int | Tuple[int, ...] | None,
+    model_kwargs: Dict[str, Any] | None = None,
+    transform_args: Dict[str, Any] | None = None,
+) -> nn.Module:
+    """Helper function to load HuggingFace models.
+
+    Args:
+        model_name: The model name to load.
+        out_indices: Whether and which multi-level patch embeddings to return.
+            Currently only out_indices=1 is supported.
+        model_kwargs: The arguments used for instantiating the model.
+        transform_args: The arguments used for instantiating the transform.
+
+    Returns: The model instance.
+    """
+    if out_indices is None:
+        tensor_transforms = transforms.ExtractCLSFeatures(**(transform_args or {}))
+    elif out_indices == 1:
+        tensor_transforms = transforms.ExtractPatchFeatures(**(transform_args or {}))
+    else:
+        raise ValueError(f"out_indices={out_indices} is currently not supported.")
+
+    return models.HuggingFaceModel(
+        model_name_or_path=model_name,
+        tensor_transforms=tensor_transforms,
+        model_kwargs=model_kwargs,
+    )

eva/vision/models/networks/backbones/pathology/__init__.py

@@ -0,0 +1,31 @@
+"""Vision Pathology Model Backbones API."""
+
+from eva.vision.models.networks.backbones.pathology.bioptimus import bioptimus_h_optimus_0
+from eva.vision.models.networks.backbones.pathology.gigapath import prov_gigapath
+from eva.vision.models.networks.backbones.pathology.histai import histai_hibou_b, histai_hibou_l
+from eva.vision.models.networks.backbones.pathology.kaiko import (
+    kaiko_vitb8,
+    kaiko_vitb16,
+    kaiko_vitl14,
+    kaiko_vits8,
+    kaiko_vits16,
+)
+from eva.vision.models.networks.backbones.pathology.lunit import lunit_vits8, lunit_vits16
+from eva.vision.models.networks.backbones.pathology.mahmood import mahmood_uni
+from eva.vision.models.networks.backbones.pathology.owkin import owkin_phikon
+
+__all__ = [
+    "kaiko_vitb16",
+    "kaiko_vitb8",
+    "kaiko_vitl14",
+    "kaiko_vits16",
+    "kaiko_vits8",
+    "owkin_phikon",
+    "lunit_vits16",
+    "lunit_vits8",
+    "mahmood_uni",
+    "bioptimus_h_optimus_0",
+    "prov_gigapath",
+    "histai_hibou_b",
+    "histai_hibou_l",
+]

eva/vision/models/networks/backbones/pathology/bioptimus.py

@@ -0,0 +1,34 @@
+"""Pathology FMs from Bioptimus."""
+
+from typing import Tuple
+
+import timm
+from torch import nn
+
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/bioptimus_h_optimus_0")
+def bioptimus_h_optimus_0(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+) -> nn.Module:
+    """Initializes the h_optimus_0 pathology FM by Bioptimus.
+
+    Args:
+        dynamic_img_size: Whether to allow the interpolation embedding
+            to be interpolated at `forward()` time when image grid changes
+            from original.
+        out_indices: Weather and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return timm.create_model(
+        model_name="hf-hub:bioptimus/H-optimus-0",
+        pretrained=True,
+        init_values=1e-5,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+        features_only=out_indices is not None,
+    )

eva/vision/models/networks/backbones/pathology/gigapath.py

@@ -0,0 +1,33 @@
+"""Pathology FMs from other/mixed entities."""
+
+from typing import Tuple
+
+import timm
+from torch import nn
+
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/prov_gigapath")
+def prov_gigapath(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+) -> nn.Module:
+    """Initializes the Prov-GigaPath pathology FM.
+
+    Args:
+        dynamic_img_size: Whether to allow the interpolation embedding
+            to be interpolated at `forward()` time when image grid changes
+            from original.
+        out_indices: Weather and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return timm.create_model(
+        model_name="hf_hub:prov-gigapath/prov-gigapath",
+        pretrained=True,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+        features_only=out_indices is not None,
+    )