feat(inference): add parallel postproc utils

Author: okunator

cellseg_models_pytorch/inference/__init__.py

@@ -1,10 +1,12 @@
+from ._base_inferer import BaseInferer
 from .folder_dataset import FolderDataset
 from .post_processor import PostProcessor
 from .predictor import Predictor
 from .resize_inferer import ResizeInferer
 from .sliding_window_inferer import SlidingWindowInferer
 
 __all__ = [
+    "BaseInferer",
     "Predictor",
     "PostProcessor",
     "ResizeInferer",

cellseg_models_pytorch/inference/_base_inferer.py

@@ -2,7 +2,7 @@
 from collections import OrderedDict
 from itertools import chain
 from pathlib import Path
-from typing import Dict, List, Tuple, Union
+from typing import Callable, Dict, List, Tuple, Union
 
 import numpy as np
 import torch
@@ -36,6 +36,8 @@ def __init__(
         save_dir: Union[Path, str] = None,
         checkpoint_path: Union[Path, str] = None,
         n_images: int = None,
+        type_post_proc: Callable = None,
+        sem_post_proc: Callable = None,
         **postproc_kwargs,
     ) -> None:
         """Inference for an image folder.
@@ -84,6 +86,12 @@ def __init__(
                 Path to the model weight checkpoints.
             n_images : int, optional
                 First n-number of images used from the `input_folder`.
+            type_post_proc : Callable, optional
+                A post-processing function for the type maps. If not None, overrides
+                the default.
+            sem_post_proc : Callable, optional
+                A post-processing function for the semantc seg maps. If not None,
+                overrides the default.
             **postproc_kwargs:
                 Arbitrary keyword arguments for the post-processing.
         """
@@ -142,6 +150,8 @@ def __init__(
             instance_postproc,
             inst_key=self.model.inst_key,
             aux_key=self.model.aux_key,
+            type_post_proc=type_post_proc,
+            sem_post_proc=sem_post_proc,
             **postproc_kwargs,
         )
 

cellseg_models_pytorch/inference/post_processor.py

@@ -1,14 +1,19 @@
-from typing import Dict, List
+from typing import Callable, Dict, List
 
 import numpy as np
 from pathos.multiprocessing import ThreadPool as Pool
-from skimage.filters import rank
-from skimage.morphology import closing, disk, opening
 from skimage.util import img_as_ubyte
 from tqdm import tqdm
 
 from ..postproc import POSTPROC_LOOKUP
-from ..utils import binarize, fill_holes_semantic, remove_debris_semantic
+from ..utils import (
+    fill_holes_semantic,
+    majority_vote_parallel,
+    majority_vote_sequential,
+    med_filt_parallel,
+    med_filt_sequential,
+    remove_debris_semantic,
+)
 
 __all__ = ["PostProcessor"]
 
@@ -28,7 +33,13 @@
 
 class PostProcessor:
     def __init__(
-        self, instance_postproc: str, inst_key: str, aux_key: str, **kwargs
+        self,
+        instance_postproc: str,
+        inst_key: str,
+        aux_key: str,
+        type_post_proc: Callable = None,
+        sem_post_proc: Callable = None,
+        **kwargs,
     ) -> None:
         """Multi-threaded post-processing.
 
@@ -42,6 +53,12 @@ def __init__(
             aux_key : Tuple[str, ...]:
                 The key/name of the model auxilliary output that is used for the
                 instance segmentation post-processing pipeline.
+            type_post_proc : Callable, optional
+                A post-processing function for the type maps. If not None, overrides
+                the default.
+            sem_post_proc : Callable, optional
+                A post-processing function for the semantc seg maps. If not None,
+                overrides the default.
             **kwargs
                 Arbitrary keyword arguments that can be used for any of the private
                 post-processing functions of this class.
@@ -57,36 +74,32 @@ def __init__(
         self.inst_key = inst_key
         self.aux_key = aux_key
         self.kwargs = kwargs
+        self.sem_post_proc = sem_post_proc
+        self.type_post_proc = type_post_proc
 
     def _get_sem_map(
         self,
         prob_map: np.ndarray,
-        use_blur: bool = False,
-        use_closing: bool = False,
-        use_opening: bool = True,
-        disk_size: int = 10,
+        parallel: bool = False,
+        kernel_width: int = 15,
         **kwargs,
     ) -> np.ndarray:
         """Run the semantic segmentation post-processing."""
-        # Median filtering to get rid of noise. Adds a lot of overhead sop optional.
-        if use_blur:
-            sem = np.zeros_like(prob_map)
-            for i in range(prob_map.shape[0]):
-                sem[i, ...] = rank.median(
-                    img_as_ubyte(prob_map[i, ...]), footprint=disk(disk_size)
-                )
-            prob_map = sem
-
-        sem = np.argmax(prob_map, axis=0)
+        sem_map = img_as_ubyte(prob_map)
 
-        if use_opening:
-            sem = opening(sem, disk(disk_size))
-
-        if use_closing:
-            sem = closing(sem, disk(disk_size))
+        if self.sem_post_proc is not None:
+            sem = self.sem_post_proc(sem_map)
+        else:
+            if parallel:
+                sem = med_filt_parallel(
+                    sem_map, kernel_size=(kernel_width, kernel_width)
+                )
+            else:
+                sem = med_filt_sequential(sem_map, kernel_width)
 
-        sem = remove_debris_semantic(sem)
-        sem = fill_holes_semantic(sem)
+            sem = np.argmax(sem, axis=0)
+            sem = remove_debris_semantic(sem)
+            sem = fill_holes_semantic(sem)
 
         return sem
 
@@ -105,31 +118,19 @@ def _get_type_map(
         self,
         prob_map: np.ndarray,
         inst_map: np.ndarray,
-        use_mask: bool = False,
+        parallel: bool = True,
         **kwargs,
     ) -> np.ndarray:
-        """Run the type map post-processing. Majority voting for each instance.
-
-        Adapted from:
-        https://github.com/vqdang/hover_net/blob/master/models/hovernet/post_proc.py
-        """
+        """Run the type map post-processing. Majority voting for each instance."""
         type_map = np.argmax(prob_map, axis=0)
-        if use_mask:
-            type_map = binarize(inst_map) * type_map
-
-        pred_id_list = np.unique(inst_map)[1:]
-        for inst_id in pred_id_list:
-            inst_type = type_map[inst_map == inst_id]
-            type_list, type_pixels = np.unique(inst_type, return_counts=True)
-            type_list = list(zip(type_list, type_pixels))
-            type_list = sorted(type_list, key=lambda x: x[1], reverse=True)
-            cell_type = type_list[0][0]
-
-            if cell_type == 0:
-                if len(type_list) > 1:
-                    cell_type = type_list[1][0]
 
-            type_map[inst_map == inst_id] = cell_type
+        if self.type_post_proc is not None:
+            type_map = self.type_post_proc(type_map, inst_map, **kwargs)
+        else:
+            if parallel:
+                type_map = majority_vote_parallel(type_map, inst_map)
+            else:
+                type_map = majority_vote_sequential(type_map, inst_map)
 
         return type_map
 
@@ -175,7 +176,6 @@ def run_parallel(
             progress_bar : bool, default=False
                 If True, a tqdm progress bar is shown.
 
-
         Returns
         -------
             List[Dict[str, np.ndarray]]:

cellseg_models_pytorch/inference/predictor.py

@@ -154,7 +154,7 @@ def classify(
         patch: torch.Tensor,
         act: Union[str, None] = "softmax",
         apply_weights: bool = False,
-    ) -> np.ndarray:
+    ) -> torch.Tensor:
         """Take in logits and output probabilities.
 
         Additionally apply a weight matrix to help with boundary artefacts.
@@ -171,7 +171,7 @@ def classify(
 
         Returns
         -------
-            np.ndarray:
+            torch.Tensor:
                 The model prediction. Same shape as input `patch`.
         """
         allowed = ("sigmoid", "softmax", "tanh", None)

cellseg_models_pytorch/inference/resize_inferer.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, Tuple, Union
+from typing import Callable, Dict, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -26,6 +26,8 @@ def __init__(
         save_dir: Union[Path, str] = None,
         checkpoint_path: Union[Path, str] = None,
         n_images: int = None,
+        type_post_proc: Callable = None,
+        sem_post_proc: Callable = None,
         **postproc_kwargs,
     ) -> None:
         """Resize inference for a folder of images.
@@ -81,6 +83,12 @@ def __init__(
                 Path to the model weight checkpoints.
             n_images : int, optional
                 First n-number of images used from the `input_folder`.
+            type_post_proc : Callable, optional
+                A post-processing function for the type maps. If not None, overrides
+                the default.
+            sem_post_proc : Callable, optional
+                A post-processing function for the semantc seg maps. If not None,
+                overrides the default.
             **postproc_kwargs:
                 Arbitrary keyword arguments for the post-processing.
         """
@@ -100,6 +108,8 @@ def __init__(
             save_dir=save_dir,
             checkpoint_path=checkpoint_path,
             n_images=n_images,
+            type_post_proc=type_post_proc,
+            sem_post_proc=sem_post_proc,
             **postproc_kwargs,
         )
 

cellseg_models_pytorch/inference/sliding_window_inferer.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, Tuple, Union
+from typing import Callable, Dict, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -29,6 +29,8 @@ def __init__(
         save_dir: Union[Path, str] = None,
         checkpoint_path: Union[Path, str] = None,
         n_images: int = None,
+        type_post_proc: Callable = None,
+        sem_post_proc: Callable = None,
         **postproc_kwargs,
     ) -> None:
         """Sliding window inference for a folder of images.
@@ -83,6 +85,12 @@ def __init__(
                 Path to the model weight checkpoints.
             n_images : int, optional
                 First n-number of images used from the `ìnput_folder`.
+            type_post_proc : Callable, optional
+                A post-processing function for the type maps. If not None, overrides
+                the default.
+            sem_post_proc : Callable, optional
+                A post-processing function for the semantc seg maps. If not None,
+                overrides the default.
             **postproc_kwargs:
                 Arbitrary keyword arguments for the post-processing.
         """
@@ -102,6 +110,8 @@ def __init__(
             save_dir=save_dir,
             checkpoint_path=checkpoint_path,
             n_images=n_images,
+            type_post_proc=type_post_proc,
+            sem_post_proc=sem_post_proc,
             **postproc_kwargs,
         )
 

cellseg_models_pytorch/inference/tests/test_inference.py

@@ -20,8 +20,7 @@ def test_slidingwin_inference(img_dir, batch_size):
         batch_size=batch_size,
         save_intermediate=False,
         device="cpu",
-        use_blur=True,
-        use_closing=True,
+        parallel=False,
     )
 
     inferer.infer()
@@ -45,8 +44,7 @@ def test_resize_inference(img_dir, batch_size):
         batch_size=batch_size,
         save_intermediate=False,
         device="cpu",
-        use_blur=True,
-        use_closing=True,
+        parallel=False,
     )
 
     inferer.infer()

cellseg_models_pytorch/utils/__init__.py

@@ -18,6 +18,10 @@
     get_inst_types,
     get_type_instances,
     label_semantic,
+    majority_vote_parallel,
+    majority_vote_sequential,
+    med_filt_parallel,
+    med_filt_sequential,
     one_hot,
     remap_label,
     remove_1px_boundary,
@@ -124,4 +128,8 @@
     "NORM_LOOKUP",
     "draw_stuff_contours",
     "draw_thing_contours",
+    "majority_vote_sequential",
+    "majority_vote_parallel",
+    "med_filt_parallel",
+    "med_filt_sequential",
 ]