🆕 Integrate Foundation Models Available VIA timm: UNI, Prov-GigaPath, 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,48 @@ 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://github.com/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":
+        # ProViT-GigaPath tile encoder: https://huggingface.co/prov-gigapath/prov-gigapath
+        # does not work with timm==0.9.8, needs timm==1.0.3: https://github.com/prov-gigapath/prov-gigapath/issues/2
+        assert (
+            timm.__version__ > "1.0.0"
+        ), "There is a bug in version `timm==0.9.8`. Tested to work from version `timm==1.0.3`"
+        feat_extract = timm.create_model(
+            "hf_hub:prov-gigapath/prov-gigapath", pretrained=True
+        )
+    else:
+        msg = f"Architecture {arch_name} not supported"
+        raise ValueError(msg)
+
+    return feat_extract
+
+
 class CNNModel(ModelABC):
     """Retrieve the model backbone and attach an extra FCN to perform classification.
 
@@ -268,3 +311,77 @@ def infer_batch(
             output = model(img_patches_device)
         # Output should be a single tensor or scalar
         return [output.cpu().numpy()]
+
+
+class TimmBackbone(ModelABC):
+    """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().__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)
+
+    @staticmethod
+    def infer_batch(
+        model: nn.Module,
+        batch_data: torch.Tensor,
+        *,
+        on_gpu: bool,
+    ) -> list[np.ndarray, ...]:
+        """Run inference on an input batch.
+
+        Contains logic for forward operation as well as i/o aggregation.
+
+        Args:
+            model (nn.Module):
+                PyTorch defined model.
+            batch_data (torch.Tensor):
+                A batch of data generated by
+                `torch.utils.data.DataLoader`.
+            on_gpu (bool):
+                Whether to run inference on a GPU.
+
+        """
+        img_patches_device = batch_data.to(select_device(on_gpu=on_gpu)).type(
+            torch.float32,
+        )  # to NCHW
+        img_patches_device = img_patches_device.permute(0, 3, 1, 2).contiguous()
+
+        # Inference mode
+        model.eval()
+        # Do not compute the gradient (not training)
+        with torch.inference_mode():
+            output = model(img_patches_device)
+        # Output should be a single tensor or scalar
+        return [output.cpu().numpy()]