feat: add CellVIT-SAM model

Author: okunator

cellseg_models_pytorch/models/__init__.py

@@ -8,6 +8,13 @@
     omnipose_base,
     omnipose_plus,
 )
+from .cellvit.cellvit import (
+    CellVitSAM,
+    cellvit_sam_base,
+    cellvit_sam_plus,
+    cellvit_sam_small,
+    cellvit_sam_small_plus,
+)
 from .hovernet.hovernet import (
     HoverNet,
     hovernet_base,
@@ -34,6 +41,10 @@
     "stardist_base": stardist_base,
     "stardist_plus": stardist_plus,
     "stardist_base_multiclass": stardist_base_multiclass,
+    "cellvit_sam_base": cellvit_sam_base,
+    "cellvit_sam_plus": cellvit_sam_plus,
+    "cellvit_sam_small": cellvit_sam_small,
+    "cellvit_sam_small_plus": cellvit_sam_small_plus,
 }
 
 
@@ -62,11 +73,11 @@ def get_model(
     if name == "stardist":
         if type == "base":
             model = MODEL_LOOKUP["stardist_base_multiclass"](
-                n_rays=32, type_classes=ntypes, **kwargs
+                type_classes=ntypes, **kwargs
             )
         elif type == "plus":
             model = MODEL_LOOKUP["stardist_plus"](
-                n_rays=32, type_classes=ntypes, sem_classes=ntissues, **kwargs
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
     elif name == "cellpose":
         if type == "base":
@@ -95,6 +106,19 @@ def get_model(
             model = MODEL_LOOKUP["hovernet_small_plus"](
                 type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
+    elif name == "cellvit":
+        if type == "base":
+            model = MODEL_LOOKUP["cellvit_sam_base"](type_classes=ntypes, **kwargs)
+        elif type == "small":
+            model = MODEL_LOOKUP["cellvit_sam_small"](type_classes=ntypes, **kwargs)
+        elif type == "plus":
+            model = MODEL_LOOKUP["cellvit_sam_plus"](
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
+            )
+        elif type == "small_plus":
+            model = MODEL_LOOKUP["cellvit_sam_small_plus"](
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
+            )
     else:
         raise ValueError("Unknown model type or name.")
 
@@ -119,4 +143,9 @@ def get_model(
     "stardist_base_multiclass",
     "MODEL_LOOKUP",
     "get_model",
+    "CellVitSAM",
+    "cellvit_sam_base",
+    "cellvit_sam_plus",
+    "cellvit_sam_small",
+    "cellvit_sam_small_plus",
 ]

cellseg_models_pytorch/models/cellvit/_conf.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, Tuple
+
+__all__ = ["_create_cellvit_args"]
+
+
+def _create_cellvit_args(
+    layer_depth: Tuple[int, ...],
+    norm: str,
+    act: str,
+    conv: str,
+    att: str,
+    preact: bool,
+    preattend: bool,
+    short_skip: str,
+    use_style: bool,
+    merge_policy: str,
+    skip_params: Dict[str, Any],
+) -> Tuple[Dict[str, Any], ...]:
+    """Create the args to build CellVit-Unet decoders."""
+    skip_params = skip_params if skip_params is not None else {"k": None}
+
+    return tuple(
+        {
+            "layer_residual": False,
+            "upsampling": "conv_transpose",
+            "merge_policy": merge_policy,
+            "short_skips": (short_skip,),
+            "block_types": (("basic",) * ld,),
+            "kernel_sizes": ((3,) * ld,),
+            "expand_ratios": ((1.0,) * ld,),
+            "groups": ((1,) * ld,),
+            "biases": ((False,) * ld,),
+            "normalizations": ((norm,) * ld,),
+            "activations": ((act,) * ld,),
+            "convolutions": ((conv,) * ld,),
+            "attentions": ((att,) + (None,) * (ld - 1),),
+            "preactivates": ((preact,) * ld,),
+            "preattends": ((preattend,) * ld,),
+            "use_styles": ((use_style,) * (ld - 1) + (False,),),
+            "skip_params": {
+                "short_skips": (short_skip,),
+                "block_types": (("basic",),),
+                "convolutions": ((conv,),),
+                "normalizations": ((norm,),),
+                "activations": ((act,),),
+                **skip_params,
+            },
+        }
+        for ld in layer_depth
+    )