fix: fix base model from_pretrained

Author: okunator

cellseg_models_pytorch/models/base/__init__.py

@@ -0,0 +1,34 @@
+__all__ = ["PRETRAINED"]
+
+PRETRAINED = {
+    "hovernet": {
+        "hgsc_v1_efficientnet_b5": {
+            "repo_id": "csmp-hub/hovernet-histo-hgsc-nuc-v1",
+            "filename": "hovernet_hgsc_v1_efficientnet_b5.safetensors",
+        },
+    },
+    "cellpose": {
+        "hgsc_v1_efficientnet_b5": {
+            "repo_id": "csmp-hub/cellpose-histo-hgsc-nuc-v1",
+            "filename": "cellpose_hgsc_v1_efficientnet_b5.safetensors",
+        },
+    },
+    "cellvit": {
+        "hgsc_v1_efficientnet_b5": {
+            "repo_id": "csmp-hub/cellvit-histo-hgsc-nuc-v1",
+            "filename": "cellvit_hgsc_v1_efficientnet_b5.safetensors",
+        },
+    },
+    "stardist": {
+        "hgsc_v1_efficientnet_b5": {
+            "repo_id": "csmp-hub/stardist-histo-hgsc-nuc-v1",
+            "filename": "stardist_hgsc_v1_efficientnet_b5.safetensors",
+        },
+    },
+    "cppnet": {
+        "hgsc_v1_efficientnet_b5": {
+            "repo_id": "csmp-hub/cppnet-histo-hgsc-nuc-v1",
+            "filename": "cppnet_hgsc_v1_efficientnet_b5.safetensors",
+        },
+    },
+}

cellseg_models_pytorch/models/base/_base_model_inst.py

@@ -3,12 +3,14 @@
 
 import numpy as np
 import torch
+from huggingface_hub import hf_hub_download
 from PIL.Image import Image
 
 from cellseg_models_pytorch.decoders.multitask_decoder import (
     SoftInstanceOutput,
     SoftSemanticOutput,
 )
+from cellseg_models_pytorch.models.base import PRETRAINED
 
 __all__ = ["BaseModelInst"]
 
@@ -24,24 +26,60 @@ def set_inference_mode(self) -> None:
     @classmethod
     def from_pretrained(
         cls,
-        weights_path: Union[str, Path],
-        n_nuc_classes: int,
-        enc_name: str = "efficientnet_b5",
-        enc_freeze: bool = False,
+        weights: Union[str, Path],
         device: torch.device = torch.device("cuda"),
         model_kwargs: Dict[str, Any] = {},
-    ) -> None:
-        """Load the model from pretrained weights."""
+    ) -> "BaseModelInst":
+        """Load the model from pretrained weights.
+
+        Parameters:
+            model_name (str):
+                Name of the pretrained model.
+            device (torch.device, default=torch.device("cuda")):
+                Device to run the model on. Default is "cuda".
+            model_kwargs (Dict[str, Any], default={}):
+                Additional arguments for the model.
+        """
+        weights_path = Path(weights)
+        if not weights_path.is_file():
+            if weights_path.as_posix() in PRETRAINED[cls.model_name].keys():
+                weights_path = Path(
+                    hf_hub_download(
+                        repo_id=PRETRAINED[cls.model_name][weights]["repo_id"],
+                        filename=PRETRAINED[cls.model_name][weights]["filename"],
+                    )
+                )
+
+            else:
+                raise ValueError(
+                    "Please provide a valid path. or a pre-trained model downloaded from the"
+                    f" csmp-hub. One of {list(PRETRAINED[cls.model_name].keys())}."
+                )
+
+        try:
+            from safetensors.torch import load_model
+        except ImportError:
+            raise ImportError(
+                "Please install `safetensors` package to load .safetensors files."
+            )
+
+        enc_name, n_nuc_classes, state_dict = cls._get_state_dict(
+            weights_path, device=device
+        )
+
         model_inst = cls(
             n_nuc_classes=n_nuc_classes,
             enc_name=enc_name,
             enc_pretrain=False,
-            enc_freeze=enc_freeze,
+            enc_freeze=False,
             device=device,
             model_kwargs=model_kwargs,
         )
-        state_dict = torch.load(weights_path, map_location=device)
-        model_inst.model.load_state_dict(state_dict, strict=True)
+
+        if weights_path.suffix == ".safetensors":
+            load_model(model_inst.model, weights_path, device.type)
+        else:
+            model_inst.model.load_state_dict(state_dict, strict=True)
 
         return model_inst
 
@@ -174,3 +212,34 @@ def post_process(
             )
 
         return x
