kaiko-eva 0.1.1__py3-none-any.whl → 0.1.5__py3-none-any.whl

eva/vision/losses/dice.py

@@ -1,4 +1,6 @@
-"""Dice loss."""
+"""Dice based loss functions."""
+
+from typing import Sequence, Tuple
 
 import torch
 from monai import losses
@@ -12,29 +14,94 @@ class DiceLoss(losses.DiceLoss):  # type: ignore
     Extends the implementation from MONAI
         - to support semantic target labels (meaning targets of shape BHW)
         - to support `ignore_index` functionality
+        - accept weight argument in list format
     """
 
-    def __init__(self, *args, ignore_index: int | None = None, **kwargs) -> None:
-        """Initialize the DiceLoss with support for ignore_index.
+    def __init__(
+        self,
+        *args,
+        ignore_index: int | None = None,
+        weight: Sequence[float] | torch.Tensor | None = None,
+        **kwargs,
+    ) -> None:
+        """Initialize the DiceLoss.
 
         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.
+            weight: A list of weights to assign to each class.
             kwargs: Key-word arguments from the base class.
         """
-        super().__init__(*args, **kwargs)
+        if weight is not None and not isinstance(weight, torch.Tensor):
+            weight = torch.tensor(weight)
+
+        super().__init__(*args, **kwargs, weight=weight)
 
         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)
+        inputs, targets = _apply_ignore_index(inputs, targets, self.ignore_index)
+        targets = _to_one_hot(targets, num_classes=inputs.shape[1])
 
         if targets.ndim == 3:
             targets = one_hot(targets[:, None, ...], num_classes=inputs.shape[1])
 
         return super().forward(inputs, targets)
