Repository: rapidsai/cucim (language: python)

Issue:
[FEA] add cucim.skimage.segmentation.expand_labels
Once #318 is merged, we can add the `expand_labels` function. I will make a PR for this soon.
Add Euclidean distance transform for images/volumes
### Overview

closes #319

This PR adds an implementation of `scipy.ndimage.distance_transform_edt`. For each foreground pixel in a binary image, this function computes the minimal Euclidean distance to reach a background pixel. This is a SciPy function rather than a scikit-image one so I have put it under `cucim.core` instead of `cucim.skimage`. Longer term we could move this upstream to CuPy, but I think we want to have a copy here so that we can use it in the near term to implement some of the missing scikit-image functionality. This function is used by the following functions in the scikit-image API, so this PR will help us implement these in future PRs:

`skimage.morphology.medial_axis` (**update**: see #342)
`skimage.segmentation.expand_labels` (**update**: see #341)
`skimage.segmentation.chan_vese` (**update**: see #343)
It is also used in examples related to `skimage.segmentation.watershed` and is needed for implementation of ITK's `SignedMaurerDistanceMapImageFilter` for [itk-cucim](https://github.com/InsightSoftwareConsortium/itk_cucim/issues/13).

The algorithm used here is adapted from the C++ code in the [PBA+ repository](https://github.com/orzzzjq/Parallel-Banding-Algorithm-plus) (MIT-licensed)

### extensions made to PBA+ kernel source
- larger sizes can be supported due to runtime code generation of the integer datatypes and encode/decode functions.

### current known limitations
- 2D and 3D only (this should cover the most common use cases)
- The `sampling` argument is not fully supported. This can likely be done with a bit of effort, but will require modifications to the CUDA kernels
- User-specified output arguments are not currently supported. We can potentially add this in the future
- The indices returned by `return_indices=True` are equally valid, but not always identical to those returned by SciPy. This is because there can be multiple indices with an identical distance, so which one gets chosen in that case is implementation-dependent.

### initial benchmarks relative to SciPy

Here `% true` is the percentage of the image that corresponds to foreground pixels. Even for fairly small images there is some benefit, with the benefit becoming two orders of magnitude at larger sizes.

shape      | % true |  cuCIM |  SciPy  | acceleration
-----------------|--------|---------|---------|--------------
(256, 256) | 5 | 0.00108 | 0.00353 | 3.25
(512, 512) | 5 | 0.00108 | 0.01552 | 14.42
(2048, 2048) | 5 | 0.00252 | 0.34434 | 136.86
(4096, 4096) | 5 | 0.00765 | 1.58948 | 207.88
(32, 32, 32) | 5 | 0.00103 | 0.00305 | 2.98
(128, 128, 128) | 5 | 0.00153 | 0.30103 | 196.26
(256, 256, 256) | 5 | 0.00763 | 3.17872 | 416.37
(384, 384, 384) | 5 | 0.02460 | 14.28779 | 580.89
(256, 256) | 50 | 0.00107 | 0.00430 | 4.01
(512, 512) | 50 | 0.00109 | 0.01878 | 17.30
(2048, 2048) | 50 | 0.00299 | 0.39304 | 131.60
(4096, 4096) | 50 | 0.00896 | 1.84686 | 206.19
(32, 32, 32) | 50 | 0.00102 | 0.00361 | 3.53
(128, 128, 128) | 50 | 0.00163 | 0.31657 | 194.66
(256, 256, 256) | 50 | 0.00914 | 3.35914 | 367.49
(384, 384, 384) | 50 | 0.03005 | 13.83219 | 460.30
(256, 256) | 95 | 0.00107 | 0.00344 | 3.22
(512, 512) | 95 | 0.00108 | 0.01638 | 15.17
(2048, 2048) | 95 | 0.00314 | 0.36996 | 117.81
(4096, 4096) | 95 | 0.00943 | 1.90475 | 202.05
(32, 32, 32) | 95 | 0.00102 | 0.00314 | 3.07
(128, 128, 128) | 95 | 0.00180 | 0.28843 | 159.80
(256, 256, 256) | 95 | 0.01073 | 3.23450 | 301.41
(384, 384, 384) | 95 | 0.03577 | 12.40526 | 346.82

### other comments
- can likely reduce memory overhead and improve performance a bit by refactoring some of the pre/post-processing code into elementwise kernels. (e.g. `encode3d`/`decode3d`, etc.)

- (JK) This may be able to leverage CuPy in the future ( https://github.com/cupy/cupy/issues/6919 )



Discussion / hints:

Fix patch:
diff --git a/python/cucim/src/cucim/skimage/segmentation/__init__.py b/python/cucim/src/cucim/skimage/segmentation/__init__.py
index b5527b8a1..154467ee7 100644
--- a/python/cucim/src/cucim/skimage/segmentation/__init__.py
+++ b/python/cucim/src/cucim/skimage/segmentation/__init__.py
@@ -1,4 +1,5 @@
 from ._clear_border import clear_border
+from ._expand_labels import expand_labels
 from ._join import join_segmentations, relabel_sequential
 from .boundaries import find_boundaries, mark_boundaries
 from .morphsnakes import (checkerboard_level_set, disk_level_set,
@@ -7,6 +8,7 @@
 from .random_walker_segmentation import random_walker
 
 __all__ = [
+    "expand_labels",
     "random_walker",
     "find_boundaries",
     "mark_boundaries",
diff --git a/python/cucim/src/cucim/skimage/segmentation/_expand_labels.py b/python/cucim/src/cucim/skimage/segmentation/_expand_labels.py
new file mode 100644
index 000000000..8ec47a7be
--- /dev/null
+++ b/python/cucim/src/cucim/skimage/segmentation/_expand_labels.py
@@ -0,0 +1,96 @@
+import cupy as cp
+
+from cucim.core.operations.morphology import distance_transform_edt
+
+
+def expand_labels(label_image, distance=1):
+    """Expand labels in label image by ``distance`` pixels without overlapping.
+
+    Given a label image, ``expand_labels`` grows label regions (connected components)
+    outwards by up to ``distance`` pixels without overflowing into neighboring regions.
+    More specifically, each background pixel that is within Euclidean distance
+    of <= ``distance`` pixels of a connected component is assigned the label of that
+    connected component.
+    Where multiple connected components are within ``distance`` pixels of a background
+    pixel, the label value of the closest connected component will be assigned (see
+    Notes for the case of multiple labels at equal distance).
+
+    Parameters
+    ----------
+    label_image : ndarray of dtype int
+        label image
+    distance : float
+        Euclidean distance in pixels by which to grow the labels. Default is one.
+
+    Returns
+    -------
+    enlarged_labels : ndarray of dtype int
+        Labeled array, where all connected regions have been enlarged
+
+    Notes
+    -----
+    Where labels are spaced more than ``distance`` pixels are apart, this is
+    equivalent to a morphological dilation with a disc or hyperball of radius ``distance``.
+    However, in contrast to a morphological dilation, ``expand_labels`` will
+    not expand a label region into a neighboring region.  
+
+    This implementation of ``expand_labels`` is derived from CellProfiler [1]_, where
+    it is known as module "IdentifySecondaryObjects (Distance-N)" [2]_.
+
+    There is an important edge case when a pixel has the same distance to
+    multiple regions, as it is not defined which region expands into that
+    space. Here, the exact behavior depends on the upstream implementation
+    of ``scipy.ndimage.distance_transform_edt``.
+
+    See Also
+    --------
+    :func:`cucim.skimage.measure.label`, :func:`cucim.skimage.morphology.dilation`  # noqa
+
+    References
+    ----------
+    .. [1] https://cellprofiler.org
+    .. [2] https://github.com/CellProfiler/CellProfiler/blob/082930ea95add7b72243a4fa3d39ae5145995e9c/cellprofiler/modules/identifysecondaryobjects.py#L559  # noqa
+
+    Examples
+    --------
+    >>> labels = np.array([0, 1, 0, 0, 0, 0, 2])
+    >>> expand_labels(labels, distance=1)
+    array([1, 1, 1, 0, 0, 2, 2])
+
+    Labels will not overwrite each other:
+
+    >>> expand_labels(labels, distance=3)
+    array([1, 1, 1, 1, 2, 2, 2])
+
+    In case of ties, behavior is undefined, but currently resolves to the
+    label closest to ``(0,) * ndim`` in lexicographical order.
+
+    >>> labels_tied = np.array([0, 1, 0, 2, 0])
+    >>> expand_labels(labels_tied, 1)
+    array([1, 1, 1, 2, 2])
+    >>> labels2d = np.array(
+    ...     [[0, 1, 0, 0],
+    ...      [2, 0, 0, 0],
+    ...      [0, 3, 0, 0]]
+    ... )
+    >>> expand_labels(labels2d, 1)
+    array([[2, 1, 1, 0],
+           [2, 2, 0, 0],
+           [2, 3, 3, 0]])
+    """
+
+    distances, nearest_label_coords = distance_transform_edt(
+        label_image == 0, return_indices=True
+    )
+    labels_out = cp.zeros_like(label_image)
+    dilate_mask = distances <= distance
+    # build the coordinates to find nearest labels,
+    # in contrast to [1] this implementation supports label arrays
+    # of any dimension
+    masked_nearest_label_coords = [
+        dimension_indices[dilate_mask]
+        for dimension_indices in nearest_label_coords
+    ]
+    nearest_labels = label_image[tuple(masked_nearest_label_coords)]
+    labels_out[dilate_mask] = nearest_labels
+    return labels_out


Test patch:
diff --git a/python/cucim/src/cucim/skimage/segmentation/tests/test_expand_labels.py b/python/cucim/src/cucim/skimage/segmentation/tests/test_expand_labels.py
new file mode 100644
index 000000000..2fe1df30c
--- /dev/null
+++ b/python/cucim/src/cucim/skimage/segmentation/tests/test_expand_labels.py
@@ -0,0 +1,226 @@
+import cupy as cp
+import pytest
+from cupy.testing import assert_array_equal
+
+from cucim.core.operations.morphology import distance_transform_edt
+from cucim.skimage import data, measure
+from cucim.skimage.segmentation import expand_labels
+
+SAMPLE1D = cp.array([0, 0, 4, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0])
+SAMPLE1D_EXPANDED_3 = cp.array(
+    [4, 4, 4, 4, 4, 4, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
+)
+
+# Some pixels are important edge cases with undefined behaviour:
+# these are the pixels that are at the same distance from
+# multiple labels. Ideally the label would be chosen at random
+# to avoid bias, but as we are relying on the index map returned
+# by the scipy.ndimage distance transform, what actually happens
+# is determined by the upstream implementation of the distance
+# tansform, thus we don't give any guarantees for the edge case pixels.
+#
+# Regardless, it seems prudent to have a test including an edge case
+# so we can detect whether future upstream changes in scipy.ndimage
+# modify the behaviour.
+
+EDGECASE1D = cp.array([0, 0, 4, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0])
+EDGECASE1D_EXPANDED_3 = cp.array(
+    [4, 4, 4, 4, 4, 4, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
+)
+
+SAMPLE2D = cp.array(
+    [
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0],
+        [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
+        [0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 1, 0, 0, 0, 0, 0, 0, 2, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    ]
+)
+
+SAMPLE2D_EXPANDED_3 = cp.array(
+    [
+        [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1, 1, 1, 0, 0, 2, 0],
+        [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
+        [1, 1, 1, 1, 1, 1, 0, 2, 2, 2, 2],
+        [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
+        [1, 1, 1, 1, 1, 0, 2, 2, 2, 2, 2],
+        [1, 1, 1, 1, 1, 0, 0, 2, 2, 2, 2],
+        [0, 0, 1, 0, 0, 0, 0, 2, 2, 2, 2],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0],
+    ]
+)
+
+# non-integer expansion
+SAMPLE2D_EXPANDED_1_5 = cp.array(
+    [
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1, 0, 0, 0, 2, 2, 2],
+        [1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2],
+        [0, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2],
+        [0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    ]
+)
+
+
+EDGECASE2D = cp.array(
+    [
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 2, 2, 0, 0, 0, 0],
+        [0, 0, 1, 1, 0, 2, 2, 0, 0, 0, 0],
+        [0, 1, 1, 1, 0, 2, 0, 0, 0, 0, 0],
+        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+    ]
+)
+
+EDGECASE2D_EXPANDED_4 = cp.array(
+    [
+        [1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 0],
+        [1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2],
+        [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
+        [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 0],
+        [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 0],
+    ]
+)
+
+SAMPLE3D = cp.array(
+    [
+        [
+            [0, 0, 0, 0],
+            [0, 3, 0, 0],
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+        ],
+
+        [
+            [0, 0, 0, 0],
+            [0, 3, 3, 0],
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+        ],
+
+        [
+            [0, 0, 0, 0],
+            [0, 3, 0, 0],
+            [0, 0, 0, 0],
+            [0, 0, 5, 0],
+        ],
+
+        [
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+            [0, 0, 5, 0],
+        ]
+    ]
+)
+
+SAMPLE3D_EXPANDED_2 = cp.array(
+    [
+        [
+            [3, 3, 3, 3],
+            [3, 3, 3, 3],
+            [3, 3, 3, 3],
+            [0, 3, 5, 0],
+        ],
+
+        [
+            [3, 3, 3, 3],
+            [3, 3, 3, 3],
+            [3, 3, 3, 3],
+            [0, 5, 5, 5],
+        ],
+
+        [
+            [3, 3, 3, 3],
+            [3, 3, 3, 3],
+            [3, 3, 5, 5],
+            [5, 5, 5, 5],
+        ],
+
+        [
+            [3, 3, 3, 0],
+            [3, 3, 3, 0],
+            [3, 3, 5, 5],
+            [5, 5, 5, 5],
+        ]
+    ]
+)
+
+SAMPLE_EDGECASE_BEHAVIOUR = cp.array(
+    [[0, 1, 0, 0], [2, 0, 0, 0], [0, 3, 0, 0]]
+)
+
+
+@pytest.mark.parametrize(
+    "input_array, expected_output, expand_distance",
+    [
+        (SAMPLE1D, SAMPLE1D_EXPANDED_3, 3),
+        (SAMPLE2D, SAMPLE2D_EXPANDED_3, 3),
+        (SAMPLE2D, SAMPLE2D_EXPANDED_1_5, 1.5),
+        (EDGECASE1D, EDGECASE1D_EXPANDED_3, 3),
+        (EDGECASE2D, EDGECASE2D_EXPANDED_4, 4),
+        (SAMPLE3D, SAMPLE3D_EXPANDED_2, 2)
+    ]
+)
+def test_expand_labels(input_array, expected_output, expand_distance):
+    if input_array.ndim == 1:
+        with pytest.raises(NotImplementedError):
+            expand_labels(input_array, expand_distance)
+    else:
+        expanded = expand_labels(input_array, expand_distance)
+        assert_array_equal(expanded, expected_output)
+
+
+@pytest.mark.parametrize('ndim', [2, 3])
+@pytest.mark.parametrize('distance', range(6))
+def test_binary_blobs(ndim, distance):
+    """Check some invariants with label expansion.
+
+    - New labels array should exactly contain the original labels array.
+    - Distance to old labels array within new labels should never exceed input
+      distance.
+    - Distance beyond the expanded labels should always exceed the input
+      distance.
+    """
+    img = data.binary_blobs(length=64, blob_size_fraction=0.05, n_dim=ndim)
+    labels = measure.label(img)
+    expanded = expand_labels(labels, distance=distance)
+    original_mask = labels != 0
+    assert_array_equal(labels[original_mask], expanded[original_mask])
+    expanded_only_mask = (expanded - labels).astype(bool)
+    distance_map = distance_transform_edt(~original_mask)
+    expanded_distances = distance_map[expanded_only_mask]
+    if expanded_distances.size > 0:
+        assert cp.all(expanded_distances <= distance)
+    beyond_expanded_distances = distance_map[~expanded.astype(bool)]
+    if beyond_expanded_distances.size > 0:
+        assert cp.all(beyond_expanded_distances > distance)
+
+
+def test_edge_case_behaviour():
+    """ Check edge case behavior to detect upstream changes
+
+    For edge cases where a pixel has the same distance to several regions,
+    lexicographical order seems to determine which region gets to expand
+    into this pixel given the current upstream behaviour in
+    scipy.ndimage.distance_map_edt.
+
+    As a result, we expect different results when transposing the array.
+    If this test fails, something has changed upstream.
+    """
+    expanded = expand_labels(SAMPLE_EDGECASE_BEHAVIOUR, 1)
+    expanded_transpose = expand_labels(SAMPLE_EDGECASE_BEHAVIOUR.T, 1)
+    assert not cp.all(expanded == expanded_transpose.T)