fix: fix inst2gdf empty masks

Author: okunator

cellseg_models_pytorch/models/base/_base_model_inst.py

@@ -56,13 +56,6 @@ def from_pretrained(
                     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
         )
@@ -77,6 +70,12 @@ def from_pretrained(
         )
 
         if weights_path.suffix == ".safetensors":
+            try:
+                from safetensors.torch import load_model
+            except ImportError:
+                raise ImportError(
+                    "Please install `safetensors` package to load .safetensors files."
+                )
             load_model(model_inst.model, weights_path, device.type)
         else:
             model_inst.model.load_state_dict(state_dict, strict=True)

cellseg_models_pytorch/models/cellpose/cellpose_unet.py

@@ -13,9 +13,7 @@
 __all__ = [
     "CellPoseUnet",
     "cellpose_nuclei",
-    "cellpose_panoptic",
     "omnipose_nuclei",
-    "omnipose_panoptic",
 ]
 
 
@@ -249,34 +247,6 @@ def cellpose_nuclei(n_nuc_classes: int, **kwargs) -> nn.Module:
     return cellpose_unet
 
 
-def cellpose_panoptic(n_nuc_classes: int, n_tissue_classes: int, **kwargs) -> nn.Module:
-    """Initialize Cellpose for panoptic segmentation.
-
-    Cellpose:
-    - https://www.nature.com/articles/s41592-020-01018-x
-
-    Parameters
-        n_nuc_classes (int):
-            Number of nuclei type classes.
-        n_tissue_classes (int):
-            Number of tissue type classes.
-        **kwargs:
-            Arbitrary key word args for the CellPoseUnet class.
-
-    Returns:
-        nn.Module: The initialized Cellpose+ U-net model.
-    """
-    cellpose_unet = CellPoseUnet(
-        decoders=("type", "tissue"),
-        heads={
-            "type": {"nuc_cellpose": 2, "nuc_type": n_nuc_classes},
-            "tissue": {"tissue_type": n_tissue_classes},
-        },
-        **kwargs,
-    )
-    return cellpose_unet
-
-
 def omnipose_nuclei(n_nuc_classes: int, **kwargs) -> nn.Module:
     """Create the baseline Omnipose U-net for nuclei segmentation.
 
@@ -300,33 +270,3 @@ def omnipose_nuclei(n_nuc_classes: int, **kwargs) -> nn.Module:
     cellpose_unet.aux_key = "omnipose"
 
     return cellpose_unet
-
-
-def omnipose_panoptic(n_nuc_classes: int, n_tissue_classes: int, **kwargs) -> nn.Module:
-    """Create the Omnipose U-net with nuclei- and tissue segmentation decoders.
-
-    Omnipose:
-    - https://www.biorxiv.org/content/10.1101/2021.11.03.467199v2
-
-    Parameters:
-        n_nuc_classes (int):
-            Number of nuclei type classes.
-        n_tissue_classes (int):
-            Number of tissue type classes.
-        **kwargs:
-            Arbitrary key word args for the CellPoseUnet class.
-
-    Returns:
-        nn.Module: The initialized Cellpose+ U-net model.
-    """
-    cellpose_unet = CellPoseUnet(
-        decoders=("type", "tissue"),
-        heads={
-            "type": {"nuc_omnipose": 2, "nuc_type": n_nuc_classes},
-            "tissue": {"tissue_type": n_tissue_classes},
-        },
-        **kwargs,
-    )
-    cellpose_unet.aux_key = "omnipose"
-
-    return cellpose_unet

cellseg_models_pytorch/models/cellvit/cellvit_unet.py

@@ -13,7 +13,6 @@
 __all__ = [
     "CellVitSamUnet",
     "cellvit_nuclei",
-    "cellvit_panoptic",
 ]
 
 
@@ -263,39 +262,3 @@ def cellvit_nuclei(enc_name: str, n_nuc_classes: int, **kwargs) -> nn.Module:
     )
 
     return cellvit_sam
-
-
-def cellvit_panoptic(
-    enc_name: str, n_nuc_classes: int, n_tissue_classes: int, **kwargs
-) -> nn.Module:
-    """Initialaize CellVit for panoptic segmentation.
-
-    CellVit:
-        - https://arxiv.org/abs/2306.15350
-
-    Parameters:
-        enc_name (str):
-            Name of the encoder. One of: "samvit_base_patch16", "samvit_base_patch16_224",
-            "samvit_huge_patch16", "samvit_large_patch16"
-        n_nuc_classes (int):
-            Number of nuclei type classes.
-        n_tissue_classes (int):
-            Number of tissue type classes.
-        **kwargs:
-            Arbitrary key word args for the CellVitSAM class.
-
-    Returns:
-        nn.Module: The initialized CellVitSAM+ model.
-    """
-    cellvit_sam = CellVitSamUnet(
-        enc_name=enc_name,
-        decoders=("hovernet", "type", "tissue"),
-        heads={
-            "hovernet": {"nuc_hovernet": 2},
-            "type": {"nuc_type": n_nuc_classes},
-            "tissue": {"tissue_type": n_tissue_classes},
-        },
-        **kwargs,
-    )
-
-    return cellvit_sam

cellseg_models_pytorch/models/cppnet/cppnet_unet.py

@@ -16,7 +16,6 @@
 __all__ = [
     "CPPNetUnet",
     "cppnet_nuclei",
-    "cppnet_panoptic",
 ]
 
 
@@ -338,38 +337,3 @@ def cppnet_nuclei(n_rays: int, n_nuc_classes: int, **kwargs) -> nn.Module:
     )
 
     return cppnet
-
-
-def cppnet_panoptic(
-    n_rays: int, n_nuc_classes: int, n_tissue_classes: int, **kwargs
-) -> nn.Module:
-    """Initialaize CPP-Net for panoptic segmentation.
-
-    CPP-Net:
-        - https://arxiv.org/abs/2102.06867
-
-    Parameters:
-        n_rays (int):
-            Number of rays predicted per each object
-        n_nuc_classes (int):
-            Number of nuclei type classes.
-        n_tissue_classes (int):
-            Number of tissue type classes.
-        **kwargs:
-            Arbitrary key word args for the CPPNet class.
-
-    Returns:
-        nn.Module: The initialized CPP-Net model.
-    """
-    cppnet = CPPNetUnet(
-        decoders=("stardist", "type", "tissue"),
-        heads={
-            "stardist": {"nuc_stardist": n_rays, "nuc_binary": 1},
-            "type": {"nuc_type": n_nuc_classes},
-            "tissue": {"tissue_type": n_tissue_classes},
-        },
-        n_rays=n_rays,
-        **kwargs,
-    )
-
-    return cppnet

cellseg_models_pytorch/models/hovernet/hovernet_unet.py

@@ -13,7 +13,6 @@
 __all__ = [
     "HoverNetUnet",
     "hovernet_nuclei",
-    "hovernet_panoptic",
 ]
 
 
@@ -244,37 +243,3 @@ def hovernet_nuclei(n_nuc_classes: int, **kwargs) -> nn.Module:
     )
 
     return hovernet
-
-
-def hovernet_panoptic(
-    n_nuc_classes: int,
-    n_tissue_classes: int,
-    **kwargs,
-) -> nn.Module:
-    """Initialaize HoverNet+ for panoptic segmentation.
-
-    HoVer-Net:
-        - https://www.sciencedirect.com/science/article/pii/S1361841519301045?via%3Dihub
-
-    Parameters:
-        n_nuc_classes (int):
-            Number of nuclei type classes.
-        n_tissue_classes (int):
-            Number of tissue type classes.
-        **kwargs:
-            Arbitrary key word args for the HoverNet class.
-
-    Returns:
-        nn.Module: The initialized HoVer-Net+ model.
-    """
-    hovernet = HoverNetUnet(
-        decoders=("hovernet", "type", "tissue"),
-        heads={
-            "hovernet": {"nuc_hovernet": 2},
-            "type": {"nuc_type": n_nuc_classes},
-            "tissue": {"tissue_type": n_tissue_classes},
-        },
-        **kwargs,
-    )
-
-    return hovernet

cellseg_models_pytorch/models/stardist/stardist_unet.py

@@ -10,7 +10,7 @@
 from cellseg_models_pytorch.encoders import Encoder
 from cellseg_models_pytorch.models.stardist._conf import _create_stardist_args
 
-__all__ = ["StarDistUnet", "stardist_nuclei", "stardist_panoptic"]
+__all__ = ["StarDistUnet", "stardist_nuclei"]
 
 
 class StarDistUnet(nn.ModuleDict):
@@ -247,40 +247,3 @@ def stardist_nuclei(n_rays: int, n_nuc_classes: int, **kwargs) -> nn.Module:
     )
 
     return stardist_unet
-
-
-def stardist_panoptic(
-    n_rays: int, n_nuc_classes: int, n_tissue_classes: int, **kwargs
-) -> nn.Module:
-    """Initialize Stardist model for panoptic segmentation.
-
-    Stardist:
-    - https://arxiv.org/abs/1806.03535
-
-    Parameters:
-        n_rays (int):
-            Number of rays predicted per each object
-        n_nuc_classes (int):
-            Number of nuclei type classes.
-        n_tissue_classes (int):
-            Number of tissue type classes.
-        **kwargs:
-            Arbitrary key word args for the StarDistUnet class.
-
-    Returns:
-        nn.Module: The initialized Panoptic Stardist model.
-    """
-    stardist_unet = StarDistUnet(
-        decoders=("stardist", "tissue"),
-        heads={
-            "stardist": {
-                "nuc_stardist": n_rays,
-                "nuc_binary": 1,
-                "nuc_type": n_nuc_classes,
-            },
-            "tissue": {"tissue_type": n_tissue_classes},
-        },
-        **kwargs,
-    )
-
-    return stardist_unet