"""Models from Paige AI company."""
from __future__ import annotations
from typing import Any
import numpy as np
import timm
import torch
from torchvision import transforms
from plismbench.models.extractor import Extractor
from plismbench.models.utils import DEFAULT_DEVICE, prepare_module
[docs]
class Virchow(Extractor):
"""Virchow model developped by Paige AI available on Hugging-Face (1).
.. note::
(1) https://huggingface.co/paige-ai/Virchow
Parameters
----------
device: int | list[int] | None = DEFAULT_DEVICE,
Compute resources to use.
If None, will use all available GPUs.
If -1, extraction will run on CPU.
mixed_precision: bool = True
Whether to use mixed_precision.
"""
def __init__(
self,
device: int | list[int] | None = DEFAULT_DEVICE,
mixed_precision: bool = False,
):
super().__init__()
self.output_dim = 2560
self.mixed_precision = mixed_precision
timm_kwargs: dict[str, Any] = {
"mlp_layer": timm.layers.SwiGLUPacked,
"act_layer": torch.nn.SiLU,
}
feature_extractor = timm.create_model(
"hf-hub:paige-ai/Virchow", pretrained=True, **timm_kwargs
)
self.feature_extractor, self.device = prepare_module(
feature_extractor,
device,
self.mixed_precision,
)
if self.device is None:
self.feature_extractor = self.feature_extractor.module
@property # type: ignore
def transform(self) -> transforms.Compose:
"""Transform method to apply element wise."""
return transforms.Compose(
[
transforms.ToTensor(), # swap axes and normalize
transforms.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
),
]
)
def __call__(self, images: torch.Tensor) -> np.ndarray:
"""Compute and return features.
Parameters
----------
images: torch.Tensor
Input of size (n_tiles, n_channels, dim_x, dim_y).
Returns
-------
torch.Tensor: Tensor of size (n_tiles, features_dim).
"""
last_hidden_state = self.feature_extractor(images.to(self.device))
class_token = last_hidden_state[:, 0]
patch_tokens = last_hidden_state[:, 1:]
features = torch.cat([class_token, patch_tokens.mean(1)], dim=-1)
return features.cpu().numpy()
[docs]
class Virchow2(Extractor):
"""Virchow2 model developped by Paige AI available on Hugging-Face (1).
You will need to be granted access to be able to use this model.
.. note::
(1) https://huggingface.co/paige-ai/Virchow2
Parameters
----------
device: int | list[int] | None = DEFAULT_DEVICE,
Compute resources to use.
If None, will use all available GPUs.
If -1, extraction will run on CPU.
mixed_precision: bool = True
Whether to use mixed_precision.
"""
def __init__(
self,
device: int | list[int] | None = DEFAULT_DEVICE,
mixed_precision: bool = False,
):
super().__init__()
self.output_dim = 2560
self.mixed_precision = mixed_precision
timm_kwargs: dict[str, Any] = {
"mlp_layer": timm.layers.SwiGLUPacked,
"act_layer": torch.nn.SiLU,
}
feature_extractor = timm.create_model(
"hf-hub:paige-ai/Virchow2", pretrained=True, **timm_kwargs
)
self.feature_extractor, self.device = prepare_module(
feature_extractor,
device,
self.mixed_precision,
)
if self.device is None:
self.feature_extractor = self.feature_extractor.module
@property # type: ignore
def transform(self) -> transforms.Compose:
"""Transform method to apply element wise."""
return transforms.Compose(
[
transforms.ToTensor(), # swap axes and normalize
transforms.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
),
]
)
def __call__(self, images: torch.Tensor) -> np.ndarray:
"""Compute and return features.
Parameters
----------
images: torch.Tensor
Input of size (n_tiles, n_channels, dim_x, dim_y).
Returns
-------
torch.Tensor: Tensor of size (n_tiles, features_dim).
"""
last_hidden_state = self.feature_extractor(images.to(self.device))
class_token = last_hidden_state[:, 0]
patch_tokens = last_hidden_state[:, 5:]
features = torch.cat([class_token, patch_tokens.mean(1)], dim=-1)
return features.cpu().numpy()
