kaiko-eva 0.1.8__py3-none-any.whl → 0.2.1__py3-none-any.whl

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

@@ -5,6 +5,9 @@ from typing import Tuple
 import timm
 from torch import 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
 
 
@@ -13,7 +16,9 @@ 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.
+    """Initializes the H-Optimus-0 pathology FM by Bioptimus.
+
+    See https://huggingface.co/bioptimus/H-optimus-0 for details.
 
     Args:
         dynamic_img_size: Whether to allow the interpolation embedding
@@ -32,3 +37,44 @@ def bioptimus_h_optimus_0(
         out_indices=out_indices,
         features_only=out_indices is not None,
     )
+
+
+@register_model("pathology/bioptimus_h0_mini")
+def bioptimus_h0_mini(
+    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 H0-mini (ViT-B) pathology FM by Bioptimus.
+
+    This model was distilled from H-Optimus-0 on 40M TCGA tiles.
+
+    See https://huggingface.co/bioptimus/H0-mini for details.
+
+    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.
+        hf_token: HuggingFace token to download the model.
+        include_patch_tokens: Whether to combine the mean aggregated patch tokens with cls token.
+
+    Returns:
+        The model instance.
+    """
+    _utils.huggingface_login(hf_token)
+    return wrappers.TimmModel(
+        model_name="hf-hub:bioptimus/H0-mini",
+        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/backbones/pathology/hkust.py

@@ -0,0 +1,69 @@
+"""Pathology FMs from Hong Kong University of Science and Technology."""
+
+import re
+from typing import Tuple
+
+import timm
+from torch import nn
+
+from eva.core.models.wrappers import _utils
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/hkust_gpfm")
+def hkust_gpfm(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+) -> nn.Module:
+    """Initializes GPFM model from Hong Kong University of Science and Technology.
+
+    Ma, J., Guo, Z., Zhou, F., Wang, Y., Xu, Y., et al. (2024).
+    Towards a generalizable pathology foundation model via unified knowledge
+    distillation (arXiv No. 2407.18449). arXiv. https://arxiv.org/abs/2407.18449
+
+    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 timm.create_model(
+        model_name="vit_large_patch14_dinov2",
+        pretrained=True,
+        pretrained_cfg={
+            "state_dict": _load_state_dict(),
+            "num_classes": 0,
+        },
+        out_indices=out_indices,
+        features_only=out_indices is not None,
+        **{
+            "img_size": 224,
+            "patch_size": 14,
+            "init_values": 1e-5,
+            "qkv_bias": True,
+            "dynamic_img_size": dynamic_img_size,
+        },
+    )
+
+
+def _load_state_dict() -> dict:
+    """Loads the state dict with model weights from github."""
+    state_dict = _utils.load_state_dict_from_url(
+        url="https://github.com/birkhoffkiki/GPFM/releases/download/ckpt/GPFM.pth",
+        md5="0dc7e345de84f385d09c8c782b4b3236",
+    )
+    return _convert_state_dict(state_dict["teacher"])
+
+
+def _convert_state_dict(state_dict: dict) -> dict:
+    """Rename state dict keys to match timm's format."""
+    state_dict = {
+        re.sub(r"blocks\.\d+\.(\d+)", r"blocks.\1", key.replace("backbone.", "")): value
+        for key, value in state_dict.items()
+    }
+    remove_keys = ["mask_token"] + [key for key in state_dict.keys() if "dino_head" in key]
+    for key in remove_keys:
+        state_dict.pop(key)
+    return state_dict

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

@@ -5,9 +5,27 @@ from typing import Tuple
 import torch
 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/kaiko_midnight_12k")
+def kaiko_midnight_12k(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the Midnight-12k pathology FM by kaiko.ai.
+
+    Args:
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return _utils.load_hugingface_model(
+        model_name="kaiko-ai/midnight",
+        out_indices=out_indices,
+        model_kwargs={"trust_remote_code": True},
+    )
+
+
 @register_model("pathology/kaiko_vits16")
 def kaiko_vits16(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None

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

@@ -1,11 +1,9 @@
 """Pathology FMs from MahmoodLab."""
 
-import os
-from pathlib import Path
 from typing import Tuple
 
-import huggingface_hub
-from loguru import logger
+import timm
+import torch
 from torch import nn
 
 from eva.vision.models import wrappers
@@ -18,7 +16,6 @@ def mahmood_uni(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,
     hf_token: str | None = None,
-    download_dir: str = os.path.join(str(Path.home()), ".cache/eva"),
 ) -> nn.Module:
     """Initializes UNI model from MahmoodLab.
 
@@ -27,29 +24,59 @@ def mahmood_uni(
             the grid size (interpolate abs and/or ROPE pos) in the forward pass.
         out_indices: Whether and which multi-level patch embeddings to return.
         hf_token: HuggingFace token to download the model.
-        download_dir: Directory to download the model checkpoint.
 
     Returns:
         The model instance.
     """
-    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)
-        _utils.huggingface_login(hf_token)
-        huggingface_hub.hf_hub_download(
-            "MahmoodLab/UNI",
-            filename="pytorch_model.bin",
-            local_dir=download_dir,
-            force_download=True,
-        )
+    _utils.huggingface_login(hf_token)
 
     return wrappers.TimmModel(
-        model_name="vit_large_patch16_224",
+        model_name="hf-hub:MahmoodLab/uni",
+        pretrained=True,
         out_indices=out_indices,
         model_kwargs={
             "init_values": 1e-5,
             "dynamic_img_size": dynamic_img_size,
         },
-        checkpoint_path=checkpoint_path,
+    )
+
+
+@register_model("pathology/mahmood_uni2_h")
+def mahmood_uni2_h(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+    hf_token: str | None = None,
+) -> nn.Module:
+    """Initializes UNI model from MahmoodLab.
+
+    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.
+        hf_token: HuggingFace token to download the model.
+
+    Returns:
+        The model instance.
+    """
+    _utils.huggingface_login(hf_token)
+
+    return wrappers.TimmModel(
+        model_name="hf-hub:MahmoodLab/UNI2-h",
+        pretrained=True,
+        out_indices=out_indices,
+        model_kwargs={
+            "img_size": 224,
+            "patch_size": 14,
+            "depth": 24,
+            "num_heads": 24,
+            "init_values": 1e-5,
+            "embed_dim": 1536,
+            "mlp_ratio": 2.66667 * 2,
+            "num_classes": 0,
+            "no_embed_class": True,
+            "mlp_layer": timm.layers.SwiGLUPacked,
+            "act_layer": torch.nn.SiLU,
+            "reg_tokens": 8,
+            "dynamic_img_size": dynamic_img_size,
+        },
     )

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

@@ -0,0 +1,11 @@
+"""Vision Radiology Model Backbones API."""
+
+from eva.vision.models.networks.backbones.radiology.swin_unetr import SwinUNETREncoder
+from eva.vision.models.networks.backbones.radiology.voco import VoCoB, VoCoH, VoCoL
+
+__all__ = [
+    "VoCoB",
+    "VoCoL",
+    "VoCoH",
+    "SwinUNETREncoder",
+]

eva/vision/models/networks/backbones/radiology/swin_unetr.py

@@ -0,0 +1,231 @@
+"""Encoder based on Swin UNETR."""
+
+from typing import List, Tuple
+
+import torch
+from monai.inferers.inferer import Inferer
+from monai.networks.blocks import unetr_block
+from monai.networks.nets import swin_unetr
+from monai.utils import misc
+from torch import nn
+
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("radiology/swin_unetr_encoder")
+class SwinUNETREncoder(nn.Module):
+    """Swin transformer encoder based on UNETR [0].
+
+    - [0] UNETR: Transformers for 3D Medical Image Segmentation
+      https://arxiv.org/pdf/2103.10504
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 1,
+        feature_size: int = 48,
+        spatial_dims: int = 3,
+        out_indices: int | None = None,
+        inferer: Inferer | None = None,
+        use_v2: bool = True,
+    ) -> None:
+        """Build the UNETR encoder.
+
+        Args:
+            in_channels: Number of input channels.
+            feature_size: The dimension of network feature size.
+            spatial_dims: Number of spatial dimensions.
+            out_indices: Number of feature outputs. If None,
+                the aggregated feature vector is returned.
+            inferer: An optional MONAI `Inferer` for efficient
+                inference during evaluation.
+            use_v2: Whether to use SwinTransformerV2.
+        """
+        super().__init__()
+
+        self._in_channels = in_channels
+        self._feature_size = feature_size
+        self._spatial_dims = spatial_dims
+        self._out_indices = out_indices
+        self._inferer = inferer
+        self._use_v2 = use_v2
+
+        self._window_size = misc.ensure_tuple_rep(7, spatial_dims)
+        self._patch_size = misc.ensure_tuple_rep(2, spatial_dims)
+
+        self.swinViT = swin_unetr.SwinTransformer(
+            in_chans=in_channels,
+            embed_dim=feature_size,
+            window_size=self._window_size,
+            patch_size=self._patch_size,
+            depths=(2, 2, 2, 2),
+            num_heads=(3, 6, 12, 24),
+            mlp_ratio=4.0,
+            qkv_bias=True,
+            drop_rate=0.0,
+            attn_drop_rate=0.0,
+            drop_path_rate=0.0,
+            norm_layer=torch.nn.LayerNorm,
+            spatial_dims=spatial_dims,
+            use_v2=use_v2,
+        )
+        self.encoder1 = unetr_block.UnetrBasicBlock(
+            spatial_dims=spatial_dims,
+            in_channels=in_channels,
+            out_channels=feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.encoder2 = unetr_block.UnetrBasicBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size,
+            out_channels=feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.encoder3 = unetr_block.UnetrBasicBlock(
+            spatial_dims=spatial_dims,
+            in_channels=2 * feature_size,
+            out_channels=2 * feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.encoder4 = unetr_block.UnetrBasicBlock(
+            spatial_dims=spatial_dims,
+            in_channels=4 * feature_size,
+            out_channels=4 * feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.encoder10 = unetr_block.UnetrBasicBlock(
+            spatial_dims=spatial_dims,
+            in_channels=16 * feature_size,
+            out_channels=16 * feature_size,
+            kernel_size=3,
+            stride=1,
+            norm_name="instance",
+            res_block=True,
+        )
+        self._pool_op = (
+            nn.AdaptiveAvgPool3d(output_size=(1, 1, 1))
+            if spatial_dims == 3
+            else nn.AdaptiveAvgPool2d(output_size=(1, 1))
+        )
+
+    def _forward_features(self, tensor: torch.Tensor) -> List[torch.Tensor]:
+        """Extracts feature maps from the Swin Transformer and encoder blocks.
+
+        Args:
+            tensor: Input tensor of shape (B, C, T, H, W).
+
+        Returns:
+            List of feature maps from encoder stages.
+        """
+        hidden_states = self.swinViT(tensor)
+        enc0 = self.encoder1(tensor)
+        enc1 = self.encoder2(hidden_states[0])
+        enc2 = self.encoder3(hidden_states[1])
+        enc3 = self.encoder4(hidden_states[2])
+        dec4 = self.encoder10(hidden_states[4])
+        return [enc0, enc1, enc2, enc3, hidden_states[3], dec4]
+
+    def forward_features(self, tensor: torch.Tensor) -> List[torch.Tensor]:
+        """Computes feature maps using either standard forward pass or inference mode.
+
+        If in inference mode (`self.training` is False) and an inference method
+        (`self._inferer`) is available,  the `_inferer` is used to extract features.
+        Otherwise, `_forward_features` is called directly.
+
+        Args:
+            tensor: Input tensor of shape (B, C, T, H, W).
+
+        Returns:
+            List of feature maps from encoder stages.
+        """
+        if not self.training and self._inferer:
+            return self._inferer(inputs=tensor, network=self._forward_features)
+
+        return self._forward_features(tensor)
+
+    def forward_encoders(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Aggregates encoder features into a single feature vector.
+
+        Args:
+            features: List of feature maps from encoder stages.
+
+        Returns:
+            Aggregated feature vector (B, C').
+        """
+        batch_size = features[0].shape[0]
+        reduced_features = []
+        for patch_features in features[:4] + features[5:]:
+            hidden_features = self._pool_op(patch_features)
+            hidden_features_reduced = hidden_features.view(batch_size, -1)
+            reduced_features.append(hidden_features_reduced)
+        return torch.cat(reduced_features, dim=1)
+
+    def forward_head(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Casts last feature map into a single feature vector.
+
+        Args:
+            features: List of encoder feature maps.
+
+        Returns:
+            Aggregated feature vector (B, C').
+        """
+        last_feature_map = features[-1]
+        pooled_features = self._pool_op(last_feature_map)
+        return torch.flatten(pooled_features, 1)
+
+    def forward_embeddings(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Computes the final aggregated feature vector.
+
+        Args:
+            tensor: Input tensor of shape (B, C, T, H, W).
+
+        Returns:
+            Aggregated feature vector of shape (B, C').
+        """
+        intermediates = self.forward_features(tensor)
+        return self.forward_encoders(intermediates)
+
+    def forward_intermediates(
+        self, tensor: torch.Tensor
+    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        """Computes encoder features and their embeddings.
+
+        Args:
+            tensor: Input tensor of shape (B, C, T, H, W).
+
+        Returns:
+            Aggregated feature vector and list of intermediate features.
+        """
+        features = self.forward_features(tensor)
+        embeddings = self.forward_encoders(features)
+        return embeddings, features
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor | List[torch.Tensor]:
+        """Forward pass through the encoder.
+
+        If `self._out_indices` is None, it returns the aggregated feature vector.
+        Otherwise, it returns the intermediate feature maps up to the specified index.
+
+        Args:
+            tensor: Input tensor of shape (B, C, T, H, W).
+
+        Returns:
+            Aggregated feature vector or intermediate features.
+        """
+        if self._out_indices is None:
+            return self.forward_embeddings(tensor)
+
+        intermediates = self.forward_features(tensor)
+        return intermediates[-1 * self._out_indices :]

eva/vision/models/networks/backbones/radiology/voco.py

@@ -0,0 +1,75 @@
+"""VoCo Self-Supervised Encoders."""
+
+from typing_extensions import override
+
+from eva.core.models.wrappers import _utils
+from eva.vision.models.networks.backbones.radiology import swin_unetr
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+class _VoCo(swin_unetr.SwinUNETREncoder):
+    """Base class for the VoCo self-supervised encoders."""
+
+    _checkpoint: str
+    """Path to the model state dict."""
+
+    _md5: str | None = None
+    """State dict MD5 validation code."""
+
+    def __init__(self, feature_size: int, out_indices: int | None = None) -> None:
+        """Initializes the model.
+
+        Args:
+            feature_size: Size of the last feature map of SwinUNETR.
+            out_indices: The number of feature maps from intermediate blocks
+                to be returned. If set to 1, only the last feature map is returned.
+        """
+        super().__init__(
+            in_channels=1,
+            feature_size=feature_size,
+            spatial_dims=3,
+            out_indices=out_indices,
+        )
+
+        self._load_checkpoint()
+
+    def _load_checkpoint(self) -> None:
+        """Loads the model checkpoint."""
+        state_dict = _utils.load_state_dict_from_url(self._checkpoint, md5=self._md5)
+        self.load_state_dict(state_dict)
+
+
+@register_model("radiology/voco_b")
+class VoCoB(_VoCo):
+    """VoCo Self-supervised pre-trained B model."""
+
+    _checkpoint = "https://huggingface.co/Luffy503/VoCo/resolve/main/VoCo_B_SSL_head.pt"
+    _md5 = "f80c4da2f81d700bdae3df188f2057eb"
+
+    @override
+    def __init__(self, out_indices: int | None = None) -> None:
+        super().__init__(feature_size=48, out_indices=out_indices)
+
+
+@register_model("radiology/voco_l")
+class VoCoL(_VoCo):
+    """VoCo Self-supervised pre-trained L model."""
+
+    _checkpoint = "https://huggingface.co/Luffy503/VoCo/resolve/main/VoCo_L_SSL_head.pt"
+    _md5 = "795095d1d43ef3808ec4c41798310136"
+
+    @override
+    def __init__(self, out_indices: int | None = None) -> None:
+        super().__init__(feature_size=96, out_indices=out_indices)
+
+
+@register_model("radiology/voco_h")
+class VoCoH(_VoCo):
+    """VoCo Self-supervised pre-trained H model."""
+
+    _checkpoint = "https://huggingface.co/Luffy503/VoCo/resolve/main/VoCo_H_SSL_head.pt"
+    _md5 = "76f95a474736b60bf5b8aad94643744d"
+
+    @override
+    def __init__(self, out_indices: int | None = None) -> None:
+        super().__init__(feature_size=192, out_indices=out_indices)

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

@@ -1,5 +1,6 @@
 """Segmentation decoder heads API."""
 
+from eva.vision.models.networks.decoders.segmentation.base import Decoder
 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 (
@@ -7,13 +8,16 @@ from eva.vision.models.networks.decoders.segmentation.semantic import (
     ConvDecoderMS,
     ConvDecoderWithImage,
     SingleLinearDecoder,
+    SwinUNETRDecoder,
 )
 
 __all__ = [
+    "Decoder",
+    "Decoder2D",
     "ConvDecoder1x1",
     "ConvDecoderMS",
-    "SingleLinearDecoder",
     "ConvDecoderWithImage",
-    "Decoder2D",
     "LinearDecoder",
+    "SingleLinearDecoder",
+    "SwinUNETRDecoder",
 ]

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

@@ -7,6 +7,7 @@ from torch import nn
 from torch.nn import functional
 
 from eva.vision.models.networks.decoders.segmentation import base
+from eva.vision.models.networks.decoders.segmentation.typings import DecoderInputs
 
 
 class LinearDecoder(base.Decoder):
@@ -104,22 +105,16 @@ class LinearDecoder(base.Decoder):
         """
         return functional.interpolate(logits, image_size, mode="bilinear")
 
-    def forward(
-        self,
-        features: List[torch.Tensor],
-        image_size: Tuple[int, int],
-    ) -> torch.Tensor:
+    def forward(self, decoder_inputs: DecoderInputs) -> torch.Tensor:
         """Maps the patch embeddings to a segmentation mask of the image size.
 
         Args:
-            features: List of multi-level image features of shape (batch_size,
-                hidden_size, n_patches_height, n_patches_width).
-            image_size: The target image size (height, width).
+            decoder_inputs: Inputs required by the decoder.
 
         Returns:
             Tensor containing scores for all of the classes with shape
             (batch_size, n_classes, image_height, image_width).
         """
-        patch_embeddings = self._forward_features(features)
+        patch_embeddings = self._forward_features(decoder_inputs.features)
         logits = self._forward_head(patch_embeddings)
-        return self._cls_seg(logits, image_size)
+        return self._cls_seg(logits, decoder_inputs.image_size)

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

@@ -5,8 +5,15 @@ from eva.vision.models.networks.decoders.segmentation.semantic.common import (
     ConvDecoderMS,
     SingleLinearDecoder,
 )
+from eva.vision.models.networks.decoders.segmentation.semantic.swin_unetr import SwinUNETRDecoder
 from eva.vision.models.networks.decoders.segmentation.semantic.with_image import (
     ConvDecoderWithImage,
 )
 
-__all__ = ["ConvDecoder1x1", "ConvDecoderMS", "SingleLinearDecoder", "ConvDecoderWithImage"]
+__all__ = [
+    "ConvDecoder1x1",
+    "ConvDecoderMS",
+    "ConvDecoderWithImage",
+    "SingleLinearDecoder",
+    "SwinUNETRDecoder",
+]