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

dcnum/segm/{segmenter_cpu.py → segmenter_mpo.py}

@@ -2,8 +2,10 @@ import abc
 import multiprocessing as mp
 import time
 import threading
+from typing import Dict
 
 import numpy as np
+import scipy.ndimage as ndi
 
 from .segmenter import Segmenter
 
@@ -13,14 +15,41 @@ from .segmenter import Segmenter
 mp_spawn = mp.get_context('spawn')
 
 
-class CPUSegmenter(Segmenter, abc.ABC):
-    def __init__(self, num_workers=None, *args, **kwargs):
-        super(CPUSegmenter, self).__init__(*args, **kwargs)
+class MPOSegmenter(Segmenter, abc.ABC):
+    hardware_processor = "cpu"
+
+    def __init__(self,
+                 *,
+                 num_workers: int = None,
+                 kwargs_mask: Dict = None,
+                 debug: bool = False,
+                 **kwargs):
+        """Segmenter with multiprocessing operation
+
+        Parameters
+        ----------
+        kwargs_mask: dict
+            Keyword arguments for mask post-processing (see `process_mask`)
+        debug: bool
+            Debugging parameters
+        kwargs:
+            Additional, optional keyword arguments for `segment_algorithm`
+            defined in the subclass.
+        """
+        super(MPOSegmenter, self).__init__(kwargs_mask=kwargs_mask,
+                                           debug=debug,
+                                           **kwargs)
         self.num_workers = num_workers or mp.cpu_count()
+        # batch input image data
         self.mp_image_raw = None
         self._mp_image_np = None
+        # batch output image data
         self.mp_labels_raw = None
         self._mp_labels_np = None
+        # batch image background offset
+        self.mp_bg_off_raw = None
+        self._mp_bg_off_np = None
+        # workers
         self._mp_workers = []
         # Image shape of the input array
         self.image_shape = None
@@ -56,6 +85,7 @@ class CPUSegmenter(Segmenter, abc.ABC):
         del state["logger"]
         del state["_mp_image_np"]
         del state["_mp_labels_np"]
+        del state["_mp_bg_off_np"]
         del state["_mp_workers"]
         return state
 
@@ -64,26 +94,26 @@ class CPUSegmenter(Segmenter, abc.ABC):
         self.__dict__.update(state)
 
     @staticmethod
-    def _create_shared_array(image_shape, batch_size, dtype):
+    def _create_shared_array(array_shape, batch_size, dtype):
         """Return raw and numpy-view on shared array
 
         Parameters
         ----------
-        image_shape: tuple of int
+        array_shape: tuple of int
             Shape of one single image in the array
         batch_size: int
             Number of images in the array
         dtype:
-            ctype, e.g. `np.ctypeslib.ctypes.c_uint8`
-            or `np.ctypeslib.ctypes.c_bool`
+            numpy dtype
         """
-        sx, sy = image_shape
-        sa_raw = mp_spawn.RawArray(dtype, int(sx * sy * batch_size))
+        ctype = np.ctypeslib.as_ctypes_type(dtype)
+        sa_raw = mp_spawn.RawArray(ctype,
+                                   int(np.prod(array_shape) * batch_size))
         # Convert the RawArray to something we can write to fast
         # (similar to memory view, but without having to cast) using
         # np.ctypeslib.as_array. See discussion in
         # https://stackoverflow.com/questions/37705974
-        sa_np = np.ctypeslib.as_array(sa_raw).reshape(batch_size, sx, sy)
+        sa_np = np.ctypeslib.as_array(sa_raw).reshape(batch_size, *array_shape)
         return sa_raw, sa_np
 
     @property
@@ -105,39 +135,49 @@ class CPUSegmenter(Segmenter, abc.ABC):
             [w.join() for w in self._mp_workers]
 
     def segment_batch(self,
-                      image_data: np.ndarray,
+                      images: np.ndarray,
                       start: int = None,
-                      stop: int = None):
-        """Perform batch segmentation of `image_data`
+                      stop: int = None,
+                      bg_off: np.ndarray = None,
+                      ):
+        """Perform batch segmentation of `images`
+
+        Before segmentation, an optional background offset correction with
+        `bg_off` is performed. After segmentation, mask postprocessing is
+        performed according to the class definition.
 
         Parameters
         ----------
