feat(utils): refactor multiprocessing codes.

Author: okunator

cellseg_models_pytorch/inference/post_processor.py

@@ -1,9 +1,7 @@
 from typing import Callable, Dict, List
 
 import numpy as np
-from pathos.multiprocessing import ThreadPool as Pool
 from skimage.util import img_as_ubyte
-from tqdm import tqdm
 
 from ..postproc import POSTPROC_LOOKUP
 from ..utils import (
@@ -13,6 +11,7 @@
     med_filt_parallel,
     med_filt_sequential,
     remove_debris_semantic,
+    run_pool,
 )
 
 __all__ = ["PostProcessor"]
@@ -166,30 +165,32 @@ def post_proc_pipeline(
         return res
 
     def run_parallel(
-        self, maps: List[Dict[str, np.ndarray]], progress_bar: bool = False
+        self,
+        maps: List[Dict[str, np.ndarray]],
+        pooltype: str = "thread",
+        maptype: str = "amap",
     ) -> List[Dict[str, np.ndarray]]:
         """Run the full post-processing pipeline in parallel for many model outputs.
 
         Parameters
         ----------
             maps : List[Dict[str, np.ndarray]]
                 The model output map dictionaries in a list.
-            progress_bar : bool, default=False
-                If True, a tqdm progress bar is shown.
+            pooltype : str, default="thread"
+                The pathos pooltype. Allowed: ("process", "thread", "serial").
+                Defaults to "thread". (Fastest in benchmarks.)
+            maptype : str, default="amap"
+                The map type of the pathos Pool object.
+                Allowed: ("map", "amap", "imap", "uimap")
+                Defaults to "amap". (Fastest in benchmarks).
 
         Returns
         -------
             List[Dict[str, np.ndarray]]:
                 The post-processed output map dictionaries in a list.
         """
-        seg_results = []
-        with Pool() as pool:
-            if progress_bar:
-                it = tqdm(pool.imap(self.post_proc_pipeline, maps), total=len(maps))
-            else:
-                it = pool.imap(self.post_proc_pipeline, maps)
-
-            for x in it:
-                seg_results.append(x)
+        seg_results = run_pool(
+            self.post_proc_pipeline, maps, ret=True, pooltype=pooltype, maptype=maptype
+        )
 
         return seg_results

cellseg_models_pytorch/utils/__init__.py

@@ -33,6 +33,7 @@
     soft_type_flatten,
     type_map_flatten,
 )
+from .multiproc import run_pool
 from .patching import (
     TilerStitcher,
     TilerStitcherTorch,
@@ -134,4 +135,5 @@
     "med_filt_parallel",
     "med_filt_sequential",
     "intersection",
+    "run_pool",
 ]

cellseg_models_pytorch/utils/file_manager.py

@@ -1,12 +1,11 @@
 import re
+from functools import partial
 from pathlib import Path
 from typing import Any, Dict, List, Tuple, Union
 
 import cv2
 import numpy as np
 import scipy.io as sio
-from pathos.multiprocessing import ThreadPool as Pool
-from tqdm import tqdm
 
 from .mask_utils import (
     bounding_box,
@@ -15,6 +14,7 @@
     get_inst_types,
     label_semantic,
 )
+from .multiproc import run_pool
 
 
 class FileHandler:
@@ -58,30 +58,34 @@ def read_mat(
         key: str = "inst_map",
         retype: bool = True,
         return_all: bool = False,
-    ) -> Union[np.ndarray, None]:
+    ) -> Union[np.ndarray, Dict[str, np.ndarray], None]:
         """Read a mask from a .mat file.
 
         If a mask is not found, return None
 
         Parameters
         ----------
             path : str or Path
-                Path to the image file.
+                Path to the .mat file.
             key : str, default="inst_map"
                 Name/key of the mask type that is being read from .mat
             retype : bool, default=True
                 Convert the matrix type.
             return_all : bool, default=False
                 Return the whole dict. Overrides the `key` arg.
 
-        Returns
-        -------
-            np.ndarray or None:
-                The mask indice matrix. Shape (H, W)
 
         Raises
         ------
             ValueError: If an illegal key is given.
+
+        Returns
+        -------
+            Union[np.ndarray, List[np.ndarray], None]:
+                if return_all == False:
+                    The instance/type/semantic labelled mask. Shape: (H, W).
+                if return_all == True:
+                    All the masks in the .mat file returned in a dictionary.
         """
         dtypes = {
             "inst_map": "int32",
@@ -468,7 +472,8 @@ def save_masks_parallel(
         classes_type: Dict[str, str] = None,
         classes_sem: Dict[str, str] = None,
         offsets: bool = False,
-        progress_bar: bool = False,
+        pooltype: str = "thread",
+        maptype: str = "amap",
         **kwargs,
     ) -> None:
         """Save the model output masks to a folder. (multi-threaded).
@@ -493,31 +498,44 @@ def save_masks_parallel(
             offsets : bool, default=False
                 If True, geojson coords are shifted by the offsets that are encoded in
                 the filenames (e.g. "x-1000_y-4000.png"). Ignored if `format` != ".json"
-            progress_bar : bool, default=False
-                If True, a tqdm progress bar is shown.
+            pooltype : str, default="thread"
+                The pathos pooltype. Allowed: ("process", "thread", "serial").
+                Defaults to "thread". (Fastest in benchmarks.)
+            maptype : str, default="amap"
+                The map type of the pathos Pool object.
+                Allowed: ("map", "amap", "imap", "uimap")
+                Defaults to "amap". (Fastest in benchmarks).
         """
-        formats = [format] * len(maps)
-        geo_formats = [geo_format] * len(maps)
-        classes_type = [classes_type] * len(maps)
-        classes_sem = [classes_sem] * len(maps)
-        offsets = [offsets] * len(maps)
-        args = tuple(
-            zip(fnames, maps, formats, geo_formats, classes_type, classes_sem, offsets)
+        func = partial(
+            FileHandler._save_masks,
+            format=format,
+            geo_format=geo_format,
+            classes_type=classes_type,
+            classes_sem=classes_sem,
+            offsets=offsets,
         )
 
-        with Pool() as pool:
-            if progress_bar:
-                it = tqdm(pool.imap(FileHandler._save_masks, args), total=len(maps))
-            else:
-                it = pool.imap(FileHandler._save_masks, args)
-
-            for _ in it:
-                pass
+        args = tuple(zip(fnames, maps))
+        run_pool(func, args, ret=False, pooltype=pooltype, maptype=maptype)
 
     @staticmethod
-    def _save_masks(args: Tuple[Dict[str, np.ndarray], str, str]) -> None:
+    def _save_masks(
+        args: Tuple[str, Dict[str, np.ndarray]],
+        format: str,
+        geo_format: str,
+        classes_type: Dict[str, str],
+        classes_sem: Dict[str, str],
+        offsets: bool,
+    ) -> None:
         """Unpacks the args for `save_mask` to enable multi-threading."""
-        return FileHandler.save_masks(*args)
+        return FileHandler.save_masks(
+            *args,
+            format=format,
+            geo_format=geo_format,
+            classes_type=classes_type,
+            classes_sem=classes_sem,
+            offsets=offsets,
+        )
 
     @staticmethod
     def get_split(string: str) -> List[str]:

cellseg_models_pytorch/utils/multiproc.py

@@ -0,0 +1,107 @@
+from typing import Any, Callable, Generator, List, Union
+
+from pathos.pools import ProcessPool, SerialPool, ThreadPool
+
+__all__ = ["run_pool"]
+
+
+def iter_pool_generator(it: Generator, res: List = None) -> Union[List[Any], None]:
+    """Iterate over a pool generator object.
+
+    Parameters
+    ----------
+        it : Generator
+            A Generator object containing results from a concurrent run.
+        res : List | None
+            An empty list, where the results from the generator will be saved.
+            If None, no results will be saved.
+
+    Returns
+    -------
+        Union[List[Any], None]:
+            A list of results or None.
+    """
+    if res is not None:
+        for x in it:
+            res.append(x)
+    else:
+        for _ in it:
+            pass
+
+    return res
+
+
+def run_pool(
+    func: Callable,
+    args: List[Any],
+    ret: bool = True,
+    pooltype: str = "thread",
+    maptype: str = "amap",
+) -> Union[List[Any], None]:
+    """Run a pathos Thread, Process or Serial pool object.
+
+    NOTE: if `ret` is set to True and `func` callable does not return anything. This
+          will return a list of None values.
+
+    Parameters
+    ----------
+        func : Callable
+            The function that will be copied to existing core and run in parallel.
+        args : List[Any]
+            A list of arguments for each of the parallelly executed functions.
+        ret : bool, default=True
+            Flag, whether to return a list of results from the pool object. Will be set
+            to False e.g. when saving data to disk in parallel etc.
+        pooltype : str, default="thread"
+            The pathos pooltype. Allowed: ("process", "thread", "serial")
+        maptype : str, default="amap"
+            The map type of the pathos Pool object.
+            Allowed: ("map", "amap", "imap", "uimap")
+
+    Raises
+    ------
+        ValueError: if illegal `pooltype` or `maptype` is given.
+
+    Returns
+    -------
+        Union[List[Any], None]:
+            A list of results or None.
+    """
+    allowed = ("process", "thread", "serial")
+    if pooltype not in allowed:
+        raise ValueError(f"Illegal `pooltype`. Got {pooltype}. Allowed: {allowed}")
+
+    allowed = ("map", "amap", "imap", "uimap")
+    if maptype not in allowed:
+        raise ValueError(f"Illegal `maptype`. Got {maptype}. Allowed: {allowed}")
+
+    Pool = None
+    if pooltype == "thread":
+        Pool = ThreadPool
+    elif pooltype == "process":
+        Pool = ProcessPool
+    else:
+        if maptype in ("amap", "uimap"):
+            raise ValueError(
+                f"`SerialPool` has only `map` & `imap` implemented. Got: {maptype}."
+            )
+        Pool = SerialPool
+
+    results = [] if ret else None
+    if maptype == "map":
+        with Pool() as pool:
+            it = pool.map(func, args)
+            results = iter_pool_generator(it, results)
+    elif maptype == "amap":
+        with Pool() as pool:
+            results = pool.amap(func, args).get()
+    elif maptype == "imap":
+        with Pool() as pool:
+            it = pool.imap(func, args)
+            results = iter_pool_generator(it, results)
+    elif maptype == "uimap":
+        with Pool() as pool:
+            it = pool.uimap(func, args)
+            results = iter_pool_generator(it, results)
+
+    return results

cellseg_models_pytorch/utils/tests/test_multiproc.py

@@ -0,0 +1,46 @@
+import pytest
+
+from cellseg_models_pytorch.utils import run_pool
+
+
+def wrap1(num):
+    return num + 2
+
+
+def wrap2(arg):
+    return
+
+
+@pytest.mark.parametrize(
+    "typesets",
+    [
+        ("thread", "imap"),
+        ("thread", "map"),
+        ("thread", "uimap"),
+        ("thread", "amap"),
+        ("process", "map"),
+        ("process", "amap"),
+        ("process", "imap"),
+        ("process", "uimap"),
+        ("serial", "map"),
+        ("serial", "imap"),
+        pytest.param(("serial", "amap"), marks=pytest.mark.xfail),
+        pytest.param(("serial", "uimap"), marks=pytest.mark.xfail),
+    ],
+)
+@pytest.mark.parametrize("func", [wrap1, wrap2])
+@pytest.mark.parametrize("ret", [True, False])
+def test_run_pool(typesets, func, ret):
+    args = [1, 2, 3, 4, 5, 6, 7, 8]
+    res = run_pool(func, args, ret=ret, pooltype=typesets[0], maptype=typesets[1])
+
+    if ret and func == wrap1:
+        assert res == [3, 4, 5, 6, 7, 8, 9, 10]
+    elif ret and func == wrap2:
+        assert res == [None, None, None, None, None, None, None, None]
+    else:
+        # with amap, we always get a return val..
+        if not ret and typesets[1] == "amap":
+            pass
+        else:
+            assert res is None

changelog.d/20221201_172540_oskari.lehtonen.md

@@ -0,0 +1,12 @@
+## Features
+
+- added more pathos.Pool options for parallel processing. Added `ThreadPool`, `ProcessPool` & `SerialPool`
+- add all the mapping methods for each Pool obj. I.e. `amap`, `imap`, `uimap` and `map`
+
+## Refactor
+
+- Refactored multiprocessing code to be reusable and moved it under `utils`.
+
+## Tests
+
+- added tests for the multiprocessing tools.