fix: minor bug fixes

Author: okunator

cellseg_models_pytorch/decoders/long_skips/cross_attn_skip.py

@@ -17,12 +17,12 @@ def __init__(
         skip_channels: Tuple[int, ...] = None,
         num_heads: int = 8,
         head_dim: int = 64,
-        n_blocks: int = 1,
-        block_types: Tuple[str, ...] = ("exact",),
-        computation_types: Tuple[str, ...] = ("basic",),
-        dropouts: Tuple[float, ...] = (0.0,),
-        biases: Tuple[bool, ...] = (False,),
-        layer_scales: Tuple[bool, ...] = (False,),
+        n_blocks: int = 2,
+        block_types: Tuple[str, ...] = ("exact", "exact"),
+        computation_types: Tuple[str, ...] = ("basic", "basic"),
+        dropouts: Tuple[float, ...] = (0.0, 0.0),
+        biases: Tuple[bool, ...] = (False, False),
+        layer_scales: Tuple[bool, ...] = (False, False),
         activation: str = "star_relu",
         mlp_ratio: int = 2,
         slice_size: int = 4,

cellseg_models_pytorch/inference/_base_inferer.py

@@ -14,7 +14,6 @@
 
 from ..utils import FileHandler, tensor_to_ndarray
 from .folder_dataset_infer import FolderDatasetInfer
-from .hdf5_dataset_infer import HDF5DatasetInfer
 from .post_processor import PostProcessor
 from .predictor import Predictor
 
@@ -126,6 +125,8 @@ def __init__(
                     " `save_dir` argument."
                 )
         elif self.path.is_file() and self.path.suffix in (".h5", ".hdf5"):
+            from .hdf5_dataset_infer import HDF5DatasetInfer
+
             ds = HDF5DatasetInfer(self.path, n_images=n_images)
         else:
             raise ValueError(
@@ -270,7 +271,9 @@ def _strip_state_dict(self, ckpt: Dict) -> OrderedDict:
         state_dict = OrderedDict()
         for k, w in ckpt["state_dict"].items():
             if "num_batches_track" not in k:
-                new_key = k.strip("model")[1:]
+                # new_key = k.strip("model")[1:]
+                spl = ["".join(kk) for kk in k.split(".")]
+                new_key = ".".join(spl[1:])
                 state_dict[new_key] = w
         ckpt["state_dict"] = state_dict
 

cellseg_models_pytorch/models/base/_timm_encoder.py

@@ -14,9 +14,25 @@ def __init__(
         pretrained: bool = True,
         in_channels: int = 3,
         depth: int = 5,
+        out_indices: List[int] = None,
         **kwargs
     ) -> None:
-        """Import any encoder from timm package."""
+        """Import any encoder from timm package.
+
+        Parameters
+        ----------
+            name : str
+                Name of the encoder.
+            pretrained : bool, optional
+                If True, load pretrained weights, by default True.
+            in_channels : int, optional
+                Number of input channels, by default 3.
+            depth : int, optional
+                Number of output features, by default 5.
+            out_indices : List[int], optional
+                Indices of the output features, by default None. If None, all the
+                features are returned.
+        """
         super().__init__()
 
         kwargs = dict(
@@ -28,9 +44,13 @@ def __init__(
 
         self.model = timm.create_model(name, **kwargs)
 
+        self.out_indices = out_indices
         self.in_channels = in_channels
         self.out_channels = tuple(self.model.feature_info.channels()[::-1])
 
+        if self.out_indices is not None:
+            self.out_channels = tuple(self.out_channels[i] for i in self.out_indices)
+
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
         """Forward pass of the encoder and return all the features."""
         features = self.model(x)

cellseg_models_pytorch/modules/mlp.py

@@ -49,6 +49,8 @@ def __init__(
         act_kwargs = act_kwargs if act_kwargs is not None else {}
         self.out_channels = in_channels if out_channels is None else out_channels
         hidden_channels = int(mlp_ratio * in_channels)
+        act_kwargs["dim_in"] = hidden_channels
+        act_kwargs["dim_out"] = hidden_channels
 
         self.fc1 = nn.Linear(in_channels, hidden_channels, bias=bias)
         self.act = Activation(activation, **act_kwargs)