-        image_data: 3d np.ndarray
+        images: 3d np.ndarray
             The time-series image data. First axis is time.
         start: int
-            First index to analyze in `image_data`
+            First index to analyze in `images`
         stop: int
-            Index after the last index to analyze in `image_data`
+            Index after the last index to analyze in `images`
+        bg_off: 1D np.ndarray
+            Optional 1D numpy array with background offset
 
         Notes
         -----
         - If the segmentation algorithm only accepts background-corrected
-          images, then `image_data` must already be background-corrected.
+          images, then `images` must already be background-corrected,
+          except for the optional `bg_off`.
         """
         if stop is None or start is None:
             start = 0
-            stop = len(image_data)
+            stop = len(images)
 
         batch_size = stop - start
-        size = np.prod(image_data.shape[1:]) * batch_size
+        size = np.prod(images.shape[1:]) * batch_size
 
         if self.image_shape is None:
-            self.image_shape = image_data[0].shape
+            self.image_shape = images[0].shape
 
         if self._mp_image_np is not None and self._mp_image_np.size != size:
             # reset image data
             self._mp_image_np = None
             self._mp_labels_np = None
+            self._mp_bg_off_np = None
             # TODO: If only the batch_size changes, don't
             #  reinitialize the workers. Otherwise, the final rest of
             #  analyzing a dataset would always take a little longer.
@@ -146,30 +186,48 @@ class CPUSegmenter(Segmenter, abc.ABC):
             self.mp_batch_index.value = -1
             self.mp_shutdown.value = 0
 
+        if bg_off is not None:
+            if not self.requires_background_correction:
+                raise ValueError(f"The segmenter {self.__class__.__name__} "
+                                 f"does not employ background correction, "
+                                 f"but the `bg_off` keyword argument was "
+                                 f"passed to `segment_chunk`. Please check "
+                                 f"your analysis pipeline.")
+            # background offset
+            if self._mp_bg_off_np is None:
+                self.mp_bg_off_raw, self._mp_bg_off_np = \
+                    self._create_shared_array(
+                        array_shape=(stop - start,),
+                        batch_size=batch_size,
+                        dtype=np.float64)
+            self._mp_bg_off_np[:] = bg_off[start:stop]
+
+        # input images
         if self._mp_image_np is None:
             self.mp_image_raw, self._mp_image_np = self._create_shared_array(
-                image_shape=self.image_shape,
+                array_shape=self.image_shape,
                 batch_size=batch_size,
-                dtype=np.ctypeslib.ctypes.c_int32,
+                dtype=images.dtype,
             )
+        self._mp_image_np[:] = images[start:stop]
 
+        # output labels
         if self._mp_labels_np is None:
             self.mp_labels_raw, self._mp_labels_np = self._create_shared_array(
-                image_shape=self.image_shape,
+                array_shape=self.image_shape,
                 batch_size=batch_size,
-                dtype=np.ctypeslib.ctypes.c_uint16,
+                dtype=np.uint16,
             )
 
-        # populate image data
-        self._mp_image_np[:] = image_data[start:stop]
-
         # Create the workers
         if self.debug:
-            worker_cls = CPUSegmenterWorkerThread
+            worker_cls = MPOSegmenterWorkerThread
             num_workers = 1
+            self.logger.debug("Running with one worker in main thread")
         else:
-            worker_cls = CPUSegmenterWorkerProcess
-            num_workers = min(self.num_workers, image_data.shape[0])
+            worker_cls = MPOSegmenterWorkerProcess
+            num_workers = min(self.num_workers, images.shape[0])
+            self.logger.debug(f"Running with {num_workers} workers")
 
         if not self._mp_workers:
             step_size = batch_size // num_workers
@@ -200,8 +258,33 @@ class CPUSegmenter(Segmenter, abc.ABC):
 
         return self._mp_labels_np
 
+    def segment_single(self, image, bg_off: float = None):
+        """Return the integer label image for an input image
+
+        Before segmentation, an optional background offset correction with
+        `bg_off` is performed. After segmentation, mask postprocessing is
+        performed according to the class definition.
+        """
+        segm_wrap = self.segment_algorithm_wrapper()
+        # optional subtraction of background offset
+        if bg_off is not None:
+            image = image - bg_off
+        # obtain mask or label
+        mol = segm_wrap(image)
+        if mol.dtype == bool:
+            # convert mask to labels
+            labels, _ = ndi.label(
+                input=mol,
+                structure=ndi.generate_binary_structure(2, 2))
+        else:
+            labels = mol
+        # optional mask/label postprocessing
+        if self.mask_postprocessing:
+            labels = self.process_mask(labels, **self.kwargs_mask)
+        return labels
 
-class CPUSegmenterWorker:
+
+class MPOSegmenterWorker:
     def __init__(self,
                  segmenter,
                  sl_start: int,
@@ -211,7 +294,7 @@ class CPUSegmenterWorker:
 
         Parameters
         ----------
-        segmenter: CPUSegmenter
+        segmenter: MPOSegmenter
             The segmentation instance
         sl_start: int
             Start of slice of input array to process
@@ -219,7 +302,7 @@ class CPUSegmenterWorker:
             Stop of slice of input array to process
         """
         # Must call super init, otherwise Thread or Process are not initialized