+
+
+class DiceCELoss(losses.dice.DiceCELoss):
+    """Combination of Dice and Cross Entropy Loss.
+
+    Extends the implementation from MONAI
+        - to support semantic target labels (meaning targets of shape BHW)
+        - to support `ignore_index` functionality
+        - accept weight argument in list format
+    """
+
+    def __init__(
+        self,
+        *args,
+        ignore_index: int | None = None,
+        weight: Sequence[float] | torch.Tensor | None = None,
+        **kwargs,
+    ) -> None:
+        """Initialize the DiceCELoss.
+
+        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.
+            weight: A list of weights to assign to each class.
+            kwargs: Key-word arguments from the base class.
+        """
+        if weight is not None and not isinstance(weight, torch.Tensor):
+            weight = torch.tensor(weight)
+
+        super().__init__(*args, **kwargs, weight=weight)
+
+        self.ignore_index = ignore_index
+
+    @override
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:  # noqa
+        inputs, targets = _apply_ignore_index(inputs, targets, self.ignore_index)
+        targets = _to_one_hot(targets, num_classes=inputs.shape[1])
+
+        return super().forward(inputs, targets)
+
+
+def _apply_ignore_index(
+    inputs: torch.Tensor, targets: torch.Tensor, ignore_index: int | None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if ignore_index is not None:
+        mask = targets != ignore_index
+        targets = targets * mask
+        inputs = torch.mul(inputs, mask.unsqueeze(1) if mask.ndim == 3 else mask)
+    return inputs, targets
+
+
+def _to_one_hot(tensor: torch.Tensor, num_classes: int) -> torch.Tensor:
+    if tensor.ndim == 3:
+        return one_hot(tensor[:, None, ...], num_classes=num_classes)
+    return tensor

eva/vision/metrics/__init__.py

@@ -0,0 +1,11 @@
+"""Default metric collections API."""
+
+from eva.vision.metrics.defaults.segmentation import MulticlassSegmentationMetrics
+from eva.vision.metrics.segmentation.generalized_dice import GeneralizedDiceScore
+from eva.vision.metrics.segmentation.mean_iou import MeanIoU
+
+__all__ = [
+    "MulticlassSegmentationMetrics",
+    "GeneralizedDiceScore",
+    "MeanIoU",
+]

eva/vision/metrics/defaults/__init__.py

@@ -0,0 +1,7 @@
+"""Default metric collections API."""
+
+from eva.vision.metrics.defaults.segmentation import MulticlassSegmentationMetrics
+
+__all__ = [
+    "MulticlassSegmentationMetrics",
+]

eva/{core → vision}/metrics/defaults/segmentation/__init__.py

@@ -1,5 +1,5 @@
 """Default segmentation metric collections API."""
 
-from eva.core.metrics.defaults.segmentation.multiclass import MulticlassSegmentationMetrics
+from eva.vision.metrics.defaults.segmentation.multiclass import MulticlassSegmentationMetrics
 
 __all__ = ["MulticlassSegmentationMetrics"]

eva/{core → vision}/metrics/defaults/segmentation/multiclass.py

@@ -1,6 +1,7 @@
 """Default metric collection for multiclass semantic segmentation tasks."""
 
-from eva.core.metrics import generalized_dice, mean_iou, structs
+from eva.core.metrics import structs
+from eva.vision.metrics.segmentation import generalized_dice, mean_iou
 
 
 class MulticlassSegmentationMetrics(structs.MetricCollection):

eva/vision/metrics/segmentation/BUILD

@@ -0,0 +1 @@
+python_sources()

eva/vision/metrics/segmentation/__init__.py

@@ -0,0 +1,9 @@
+"""Segmentation metrics API."""
+
+from eva.vision.metrics.segmentation.generalized_dice import GeneralizedDiceScore
+from eva.vision.metrics.segmentation.mean_iou import MeanIoU
+
+__all__ = [
+    "GeneralizedDiceScore",
+    "MeanIoU",
+]

eva/vision/metrics/segmentation/_utils.py

@@ -0,0 +1,69 @@
+"""Utils for segmentation metric collections."""
+
+from typing import Tuple
+
+import torch
+
+
+def apply_ignore_index(
+    preds: torch.Tensor, target: torch.Tensor, ignore_index: int, num_classes: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Applies the ignore index to the predictions and target tensors.
+
+    1. Masks the values in the target tensor that correspond to the ignored index.
+    2. Remove the channel corresponding to the ignored index from both tensors.
+
+    Args:
+        preds: The predictions tensor. Expected to be of shape `(N,C,...)`.
+        target: The target tensor. Expected to be of shape `(N,C,...)`.
+        ignore_index: The index to ignore.
+        num_classes: The number of classes.
+
+    Returns:
+        The modified predictions and target tensors of shape `(N,C-1,...)`.
+    """
+    if ignore_index < 0:
+        raise ValueError("ignore_index must be a non-negative integer")
+
+    ignore_mask = preds[:, ignore_index] == 1
+    target = target * (~ignore_mask.unsqueeze(1))
+
+    preds = _ignore_tensor_channel(preds, ignore_index)
+    target = _ignore_tensor_channel(target, ignore_index)
+
+    return preds, target
+
+
+def index_to_one_hot(tensor: torch.Tensor, num_classes: int) -> torch.Tensor:
+    """Converts an index tensor to a one-hot tensor.
+
+    Args:
+        tensor: The index tensor to convert. Expected to be of shape `(N,...)`.
+        num_classes: The number of classes to one-hot encode.
+
+    Returns:
+        A one-hot tensor of shape `(N,C,...)`.
+    """
+    if not _is_one_hot(tensor):
+        tensor = torch.nn.functional.one_hot(tensor.long(), num_classes=num_classes).movedim(-1, 1)
+    return tensor
+
+
+def _ignore_tensor_channel(tensor: torch.Tensor, ignore_index: int) -> torch.Tensor:
+    """Removes the channel corresponding to the specified ignore index.
+
+    Args:
+        tensor: The tensor to remove the channel from. Expected to be of shape `(N,C,...)`.
+        ignore_index: The index of the channel dimension (C) to remove.
+
+    Returns:
+        A tensor without the specified channel `(N,C-1,...)`.
+    """
+    if ignore_index < 0:
+        raise ValueError("ignore_index must be a non-negative integer")
+    return torch.cat([tensor[:, :ignore_index], tensor[:, ignore_index + 1 :]], dim=1)
+
+
+def _is_one_hot(tensor: torch.Tensor, expected_dim: int = 4) -> bool:
+    """Checks if the tensor is a one-hot tensor."""
+    return bool((tensor.bool() == tensor).all()) and tensor.ndim == expected_dim

eva/{core/metrics → vision/metrics/segmentation}/generalized_dice.py

@@ -6,6 +6,8 @@ import torch
 from torchmetrics import segmentation
 from typing_extensions import override
 
+from eva.vision.metrics.segmentation import _utils
+
 
 class GeneralizedDiceScore(segmentation.GeneralizedDiceScore):
     """Defines the Generalized Dice Score.
@@ -30,8 +32,6 @@ class GeneralizedDiceScore(segmentation.GeneralizedDiceScore):
             include_background: Whether to include the background class in the computation
             weight_type: The type of weight to apply to each class. Can be one of `"square"`,
                 `"simple"`, or `"linear"`.
-            input_format: What kind of input the function receives. Choose between ``"one-hot"``
-                for one-hot encoded tensors or ``"index"`` for index tensors.
             ignore_index: Integer specifying a target class to ignore. If given, this class
                 index does not contribute to the returned score, regardless of reduction method.
             per_class: Whether to compute the IoU for each class separately. If set to ``False``,
@@ -39,21 +39,23 @@ class GeneralizedDiceScore(segmentation.GeneralizedDiceScore):
             kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
         """
         super().__init__(
-            num_classes=num_classes,
+            num_classes=num_classes
+            - (ignore_index is not None)
+            + (ignore_index == 0 and not include_background),
             include_background=include_background,
             weight_type=weight_type,
             per_class=per_class,
             **kwargs,
         )
-
+        self.orig_num_classes = num_classes
         self.ignore_index = ignore_index
 
     @override
     def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        preds = _utils.index_to_one_hot(preds, num_classes=self.orig_num_classes)
+        target = _utils.index_to_one_hot(target, num_classes=self.orig_num_classes)
         if self.ignore_index is not None:
-            mask = target != self.ignore_index
-            mask = mask.all(dim=-1, keepdim=True)
-            preds = preds * mask
-            target = target * mask
-
-        super().update(preds=preds, target=target)
+            preds, target = _utils.apply_ignore_index(
+                preds, target, self.ignore_index, self.num_classes
+            )
+        super().update(preds=preds.long(), target=target.long())

eva/vision/metrics/segmentation/mean_iou.py

@@ -0,0 +1,57 @@
+"""MeanIoU metric for semantic segmentation."""
+
+from typing import Any
+
+import torch
+from torchmetrics import segmentation
+from typing_extensions import override
+
+from eva.vision.metrics.segmentation import _utils
+
+
+class MeanIoU(segmentation.MeanIoU):
+    """MeanIoU (mIOU) metric for semantic segmentation.
+
+    It expands the `torchmetrics` class by including an `ignore_index`
+    functionality.
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        include_background: bool = True,
+        ignore_index: int | None = None,
+        per_class: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes the metric.
+
+        Args:
+            num_classes: The number of classes in the segmentation problem.
+            include_background: Whether to include the background class in the computation
+            ignore_index: Integer specifying a target class to ignore. If given, this class
+                index does not contribute to the returned score, regardless of reduction method.
+            per_class: Whether to compute the IoU for each class separately. If set to ``False``,
+                the metric will compute the mean IoU over all classes.
+            kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+        """
+        super().__init__(
+            num_classes=num_classes
+            - (ignore_index is not None)
+            + (ignore_index == 0 and not include_background),
+            include_background=include_background,
+            per_class=per_class,
+            **kwargs,
+        )
+        self.orig_num_classes = num_classes
+        self.ignore_index = ignore_index
+
+    @override
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        preds = _utils.index_to_one_hot(preds, num_classes=self.orig_num_classes)
+        target = _utils.index_to_one_hot(target, num_classes=self.orig_num_classes)
+        if self.ignore_index is not None:
+            preds, target = _utils.apply_ignore_index(
+                preds, target, self.ignore_index, self.num_classes
+            )
+        super().update(preds=preds.long(), target=target.long())

eva/vision/models/modules/semantic_segmentation.py

@@ -15,6 +15,7 @@ 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
+from eva.vision.models.networks.decoders.segmentation.typings import DecoderInputs
 
 
 class SemanticSegmentationModule(module.ModelModule):
@@ -101,9 +102,9 @@ class SemanticSegmentationModule(module.ModelModule):
                 "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:])
+        features = self.encoder(inputs) if self.encoder else inputs
+        decoder_inputs = DecoderInputs(features, to_size or inputs.shape[-2:], inputs)  # type: ignore
+        return self.decoder(decoder_inputs)
 
     @override
     def training_step(self, batch: INPUT_TENSOR_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:

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

@@ -1,7 +1,9 @@
 """Utilis for backbone networks."""
 
+import os
 from typing import Any, Dict, Tuple
 
+import huggingface_hub
 from torch import nn
 
 from eva import models
@@ -37,3 +39,13 @@ def load_hugingface_model(
         tensor_transforms=tensor_transforms,
         model_kwargs=model_kwargs,
     )
+
+
+def huggingface_login(hf_token: str | None = None):
+    token = hf_token or os.environ.get("HF_TOKEN")
+    if not token:
+        raise ValueError(
+            "Please provide a HuggingFace token to download the model. "
+            "You can either pass it as an argument or set the env variable HF_TOKEN."
+        )
+    huggingface_hub.login(token=token)

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

@@ -12,7 +12,8 @@ from eva.vision.models.networks.backbones.pathology.kaiko import (
 )
 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
+from eva.vision.models.networks.backbones.pathology.owkin import owkin_phikon, owkin_phikon_v2
+from eva.vision.models.networks.backbones.pathology.paige import paige_virchow2
 
 __all__ = [
     "kaiko_vitb16",
@@ -21,6 +22,7 @@ __all__ = [
     "kaiko_vits16",
     "kaiko_vits8",
     "owkin_phikon",
+    "owkin_phikon_v2",
     "lunit_vits16",
     "lunit_vits8",
     "mahmood_uni",
@@ -28,4 +30,5 @@ __all__ = [
     "prov_gigapath",
     "histai_hibou_b",
     "histai_hibou_l",
+    "paige_virchow2",
 ]

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

@@ -12,6 +12,9 @@ from eva.vision.models.networks.backbones.registry import register_model
 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).
 
+    Uses a customized implementation of the DINOv2 architecture from the transformers
+    library to add support for registers, which requires the trust_remote_code=True flag.
+
     Args:
         out_indices: Whether and which multi-level patch embeddings to return.
             Currently only out_indices=1 is supported.
@@ -23,7 +26,7 @@ def histai_hibou_b(out_indices: int | Tuple[int, ...] | None = None) -> nn.Modul
         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,
+        transform_args={"num_register_tokens": 4} if out_indices is not None else None,
     )
 
 
@@ -31,6 +34,9 @@ def histai_hibou_b(out_indices: int | Tuple[int, ...] | None = None) -> nn.Modul
 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).
 
+    Uses a customized implementation of the DINOv2 architecture from the transformers
+    library to add support for registers, which requires the trust_remote_code=True flag.
+
     Args:
         out_indices: Whether and which multi-level patch embeddings to return.
             Currently only out_indices=1 is supported.
@@ -42,5 +48,5 @@ def histai_hibou_l(out_indices: int | Tuple[int, ...] | None = None) -> nn.Modul
         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,
+        transform_args={"num_register_tokens": 4} if out_indices is not None else None,
     )

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

@@ -9,6 +9,7 @@ from loguru import logger
 from torch import nn
 
 from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones import _utils
 from eva.vision.models.networks.backbones.registry import register_model
 
 
@@ -31,19 +32,11 @@ def mahmood_uni(
     Returns:
         The model instance.
     """
-    token = hf_token or os.environ.get("HF_TOKEN")
-    if not token:
-        raise ValueError(
-            "Please provide a HuggingFace token to download the model. "
-            "You can either pass it as an argument or set the env variable HF_TOKEN."
-        )
-
     checkpoint_path = os.path.join(download_dir, "pytorch_model.bin")
-
     if not os.path.exists(checkpoint_path):
         logger.info(f"Downloading the model checkpoint to {download_dir} ...")
         os.makedirs(download_dir, exist_ok=True)
-        huggingface_hub.login(token=token)
+        _utils.huggingface_login(hf_token)
         huggingface_hub.hf_hub_download(
             "MahmoodLab/UNI",
             filename="pytorch_model.bin",

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

@@ -20,3 +20,17 @@ def owkin_phikon(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
         The model instance.
     """
     return _utils.load_hugingface_model(model_name="owkin/phikon", out_indices=out_indices)
+
+
+@register_model("pathology/owkin_phikon_v2")
+def owkin_phikon_v2(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the phikon-v2 pathology FM by owkin (https://huggingface.co/owkin/phikon-v2).
+
+    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="owkin/phikon-v2", out_indices=out_indices)

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

@@ -0,0 +1,51 @@
+"""Pathology FMs from paige.ai.
+
+Source: https://huggingface.co/paige-ai/
+"""
+
+from typing import Tuple
+
+import timm
+import torch.nn as nn
+
+from eva.core.models import transforms
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones import _utils
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/paige_virchow2")
+def paige_virchow2(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+    hf_token: str | None = None,
+    include_patch_tokens: bool = False,
+) -> nn.Module:
+    """Initializes the Virchow2 pathology FM by paige.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.
+        include_patch_tokens: Whether to combine the mean aggregated patch tokens with cls token.
+        hf_token: HuggingFace token to download the model.
+
+    Returns:
+        The model instance.
+    """
+    _utils.huggingface_login(hf_token)
+    return wrappers.TimmModel(
+        model_name="hf-hub:paige-ai/Virchow2",
+        out_indices=out_indices,
+        pretrained=True,
+        model_kwargs={
+            "dynamic_img_size": dynamic_img_size,
+            "mlp_layer": timm.layers.SwiGLUPacked,
+            "act_layer": nn.SiLU,
+        },
+        tensor_transforms=(
+            transforms.ExtractCLSFeatures(include_patch_tokens=include_patch_tokens)
+            if out_indices is None
+            else None
+        ),
+    )

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

@@ -1,6 +1,6 @@
 """Decoder heads API."""
 
 from eva.vision.models.networks.decoders import segmentation
-from eva.vision.models.networks.decoders.decoder import Decoder
+from eva.vision.models.networks.decoders.segmentation.base import Decoder
 
 __all__ = ["segmentation", "Decoder"]

eva/vision/models/networks/decoders/segmentation/__init__.py

@@ -1,11 +1,19 @@
 """Segmentation decoder heads API."""
 
-from eva.vision.models.networks.decoders.segmentation.common import (
+from eva.vision.models.networks.decoders.segmentation.decoder2d import Decoder2D
+from eva.vision.models.networks.decoders.segmentation.linear import LinearDecoder
+from eva.vision.models.networks.decoders.segmentation.semantic import (
     ConvDecoder1x1,
     ConvDecoderMS,
+    ConvDecoderWithImage,
     SingleLinearDecoder,
 )
-from eva.vision.models.networks.decoders.segmentation.conv2d import ConvDecoder
-from eva.vision.models.networks.decoders.segmentation.linear import LinearDecoder
 
-__all__ = ["ConvDecoder1x1", "ConvDecoderMS", "SingleLinearDecoder", "ConvDecoder", "LinearDecoder"]
+__all__ = [
+    "ConvDecoder1x1",
+    "ConvDecoderMS",
+    "SingleLinearDecoder",
+    "ConvDecoderWithImage",
+    "Decoder2D",
+    "LinearDecoder",
+]

eva/vision/models/networks/decoders/segmentation/base.py

@@ -0,0 +1,16 @@
+"""Semantic segmentation decoder base class."""
+
+import abc
+
+import torch
+from torch import nn
+
+from eva.vision.models.networks.decoders.segmentation.typings import DecoderInputs
+
+
+class Decoder(nn.Module, abc.ABC):
+    """Abstract base class for segmentation decoders."""
+
+    @abc.abstractmethod
+    def forward(self, decoder_inputs: DecoderInputs) -> torch.Tensor:
+        """Forward pass of the decoder."""