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

dcnum/read/mapped.py

@@ -0,0 +1,87 @@
+import functools
+
+import numbers
+
+import h5py
+import numpy as np
+
+
+class MappedHDF5Dataset:
+    def __init__(self,
+                 h5ds: h5py.Dataset,
+                 mapping_indices: np.ndarray):
+        """An index-mapped object for accessing an HDF5 dataset
+
+        Parameters
+        ----------
+        h5ds: h5py.Dataset
+            HDF5 dataset from which to map data
+        mapping_indices: np.ndarray
+            numpy indexing array containing integer indices
+        """
+        self.h5ds = h5ds
+        self.mapping_indices = mapping_indices
+        self.shape = (mapping_indices.size,) + h5ds.shape[1:]
+
+    def __getitem__(self, idx):
+        if isinstance(idx, numbers.Integral):
+            return self.h5ds[self.mapping_indices[idx]]
+        else:
+            idx_mapped = self.mapping_indices[idx]
+            return self.h5ds[idx_mapped]
+
+
+def get_mapping_indices(
+        index_mapping: numbers.Integral | slice | list | np.ndarray
+        ):
+    """Return integer numpy array with mapping indices for a range
+
+    Parameters
+    ----------
+    index_mapping: numbers.Integral | slice | list | np.ndarray
+        Several options you have here:
+        - integer: results in np.arrange(integer)
+        - slice: results in np.arrange(slice.start, slice.stop, slice.step)
+        - list or np.ndarray: returns the input as  unit32 array
+    """
+    if isinstance(index_mapping, numbers.Integral):
+        return _get_mapping_indices_cached(index_mapping)
+    elif isinstance(index_mapping, slice):
+        return _get_mapping_indices_cached(
+            (index_mapping.start, index_mapping.stop, index_mapping.step))
+    elif isinstance(index_mapping, (np.ndarray, list)):
+        return np.array(index_mapping, dtype=np.uint32)
+    else:
+        raise ValueError(f"Invalid type for `index_mapping`: "
+                         f"{type(index_mapping)} ({index_mapping})")
+
+
+@functools.lru_cache(maxsize=100)
+def _get_mapping_indices_cached(
+        index_mapping: numbers.Integral | tuple
+        ):
+    if isinstance(index_mapping, numbers.Integral):
+        return np.arange(index_mapping)
+    elif isinstance(index_mapping, tuple):
+        im_slice = slice(*index_mapping)
+        if im_slice.stop is None or im_slice.start is None:
+            raise NotImplementedError(
+                "Slices must have start and stop defined")
+        return np.arange(im_slice.start, im_slice.stop, im_slice.step)
+    elif isinstance(index_mapping, list):
+        return np.array(index_mapping, dtype=np.uint32)
+    else:
+        raise ValueError(f"Invalid type for cached `index_mapping`: "
+                         f"{type(index_mapping)} ({index_mapping})")
+
+
+def get_mapped_object(obj, index_mapping=None):
+    if index_mapping is None:
+        return obj
+    elif isinstance(obj, h5py.Dataset):
+        return MappedHDF5Dataset(
+            obj,
+            mapping_indices=get_mapping_indices(index_mapping))
+    else:
+        raise ValueError(f"No recipe to convert object of type {type(obj)} "
+                         f"({obj}) to an index-mapped object")

dcnum/segm/__init__.py

@@ -1,6 +1,9 @@
 # flake8: noqa: F401
-from .segmenter import Segmenter, get_available_segmenters
-from .segmenter_cpu import CPUSegmenter
-from .segmenter_gpu import GPUSegmenter
+from .segmenter import (
+    Segmenter, SegmenterNotApplicableError, get_available_segmenters
+)
+from .segmenter_mpo import MPOSegmenter
+from .segmenter_sto import STOSegmenter
 from .segmenter_manager_thread import SegmenterManagerThread
 from . import segm_thresh
+from . import segm_torch

dcnum/segm/segm_thresh.py

@@ -1,7 +1,7 @@
-from .segmenter_cpu import CPUSegmenter
+from .segmenter_mpo import MPOSegmenter
 
 
-class SegmentThresh(CPUSegmenter):
+class SegmentThresh(MPOSegmenter):
     mask_postprocessing = True
     mask_default_kwargs = {
         "clear_border": True,
@@ -10,22 +10,10 @@ class SegmentThresh(CPUSegmenter):
     }
     requires_background_correction = True
 
