dcnum 0.13.2__py3-none-any.whl → 0.23.1__py3-none-any.whl

dcnum/segm/segmenter.py

@@ -1,43 +1,69 @@
 import abc
+import copy
 import functools
 import inspect
 import logging
+from typing import Dict
 
 import cv2
 import numpy as np
 import scipy.ndimage as ndi
 from skimage import morphology
 
-from ..meta.ppid import kwargs_to_ppid
+from ..meta.ppid import kwargs_to_ppid, ppid_to_kwargs
+
+
+class SegmenterNotApplicableError(BaseException):
+    """Used to indicate when a dataset cannot be segmented with a segmenter"""
+    def __init__(self, segmenter_class, reasons_list):
+        super(SegmenterNotApplicableError, self).__init__(
+            f"The dataset cannot be segmented with the "
+            f"'{segmenter_class.get_ppid_code()}' segmenter: "
+            f"{', '.join(reasons_list)}"
+        )
+        self.reasons_list = reasons_list
+        self.segmenter_class = segmenter_class
 
 
 class Segmenter(abc.ABC):
+    #: Required hardware ("cpu" or "gpu") defined in first-level subclass.
+    hardware_processor = "none"
     #: Whether to enable mask post-processing. If disabled, you should
     #: make sure that your mask is properly defined and cleaned or you
-    #: have to call `process_mask` in your `segment_approach` implementation.
-    mask_postprocessing = False
+    #: have to call `process_mask` in your `segment_algorithm` implementation.
+    mask_postprocessing = True
     #: Default keyword arguments for mask post-processing. See `process_mask`
     #: for available options.
     mask_default_kwargs = {}
     #: If the segmenter requires a background-corrected image, set this to True
     requires_background_correction = False
 
-    def __init__(self, kwargs_mask=None, debug=False, **kwargs):
-        """Base segemnter
+    def __init__(self,
+                 *,
+                 kwargs_mask: Dict = None,
+                 debug: bool = False,
+                 **kwargs):
+        """Base segmenter class
+
+        This is the base segmenter class for the multiprocessing operation
+        segmenter :class:`.MPOSegmenter` (multiple subprocesses are spawned
+        and each of them works on a queue of images) and the single-threaded
+        operation segmenter :class:`.STOSegmenter` (e.g. for batch
+        segmentation on a GPU).
 
         Parameters
         ----------
-        data: HDF5Data
-            Instance containing the raw data. Requires at least the
-            `image` and `image_bg` attributes. Some segemnters require
-            more properties, so make sure to use :class:`.HDF5Data`.
         kwargs_mask: dict
             Keyword arguments for mask post-processing (see `process_mask`)
+        debug: bool
+            Enable debugging mode (e.g. CPU segmenter runs in one thread)
+        kwargs:
+            Additional, optional keyword arguments for `segment_batch`.
         """
         self.debug = debug
         self.logger = logging.getLogger(__name__).getChild(
             self.__class__.__name__)
-        spec = inspect.getfullargspec(self.segment_approach)
+        spec = inspect.getfullargspec(self.segment_algorithm)
         #: custom keyword arguments for the subclassing segmenter
         self.kwargs = spec.kwonlydefaults or {}
         self.kwargs.update(kwargs)
@@ -55,19 +81,26 @@ class Segmenter(abc.ABC):
                 "`kwargs_mask` has been specified, but mask post-processing "
                 f"is disabled for segmenter {self.__class__}")
 
-    @classmethod
-    def key(cls):
-        """The unique key/name of this segmenter class"""
-        key = cls.__name__.lower()
-        if key.startswith("segment"):
-            key = key[7:]
-        return key
+    @staticmethod
+    @functools.cache
+    def get_border(shape):
+        """Cached boolean image with outer pixels set to True"""
+        border = np.zeros(shape, dtype=bool)
+        border[[0, -1], :] = True
+        border[:, [0, -1]] = True
+        return border
+
+    @staticmethod
+    @functools.cache
+    def get_disk(radius):
+        """Cached `skimage.morphology.disk(radius)`"""
+        return morphology.disk(radius)
 
     def get_ppid(self):
         """Return a unique segmentation pipeline identifier
 
         The pipeline identifier is universally applicable and must
