Add GenBio-PathFM patch encoder (ViT-Giant, 4608-dim, JEDI training) (#9)

Author: saarthak-kapse

README.md

@@ -62,6 +62,7 @@ An API and command line interface (CLI) are provided to allow users to download
 * Virchow2: https://huggingface.co/paige-ai/Virchow2
 * H-optimus-0: https://huggingface.co/bioptimus/H-optimus-0
 * H-optimus-1: https://huggingface.co/bioptimus/H-optimus-1
+* GenBio-PathFM: https://huggingface.co/genbio-ai/genbio-pathfm
 * CONCH: https://huggingface.co/MahmoodLab/CONCH
 * TITAN/CONCHv1.5: https://huggingface.co/MahmoodLab/TITAN
 * Phikon: https://huggingface.co/owkin/phikon

src/thunder/benchmark.py

@@ -20,7 +20,7 @@ def benchmark(
 
     where options are:
         - dataset: *bach*, *bracs*, *break_his*, *ccrcc*, *crc*, *esca*, *mhist*, *ocelot*, *pannuke*, *patch_camelyon*, *segpath_epithelial*, *segpath_lymphocytes*, *tcga_crc_msi*, *tcga_tils*, *tcga_uniform*, *wilds*
-        - model: *hiboub*, *hiboul*, *hoptimus0*, *hoptimus1*, *midnight*, *phikon*, *phikon2*, *uni*, *uni2h*, *virchow*, *virchow2*, *conch*, *titan*, *keep*, *musk*, *plip*, *quiltnetb32*, *dinov2base*, *dinov2large*, *vitbasepatch16224in21k*, *vitlargepatch16224in21k*, *clipvitbasepatch32*, *clipvitlargepatch14*
+        - model: *hiboub*, *hiboul*, *hoptimus0*, *hoptimus1*, *genbio-pathfm*, *midnight*, *phikon*, *phikon2*, *uni*, *uni2h*, *virchow*, *virchow2*, *conch*, *titan*, *keep*, *musk*, *plip*, *quiltnetb32*, *dinov2base*, *dinov2large*, *vitbasepatch16224in21k*, *vitlargepatch16224in21k*, *clipvitbasepatch32*, *clipvitlargepatch14*
         - task: *adversarial_attack*, *alignment_scoring*, *image_retrieval*, *knn*, *linear_probing*, *pre_computing_embeddings*, *segmentation*, *simple_shot*, *transformation_invariance*, *zero_shot_vlm*
         - loading_mode: *online_loading*, *image_pre_loading*, *embedding_pre_loading*
 

src/thunder/config/pretrained_model/genbio-pathfm.yaml

@@ -0,0 +1,6 @@
+model_name: genbio-pathfm
+type: safetensors
+vlm: false
+emb_dim: 4608
+ckpt_path: ${oc.env:THUNDER_BASE_DATA_FOLDER}/pretrained_ckpts/genbiopathfm/model.pth
+hf_tag: hf-hub:genbio-ai/genbio-pathfm

src/thunder/models/download.py

@@ -38,6 +38,10 @@
         "bioptimus/H-optimus-1",
         "pytorch_model.bin",
     ),  # H-optimus 1 (https://huggingface.co/bioptimus/H-optimus-1)
+    "genbio-pathfm": (
+        "genbio-ai/genbio-pathfm",
+        "model.pth",
+    ),  # GenBio-PathFM (https://huggingface.co/genbio-ai/genbio-pathfm)
     "provgigapath": (
         "prov-gigapath/prov-gigapath",
         "pytorch_model.bin",
@@ -156,6 +160,7 @@ def download_models(models: Union[List[str], str]) -> None:
         * hoptimus0
         * h0mini
         * hoptimus1
+        * genbio-pathfm
         * provgigapath
         * conch
         * titan

src/thunder/models/pretrained_models.py

@@ -126,6 +126,8 @@ def get_model(model_cfg: dict, device: str):
             model, transform = get_titan(model_cfg.ckpt_path)
         elif model_cfg.model_name == "midnight":
             model, transform = get_midnight(model_cfg.ckpt_path)
+        elif model_cfg.model_name == "genbio-pathfm":
+            model, transform = get_genbio_pathfm(model_cfg.ckpt_path)
         else:
             model, transform, tokenizer = get_from_safetensors(
                 model_cfg.ckpt_path, use_fast="dinov3" in model_cfg.model_name
@@ -366,6 +368,8 @@ def extract_embedding(src, pretrained_model, task_type="linear_probing"):
                     emb = pretrained_model.trunk(src, return_all_tokens=True)[:, 1:]
                 elif model_cfg.model_name == "openmidnight":
                     emb = pretrained_model.get_intermediate_layers(src)[0]
+                elif model_cfg.model_name == "genbio-pathfm":
+                    emb = pretrained_model.forward_with_patches(src)[1]
                 else:
                     emb = pretrained_model.forward_features(src)[:, 1:]
 
@@ -394,6 +398,7 @@ def get_model_from_name(model_name: str, device: str):
         * hoptimus0
         * h0mini
         * hoptimus1
+        * genbio-pathfm
         * provgigapath
         * conch
         * titan
@@ -756,3 +761,36 @@ def get_titan(ckpt_path: str):
     model, transform = titan.return_conch()
 
     return model, transform
+
+
+def get_genbio_pathfm(ckpt_path: str):
+    """
+    Adapted from:
+    - https://github.com/genbio-ai/genbio-pathfm
+
+    :param ckpt_path: path to the stored checkpoint.
+    """
+    try:
+        from genbio_pathfm.model import GenBio_PathFM_Inference
+    except ImportError:
+        raise ImportError(
+            "In order to use GenBio-PathFM, please run the following: 'pip install git+https://github.com/genbio-ai/genbio-pathfm.git'"
+        )
+
+    from torchvision import transforms
+
+    # Model    
+    model = GenBio_PathFM_Inference(ckpt_path, device="cpu")
+
+    # Transform
+    transform = transforms.Compose([
+        transforms.Resize((224,224)),  
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=(0.697, 0.575, 0.728), 
+            std=(0.188, 0.240, 0.187)
+        ),
+    ])
+
+    return model, transform
+

src/thunder/utils/constants.py

@@ -9,6 +9,7 @@ class ModelConstants(Enum):
         "hoptimus0",
         "h0mini",
         "hoptimus1",
+        "genbio-pathfm",
         "kaiko_vits8",
         "kaiko_vits16",
         "kaiko_vitb8",