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

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

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

@@ -0,0 +1,62 @@
+"""Pathology FMs from MahmoodLab."""
+
+import os
+from pathlib import Path
+from typing import Tuple
+
+import huggingface_hub
+from loguru import logger
+from torch import nn
+
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/mahmood_uni")
+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.
+
+    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.
+        download_dir: Directory to download the model checkpoint.
+
+    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)
+        huggingface_hub.hf_hub_download(
+            "MahmoodLab/UNI",
+            filename="pytorch_model.bin",
+            local_dir=download_dir,
+            force_download=True,
+        )
+
+    return wrappers.TimmModel(
+        model_name="vit_large_patch16_224",
+        out_indices=out_indices,
+        model_kwargs={
+            "init_values": 1e-5,
+            "dynamic_img_size": dynamic_img_size,
+        },
+        checkpoint_path=checkpoint_path,
+    )

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

@@ -0,0 +1,22 @@
+"""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/owkin_phikon")
+def owkin_phikon(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the phikon pathology FM by owkin (https://huggingface.co/owkin/phikon).
+
+    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", out_indices=out_indices)

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

@@ -0,0 +1,47 @@
+"""Backbone Model Registry."""
+
+from typing import Any, Callable, Dict, List
+
+import torch.nn as nn
+
+
+class BackboneModelRegistry:
+    """A model registry for accessing backbone models by name."""
+
+    _registry: Dict[str, Callable[..., nn.Module]] = {}
+
+    @classmethod
+    def register(cls, name: str) -> Callable:
+        """Decorator to register a new model."""
+
+        def decorator(model_fn: Callable[..., nn.Module]) -> Callable[..., nn.Module]:
+            if name in cls._registry:
+                raise ValueError(f"Model {name} is already registered.")
+            cls._registry[name] = model_fn
+            return model_fn
+
+        return decorator
+
+    @classmethod
+    def get(cls, model_name: str) -> Callable[..., nn.Module]:
+        """Gets a model function from the registry."""
+        if model_name not in cls._registry:
+            raise ValueError(f"Model {model_name} not found in the registry.")
+        return cls._registry[model_name]
+
+    @classmethod
+    def load_model(cls, model_name: str, model_kwargs: Dict[str, Any] | None = None) -> nn.Module:
+        """Loads & initializes a model class from the registry."""
+        model_fn = cls.get(model_name)
+        return model_fn(**(model_kwargs or {}))
+
+    @classmethod
+    def list_models(cls) -> List[str]:
+        """List all models in the registry."""
+        register_models = [name for name in cls._registry.keys() if not name.startswith("timm")]
+        return register_models + ["timm/<model_name>"]
+
+
+def register_model(name: str) -> Callable:
+    """Simple decorator to register a model."""
+    return BackboneModelRegistry.register(name)

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

@@ -0,0 +1,5 @@
+"""timm backbones API."""
+
+from eva.vision.models.networks.backbones.timm.backbones import timm_model
+
+__all__ = ["timm_model"]

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

@@ -0,0 +1,54 @@
+"""timm backbones."""
+
+import functools
+from typing import Tuple
+
+import timm
+from loguru import logger
+from torch import nn
+
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones.registry import BackboneModelRegistry
+
+
+def timm_model(
+    model_name: str,
+    checkpoint_path: str | None = None,
+    pretrained: bool = False,
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+    **kwargs,
+) -> nn.Module:
+    """Initializes any ViT model from timm with weights from a specified checkpoint.
+
+    Args:
+        model_name: The name of the model to load.
+        checkpoint_path: The path to the checkpoint file.
+        pretrained: If set to `True`, load pretrained ImageNet-1k weights.
+        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.
+        **kwargs: Additional arguments to pass to the model
+
+    Returns:
+        The VIT model instance.
+    """
+    logger.info(
+        f"Loading timm model {model_name}"
+        + (f"using checkpoint {checkpoint_path}" if checkpoint_path else "")
+    )
+    return wrappers.TimmModel(
+        model_name=model_name,
+        checkpoint_path=checkpoint_path or "",
+        pretrained=pretrained,
+        out_indices=out_indices,
+        model_kwargs=kwargs,
+    )
+
+
+BackboneModelRegistry._registry.update(
+    {
+        f"timm/{model_name}": functools.partial(timm_model, model_name=model_name)
+        for model_name in timm.list_models()
+    }
+)

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

@@ -0,0 +1,8 @@
+"""Universal Vision Model Backbones API."""
+
+from eva.vision.models.networks.backbones.universal.vit import (
+    vit_small_patch16_224_dino,
+    vit_small_patch16_224_random,
+)
+
+__all__ = ["vit_small_patch16_224_dino", "vit_small_patch16_224_random"]

eva/vision/models/networks/backbones/universal/vit.py

@@ -0,0 +1,54 @@
+"""Vision Transformers base universal backbones."""
+
+from typing import Tuple
+
+import timm
+from torch import nn
+
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("universal/vit_small_patch16_224_random")
+def vit_small_patch16_224_random(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViTS-16 baseline model with random weights.
+
+    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 torch ViTS-16 based foundation model.
+    """
+    return timm.create_model(
+        model_name="vit_small_patch16_224.dino",
+        pretrained=False,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )
+
+
+@register_model("universal/vit_small_patch16_224_dino")
+def vit_small_patch16_224_dino(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViTS-16 baseline model pretrained w/ DINO.
+
+    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 torch ViTS-16 based foundation model.
+    """
+    return timm.create_model(
+        model_name="vit_small_patch16_224.dino",
+        pretrained=True,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )

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

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

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

@@ -0,0 +1,7 @@
+"""Semantic segmentation decoder base class."""
+
+from torch import nn
+
+
+class Decoder(nn.Module):
+    """Semantic segmentation decoder base class."""

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

@@ -0,0 +1,11 @@
+"""Segmentation decoder heads API."""
+
+from eva.vision.models.networks.decoders.segmentation.common import (
+    ConvDecoder1x1,
+    ConvDecoderMS,
+    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"]

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

@@ -0,0 +1,74 @@
+"""Common semantic segmentation decoders.
+
+This module contains implementations of different types of decoder models 
+used in semantic segmentation. These decoders convert the high-level features 
+output by an encoder into pixel-wise predictions for segmentation tasks.
+"""
+
+from torch import nn
+
+from eva.vision.models.networks.decoders.segmentation import conv2d, linear
+
+
+class ConvDecoder1x1(conv2d.ConvDecoder):
+    """A convolutional decoder with a single 1x1 convolutional layer."""
+
+    def __init__(self, in_features: int, num_classes: int) -> None:
+        """Initializes the decoder.
+
+        Args:
+            in_features: The hidden dimension size of the embeddings.
+            num_classes: Number of output classes as channels.
+        """
+        super().__init__(
+            layers=nn.Conv2d(
+                in_channels=in_features,
+                out_channels=num_classes,
+                kernel_size=(1, 1),
+            ),
+        )
+
+
+class ConvDecoderMS(conv2d.ConvDecoder):
+    """A multi-stage convolutional decoder with upsampling and convolutional layers.
+
+    This decoder applies a series of upsampling and convolutional layers to transform
+    the input features into output predictions with the desired spatial resolution.
+
+    This decoder is based on the `+ms` segmentation decoder from DINOv2
+    (https://arxiv.org/pdf/2304.07193)
+    """
+
+    def __init__(self, in_features: int, num_classes: int) -> None:
+        """Initializes the decoder.
+
+        Args:
+            in_features: The hidden dimension size of the embeddings.
+            num_classes: Number of output classes as channels.
+        """
+        super().__init__(
+            layers=nn.Sequential(
+                nn.Upsample(scale_factor=2),
+                nn.Conv2d(in_features, 64, kernel_size=(3, 3), padding=(1, 1)),
+                nn.Upsample(scale_factor=2),
+                nn.Conv2d(64, num_classes, kernel_size=(3, 3), padding=(1, 1)),
+            ),
+        )
+
+
+class SingleLinearDecoder(linear.LinearDecoder):
+    """A simple linear decoder with a single fully connected layer."""
+
+    def __init__(self, in_features: int, num_classes: int) -> None:
+        """Initializes the decoder.
+
+        Args:
+            in_features: The hidden dimension size of the embeddings.
+            num_classes: Number of output classes as channels.
+        """
+        super().__init__(
+            layers=nn.Linear(
+                in_features=in_features,
+                out_features=num_classes,
+            ),
+        )