docs: docstr fixes

Author: okunator

cellseg_models_pytorch/models/__init__.py

@@ -1,3 +1,5 @@
+import torch.nn as nn
+
 from .base._multitask_unet import MultiTaskUnet
 from .cellpose.cellpose import (
     CellPoseUnet,
@@ -26,6 +28,7 @@
     "omnipose_base": omnipose_base,
     "omnipose_plus": omnipose_plus,
     "hovernet_base": hovernet_base,
+    "hovernet_plus": hovernet_plus,
     "hovernet_small": hovernet_small,
     "hovernet_small_plus": hovernet_small_plus,
     "stardist_base": stardist_base,
@@ -34,43 +37,63 @@
 }
 
 
-def get_model(name: str, type: str, ntypes: int = None, ntissues: int = None):
-    """Get the corect model at hand given name and type."""
+def get_model(
+    name: str, type: str, ntypes: int = None, ntissues: int = None, **kwargs
+) -> nn.Module:
+    """Get the corect model at hand given name and type.
+
+    Parameters
+    ----------
+        name : str
+            Name of the model.
+        type : str
+            Type of the model. One of "base", "plus", "small", "small_plus".
+        ntypes : int
+            Number of cell types to segment.
+        ntissues : int
+            Number of tissue types to segment.
+        **kwargs : dict
+            Additional keyword arguments.
+
+    Returns
+    -------
+        nn.Module: The specified model.
+    """
     if name == "stardist":
         if type == "base":
             model = MODEL_LOOKUP["stardist_base_multiclass"](
-                n_rays=32, type_classes=ntypes
+                n_rays=32, type_classes=ntypes, **kwargs
             )
         elif type == "plus":
             model = MODEL_LOOKUP["stardist_plus"](
-                n_rays=32, type_classes=ntypes, sem_classes=ntissues
+                n_rays=32, type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
     elif name == "cellpose":
         if type == "base":
-            model = MODEL_LOOKUP["cellpose_base"](type_classes=ntypes)
+            model = MODEL_LOOKUP["cellpose_base"](type_classes=ntypes, **kwargs)
         elif type == "plus":
             model = MODEL_LOOKUP["cellpose_plus"](
-                type_classes=ntypes, sem_classes=ntissues
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
     elif name == "omnipose":
         if type == "base":
-            model = MODEL_LOOKUP["omnipose_base"](type_classes=ntypes)
+            model = MODEL_LOOKUP["omnipose_base"](type_classes=ntypes, **kwargs)
         elif type == "plus":
             model = MODEL_LOOKUP["omnipose_plus"](
-                type_classes=ntypes, sem_classes=ntissues
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
     elif name == "hovernet":
         if type == "base":
-            model = MODEL_LOOKUP["hovernet_base"](type_classes=ntypes)
+            model = MODEL_LOOKUP["hovernet_base"](type_classes=ntypes, **kwargs)
         elif type == "small":
-            model = MODEL_LOOKUP["hovernet_small"](type_classes=ntypes)
+            model = MODEL_LOOKUP["hovernet_small"](type_classes=ntypes, **kwargs)
         elif type == "plus":
             model = MODEL_LOOKUP["hovernet_plus"](
-                type_classes=ntypes, sem_classes=ntissues
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
         elif type == "small_plus":
             model = MODEL_LOOKUP["hovernet_small_plus"](
-                type_classes=ntypes, sem_classes=ntissues
+                type_classes=ntypes, sem_classes=ntissues, **kwargs
             )
     else:
         raise ValueError("Unknown model type or name.")

cellseg_models_pytorch/modules/transformers.py

@@ -243,7 +243,6 @@ def __init__(
                 f"Illegal args: {illegal_args}"
             )
 
-        # self.tr_blocks = nn.ModuleDict()
         self.tr_blocks = nn.ModuleList()
         self.layer_scales = nn.ModuleList()
         blocks = list(range(n_blocks))

cellseg_models_pytorch/training/lit/lightning_experiment.py

@@ -49,7 +49,7 @@ def __init__(
                 Use underscores to create joint loss functions. e.g. "dice_ce_tversky".
             branch_loss_params : Dict[str, Dict[str, Any]], optional
                 Params for the different losses at different branches. For example to
-                use LS or class weighting when computing the losses.
+                use label smoothing or class weighting when computing the losses.
                 E.g. {"inst": {"apply_ls": True}, "sem": {"edge_weight": False}}
             branch_metrics : Dict[str, List[str]], optional
                 A Dict of branch names mapped to a list of strings specifying a metrics.
@@ -65,7 +65,7 @@ def __init__(
             optim_params : Dict[str, Dict[str, Any]]
                 optim paramas like learning rates, weight decays etc for diff parts of
                 the network.
-                E.g. {"encoder": {"weight_decay: 0.1, "lr": 0.1}, "sem": {"lr": 0.01}}
+                E.g. {"encoder": {"weight_decay": 0.1, "lr": 0.1}, "sem": {"lr": 0.01}}
                 or {"learning_rate": 0.005, "weight_decay": 0.03}
             lookahead : bool, default=False
                 Flag whether the optimizer uses lookahead.

cellseg_models_pytorch/utils/file_manager.py

@@ -469,8 +469,8 @@ def write_gson(
             import geojson
         except ModuleNotFoundError:
             raise ModuleNotFoundError(
-                "To use the `FileHandler.mask2geojson`, pytorch-lightning is required. "
-                "Install with `pip install pytorch-lightning`"
+                "To use the `FileHandler.mask2geojson`, geojson is required. "
+                "Install with `pip install geojson`"
             )
 
         fname = Path(fname)