fix(nodels): Enable Dippa cellseg model conversion

Author: okunator

cellseg_models_pytorch/inference/post_processor.py

@@ -161,6 +161,7 @@ def post_proc_pipeline(
 
         if "type" in maps.keys():
             res["type"] = self._get_type_map(maps["type"], res["inst"], **self.kwargs)
+            res["inst"] *= res["type"] > 0
 
         return res
 

cellseg_models_pytorch/models/base/_base_model.py

@@ -86,7 +86,11 @@ def _check_decoder_args(
 
     def _get_inner_keys(self, d: Dict[str, Dict[str, Any]]) -> List[str]:
         """Get the inner dict keys from a nested dict."""
-        return chain.from_iterable(list(d[k].keys()) for k in d.keys())
+        return list(chain.from_iterable(list(d[k].keys()) for k in d.keys()))
+
+    def _flatten_inner_dicts(self, d: Dict[str, Dict[str, Any]]) -> List[str]:
+        """Get the inner dicts as one dict from a nested dict."""
+        return dict(chain.from_iterable(list(d[k].items()) for k in d.keys()))
 
     def _check_head_args(
         self, heads: Dict[str, int], decoders: Tuple[str, ...]

cellseg_models_pytorch/models/base/_multitask_unet.py

@@ -26,6 +26,8 @@ def __init__(
         enc_name: str = "resnet50",
         enc_pretrain: bool = True,
         enc_freeze: bool = False,
+        inst_key: str = None,
+        aux_key: str = None,
     ) -> None:
         """Create a universal multi-task (2D) unet.
 
@@ -63,11 +65,19 @@ def __init__(
                 Whether to use imagenet pretrained weights in the encoder.
             enc_freeze : bool, default=False
                 Freeze encoder weights for training.
+            inst_key : str, optional
+                The key for the model output that will be used in the instance
+                segmentation post-processing pipeline as the binary segmentation result.
+            aux_key : str, optional
+                The key for the model output that will be used in the instance
+                segmentation post-processing pipeline as the auxilliary map.
         """
         super().__init__()
         self.enc_freeze = enc_freeze
         use_style = style_channels is not None
         self.heads = heads
+        self.inst_key = inst_key
+        self.aux_key = aux_key
 
         # set timm encoder
         self.encoder = TimmEncoder(

cellseg_models_pytorch/modules/conv/ws_conv.py

@@ -41,7 +41,7 @@ def __init__(
         )
         self.eps = eps
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Weight standardized convolution forward pass."""
         weight = self.weight
 

cellseg_models_pytorch/modules/conv_base.py

@@ -12,6 +12,7 @@
     "InvertedBottleneckConv",
     "FusedMobileInvertedConv",
     "HoverNetDenseConv",
+    "BasicConvOld",
 ]
 
 
@@ -829,3 +830,135 @@ def forward_features_preact(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv(x)
 
         return x
+
+
+class BasicConvOld(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        same_padding: bool = True,
+        normalization: str = "bn",
+        activation: str = "relu",
+        convolution: str = "conv",
+        preactivate: bool = False,
+        kernel_size=3,
+        groups: int = 1,
+        bias: bool = False,
+        attention: str = None,
+        preattend: bool = False,
+        **kwargs
+    ) -> None:
+        """Conv-block (basic) parent class.
+
+        Parameters
+        ----------
+            in_channels : int
+                Number of input channels.
+            out_channels : int
+                Number of output channels.
+            same_padding : bool, default=True
+                if True, performs same-covolution.
+            normalization : str, default="bn":
+                Normalization method.
+                One of: "bn", "bcn", "gn", "in", "ln", None
+            activation : str, default="relu"
+                Activation method.
+                One of: "mish", "swish", "relu", "relu6", "rrelu", "selu",
+                "celu", "gelu", "glu", "tanh", "sigmoid", "silu", "prelu",
+                "leaky-relu", "elu", "hardshrink", "tanhshrink", "hardsigmoid"
+            convolution : str, default="conv"
+                The convolution method. One of: "conv", "wsconv", "scaled_wsconv"
+            preactivate : bool, default=False
+                If True, normalization will be applied before convolution.
+            kernel_size : int, default=3
+                The size of the convolution kernel.
+            groups : int, default=1
+                Number of groups the kernels are divided into. If `groups == 1`
+                normal convolution is applied. If `groups = in_channels`
+                depthwise convolution is applied.
+            bias : bool, default=False,
+                Include bias term in the convolution.
+            attention : str, default=None
+                Attention method. One of: "se", "scse", "gc", "eca", None
+            preattend : bool, default=False
+                If True, Attention is applied at the beginning of forward pass.
+        """
+        super().__init__()
+        self.conv_choice = convolution
+        self.out_channels = out_channels
+        self.preattend = preattend
+        self.preactivate = preactivate
+
+        # set norm channel number for preactivation or normal
+        norm_channels = in_channels if preactivate else self.out_channels
+
+        # set padding. Works if dilation or stride are not adjusted
+        padding = (kernel_size - 1) // 2 if same_padding else 0
+
+        self.conv = Conv(
+            name=self.conv_choice,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            groups=groups,
+            padding=padding,
+            bias=bias,
+        )
+
+        self.norm = Norm(normalization, num_features=norm_channels)
+        self.act = Activation(activation)
+
+        # set attention channels
+        att_channels = in_channels if preattend else self.out_channels
+        self.att = Attention(attention, in_channels=att_channels)
+
+        self.downsample = None
+        if in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                Conv(
+                    self.conv_choice,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    bias=False,
+                    kernel_size=1,
+                    padding=0,
+                ),
+                Norm(normalization, num_features=out_channels),
+            )
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        identity = x
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        x = self.att(x)
+
+        # residual
+        x = self.conv(x)
+        x = self.norm(x)
+
+        x += identity
+        x = self.act(x)
+
+        return x
+
+    def forward_features_preact(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with pre-activation."""
+        identity = x
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        # pre-attention
+        x = self.att(x)
+
+        # preact residual
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.conv(x)
+
+        x += identity
+        x = self.act(x)
+
+        return x

cellseg_models_pytorch/modules/conv_block.py

@@ -3,6 +3,7 @@
 
 from .conv_base import (
     BasicConv,
+    BasicConvOld,
     BottleneckConv,
     DepthWiseSeparableConv,
     FusedMobileInvertedConv,
@@ -20,6 +21,7 @@
     "mbconv": InvertedBottleneckConv,
     "fmbconv": FusedMobileInvertedConv,
     "hover_dense": HoverNetDenseConv,
+    "basic_old": BasicConvOld,
 }
 
 

changelog.d/20221006_114649_oskari.lehtonen.md

@@ -0,0 +1,5 @@
+## Fixes
+
+- Add a conv block `BasicConvOld` to enable `Dippa` to cellseg conversion of models.
+- Fix `inst_key`, `aux_key` bug in `MultiTaskUnet`
+- Add a type_map > 0 masking for the `inst_map`s in post-processing