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

dcnum/segm/segm_torch/torch_preproc.py

@@ -0,0 +1,114 @@
+from typing import Tuple
+
+import numpy as np
+
+
+def preprocess_images(images: np.ndarray,
+                      norm_mean: float | None,
+                      norm_std: float | None,
+                      image_shape: Tuple[int, int] = None,
+                      ):
+    """Transform image data to something torch models expect
+
+    The transformation includes:
+    - normalization (division by 255, subtraction of mean, division by std)
+    - cropping and padding of the input images to `image_shape`. For padding,
+      the median of each *individual* image is used.
+    - casting the input images to four dimensions
+      (batch_size, 1, height, width) where the second axis is "channels"
+
+    Parameters
+    ----------
+    images:
+        Input image array (batch_size, height_in, width_in). If this is a
+        2D image, it will be reshaped to a 3D image with a batch_size of 1.
+    norm_mean:
+        Mean value used for standard score data normalization, i.e.
+        `normalized = `(images / 255 - norm_mean) / norm_std`; Set
+        to None to disable normalization.
+    norm_std:
+        Standard deviation used for standard score data normalization;
+        Set to None to disable normalization (see above).
+    image_shape
+        Image shape for which the model was created (height, width).
+        If the image shape does not match the input image shape, then
+        the input images are padded/cropped to fit the image shape of
+        the model.
+
+    Returns
+    -------
+    image_proc:
+        3D array with preprocessed image data of shape
+        (batch_size, 1, height, width)
+    """
+    if len(images.shape) == 2:
+        # Insert indexing axis (batch dimension)
+        images = images[np.newaxis, :, :]
+
+    batch_size = images.shape[0]
+
+    # crop and pad the images based on what the model expects
+    image_shape_act = images.shape[1:]
+    if image_shape is None:
+        # model fits perfectly to input data
+        image_shape = image_shape_act
+
+    # height
+    hdiff = image_shape_act[0] - image_shape[0]
+    ht = abs(hdiff) // 2
+    hb = abs(hdiff) - ht
+    # width
+    wdiff = image_shape_act[1] - image_shape[1]
+    wl = abs(wdiff) // 2
+    wr = abs(wdiff) - wl
+    # helper variables
+    wpad = wdiff < 0
+    wcrp = wdiff > 0
+    hpad = hdiff < 0
+    hcrp = hdiff > 0
+
+    # The easy part is the cropping
+    if hcrp or wcrp:
+        # define slices for width and height
+        slice_hc = slice(ht, -hb) if hcrp else slice(None, None)
+        slice_wc = slice(wl, -wr) if wcrp else slice(None, None)
+        img_proc = images[:, slice_hc, slice_wc]
+    else:
+        img_proc = images
+
+    # The hard part is the padding
+    if hpad or wpad:
+        # compute median for each original input image
+        img_med = np.median(images, axis=(1, 2))
+        # broadcast the median array from 1D to 3D
+        img_med = img_med[:, None, None]
+
+        # define slices for width and height
+        slice_hp = slice(ht, -hb) if hpad else slice(None, None)
+        slice_wp = slice(wl, -wr) if wpad else slice(None, None)
+
+        # empty padding image stack with the shape required for the model
+        img_pad = np.empty(shape=(batch_size, image_shape[0], image_shape[1]),
+                           dtype=np.float32)
+        # fill in original data
+        img_pad[:, slice_hp, slice_wp] = img_proc
+        # fill in background data for height
+        if hpad:
+            img_pad[:, :ht, :] = img_med
+            img_pad[:, -hb:, :] = img_med
+        # fill in background data for width
+        if wpad:
+            img_pad[:, :, :wl] = img_med
+            img_pad[:, :, -wr:] = img_med
+        # Replace img_norm
+        img_proc = img_pad
+
+    if norm_mean is None or norm_std is None:
+        # convert to float32
+        img_norm = img_proc.astype(np.float32)
+    else:
+        # normalize images
+        img_norm = (img_proc.astype(np.float32) / 255 - norm_mean) / norm_std
+
+    # Add a "channels" axis for the ML models.
+    return img_norm[:, np.newaxis, :, :]

dcnum/segm/segmenter.py

@@ -13,13 +13,25 @@ from skimage import morphology
 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 = {}
@@ -31,21 +43,27 @@ class Segmenter(abc.ABC):
                  kwargs_mask: Dict = None,
                  debug: bool = False,
                  **kwargs):
-        """Base segmenter
+        """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
         ----------
         kwargs_mask: dict
             Keyword arguments for mask post-processing (see `process_mask`)
         debug: bool
-            Debugging parameters
+            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)
@@ -68,10 +86,8 @@ class Segmenter(abc.ABC):
     def get_border(shape):
         """Cached boolean image with outer pixels set to True"""
         border = np.zeros(shape, dtype=bool)
-        border[0, :] = True
-        border[-1, :] = True
-        border[:, 0] = True
-        border[:, -1] = True
+        border[[0, -1], :] = True
+        border[:, [0, -1]] = True
         return border
 
     @staticmethod
@@ -84,7 +100,7 @@ class Segmenter(abc.ABC):
         """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::
@@ -92,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
 
@@ -121,25 +137,40 @@ class Segmenter(abc.ABC):
         get_ppid: Same method for class instances
         """
         kwargs = copy.deepcopy(kwargs)
-        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)
-        code = cls.get_ppid_code()
-        csegm = kwargs_to_ppid(cls, "segment_approach", kwargs)
-        cmask = kwargs_to_ppid(cls, "process_mask", kwargs_mask_used)
-        return ":".join([code, csegm, cmask])
+        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
     def get_ppkw_from_ppid(segm_ppid):
         """Return keyword arguments for this pipeline identifier"""
-        code, pp_kwargs, pp_kwargs_mask = segm_ppid.split(":")
+        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]
@@ -148,11 +179,13 @@ class Segmenter(abc.ABC):
             raise ValueError(
                 f"Could not find segmenter '{code}'!")
         kwargs = ppid_to_kwargs(cls=cls,
-                                method="segment_approach",
+                                method="segment_algorithm",
                                 ppid=pp_kwargs)
-        kwargs["kwargs_mask"] = ppid_to_kwargs(cls=cls,
-                                               method="process_mask",
-                                               ppid=pp_kwargs_mask)
+        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
@@ -169,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`
@@ -181,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
@@ -196,25 +235,6 @@ class Segmenter(abc.ABC):
                         continue
                     labels[labels == li] = 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))
-
         if fill_holes:
             # Floodfill only works with uint8 (too small) or int32
             if labels.dtype != np.int32:
@@ -226,58 +246,139 @@ 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 the approach the subclasses implement.
+        This is implemented in the MPO and STO segmenters.
         """
 
+    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

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,40 +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():
-            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":
-                    # It's the segmenters turn. Note that we use '!= "e"',
-                    # because the initial value is "\x00".
-                    break
-                else:
-                    # Try another slot.
-                    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
@@ -116,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()