🆕 Integrate Foundation Models Available VIA timm: UNI, Virchow, Hibou, H-optimus-0

tiatoolbox/models/architecture/vanilla.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+import timm
 import torch
 import torchvision.models as torch_models
 from torch import nn
@@ -79,6 +80,45 @@ def _get_architecture(
     return model.features
 
 
+def _get_timm_architecture(
+    arch_name: str,
+    weights: str or WeightsEnum = "DEFAULT",
+    **kwargs: dict,
+) -> list[nn.Sequential, ...] | nn.Sequential:
+    """Get architecture and weights for pathology-specific timm models.
+
+    Args:
+        arch_name (str):
+            Architecture name.
+        weights:
+            path to pretrained weights
+        kwargs (dict):
+            Key-word arguments.
+
+    Returns:
+        A ready-to-use timm model.
+    """
+    if arch_name == "uni_v1":
+        # UNI tile encoder: https://huggingface.co/MahmoodLab/UNI
+        feat_extract = timm.create_model(
+            "hf-hub:MahmoodLab/UNI",
+            pretrained=True,
+            init_values=1e-5,
+            dynamic_img_size=True,
+        )
+    elif arch_name == "prov-gigapath" and timm.__version__ > "1.0.0":
+        # ProViT-GigaPath tile encoder: https://huggingface.co/prov-gigapath/prov-gigapath
+        # There is a bug in timm version 0.9.8 (https://github.com/prov-gigapath/prov-gigapath/issues/2). Version 1.0.3 has been tested successfully.
+        feat_extract = timm.create_model(
+            "hf_hub:prov-gigapath/prov-gigapath", pretrained=True
+        )
+    else:
+        msg = f"Architecture {arch_name} not supported. If you are loading timm models, only timm version > `1.0.3` are supported."
+        raise ValueError(msg)
+
+    return feat_extract
+
+
 class CNNModel(ModelABC):
     """Retrieve the model backbone and attach an extra FCN to perform classification.
 
@@ -268,3 +308,43 @@ def infer_batch(
             output = model(img_patches_device)
         # Output should be a single tensor or scalar
         return [output.cpu().numpy()]
+
+
+class TimmBackbone(CNNBackbone):
+    """Retrieve the pathology-specific tile encoder from timm.
+
+    This is a wrapper for pretrained models within timm.
+
+    Args:
+        backbone (str):
+            Model name. Currently, the tool supports following
+             model names and their default associated weights from timm.
+             - "uni_v1"
+             - "prov-gigapath"
+
+    Examples:
+        >>> # Creating UNI tile encoder
+        >>> model = TimmBackbone(backbone="uni_v1")
+        >>> model.eval()  # set to evaluation mode
+        >>> # dummy sample in NHWC form
+        >>> samples = torch.rand(4, 3, 224, 224)
+        >>> features = model(samples)
+        >>> features.shape  # feature vector
+        torch.Size([4, 1024])
+    """
+
+    def __init__(self: TimmBackbone, backbone: str) -> None:
+        """Initialize :class:`TimmBackbone`."""
+        super(CNNBackbone, self).__init__()
+        self.feat_extract = _get_timm_architecture(backbone)
+
+    def forward(self: TimmBackbone, imgs: torch.Tensor) -> torch.Tensor:
+        """Pass input data through the model.
+
+        Args:
+            imgs (torch.Tensor):
+                Model input.
+
+        """
+        feats = self.feat_extract(imgs)
+        return torch.flatten(feats, 1)