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

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,
+    )

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

@@ -0,0 +1,46 @@
+"""Pathology FMs from owkin."""
+
+from typing import Tuple
+
+from torch import nn
+
+from eva.vision.models.networks.backbones import _utils
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/histai_hibou_b")
+def histai_hibou_b(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the hibou-B pathology FM by hist.ai (https://huggingface.co/histai/hibou-B).
+
+    Args:
+        out_indices: Whether and which multi-level patch embeddings to return.
+            Currently only out_indices=1 is supported.
+
+    Returns:
+        The model instance.
+    """
+    return _utils.load_hugingface_model(
+        model_name="histai/hibou-B",
+        out_indices=out_indices,
+        model_kwargs={"trust_remote_code": True},
+        transform_args={"ignore_remaining_dims": True} if out_indices is not None else None,
+    )
+
+
+@register_model("pathology/histai_hibou_l")
+def histai_hibou_l(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the hibou-L pathology FM by hist.ai (https://huggingface.co/histai/hibou-L).
+
+    Args:
+        out_indices: Whether and which multi-level patch embeddings to return.
+            Currently only out_indices=1 is supported.
+
+    Returns:
+        The model instance.
+    """
+    return _utils.load_hugingface_model(
+        model_name="histai/hibou-L",
+        out_indices=out_indices,
+        model_kwargs={"trust_remote_code": True},
+        transform_args={"ignore_remaining_dims": True} if out_indices is not None else None,
+    )

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

@@ -0,0 +1,123 @@
+"""Pathology FMs from kaiko.ai."""
+
+from typing import Tuple
+
+import torch
+from torch import nn
+
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/kaiko_vits16")
+def kaiko_vits16(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTS-16 pathology FM by kaiko.ai.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return torch.hub.load(  # type: ignore
+        repo_or_dir="kaiko-ai/towards_large_pathology_fms",
+        model="vits16",
+        trust_repo=True,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+    )
+
+
+@register_model("pathology/kaiko_vits8")
+def kaiko_vits8(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTS-8 pathology FM by kaiko.ai.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return torch.hub.load(  # type: ignore
+        repo_or_dir="kaiko-ai/towards_large_pathology_fms",
+        model="vits8",
+        trust_repo=True,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+    )
+
+
+@register_model("pathology/kaiko_vitb16")
+def kaiko_vitb16(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTB-16 pathology FM by kaiko.ai.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return torch.hub.load(  # type: ignore
+        repo_or_dir="kaiko-ai/towards_large_pathology_fms",
+        model="vitb16",
+        trust_repo=True,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+    )
+
+
+@register_model("pathology/kaiko_vitb8")
+def kaiko_vitb8(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTB-8 pathology FM by kaiko.ai.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return torch.hub.load(  # type: ignore
+        repo_or_dir="kaiko-ai/towards_large_pathology_fms",
+        model="vitb8",
+        trust_repo=True,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+    )
+
+
+@register_model("pathology/kaiko_vitl14")
+def kaiko_vitl14(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTL-14 pathology FM by kaiko.ai.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return torch.hub.load(  # type: ignore
+        repo_or_dir="kaiko-ai/towards_large_pathology_fms",
+        model="vitl14",
+        trust_repo=True,
+        dynamic_img_size=dynamic_img_size,
+        out_indices=out_indices,
+    )

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

@@ -0,0 +1,68 @@
+"""Pathology FMs from Lunit.
+
+Source: https://github.com/lunit-io/benchmark-ssl-pathology/releases
+
+For training the vit-s models the following standardization parameters were used:
+
+mean: [ 0.70322989, 0.53606487, 0.66096631 ]
+std: [ 0.21716536, 0.26081574, 0.20723464 ]
+"""
+
+from typing import Tuple
+
+from torch import nn
+
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones.registry import register_model
+
+VITS_URL_PREFIX = (
+    "https://github.com/lunit-io/benchmark-ssl-pathology/releases/download/pretrained-weights"
+)
+
+
+@register_model("pathology/lunit_vits16")
+def lunit_vits16(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTS-16 pathology FM by lunit.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return wrappers.TimmModel(
+        model_name="vit_small_patch16_224.dino",
+        out_indices=out_indices,
+        model_kwargs={
+            "dynamic_img_size": dynamic_img_size,
+        },
+        checkpoint_path=f"{VITS_URL_PREFIX}/dino_vit_small_patch16_ep200.torch",
+    )
+
+
+@register_model("pathology/lunit_vits8")
+def lunit_vits8(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes the ViTS-8 pathology FM by lunit.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return wrappers.TimmModel(
+        model_name="vit_small_patch8_224.dino",
+        out_indices=out_indices,
+        model_kwargs={
+            "dynamic_img_size": dynamic_img_size,
+        },
+        checkpoint_path=f"{VITS_URL_PREFIX}/dino_vit_small_patch8_ep200.torch",
+    )