+
+    @staticmethod
+    def _get_state_dict(
+        weights_path: Union[str, Path], device: torch.device = torch.device("cuda")
+    ) -> None:
+        """Load the model from pretrained weights."""
+        weights_path = Path(weights_path)
+        if not weights_path.exists():
+            raise ValueError(f"Model weights not found at {weights_path}")
+        if weights_path.suffix == ".safetensors":
+            try:
+                from safetensors.torch import load_file
+            except ImportError:
+                raise ImportError(
+                    "Please install `safetensors` package to load .safetensors files."
+                )
+            state_dict = load_file(weights_path, device=device.type)
+        else:
+            state_dict = torch.load(weights_path, map_location=device)
+
+        # infer encoder name and number of classes from state_dict
+        enc_keys = [key for key in state_dict.keys() if "encoder." in key]
+        enc_name = enc_keys[0].split(".")[0] if enc_keys else None
+        nuc_type_head_key = next(
+            key
+            for key in state_dict.keys()
+            if "nuc_type_head.head" in key and "weight" in key
+        )
+        n_nuc_classes = state_dict[nuc_type_head_key].shape[0]
+
+        return enc_name, n_nuc_classes, state_dict

cellseg_models_pytorch/models/cellpose/cellpose.py

@@ -11,6 +11,8 @@
 
 
 class CellPose(BaseModelInst):
+    model_name = "cellpose"
+
     def __init__(
         self,
         n_nuc_classes: int,

cellseg_models_pytorch/models/cellpose/cellpose_unet.py

@@ -19,7 +19,7 @@
 ]
 
 
