feat: add enc & dec feat return option to forward

Author: okunator

cellseg_models_pytorch/models/base/_base_model.py

@@ -10,6 +10,19 @@
 
 
 class BaseMultiTaskSegModel(nn.ModuleDict):
+    def forward_features(
+        self, x: torch.Tensor
+    ) -> Tuple[List[torch.Tensor], Dict[str, torch.Tensor]]:
+        """Forward pass for encoder, style and decoders.
+
+        NOTE: Returns both encoder and decoder features, not style.
+        """
+        feats = self.forward_encoder(x)
+        style = self.forward_style(feats[0])
+        dec_feats = self.forward_dec_features(feats, style)
+
+        return feats, dec_feats
+
     def forward_encoder(self, x: torch.Tensor) -> List[torch.Tensor]:
         """Forward the model encoder."""
         self._check_input_shape(x)

cellseg_models_pytorch/models/cellpose/cellpose.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -199,11 +199,42 @@ def __init__(
         if enc_freeze:
             self.freeze_encoder()
 
-    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
-        """Forward pass of Cellpose U-net."""
-        feats = self.forward_encoder(x)
-        style = self.forward_style(feats[0])
-        dec_feats = self.forward_dec_features(feats, style)
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_feats: bool = False,
+    ) -> Union[
+        Dict[str, torch.Tensor],
+        Tuple[
+            List[torch.Tensor],
+            Dict[str, torch.Tensor],
+            Dict[str, torch.Tensor],
+        ],
+    ]:
+        """Forward pass of Cellpose U-net.
+
+        Parameters
+        ----------
+            x : torch.Tensor
+                Input image batch. Shape: (B, C, H, W).
+            return_feats : bool, default=False
+                If True, encoder, decoder, and head outputs will all be returned
+
+        Returns
+        -------
+        Union[
+            Dict[str, torch.Tensor],
+            Tuple[
+                List[torch.Tensor],
+                Dict[str, torch.Tensor],
+                Dict[str, torch.Tensor],
+            ],
+        ]:
+            Dictionary mapping of output names to outputs or if `return_feats == True`
+            returns also the encoder features in a list, decoder features as a dict
+            mapping decoder names to outputs and the final head outputs dict.
+        """
+        feats, dec_feats = self.forward_features(x)
 
         for decoder_name in self.heads.keys():
             for head_name in self.heads[decoder_name].keys():
@@ -212,6 +243,9 @@ def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
 
         out = self.forward_heads(dec_feats)
 
+        if return_feats:
+            return feats, dec_feats, out
+
         return out
 
 

cellseg_models_pytorch/models/hovernet/hovernet.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -192,14 +192,47 @@ def __init__(
         if enc_freeze:
             self.freeze_encoder()
 
-    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
-        """Forward pass of HoVer-Net."""
-        feats = self.forward_encoder(x)
-        style = self.forward_style(feats[0])
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_feats: bool = False,
+    ) -> Union[
+        Dict[str, torch.Tensor],
+        Tuple[
+            List[torch.Tensor],
+            Dict[str, torch.Tensor],
+            Dict[str, torch.Tensor],
+        ],
+    ]:
+        """Forward pass of HoVer-Net.
 
-        dec_feats = self.forward_dec_features(feats, style)
+        Parameters
+        ----------
+            x : torch.Tensor
+                Input image batch. Shape: (B, C, H, W).
+            return_feats : bool, default=False
+                If True, encoder, decoder, and head outputs will all be returned
+
+        Returns
+        -------
+        Union[
+            Dict[str, torch.Tensor],
+            Tuple[
+                List[torch.Tensor],
+                Dict[str, torch.Tensor],
+                Dict[str, torch.Tensor],
+            ],
+        ]:
+            Dictionary mapping of output names to outputs or if `return_feats == True`
+            returns also the encoder features in a list, decoder features as a dict
+            mapping decoder names to outputs and the final head outputs dict.
+        """
+        feats, dec_feats = self.forward_features(x)
         out = self.forward_heads(dec_feats)
 
+        if return_feats:
+            return feats, dec_feats, out
+
         return out
 
 

cellseg_models_pytorch/models/stardist/stardist.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -207,12 +207,43 @@ def __init__(
         if enc_freeze:
             self.freeze_encoder()
 
-    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
-        """Forward pass of Stardist."""
-        feats = self.forward_encoder(x)
-        style = self.forward_style(feats[0])
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_feats: bool = False,
+    ) -> Union[
+        Dict[str, torch.Tensor],
+        Tuple[
+            List[torch.Tensor],
+            Dict[str, torch.Tensor],
+            Dict[str, torch.Tensor],
+        ],
+    ]:
+        """Forward pass of Stardist.
+
+        Parameters
+        ----------
+            x : torch.Tensor
+                Input image batch. Shape: (B, C, H, W).
+            return_feats : bool, default=False
+                If True, encoder, decoder, and head outputs will all be returned
+
+        Returns
+        -------
+        Union[
+            Dict[str, torch.Tensor],
+            Tuple[
+                List[torch.Tensor],
+                Dict[str, torch.Tensor],
+                Dict[str, torch.Tensor],
+            ],
+        ]:
+            Dictionary mapping of output names to outputs or if `return_feats == True`
+            returns also the encoder features in a list, decoder features as a dict
+            mapping decoder names to outputs and the final head outputs dict.
+        """
+        feats, dec_feats = self.forward_features(x)
 
-        dec_feats = self.forward_dec_features(feats, style)
         # Extra convs after decoders
         for e in self.extra_convs.keys():
             for extra_conv in self.extra_convs[e].keys():
@@ -230,6 +261,9 @@ def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
 
         out = self.forward_heads(dec_feats)
 
+        if return_feats:
+            return feats, dec_feats, out
+
         return out
 
 

cellseg_models_pytorch/postproc/functional/stardist/stardist.py

@@ -204,7 +204,7 @@ def post_proc_stardist(
     nms are accelerated with `numba` and `scipy.spatial.KDtree`.
 
     NOTE:
-    This implementaiton of the stardist post-processing is actually nearly twice
+    This implementaiton of the stardist post-processing is actually nearly 2x
     faster than the original version if `trim_bboxes` is set to True. The resulting
     segmentation is not an exact match but the differences are mostly neglible.
 

changelog.d/20221125_124720_oskari.lehtonen.md

@@ -0,0 +1,3 @@
+## Features
+
+- Add option to return encoder features, and decoder features along the outputs in the forward pass of any model.