-    def __init__(self, thresh=-6, *args, **kwargs):
-        """Simple image thresholding segmentation
-
-        Parameters
-        ----------
-        thresh: int
-            grayscale threhold value for creating the mask image;
-            For a background-corrected image, pixels with values below
-            this value are considered to be part of the mask.
-        """
-        super(SegmentThresh, self).__init__(thresh=thresh, *args, **kwargs)
-
     @staticmethod
-    def segment_approach(image, *,
-                         thresh: float = -6):
-        """Mask retrieval as it is done in Shape-In
+    def segment_algorithm(image, *,
+                          thresh: float = -6):
+        """Mask retrieval using basic thresholding
 
         Parameters
         ----------
@@ -39,7 +27,7 @@ class SegmentThresh(CPUSegmenter):
         Returns
         -------
         mask: 2d boolean ndarray
-            Mask image for the give index
+            Mask image for the given index
         """
         assert thresh < 0, "threshold values above zero not supported!"
         return image < thresh

dcnum/segm/segm_torch/__init__.py

@@ -0,0 +1,19 @@
+import importlib
+
+try:
+    torch = importlib.import_module("torch")
+    req_maj = 2
+    req_min = 3
+    ver_tuple = torch.__version__.split(".")
+    act_maj = int(ver_tuple[0])
+    act_min = int(ver_tuple[1])
+    if act_maj < req_maj or (act_maj == req_maj and act_min < req_min):
+        raise ValueError(f"Your PyTorch version {act_maj}.{act_min} is not "
+                         f"supported, please update to at least "
+                         f"{req_maj}.{req_min}")
+except ImportError:
+    pass
+else:
+    from .segm_torch_mpo import SegmentTorchMPO  # noqa: F401
+    if torch.cuda.is_available():
+        from .segm_torch_sto import SegmentTorchSTO  # noqa: F401

dcnum/segm/segm_torch/segm_torch_base.py

@@ -0,0 +1,125 @@
+import functools
+import pathlib
+import re
+from typing import Dict
+
+from ...meta import paths
+
+from ..segmenter import Segmenter, SegmenterNotApplicableError
+
+from .torch_model import load_model
+
+
+class TorchSegmenterBase(Segmenter):
+    """Torch segmenters that use a pretrained model for segmentation"""
+    requires_background_correction = False
+    mask_postprocessing = True
+    mask_default_kwargs = {
+        "clear_border": True,
+        "fill_holes": True,
+        "closing_disk": 0,
+    }
+
+    @classmethod
+    def get_ppid_from_ppkw(cls, kwargs, kwargs_mask=None):
+        kwargs_new = kwargs.copy()
+        # Make sure that the `model_file` kwarg is actually just a filename
+        # so that the pipeline identifier only contains the name, but not
+        # the full path.
+        if "model_file" in kwargs:
+            model_file = kwargs["model_file"]
+            mpath = pathlib.Path(model_file)
+            if mpath.exists():
+                # register the location of the file in the search path
+                # registry so other threads/processes will find it.
+                paths.register_search_path("torch_model_files", mpath.parent)
+                kwargs_new["model_file"] = mpath.name
+        return super(TorchSegmenterBase, cls).get_ppid_from_ppkw(kwargs_new,
+                                                                 kwargs_mask)
+
+    @classmethod
+    def validate_applicability(cls,
+                               segmenter_kwargs: Dict,
+                               meta: Dict = None,
+                               logs: Dict = None):
+        """Validate the applicability of this segmenter for a dataset
+
+        The applicability is defined by the metadata in the segmentation
+        model.
+
+        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
+        ------
+        SegmenterNotApplicable
+            If the segmenter is not applicable to the dataset
+        """
+        if "model_file" not in segmenter_kwargs:
+            raise ValueError("A `model_file` must be provided in the "
+                             "`segmenter_kwargs` to validate applicability")
+
+        model_file = segmenter_kwargs["model_file"]
+        _, model_meta = load_model(model_file, device="cpu")
+
+        reasons_list = []
+        validators = {
+            "meta": functools.partial(
+                cls._validate_applicability_item,
+                data_dict=meta,
+                reasons_list=reasons_list),
+            "logs": functools.partial(
+                cls._validate_applicability_item,
+                # convert logs to strings
+                data_dict={key: "\n".join(val) for key, val in logs.items()},
+                reasons_list=reasons_list)
+        }
+        for item in model_meta.get("validation", []):
+            it = item["type"]
+            if it in validators:
+                validators[it](item)
+            else:
+                reasons_list.append(
+                    f"invalid validation type {it} in {model_file}")
+
+        if reasons_list:
+            raise SegmenterNotApplicableError(segmenter_class=cls,
+                                              reasons_list=reasons_list)
+
+        return True
+
+    @staticmethod
+    def _validate_applicability_item(item, data_dict, reasons_list):
+        """Populate `reasons_list` with invalid entries
+
+        Example `data_dict`::
+
+            {"type": "meta",
+             "key": "setup:region",
+             "allow-missing-key": False,
+             "regexp": "^channel$",
+             "regexp-negate": False,
+             "reason": "only channel region supported",
+             }
+        """
+        key = item["key"]
+        if key in data_dict:
+            regexp = re.compile(item["regexp"])
+            matched = bool(regexp.match(data_dict[key]))
+            negate = item.get("regexp-negate", False)
+            valid = matched if not negate else not matched
+            if not valid:
+                reasons_list.append(item.get("reason", "unknown reason"))
+        elif not item.get("allow-missing-key", False):
+            reasons_list.append(f"Key '{key}' missing in {item['type']}")

dcnum/segm/segm_torch/segm_torch_mpo.py

@@ -0,0 +1,71 @@
+import numpy as np
+import torch
+
+from ..segmenter_mpo import MPOSegmenter
+
+from .segm_torch_base import TorchSegmenterBase
+from .torch_model import load_model
+from .torch_preproc import preprocess_images
+from .torch_postproc import postprocess_masks
+
+
+class SegmentTorchMPO(TorchSegmenterBase, MPOSegmenter):
+    """PyTorch segmentation (multiprocessing version)"""
+
+    @staticmethod
+    def segment_algorithm(image, *,
+                          model_file: str = None):
+        """
+        Parameters
+        ----------
+        image: 2d ndarray
+            event image
+        model_file: str
+            path to or name of a dcnum model file (.dcnm); if only a
+            name is provided, then the "torch_model_files" directory
+            paths are searched for the file name
+
+        Returns
+        -------
+        mask: 2d boolean or integer ndarray
+            mask or labeling image for the give index
+        """
+        if model_file is None:
+            raise ValueError("Please specify a .dcnm model file!")
+
+        # Set number of pytorch threads to 1, because dcnum is doing
+        # all the multiprocessing.
+        # https://pytorch.org/docs/stable/generated/torch.set_num_threads.html#torch.set_num_threads
+        torch.set_num_threads(1)
+        device = torch.device("cpu")
+
+        # Load model and metadata
+        model, model_meta = load_model(model_file, device)
+
+        image_preproc = preprocess_images(image[np.newaxis, :, :],
+                                          **model_meta["preprocessing"])
+
+        image_ten = torch.from_numpy(image_preproc)
+
+        # Move image tensors to device
+        image_ten_on_device = image_ten.to(device)
+        # Model inference
+        pred_tensor = model(image_ten_on_device)
+
+        # Convert cuda-tensor into numpy mask array. The `pred_tensor`
+        # array is still of the shape (1, 1, H, W). The `masks`
+        # array is of shape (1, H, W). We can optionally label it
+        # here (we have to if the shapes don't match) or do it in
+        # postprocessing.
+        masks = pred_tensor.detach().cpu().numpy()[0] >= 0.5
+
+        # Perform postprocessing in cases where the image shapes don't match
+        assert len(masks[0].shape) == len(image.shape), "sanity check"
+        if masks[0].shape != image.shape:
+            labels = postprocess_masks(
+                masks=masks,
+                original_image_shape=image.shape,
+            )
+            return labels[0]
+        else:
+            return masks[0]

dcnum/segm/segm_torch/segm_torch_sto.py

@@ -0,0 +1,88 @@
+from dcnum.segm import STOSegmenter
+import numpy as np
+import torch
+
+from .segm_torch_base import TorchSegmenterBase
+from .torch_model import load_model
+from .torch_preproc import preprocess_images
+from .torch_postproc import postprocess_masks
+
+
+class SegmentTorchSTO(TorchSegmenterBase, STOSegmenter):
+    """PyTorch segmentation (GPU version)"""
+
+    @staticmethod
+    def _segment_in_batches(imgs_t, model, batch_size, device):
+        """Segment image data in batches"""
+        size = len(imgs_t)
+        # Create empty array to fill up with segmented batches
+        masks = np.empty((len(imgs_t), *imgs_t[0].shape[-2:]),
+                         dtype=bool)
+
+        for start_idx in range(0, size, batch_size):
+            batch = imgs_t[start_idx:start_idx + batch_size]
+            # Move image tensors to cuda
+            batch = torch.tensor(batch, device=device)
+            # Model inference
+            batch_seg = model(batch)
+            # Remove extra dim [B, C, H, W] --> [B, H, W]
+            batch_seg = batch_seg.squeeze(1)
+            # Convert cuda-tensor into numpy arrays
+            batch_seg_np = batch_seg.detach().cpu().numpy()
+            # Fill empty array with segmented batch
+            masks[start_idx:start_idx + batch_size] = batch_seg_np >= 0.5
+
+        return masks
+
+    @staticmethod
+    def segment_algorithm(images, gpu_id=None, batch_size=50, *,
+                          model_file: str = None):
+        """
+        Parameters
+        ----------
+        images: 3d ndarray
+            array of N event images of shape (N, H, W)
+        gpu_id: str
+            optional argument specifying the GPU to use
+        batch_size: int
+            number of images to process in one batch
+        model_file: str
+            path to or name of a dcnum model file (.dcnm); if only a
+            name is provided, then the "torch_model_files" directory
+            paths are searched for the file name
+
+        Returns
+        -------
+        mask: 2d boolean or integer ndarray
+            mask or label images of shape (N, H, W)
+        """
+        if model_file is None:
+            raise ValueError("Please specify a model file!")
+
+        # Determine device to use
+        device = torch.device(gpu_id if gpu_id is not None else "cuda")
+
+        # Load model and metadata
+        model, model_meta = load_model(model_file, device)
+
+        # Preprocess the images
+        image_preproc = preprocess_images(images,
+                                          **model_meta["preprocessing"])
+        # Model inference
+        # The `masks` array has the shape (len(images), H, W), where
+        # H and W may be different from the corresponding axes in `images`.
+        masks = SegmentTorchSTO._segment_in_batches(image_preproc,
+                                                    model,
+                                                    batch_size,
+                                                    device
+                                                    )
+
+        # Perform postprocessing in cases where the image shapes don't match
+        assert len(masks.shape[1:]) == len(images.shape[1:]), "sanity check"
+        if masks.shape[1:] != images.shape[1:]:
+            labels = postprocess_masks(
+                masks=masks,
+                original_image_shape=images.shape[1:])
+            return labels
+        else:
+            return masks

dcnum/segm/segm_torch/torch_model.py

@@ -0,0 +1,95 @@
+import errno
+import functools
+import hashlib
+import json
+import logging
+import os
+import pathlib
+
+import torch
+
+from ...meta import paths
+
+
+logger = logging.getLogger(__name__)
+
+
+def check_md5sum(path):
+    """Verify the last five characters of the file stem with its MD5 hash"""
+    md5 = hashlib.md5(path.read_bytes()).hexdigest()
+    if md5[:5] != path.stem.split("_")[-1]:
+        raise ValueError(f"MD5 mismatch for {path} ({md5})! Expected the "
+                         f"input file to end with '{md5[:5]}{path.suffix}'.")
+
+
+@functools.cache
+def load_model(path_or_name, device):
+    """Load a PyTorch model + metadata from a TorchScript jit checkpoint
+
+    Parameters
+    ----------
+    path_or_name: str or pathlib.Path
+        jit checkpoint file; For dcnum, these files have the suffix .dcnm
+        and contain a special `_extra_files["dcnum_meta.json"]` extra
+        file that can be loaded via `torch.jit.load` (see below).
+    device: str or torch.device
+        device on which to run the model
+
+    Returns
+    -------
+    model_jit: torch.jit.ScriptModule
+        loaded PyTorch model stored as a TorchScript module
+    model_meta: dict
+        metadata associated with the loaded model
+    """
+    model_path = retrieve_model_file(path_or_name)
+    # define an extra files mapping dictionary that loads the model's metadata
+    extra_files = {"dcnum_meta.json": ""}
+    # load model
+    model_jit = torch.jit.load(model_path,
+                               _extra_files=extra_files,
+                               map_location=device)
+    # load model metadata
+    model_meta = json.loads(extra_files["dcnum_meta.json"])
+    # set model to evaluation mode
+    model_jit.eval()
+    # optimize for inference on device
+    model_jit = torch.jit.optimize_for_inference(model_jit)
+    return model_jit, model_meta
+
+
+@functools.cache
+def retrieve_model_file(path_or_name):
+    """Retrieve a dcnum torch model file
+
+    If a path to a model is given, then this path is returned directly.
+    If a file name is given, then look for the file with
+    :func:`dcnum.meta.paths.find_file` using the "torch_model_file"
+    topic.
+    """
+    # Did the user already pass a path?
+    if isinstance(path_or_name, pathlib.Path):
+        if path_or_name.exists():
+            path = path_or_name
+        else:
+            try:
+                return retrieve_model_file(path_or_name.name)
+            except BaseException:
+                raise FileNotFoundError(errno.ENOENT,
+                                        os.strerror(errno.ENOENT),
+                                        str(path_or_name))
+    elif isinstance(path_or_name, str):
+        name = path_or_name.strip()
+        # We now have a string for a filename, and we have to figure out what
+        # the path is. There are several options, including cached files.
+        if pathlib.Path(name).exists():
+            path = pathlib.Path(name)
+        else:
+            path = paths.find_file("torch_model_files", name)
+    else:
+        raise ValueError(
+            f"Please pass a string or a path, got {type(path_or_name)}!")
+
+    logger.info(f"Found dcnum model file {path}")
+    check_md5sum(path)
+    return path

dcnum/segm/segm_torch/torch_postproc.py

@@ -0,0 +1,93 @@
+from typing import Tuple
+
+from ..segmenter import Segmenter
+
+import numpy as np
+from scipy import ndimage as ndi
+
+
+def postprocess_masks(masks,
+                      original_image_shape: Tuple[int, int]):
+    """Postprocess mask images from ML segmenters
+
+    The transformation includes:
+    - Revert the cropping and padding operations done in
+      :func:`.preprocess_images` by padding with zeros and cropping.
+    - If the original image shape is larger than the mask image shape,
+      also clear borders in an intermediate step
+      (maks postprocessing using :func:`Segmenter.process_mask`).
+
+    Parameters
+    ----------
+    masks: 3d or 4d ndarray
+        Mask data in shape (batch_size, 1, imagex_size, imagey_size)
+        or (batch_size, imagex_size, imagey_size).
+    original_image_shape: tuple of (int, int)
+        The required output mask shape for one event. This required for
+        doing the inverse of what is done in :func:`.preprocess_images`.
+
+    Returns
+    -------
+    labels_proc: np.ndarray
+        An integer array with the same dimensions as the original image
+        data passed to :func:`.preprocess_images`. The shape of this array
+        is (batch_size, original_image_shape[0], original_image_shape[1]).
+    """
+    # If output of model is 4d, remove channel axis
+    if len(masks.shape) == 4:
+        masks = masks[:, 0, :, :]
+
+    # Label the mask image
+    labels = np.empty(masks.shape, dtype=np.uint16)
+    label_struct = ndi.generate_binary_structure(2, 2)
+    for ii in range(masks.shape[0]):
+        ndi.label(
+            input=masks[ii],
+            output=labels[ii],
+            structure=label_struct)
+
+    batch_size = labels.shape[0]
+
+    # Revert padding and cropping from preprocessing
+    mask_shape_ret = labels.shape[1:]
+    # height
+    s0diff = original_image_shape[0] - mask_shape_ret[0]
+    s0t = abs(s0diff) // 2
+    s0b = abs(s0diff) - s0t
+    # width
+    s1diff = original_image_shape[1] - mask_shape_ret[1]
+    s1l = abs(s1diff) // 2
+    s1r = abs(s1diff) - s1l
+
+    if s0diff > 0 or s1diff > 0:
+        # The masks that we have must be padded. Before we do that, we have
+        # to remove events on the edges, otherwise we will have half-segmented
+        # cell events in the output array.
+        for ii in range(batch_size):
+            labels[ii] = Segmenter.process_mask(labels[ii],
+                                                clear_border=True,
+                                                fill_holes=False,
+                                                closing_disk=0)
+
+    # Crop first, only then pad.
+    if s1diff > 0:
+        labels_pad = np.zeros((batch_size,
+                              labels.shape[1],
+                              original_image_shape[1]),
+                              dtype=np.uint16)
+        labels_pad[:, :, s1l:-s1r] = labels
+        labels = labels_pad
+    elif s1diff < 0:
+        labels = labels[:, :, s1l:-s1r]
+
+    if s0diff > 0:
+        labels_pad = np.zeros((batch_size,
+                              original_image_shape[0],
+                              original_image_shape[1]),
+                              dtype=np.uint16)
+        labels_pad[:, s0t:-s0b, :] = labels
+        labels = labels_pad
+    elif s0diff < 0:
+        labels = labels[:, s0t:-s0b, :]
+
+    return labels