-        be backwards-compatible (future versions of dcevent will
+        be backwards-compatible (future versions of dcnum will
         correctly acknowledge the ID).
 
         The segmenter pipeline ID is defined as::
@@ -75,7 +108,7 @@ class Segmenter(abc.ABC):
             KEY:KW_APPROACH:KW_MASK
 
         Where KEY is e.g. "legacy" or "watershed", and KW_APPROACH is a
-        list of keyword arguments for `segment_approach`, e.g.::
+        list of keyword arguments for `segment_algorithm`, e.g.::
 
             thresh=-6^blur=0
 
@@ -85,36 +118,75 @@ class Segmenter(abc.ABC):
 
         KW_MASK represents keyword arguments for `process_mask`.
         """
-        return self.get_ppid_from_kwargs(self.kwargs, self.kwargs_mask)
+        return self.get_ppid_from_ppkw(self.kwargs, self.kwargs_mask)
+
+    @classmethod
+    def get_ppid_code(cls):
+        """The unique code/name of this segmenter class"""
+        code = cls.__name__.lower()
+        if code.startswith("segment"):
+            code = code[7:]
+        return code
 
     @classmethod
-    def get_ppid_from_kwargs(cls, kwargs, kwargs_mask=None):
+    def get_ppid_from_ppkw(cls, kwargs, kwargs_mask=None):
         """Return the pipeline ID from given keyword arguments
 
         See Also
         --------
         get_ppid: Same method for class instances
         """
-        if kwargs_mask is None and kwargs.get("kwargs_mask", None) is None:
-            raise KeyError("`kwargs_mask` must be either specified as "
-                           "keyword argument to this method or as a key "
-                           "in `kwargs`!")
-        if kwargs_mask is None:
-            # see check above (kwargs_mask may also be {})
-            kwargs_mask = kwargs["kwargs_mask"]
-        # Start with the default mask kwargs defined for this subclass
-        kwargs_mask_used = cls.mask_default_kwargs
-        kwargs_mask_used.update(kwargs_mask)
-        key = cls.key()
-        csegm = kwargs_to_ppid(cls, "segment_approach", kwargs)
-        cmask = kwargs_to_ppid(cls, "process_mask", kwargs_mask_used)
-        return ":".join([key, csegm, cmask])
+        kwargs = copy.deepcopy(kwargs)
+        if cls.mask_postprocessing:
+            if kwargs_mask is None and kwargs.get("kwargs_mask", None) is None:
+                raise KeyError("`kwargs_mask` must be either specified as "
+                               "keyword argument to this method or as a key "
+                               "in `kwargs`!")
+            if kwargs_mask is None:
+                # see check above (kwargs_mask may also be {})
+                kwargs_mask = kwargs.pop("kwargs_mask")
+            # Start with the default mask kwargs defined for this subclass
+            kwargs_mask_used = copy.deepcopy(cls.mask_default_kwargs)
+            kwargs_mask_used.update(kwargs_mask)
+        elif kwargs_mask:
+            raise ValueError(f"The segmenter '{cls.__name__}' does not "
+                             f"support mask postprocessing, but 'kwargs_mask' "
+                             f"was provided: {kwargs_mask}")
+
+        ppid_parts = [
+            cls.get_ppid_code(),
+            kwargs_to_ppid(cls, "segment_algorithm", kwargs),
+            ]
+
+        if cls.mask_postprocessing:
+            ppid_parts.append(
+                kwargs_to_ppid(cls, "process_mask", kwargs_mask_used))
+
+        return ":".join(ppid_parts)
 
     @staticmethod
-    @functools.cache
-    def get_disk(radius):
-        """Cached `skimage.morphology.disk(radius)`"""
-        return morphology.disk(radius)
+    def get_ppkw_from_ppid(segm_ppid):
+        """Return keyword arguments for this pipeline identifier"""
+        ppid_parts = segm_ppid.split(":")
+        code = ppid_parts[0]
+        pp_kwargs = ppid_parts[1]
+
+        for cls_code in get_available_segmenters():
+            if cls_code == code:
+                cls = get_available_segmenters()[cls_code]
+                break
+        else:
+            raise ValueError(
+                f"Could not find segmenter '{code}'!")
+        kwargs = ppid_to_kwargs(cls=cls,
+                                method="segment_algorithm",
+                                ppid=pp_kwargs)
+        if cls.mask_postprocessing:
+            pp_kwargs_mask = ppid_parts[2]
+            kwargs["kwargs_mask"] = ppid_to_kwargs(cls=cls,
+                                                   method="process_mask",
+                                                   ppid=pp_kwargs_mask)
+        return kwargs
 
     @staticmethod
     def process_mask(labels, *,
@@ -130,8 +202,8 @@ class Segmenter(abc.ABC):
 
         Parameters
         ----------
-        labels: 2d integer ndarray
-            Labeled input (contains blobs with same number)
+        labels: 2d integer or boolean ndarray
+            Labeled input (contains blobs consisting of unique numbers)
         clear_border: bool
             clear the image boarder using
             :func:`skimage.segmentation.clear_border`
@@ -142,6 +214,12 @@ class Segmenter(abc.ABC):
             if > 0, perform a binary closing with a disk
             of that radius in pixels
         """
+        if labels.dtype == bool:
+            # Convert mask image to labels
+            labels, _ = ndi.label(
+                input=labels,
+                structure=ndi.generate_binary_structure(2, 2))
+
         if clear_border:
             #
             # from skimage import segmentation
@@ -149,33 +227,13 @@ class Segmenter(abc.ABC):
             #
             if (labels[0, :].sum() or labels[-1, :].sum()
                     or labels[:, 0].sum() or labels[:, -1].sum()):
-                border = np.zeros_like(labels, dtype=bool)
-                border[0] = True
-                border[-1] = True
-                border[:, 0] = True
-                border[:, -1] = True
+                border = Segmenter.get_border(labels.shape)
                 indices = sorted(np.unique(labels[border]))
-                for ii in indices[1:]:
-                    labels[labels == ii] = 0
-
-        # scikit-image is too slow for us here. So we use OpenCV.
-        # https://github.com/scikit-image/scikit-image/issues/1190
-
-        if closing_disk:
-            #
-            # from skimage import morphology
-            # morphology.binary_closing(
-            #    mask,
-            #    footprint=morphology.disk(closing_disk),
-            #    out=mask)
-            #
-            element = Segmenter.get_disk(closing_disk)
-            labels_uint8 = np.array(labels, dtype=np.uint8)
-            labels_dilated = cv2.dilate(labels_uint8, element)
-            labels_eroded = cv2.erode(labels_dilated, element)
-            labels, _ = ndi.label(
-                input=labels_eroded > 0,
-                structure=ndi.generate_binary_structure(2, 2))
+                for li in indices:
+                    if li == 0:
+                        # ignore background values
+                        continue
+                    labels[labels == li] = 0
 
         if fill_holes:
             # Floodfill only works with uint8 (too small) or int32
@@ -188,55 +246,146 @@ class Segmenter(abc.ABC):
             # Floodfill algorithm fills the background image and
             # the resulting inversion is the image with holes filled.
             # This will destroy labels (adding 2,147,483,647 to background)
+            # Since floodfill will use the upper left corner of the image as
+            # a seed, we have to make sure it is set to background. We set
+            # a line of pixels in the upper channel wall to zero to be sure.
+            labels[0, :] = 0
+            # ...and a 4x4 pixel region in the top left corner.
+            labels[1, :2] = 0
             cv2.floodFill(labels, None, (0, 0), 2147483647)
             mask = labels != 2147483647
             labels, _ = ndi.label(
                 input=mask,
                 structure=ndi.generate_binary_structure(2, 2))
 
-        return labels
-
-    def segment_chunk(self, image_data, chunk):
-        """Return the integer labels for one `image_data` chunk"""
-        data = image_data.get_chunk(chunk)
-        return self.segment_batch(data)
-
-    def segment_frame(self, image):
-        """Return the integer label image for `index`"""
-        segm_wrap = self.segment_frame_wrapper()
-        # obtain mask or label
-        mol = segm_wrap(image)
-        if mol.dtype == bool:
-            # convert mask to label
+        if closing_disk:
+            # scikit-image is too slow for us here. So we use OpenCV.
+            # https://github.com/scikit-image/scikit-image/issues/1190
+            #
+            # from skimage import morphology
+            # morphology.binary_closing(
+            #    mask,
+            #    footprint=morphology.disk(closing_disk),
+            #    out=mask)
+            #
+            element = Segmenter.get_disk(closing_disk)
+            # Note: erode/dilate not implemented for int32
+            labels_uint8 = np.array(labels, dtype=np.uint8)
+            # Historically, we would like to do a closing (dilation followed
+            # by erosion) on the image data where lower brightness values
+            # meant "we have an event". However, since we are now working
+            # with labels instead of image data (0 is background and labels
+            # are enumerated with integers), high "brightness" values are
+            # actually the event. Thus, we have to perform an opening
+            # (erosion followed by dilation) of the label image.
+            labels_eroded = cv2.erode(labels_uint8, element)
+            labels_dilated = cv2.dilate(labels_eroded, element)
             labels, _ = ndi.label(
-                input=mol,
+                input=labels_dilated > 0,
                 structure=ndi.generate_binary_structure(2, 2))
-        else:
-            labels = mol
-        # optional postprocessing
-        if self.mask_postprocessing:
-            labels = self.process_mask(labels, **self.kwargs_mask)
+
         return labels
 
+    @staticmethod
+    @abc.abstractmethod
+    def segment_algorithm(image):
+        """The segmentation algorithm implemented in the subclass
+
+        Perform segmentation and return integer label or binary mask image
+        """
+
     @functools.cache
-    def segment_frame_wrapper(self):
+    def segment_algorithm_wrapper(self):
+        """Wraps `self.segment_algorithm` to only accept an image
+
+        The static method `self.segment_algorithm` may optionally accept
+        keyword arguments `self.kwargs`. This wrapper returns the
+        wrapped method that only accepts the image as an argument. This
+        makes sense if you want to unify
+        """
         if self.kwargs:
             # For segmenters that accept keyword arguments.
-            segm_wrap = functools.partial(self.segment_approach,
+            segm_wrap = functools.partial(self.segment_algorithm,
                                           **self.kwargs)
         else:
             # For segmenters that don't accept keyword arguments.
-            segm_wrap = self.segment_approach
+            segm_wrap = self.segment_algorithm
         return segm_wrap
 
-    @staticmethod
     @abc.abstractmethod
-    def segment_approach(image):
-        """Perform segmentation and return integer label or binary mask image
+    def segment_batch(self, images, start=None, stop=None, bg_off=None):
+        """Return the integer labels for an entire batch
+
+        This is implemented in the MPO and STO segmenters.
+        """
 
-        This is the approach the subclasses implement.
+    def segment_chunk(self, image_data, chunk, bg_off=None):
+        """Return the integer labels for one `image_data` chunk
+
+        This is a wrapper for `segment_batch`.
+
+        Parameters
+        ----------
+        image_data:
+            Instance of dcnum's :class:`.BaseImageChunkCache` with
+            the methods `get_chunk` and `get_chunk_slice`.
+        chunk: int
+            Integer identifying the chunk in `image_data` to segment
+        bg_off: ndarray
+            Optional 1D array with same length as `image_data` that holds
+            additional background offset values that should be subtracted
+            from the image data before segmentation. Should only be
+            used in combination with segmenters that have
+            `requires_background_correction` set to True.
         """
+        images = image_data.get_chunk(chunk)
+        if bg_off is not None:
+            bg_off_chunk = bg_off[image_data.get_chunk_slice(chunk)]
+        else:
+            bg_off_chunk = None
+        return self.segment_batch(images, bg_off=bg_off_chunk)
 
     @abc.abstractmethod
-    def segment_batch(self, data, start=None, stop=None):
-        """Return the integer labels for an entire batch"""
+    def segment_single(self, image):
+        """Return the integer label for one image
+
+        This is implemented in the MPO and STO segmenters.
+        """
+
+    @classmethod
+    def validate_applicability(cls,
+                               segmenter_kwargs: Dict,
+                               meta: Dict = None,
+                               logs: Dict = None):
+        """Validate the applicability of this segmenter for a dataset
+
+        Parameters
+        ----------
+        segmenter_kwargs: dict
+            Keyword arguments for the segmenter
+        meta: dict
+            Dictionary of metadata from an :class:`HDF5Data` instance
+        logs: dict
+            Dictionary of logs from an :class:`HDF5Data` instance
+
+        Returns
+        -------
+        applicable: bool
+            True if the segmenter is applicable to the dataset
+
+        Raises
+        ------
+        SegmenterNotApplicableError
+            If the segmenter is not applicable to the dataset
+        """
+        return True
+
+
+@functools.cache
+def get_available_segmenters():
+    """Return dictionary of available segmenters"""
+    segmenters = {}
+    for scls in Segmenter.__subclasses__():
+        for cls in scls.__subclasses__():
+            segmenters[cls.get_ppid_code()] = cls
+    return segmenters

dcnum/segm/segmenter_manager_thread.py

@@ -8,7 +8,7 @@ import numpy as np
 from ..read.cache import HDF5ImageCache, ImageCorrCache
 
 from .segmenter import Segmenter
-from .segmenter_cpu import CPUSegmenter
+from .segmenter_mpo import MPOSegmenter
 
 
 class SegmenterManagerThread(threading.Thread):
@@ -17,7 +17,7 @@ class SegmenterManagerThread(threading.Thread):
                  image_data: HDF5ImageCache | ImageCorrCache,
                  slot_states: mp.Array,
                  slot_chunks: mp.Array,
-                 debug: bool = False,
+                 bg_off: np.ndarray = None,
                  *args, **kwargs):
         """Manage the segmentation of image data
 
@@ -38,10 +38,10 @@ class SegmenterManagerThread(threading.Thread):
         slot_chunks:
             For each slot in `slot_states`, this shared array defines
             on which chunk in `image_data` the segmentation took place.
-        debug:
-            Whether to run in debugging mode (more verbose messages and
-            CPU-based segmentation is done in one single thread instead
-            of in multiple subprocesses).
+        bg_off:
+            1d array containing additional background image offset values
+            that are added to each background image before subtraction
+            from the input image
 
         Notes
         -----
@@ -65,6 +65,9 @@ class SegmenterManagerThread(threading.Thread):
         self.segmenter = segmenter
         #: Image data which is being segmented
         self.image_data = image_data
+        #: Additional, optional background offset
+        self.bg_off = (
+            bg_off if self.segmenter.requires_background_correction else None)
         #: Slot states
         self.slot_states = slot_states
         #: Current slot chunk index for the slot states
@@ -73,37 +76,45 @@ class SegmenterManagerThread(threading.Thread):
         self.labels_list = [None] * len(self.slot_states)
         #: Time counter for segmentation
         self.t_count = 0
-        #: Whether running in debugging mode
-        self.debug = debug
 
     def run(self):
+        num_slots = len(self.slot_states)
         # We iterate over all the chunks of the image data.
         for chunk in self.image_data.iter_chunks():
-            num_slots = len(self.slot_states)
-            cur_slot = 0
-            empty_slots = 0
+            unavailable_slots = 0
+            found_free_slot = False
             # Wait for a free slot to perform segmentation (compute labels)
-            while True:
-                # - "e" there is data from the segmenter (the extractor
-                #   can take it and process it)
-                # - "s" the extractor processed the data and is waiting
-                #   for the segmenter
-                if self.slot_states[cur_slot] != "e":
-                    break
-                else:
-                    empty_slots += 1
-                    cur_slot = (cur_slot + 1) % num_slots
-                if empty_slots >= num_slots:
-                    # There is nothing to do, try to avoid 100% CPU
-                    empty_slots = 0
-                    time.sleep(.01)
+            while not found_free_slot:
+                # We sort the slots according to the slot chunks so that we
+                # always process the slot with the smallest slot chunk number
+                # first. Initially, the slot_chunks array is filled with
+                # zeros, but we populate it here.
+                for cur_slot in np.argsort(self.slot_chunks):
+                    # - "e" there is data from the segmenter (the extractor
+                    #   can take it and process it)
+                    # - "s" the extractor processed the data and is waiting
+                    #   for the segmenter
+                    if self.slot_states[cur_slot] != "e":
+                        # It's the segmenter's turn. Note that we use '!= "e"',
+                        # because the initial value is "\x00".
+                        found_free_slot = True
+                        break
+                    else:
+                        # Try another slot.
+                        unavailable_slots += 1
+                    if unavailable_slots >= num_slots:
+                        # There is nothing to do, try to avoid 100% CPU
+                        unavailable_slots = 0
+                        time.sleep(.1)
 
             t1 = time.monotonic()
 
             # We have a free slot to compute the segmentation
             labels = self.segmenter.segment_chunk(
                 image_data=self.image_data,
-                chunk=chunk)
+                chunk=chunk,
+                bg_off=self.bg_off,
+            )
 
             # TODO: make this more memory efficient (pre-shared mp.Arrays?)
             # Store labels in a list accessible by the main thread
@@ -113,12 +124,12 @@ class SegmenterManagerThread(threading.Thread):
             # This must be done last: Let the extractor know that this
             # slot is ready for processing.
             self.slot_states[cur_slot] = "e"
-            self.logger.debug(f"Segmented one chunk: {chunk}")
+            self.logger.debug(f"Segmented chunk {chunk} in slot {cur_slot}")
 
             self.t_count += time.monotonic() - t1
 
         # Cleanup
-        if isinstance(self.segmenter, CPUSegmenter):
+        if isinstance(self.segmenter, MPOSegmenter):
             # Join the segmentation workers.
             self.segmenter.join_workers()