-        super(CPUSegmenterWorker, self).__init__()
+        super(MPOSegmenterWorker, self).__init__()
         self.segmenter = segmenter
         # Value incrementing the batch index. Starts with 0 and is
         # incremented every time :func:`Segmenter.segment_batch` is
@@ -231,8 +314,10 @@ class CPUSegmenterWorker:
         # Shutdown bit tells workers to stop when set to != 0
         self.shutdown = segmenter.mp_shutdown
         # The image data for segmentation
-        self.image_data_raw = segmenter.mp_image_raw
-        # Boolean mask array
+        self.image_arr_raw = segmenter.mp_image_raw
+        # Background data offset
+        self.bg_off = segmenter.mp_bg_off_raw
+        # Integer output label array
         self.labels_data_raw = segmenter.mp_labels_raw
         # The shape of one image
         self.image_shape = segmenter.image_shape
@@ -244,10 +329,14 @@ class CPUSegmenterWorker:
         # We have to create the numpy-versions of the mp.RawArrays here,
         # otherwise we only get some kind of copy in the new process
         # when we use "spawn" instead of "fork".
-        labels_data = np.ctypeslib.as_array(self.labels_data_raw).reshape(
+        labels_arr = np.ctypeslib.as_array(self.labels_data_raw).reshape(
             -1, self.image_shape[0], self.image_shape[1])
-        image_data = np.ctypeslib.as_array(self.image_data_raw).reshape(
+        image_arr = np.ctypeslib.as_array(self.image_arr_raw).reshape(
             -1, self.image_shape[0], self.image_shape[1])
+        if self.bg_off is not None:
+            bg_off_data = np.ctypeslib.as_array(self.bg_off)
+        else:
+            bg_off_data = None
 
         idx = self.sl_start
         itr = 0  # current iteration (incremented when we reach self.sl_stop)
@@ -261,8 +350,13 @@ class CPUSegmenterWorker:
                     with self.batch_worker:
                         self.batch_worker.value += 1
                 else:
-                    labels_data[idx, :, :] = self.segmenter.segment_frame(
-                        image_data[idx])
+                    if bg_off_data is None:
+                        bg_off = None
+                    else:
+                        bg_off = bg_off_data[idx]
+
+                    labels_arr[idx, :, :] = self.segmenter.segment_single(
+                        image=image_arr[idx], bg_off=bg_off)
                     idx += 1
             elif self.shutdown.value:
                 break
@@ -271,11 +365,11 @@ class CPUSegmenterWorker:
                 time.sleep(.01)
 
 
-class CPUSegmenterWorkerProcess(CPUSegmenterWorker, mp_spawn.Process):
+class MPOSegmenterWorkerProcess(MPOSegmenterWorker, mp_spawn.Process):
     def __init__(self, *args, **kwargs):
-        super(CPUSegmenterWorkerProcess, self).__init__(*args, **kwargs)
+        super(MPOSegmenterWorkerProcess, self).__init__(*args, **kwargs)
 
 
-class CPUSegmenterWorkerThread(CPUSegmenterWorker, threading.Thread):
+class MPOSegmenterWorkerThread(MPOSegmenterWorker, threading.Thread):
     def __init__(self, *args, **kwargs):
-        super(CPUSegmenterWorkerThread, self).__init__(*args, **kwargs)
+        super(MPOSegmenterWorkerThread, self).__init__(*args, **kwargs)

dcnum/segm/segmenter_sto.py

@@ -0,0 +1,110 @@
+import abc
+from typing import Dict
+
+import numpy as np
+import scipy.ndimage as ndi
+
+
+from .segmenter import Segmenter
+
+
+class STOSegmenter(Segmenter, abc.ABC):
+    hardware_processor = "gpu"
+
+    def __init__(self,
+                 *,
+                 num_workers: int = None,
+                 kwargs_mask: Dict = None,
+                 debug: bool = False,
+                 **kwargs
+                 ):
+        """Segmenter with single thread operation
+
+        Parameters
+        ----------
+        kwargs_mask: dict
+            Keyword arguments for mask post-processing (see `process_mask`)
+        debug: bool
+            Debugging parameters
+        kwargs:
+            Additional, optional keyword arguments for `segment_algorithm`
+            defined in the subclass.
+        """
+        if num_workers not in [None, 1]:
+            raise ValueError(f"Number of workers must not be larger than 1 "
+                             f"for GPU segmenter, got '{num_workers}'!")
+        super(STOSegmenter, self).__init__(kwargs_mask=kwargs_mask,
+                                           debug=debug,
+                                           **kwargs)
+
+    def segment_batch(self,
+                      images: np.ndarray,
+                      start: int = None,
+                      stop: int = None,
+                      bg_off: np.ndarray = None,
+                      ):
+        """Perform batch segmentation of `images`
+
+        Before segmentation, an optional background offset correction with
+        `bg_off` is performed. After segmentation, mask postprocessing is
+        performed according to the class definition.
+
+        Parameters
+        ----------
+        images: 3d np.ndarray
+            The time-series image data. First axis is time.
+        start: int
+            First index to analyze in `images`
+        stop: int
+            Index after the last index to analyze in `images`
+        bg_off: 1D np.ndarray
+            Optional 1D numpy array with background offset
+
+        Notes
+        -----
+        - If the segmentation algorithm only accepts background-corrected
+          images, then `images` must already be background-corrected,
+          except for the optional `bg_off`.
+        """
+        if stop is None or start is None:
+            start = 0
+            stop = len(images)
+
+        image_slice = images[start:stop]
+        segm = self.segment_algorithm_wrapper()
+
+        if bg_off is not None:
+            if not self.requires_background_correction:
+                raise ValueError(f"The segmenter {self.__class__.__name__} "
+                                 f"does not employ background correction, "
+                                 f"but the `bg_off` keyword argument was "
+                                 f"passed to `segment_chunk`. Please check "
+                                 f"your analysis pipeline.")
+            image_slice = image_slice - bg_off.reshape(-1, 1, 1)
+        labels = segm(image_slice)
+
+        # Make sure we have integer labels and perform mask postprocessing
+        if labels.dtype == bool:
+            new_labels = np.zeros_like(labels, dtype=np.uint16)
+            for ii in range(len(labels)):
+                ndi.label(
+                    input=labels[ii],
+                    output=new_labels[ii],
+                    structure=ndi.generate_binary_structure(2, 2))
+            labels = new_labels
+
+        # Perform mask postprocessing
+        if self.mask_postprocessing:
+            for ii in range(len(labels)):
+                labels[ii] = self.process_mask(labels[ii], **self.kwargs_mask)
+
+        return labels
+
+    def segment_single(self, image, bg_off: float = None):
+        """This is a convenience-wrapper around `segment_batch`"""
+        if bg_off is None:
+            bg_off_batch = None
+        else:
+            bg_off_batch = np.atleast_1d(bg_off)
+        images = image[np.newaxis]
+        return self.segment_batch(images, bg_off=bg_off_batch)[0]

dcnum/write/__init__.py

@@ -1,4 +1,6 @@
 # flake8: noqa: F401
 from .deque_writer_thread import DequeWriterThread
 from .queue_collector_thread import EventStash, QueueCollectorThread
-from .writer import HDF5Writer, copy_metadata, create_with_basins
+from .writer import (
+    HDF5Writer, copy_features, copy_metadata, create_with_basins,
+    set_default_filter_kwargs)

dcnum/write/deque_writer_thread.py

@@ -1,14 +1,17 @@
 import collections
+import logging
 import pathlib
 import threading
 import time
 
+import h5py
+
 from .writer import HDF5Writer
 
 
 class DequeWriterThread(threading.Thread):
     def __init__(self,
-                 path_out: pathlib.Path,
+                 path_out: pathlib.Path | h5py.File,
                  dq: collections.deque,
                  ds_kwds: dict = None,
                  mode: str = "a",
@@ -19,11 +22,12 @@ class DequeWriterThread(threading.Thread):
         ----------
         path_out:
             Path to the output HDF5 file
-        dq:
+        dq: collections.deque
             `collections.deque` object from which data are taken
             using `popleft()`.
         """
         super(DequeWriterThread, self).__init__(*args, **kwargs)
+        self.logger = logging.getLogger("dcnum.write.DequeWriterThread")
         if mode == "w":
             path_out.unlink(missing_ok=True)
         self.writer = HDF5Writer(path_out, mode=mode, ds_kwds=ds_kwds)
@@ -40,15 +44,21 @@ class DequeWriterThread(threading.Thread):
         self.may_stop_loop = True
 
     def run(self):
+        time_tot = 0
         while True:
+            ldq = len(self.dq)
             if self.must_stop_loop:
                 break
-            elif len(self.dq):
-                feat, data = self.dq.popleft()
-                self.writer.store_feature_chunk(feat=feat, data=data)
+            elif ldq:
+                t0 = time.perf_counter()
+                for _ in range(ldq):
+                    feat, data = self.dq.popleft()
+                    self.writer.store_feature_chunk(feat=feat, data=data)
+                time_tot += time.perf_counter() - t0
             elif self.may_stop_loop:
                 break
             else:
                 # wait for the next item to arrive
                 time.sleep(.1)
+        self.logger.info(f"Disk time: {time_tot:.1f}s")
         self.writer.close()

dcnum/write/queue_collector_thread.py

@@ -171,7 +171,7 @@ class QueueCollectorThread(threading.Thread):
         self.event_queue.cancel_join_thread()
         # Indexes the current frame in `self.data`.
         last_idx = 0
-        self.logger.debug("Started collector thread.")
+        self.logger.debug("Started collector thread")
         while True:
             # Slice of the shared nevents array. If it contains -1 values,
             # this means that some of the frames have not yet been processed.
@@ -184,10 +184,10 @@ class QueueCollectorThread(threading.Thread):
 
             if len(cur_nevents) == 0:
                 self.logger.info(
-                    "Reached the end of the current dataset (frame "
+                    "Reached dataset end (frame "
                     # `last_idx` is the size of the dataset in the end,
                     # because `len(cur_nevents)` is always added to it.
-                    f"{last_idx} of {len(self.feat_nevents)}).")
+                    f"{last_idx} of {len(self.feat_nevents)})")
                 break
 
             # We have reached the writer threshold. This means the extractor
@@ -245,20 +245,14 @@ class QueueCollectorThread(threading.Thread):
             # the events that we just saved.
             indices = stash.indices_for_data
 
-            # Write all the scalar features.
-            for feat in self.data.features_scalar_frame:
-                self.writer_dq.append((feat, self.data[feat][indices]))
-
-            # Write the image and background data.
-            imdat = np.zeros((stash.size,) + self.data.image.image_shape,
-                             dtype=np.uint8)
-            bgdat = np.zeros((stash.size,) + self.data.image.image_shape,
-                             dtype=np.uint8)
-            for ii, idx in enumerate(indices):
-                imdat[ii] = self.data.image[idx]
-                bgdat[ii] = self.data.image_bg[idx]
-            self.writer_dq.append(("image", imdat))
-            self.writer_dq.append(("image_bg", bgdat))
+            # This is the unmapped index from the input HDF5Data instance.
+            # Unmapped means that this only enumerates HDF5Data, but since
+            # HDF5Data can be mapped, the index does not necessarily enumerate
+            # the underlying HDF5 file. Later on, we will have to convert this
+            # to the correct "basinmap0" feature
+            # (see `DCNumJobRunner.task_enforce_basin_strategy`)
+            self.writer_dq.append(("index_unmapped",
+                                   np.array(indices, dtype=np.uint32)))
 
             # Write the number of events.
             self.writer_dq.append(("nevents",
@@ -273,3 +267,6 @@ class QueueCollectorThread(threading.Thread):
 
             # Increment current frame index.
             last_idx += len(cur_nevents)
+
+        self.logger.info(f"Counted {self.written_events} events")
+        self.logger.debug(f"Counted {self.written_frames} frames")