feat(models): Add multi-task U-net for experiments

Author: okunator

cellseg_models_pytorch/models/__init__.py

@@ -1,3 +1,4 @@
+from .base._multitask_unet import MultiTaskUnet
 from .cellpose.cellpose import (
     CellPoseUnet,
     cellpose_base,
@@ -20,6 +21,7 @@
 )
 
 __all__ = [
+    "MultiTaskUnet",
     "HoverNet",
     "hovernet_base",
     "hovernet_plus",

cellseg_models_pytorch/models/base/_multitask_unet.py

@@ -0,0 +1,135 @@
+from typing import Any, Dict, Tuple
+
+import torch
+
+from ...decoders import Decoder
+from ...modules.misc_modules import StyleReshape
+from ._base_model import BaseMultiTaskSegModel
+from ._seg_head import SegHead
+from ._timm_encoder import TimmEncoder
+
+__all__ = ["MultiTaskUnet"]
+
+
+class MultiTaskUnet(BaseMultiTaskSegModel):
+    def __init__(
+        self,
+        decoders: Tuple[str, ...],
+        heads: Dict[str, Dict[str, int]],
+        n_layers: Dict[str, Tuple[int, ...]],
+        n_blocks: Dict[str, Tuple[Tuple[int, ...], ...]],
+        out_channels: Dict[str, Tuple[int, ...]],
+        long_skips: Dict[str, str],
+        dec_params: Dict[str, Tuple[Dict[str, Any], ...]],
+        depth: int = 4,
+        style_channels: int = 256,
+        enc_name: str = "resnet50",
+        enc_pretrain: bool = True,
+        enc_freeze: bool = False,
+    ) -> None:
+        """Create a universal multi-task (2D) unet.
+
+        NOTE: For experimental purposes.
+
+        Parameters
+        ----------
+            decoders : Tuple[str, ...]
+                Names of the decoder branches of this network. E.g. ("cellpose", "sem")
+            heads : Dict[str, Dict[str, int]]
+                Names of the decoder branches (has to match `decoders`) mapped to dicts
+                 of output name - number of output classes. E.g.
+                {"cellpose": {"type": 4, "cellpose": 2}, "sem": {"sem": 5}}
+            n_layers : Dict[str, Tuple[int, ...]]
+                The number of conv layers inside each of the decoder stages.
+            n_blocks : Dict[str, Tuple[Tuple[int, ...], ...]]
+                The number of blocks inside each conv-layer in each decoder stage.
+            out_channels : Tuple[int, ...]
+                Out channels for each decoder stage.
+            long_skips : Dict[str, str]
+                long skip method to be used. One of: "unet", "unetpp", "unet3p",
+                "unet3p-lite", None
+            dec_params : Dict[str, Tuple[Dict[str, Any], ...]])
+                The keyword args for each of the distinct decoder stages. Incudes the
+                parameters for the long skip connections and convolutional layers of the
+                decoder itself. See the `DecoderStage` documentation for more info.
+            depth : int, default=4
+                The depth of the encoder. I.e. Number of returned feature maps from
+                the encoder. Maximum depth = 5.
+            style_channels : int, default=256
+                Number of style vector channels. If None, style vectors are ignored.
+            enc_name : str, default="resnet50"
+                Name of the encoder. See timm docs for more info.
+            enc_pretrain : bool, default=True
+                Whether to use imagenet pretrained weights in the encoder.
+            enc_freeze : bool, default=False
+                Freeze encoder weights for training.
+        """
+        super().__init__()
+        self.enc_freeze = enc_freeze
+        use_style = style_channels is not None
+        self.heads = heads
+
+        # set timm encoder
+        self.encoder = TimmEncoder(
+            enc_name,
+            depth=depth,
+            pretrained=enc_pretrain,
+        )
+
+        # style
+        self.make_style = None
+        if use_style:
+            self.make_style = StyleReshape(self.encoder.out_channels[0], style_channels)
+
+            # set decoders
+        for decoder_name in decoders:
+            decoder = Decoder(
+                enc_channels=list(self.encoder.out_channels),
+                style_channels=style_channels,
+                out_channels=out_channels[decoder_name],
+                long_skip=long_skips[decoder_name],
+                n_layers=n_layers[decoder_name],
+                n_blocks=n_blocks[decoder_name],
+                stage_params=dec_params[decoder_name],
+            )
+            self.add_module(f"{decoder_name}_decoder", decoder)
+
+        # set heads
+        for decoder_name in heads.keys():
+            for output_name, n_classes in heads[decoder_name].items():
+                seg_head = SegHead(
+                    in_channels=decoder.out_channels,
+                    out_channels=n_classes,
+                    kernel_size=1,
+                )
+                self.add_module(f"{output_name}_seg_head", seg_head)
+
+        self.name = f"MultiTaskUnet-{enc_name}"
+
+        # init decoder weights
+        self.initialize()
+
+        # freeze encoder if specified
+        if enc_freeze:
+            self.freeze_encoder()
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Forward pass of Multi-task U-net."""
+        self._check_input_shape(x)
+
+        feats = self.encoder(x)
+
+        style = None
+        if self.make_style is not None:
+            style = self.make_style(feats[0])
+
+        dec_feats = self.forward_dec_features(feats, style)
+
+        for decoder_name in self.heads.keys():
+            for head_name in self.heads[decoder_name].keys():
+                k = self.aux_key if head_name not in dec_feats.keys() else head_name
+                dec_feats[head_name] = dec_feats[k]
+
+        out = self.forward_heads(dec_feats)
+
+        return out

cellseg_models_pytorch/models/cellpose/_conf.py

@@ -15,7 +15,7 @@ def _create_cellpose_args(
     use_style: bool,
     merge_policy: str,
     skip_params: Dict[str, Any],
-) -> Dict[str, Any]:
+) -> Tuple[Dict[str, Any], ...]:
     """Create the args to build CellPose-Unet architecture."""
     skip_params = skip_params if skip_params is not None else {"k": None}
 

cellseg_models_pytorch/models/cellpose/cellpose.py

@@ -22,7 +22,7 @@
 class CellPoseUnet(BaseMultiTaskSegModel):
     def __init__(
         self,
-        decoders: Tuple[str],
+        decoders: Tuple[str, ...],
         heads: Dict[str, Dict[str, int]],
         inst_key: str = "type",
         depth: int = 4,
@@ -65,7 +65,7 @@ def __init__(
 
         Parameters
         ----------
-            decoders : Tuple[str]
+            decoders : Tuple[str, ...]
                 Names of the decoder branches of this network. E.g. ("cellpose", "sem")
             heads : Dict[str, Dict[str, int]]
                 Names of the decoder branches (has to match `decoders`) mapped to dicts

cellseg_models_pytorch/models/hovernet/_conf.py

@@ -15,7 +15,7 @@ def _create_hovernet_args(
     use_style: bool,
     merge_policy: str,
     skip_params: Dict[str, Any],
-) -> Dict[str, Any]:
+) -> Tuple[Dict[str, Any], ...]:
     """Create the correct args to build HoVerNet architecture."""
     skip_params = skip_params if skip_params is not None else {"k": None}
 

cellseg_models_pytorch/models/hovernet/hovernet.py

@@ -62,7 +62,7 @@ def __init__(
 
         Parameters
         ----------
-            decoders : Tuple[str]
+            decoders : Tuple[str, ...]
                 Names of the decoder branches of this network. E.g. ("hovernet", "sem")
             heads : Dict[str, Dict[str, int]]
                 The segmentation heads of the architecture. I.e. Names of the decoder

cellseg_models_pytorch/models/stardist/_conf.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Tuple
 
 __all__ = ["_create_stardist_args"]
 
@@ -16,7 +16,7 @@ def _create_stardist_args(
     block_type: str,
     merge_policy: str,
     skip_params: Dict[str, Any],
-) -> Dict[str, Any]:
+) -> Tuple[Dict[str, Any], ...]:
     """Create the args to build CellPose-Unet architecture."""
     skip_params = skip_params if skip_params is not None else {"k": None}
 

cellseg_models_pytorch/models/stardist/stardist.py

@@ -57,7 +57,7 @@ def __init__(
 
         Parameters
         ----------
-            decoders : Tuple[str]
+            decoders : Tuple[str, ...]
                 Names of the decoder branches of this network. E.g. ("stardist", "sem")
             extra_convs : Dict[str, Dict[str, int]]
                 The extra conv blocks before segmentation heads of the architecture.

cellseg_models_pytorch/models/tests/test_models.py

@@ -2,6 +2,7 @@
 import torch
 
 from cellseg_models_pytorch.models import (
+    MultiTaskUnet,
     cellpose_base,
     cellpose_plus,
     hovernet_base,
@@ -76,3 +77,20 @@ def test_omnipose_fwdbwd(model):
         assert y["sem"].shape == torch.Size([1, 3, 64, 64])
 
     assert y["omnipose"].shape == torch.Size([1, 2, 64, 64])
+
+
+def test_multitaskunet_fwdbwd():
+    x = torch.rand([1, 3, 64, 64])
+    m = MultiTaskUnet(
+        decoders=("sem",),
+        heads={"sem": {"sem": 3}},
+        n_layers={"sem": (1, 1, 1, 1)},
+        n_blocks={"sem": ((2,), (2,), (2,), (2,))},
+        out_channels={"sem": (128, 64, 32, 16)},
+        long_skips={"sem": "unet"},
+        dec_params={"sem": None},
+    )
+    y = m(x)
+    y["sem"].mean().backward()
+
+    assert y["sem"].shape == torch.Size([1, 3, 64, 64])

changelog.d/20220908_144856_oskari.lehtonen.md

@@ -0,0 +1,11 @@
+### Features
+
+- **models**: Add a universal multi-task U-net model builder (experimental)
+
+### Type Hints
+
+- **models**: Fix incorrect type hints.
+
+### Test
+
+- **models**: Update tests for multi-task U-Net