-class CellPoseUnet(nn.Module):
+class CellPoseUnet(nn.ModuleDict):
     def __init__(
         self,
         decoders: Tuple[str, ...],
@@ -127,6 +127,7 @@ def __init__(
         super().__init__()
         self.inst_key = inst_key
         self.aux_key = "cellpose"
+        self.enc_name = enc_name
 
         if enc_out_indices is None:
             enc_out_indices = tuple(range(depth))
@@ -155,18 +156,21 @@ def __init__(
         )
 
         # set encoder
-        self.encoder = Encoder(
-            timm_encoder_name=enc_name,
-            timm_encoder_out_indices=enc_out_indices,
-            timm_encoder_pretrained=enc_pretrain,
-            timm_extra_kwargs=encoder_kws,
+        self.add_module(
+            self.enc_name,
+            Encoder(
+                timm_encoder_name=enc_name,
+                timm_encoder_out_indices=enc_out_indices,
+                timm_encoder_pretrained=enc_pretrain,
+                timm_extra_kwargs=encoder_kws,
+            ),
         )
 
         self.decoder = MultiTaskDecoder(
             decoders=decoders,
             heads=heads,
             out_channels=out_channels,
-            enc_feature_info=self.encoder.feature_info,
+            enc_feature_info=self[self.enc_name].feature_info,
             n_layers=n_layers,
             n_blocks=n_blocks,
             stage_kws=stage_kws,
@@ -181,7 +185,7 @@ def __init__(
 
         # freeze encoder if specified
         if enc_freeze:
-            self.encoder.freeze_encoder()
+            self[self.enc_name].freeze_encoder()
 
         self.name = f"CellPoseUnet-{enc_name}"
 
@@ -204,7 +208,7 @@ def forward(self, x: torch.Tensor, return_pred_only: bool = True) -> Dict[str, A
                     - "dec_feats": Dict[str, List[torch.Tensor]].
                     - "enc_out": torch.Tensor.
         """
-        enc_output, feats = self.encoder(x)
+        enc_output, feats = self[self.enc_name](x)
         dec_out: DecoderSoftOutput = self.decoder(feats, x)
 
         res = {

cellseg_models_pytorch/models/cellvit/cellvit.py

@@ -11,6 +11,8 @@
 
 
 class CellVit(BaseModelInst):
+    model_name = "cellvit"
+
     def __init__(
         self,
         n_nuc_classes: int,

cellseg_models_pytorch/models/cellvit/cellvit_unet.py

@@ -17,7 +17,7 @@
 ]
 
 
-class CellVitSamUnet(nn.Module):
+class CellVitSamUnet(nn.ModuleDict):
     def __init__(
         self,
         decoders: Tuple[str, ...],
@@ -119,9 +119,22 @@ def __init__(
                 segmentation post-processing pipeline as the binary segmentation result.
         """
         super().__init__()
+        allowed = (
+            "samvit_base_patch16",
+            "samvit_base_patch16_224",
+            "samvit_huge_patch16",
+            "samvit_large_patch16",
+        )
+        if enc_name not in allowed:
+            raise ValueError(
+                f"Wrong encoder name. Got: {enc_name}. "
+                f"Allowed encoder for CellVit: {allowed}"
+            )
+
         self.inst_key = inst_key
         self.aux_key = "hovernet"
         self.depth = len(layer_depths)
+        self.enc_name = enc_name
 
         if enc_out_indices is None:
             enc_out_indices = tuple(range(self.depth))
@@ -151,31 +164,22 @@ def __init__(
             skip_kws,
         )
 
-        allowed = (
-            "samvit_base_patch16",
-            "samvit_base_patch16_224",
-            "samvit_huge_patch16",
-            "samvit_large_patch16",
-        )
-        if enc_name not in allowed:
-            raise ValueError(
-                f"Wrong encoder name. Got: {enc_name}. "
-                f"Allowed encoder for CellVit: {allowed}"
-            )
-
         # set encoders
-        self.encoder = Encoder(
-            timm_encoder_name=enc_name,
-            timm_encoder_out_indices=enc_out_indices,
-            timm_encoder_pretrained=enc_pretrain,
-            timm_extra_kwargs=encoder_kws,
+        self.add_module(
+            self.enc_name,
+            Encoder(
+                timm_encoder_name=enc_name,
+                timm_encoder_out_indices=enc_out_indices,
+                timm_encoder_pretrained=enc_pretrain,
+                timm_extra_kwargs=encoder_kws,
+            ),
         )
 
         self.decoder = MultiTaskDecoder(
             decoders=decoders,
             heads=heads,
             out_channels=out_channels,
-            enc_feature_info=self.encoder.feature_info,
+            enc_feature_info=self[self.enc_name].feature_info,
             n_layers=n_layers,
             n_blocks=n_blocks,
             stage_kws=stage_kws,
@@ -190,7 +194,7 @@ def __init__(
 
         # freeze encoder if specified
         if enc_freeze:
-            self.encoder.freeze_encoder()
+            self[self.enc_name].freeze_encoder()
 
         self.name = f"CellVit-{enc_name}"
 
@@ -213,7 +217,7 @@ def forward(self, x: torch.Tensor, return_pred_only: bool = True) -> Dict[str, A
                     - "dec_feats": Dict[str, List[torch.Tensor]].
                     - "enc_out": torch.Tensor.
         """
-        enc_output, feats = self.encoder(x)
+        enc_output, feats = self[self.enc_name](x)
         dec_out: DecoderSoftOutput = self.decoder(feats, x)
 
         res = {

cellseg_models_pytorch/models/cppnet/cppnet.py

@@ -11,6 +11,8 @@
 
 
 class CPPNet(BaseModelInst):
+    model_name = "cppnet"
+
     def __init__(
         self,
         n_nuc_classes: int,

cellseg_models_pytorch/models/cppnet/cppnet_unet.py

@@ -87,7 +87,7 @@ def forward(
         return ray_refined, confidence_refined
 
 
-class CPPNetUnet(nn.Module):
+class CPPNetUnet(nn.ModuleDict):
     def __init__(
         self,
         decoders: Tuple[str, ...],
@@ -194,6 +194,7 @@ def __init__(
         self.inst_key = inst_key
         self.aux_key = "stardist"
         self.n_rays = n_rays
+        self.enc_name = enc_name
 
         if enc_out_indices is None:
             enc_out_indices = tuple(range(depth))
@@ -221,18 +222,21 @@ def __init__(
         )
 
         # set encoder
-        self.encoder = Encoder(
-            timm_encoder_name=enc_name,
-            timm_encoder_out_indices=enc_out_indices,
-            timm_encoder_pretrained=enc_pretrain,
-            timm_extra_kwargs=encoder_kws,
+        self.add_module(
+            self.enc_name,
+            Encoder(
+                timm_encoder_name=enc_name,
+                timm_encoder_out_indices=enc_out_indices,
+                timm_encoder_pretrained=enc_pretrain,
+                timm_extra_kwargs=encoder_kws,
+            ),
         )
 
         self.decoder = MultiTaskDecoder(
             decoders=decoders,
             heads=heads,
             out_channels=out_channels,
-            enc_feature_info=self.encoder.feature_info,
+            enc_feature_info=self[self.enc_name].feature_info,
             n_layers=n_layers,
             n_blocks=n_blocks,
             stage_kws=stage_kws,
@@ -255,7 +259,7 @@ def __init__(
 
         # freeze encoder if specified
         if enc_freeze:
-            self.encoder.freeze_encoder()
+            self[self.enc_name].freeze_encoder()
 
     def forward(self, x: torch.Tensor, return_pred_only: bool = True) -> Dict[str, Any]:
         """Forward pass of Cellpose U-net.
@@ -276,7 +280,7 @@ def forward(self, x: torch.Tensor, return_pred_only: bool = True) -> Dict[str, A
                     - "dec_feats": Dict[str, List[torch.Tensor]].
                     - "enc_out": torch.Tensor.
         """
-        enc_output, feats = self.encoder(x)
+        enc_output, feats = self[self.enc_name](x)
         dec_out: DecoderSoftOutput = self.decoder(feats, x)
         if dec_out.nuc_map is not None:
             dec_name = dec_out.nuc_map.parents["aux_map"][0]