Repository: cupy/cupy

Issue:
Support `scipy.ndimage.label`
Per request:

https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.label.html

@anaruse 


Discussion:
I would like to implement this.

Will be raising a PR soon.
Do I need to implement https://github.com/scipy/scipy/blob/e7d5dd4a77a2c330a467fde773e267d5cbb90f13/scipy/ndimage/src/_ni_label.pyx#L204 here or I should directly use `scipy.ndimage.ni_label._label`? @emcastillo 
You have to re-implement it using cuda kernels and/or cupy arrays.
Agree this would be a nice feature to have 😄

cc @jni @GenevieveBuckley
@rushabh-v, I suppose you have already started working on this, but let me do this one. This is needed for us immediately. Thanks.
Okay.
Even I will be busier for a few more days, so you are welcome to go ahead with this 👍
Thanks!
Thanks for working on this Akira! 😄 Would be interested in looking at a PR once it is ready 😉

Resolving patch:
diff --git a/cupyx/scipy/ndimage/__init__.py b/cupyx/scipy/ndimage/__init__.py
--- a/cupyx/scipy/ndimage/__init__.py
+++ b/cupyx/scipy/ndimage/__init__.py
@@ -6,3 +6,5 @@
 from cupyx.scipy.ndimage.interpolation import rotate  # NOQA
 from cupyx.scipy.ndimage.interpolation import shift  # NOQA
 from cupyx.scipy.ndimage.interpolation import zoom  # NOQA
+
+from cupyx.scipy.ndimage.measurements import label  # NOQA
diff --git a/cupyx/scipy/ndimage/measurements.py b/cupyx/scipy/ndimage/measurements.py
new file mode 100644
--- /dev/null
+++ b/cupyx/scipy/ndimage/measurements.py
@@ -0,0 +1,202 @@
+import numpy
+
+import cupy
+
+
+def label(input, structure=None, output=None):
+    """Labels features in an array
+
+    Args:
+        input (cupy.ndarray): The input array.
+        structure (array_like or None): A structuring element that defines
+            feature connections. ```structure``` must be centersymmetric. If
+            None, structure is automatically generated with a squared
+            connectivity equal to one.
+        output (cupy.ndarray, dtype or None): The array in which to place the
+            output.
+
+    Returns:
+        label (cupy.ndarray): An integer array where each unique feature in
+            ```input``` has a unique label in the array.
+        num_features (int): Number of features found.
+
+    .. warning::
+
+        This function may synchronize the device.
+
+    .. seealso:: :func:`scipy.ndimage.label`
+    """
+    if not isinstance(input, cupy.ndarray):
+        raise TypeError('input must be cupy.ndarray')
+    if input.dtype.char in 'FD':
+        raise TypeError('Complex type not supported')
+    if structure is None:
+        structure = _generate_binary_structure(input.ndim, 1)
+    elif isinstance(structure, cupy.ndarray):
+        structure = cupy.asnumpy(structure)
+    structure = numpy.array(structure, dtype=bool)
+    if structure.ndim != input.ndim:
+        raise RuntimeError('structure and input must have equal rank')
+    for i in structure.shape:
+        if i != 3:
+            raise ValueError('structure dimensions must be equal to 3')
+
+    if isinstance(output, cupy.ndarray):
+        if output.shape != input.shape:
+            raise ValueError("output shape not correct")
+        caller_provided_output = True
+    else:
+        caller_provided_output = False
+        if output is None:
+            output = cupy.empty(input.shape, numpy.int32)
+        else:
+            output = cupy.empty(input.shape, output)
+
+    if input.size == 0:
+        # empty
+        maxlabel = 0
+    elif input.ndim == 0:
+        # 0-dim array
+        maxlabel = 0 if input.item() == 0 else 1
+        output[...] = maxlabel
+    else:
+        if output.dtype != numpy.int32:
+            y = cupy.empty(input.shape, numpy.int32)
+        else:
+            y = output
+        maxlabel = _label(input, structure, y)
+        if output.dtype != numpy.int32:
+            output[...] = y[...]
+
+    if caller_provided_output:
+        return maxlabel
+    else:
+        return output, maxlabel
+
+
+def _generate_binary_structure(rank, connectivity):
+    if connectivity < 1:
+        connectivity = 1
+    if rank < 1:
+        return numpy.array(True, dtype=bool)
+    output = numpy.fabs(numpy.indices([3] * rank) - 1)
+    output = numpy.add.reduce(output, 0)
+    return output <= connectivity
+
+
+def _label(x, structure, y):
+    elems = numpy.where(structure != 0)
+    vecs = [elems[dm] - 1 for dm in range(x.ndim)]
+    offset = vecs[0]
+    for dm in range(1, x.ndim):
+        offset = offset * 3 + vecs[dm]
+    indxs = numpy.where(offset < 0)[0]
+    dirs = [[vecs[dm][dr] for dm in range(x.ndim)] for dr in indxs]
+    dirs = cupy.array(dirs, dtype=numpy.int32)
+    ndirs = indxs.shape[0]
+    y_shape = cupy.array(y.shape, dtype=numpy.int32)
+    count = cupy.zeros(2, dtype=numpy.int32)
+    _kernel_init()(x, y)
+    _kernel_connect()(y_shape, dirs, ndirs, x.ndim, y, size=y.size)
+    _kernel_count()(y, count, size=y.size)
+    maxlabel = int(count[0])
+    labels = cupy.empty(maxlabel, dtype=numpy.int32)
+    _kernel_labels()(y, count, labels, size=y.size)
+    _kernel_finalize()(maxlabel, cupy.sort(labels), y, size=y.size)
+    return maxlabel
+
+
+def _kernel_init():
+    return cupy.ElementwiseKernel(
+        'X x', 'Y y', 'if (x == 0) { y = -1; } else { y = i; }',
+        'cupyx_nd_label_init')
+
+
+def _kernel_connect():
+    return cupy.ElementwiseKernel(
+        'raw int32 shape, raw int32 dirs, int32 ndirs, int32 ndim',
+        'raw Y y',
+        '''
+        if (y[i] < 0) continue;
+        for (int dr = 0; dr < ndirs; dr++) {
+            int j = i;
+            int rest = j;
+            int stride = 1;
+            int k = 0;
+            for (int dm = ndim-1; dm >= 0; dm--) {
+                int pos = rest % shape[dm] + dirs[dm + dr * ndim];
+                if (pos < 0 || pos >= shape[dm]) {
+                    k = -1;
+                    break;
+                }
+                k += pos * stride;
+                rest /= shape[dm];
+                stride *= shape[dm];
+            }
+            if (k < 0) continue;
+            if (y[k] < 0) continue;
+            while (1) {
+                while (j != y[j]) { j = y[j]; }
+                while (k != y[k]) { k = y[k]; }
+                if (j == k) break;
+                if (j < k) {
+                    int old = atomicCAS( &y[k], k, j );
+                    if (old == k) break;
+                    k = old;
+                }
+                else {
+                    int old = atomicCAS( &y[j], j, k );
+                    if (old == j) break;
+                    j = old;
+                }
+            }
+        }
+        ''',
+        'cupyx_nd_label_connect')
+
+
+def _kernel_count():
+    return cupy.ElementwiseKernel(
+        '', 'raw Y y, raw int32 count',
+        '''
+        if (y[i] < 0) continue;
+        int j = i;
+        while (j != y[j]) { j = y[j]; }
+        if (j != i) y[i] = j;
+        else atomicAdd(&count[0], 1);
+        ''',
+        'cupyx_nd_label_count')
+
+
+def _kernel_labels():
+    return cupy.ElementwiseKernel(
+        '', 'raw Y y, raw int32 count, raw int32 labels',
+        '''
+        if (y[i] != i) continue;
+        int j = atomicAdd(&count[1], 1);
+        labels[j] = i;
+        ''',
+        'cupyx_nd_label_labels')
+
+
+def _kernel_finalize():
+    return cupy.ElementwiseKernel(
+        'int32 maxlabel', 'raw int32 labels, raw Y y',
+        '''
+        if (y[i] < 0) {
+            y[i] = 0;
+            continue;
+        }
+        int yi = y[i];
+        int j_min = 0;
+        int j_max = maxlabel - 1;
+        int j = (j_min + j_max) / 2;
+        while (j_min < j_max) {
+            if (yi == labels[j]) break;
+            if (yi < labels[j]) j_max = j - 1;
+            else j_min = j + 1;
+            j = (j_min + j_max) / 2;
+        }
+        y[i] = j + 1;
+        ''',
+        'cupyx_nd_label_finalize')


Test patch:
diff --git a/tests/cupyx_tests/scipy_tests/ndimage_tests/test_measurements.py b/tests/cupyx_tests/scipy_tests/ndimage_tests/test_measurements.py
new file mode 100644
--- /dev/null
+++ b/tests/cupyx_tests/scipy_tests/ndimage_tests/test_measurements.py
@@ -0,0 +1,98 @@
+import unittest
+
+import numpy
+
+from cupy import testing
+import cupyx.scipy.ndimage  # NOQA
+
+try:
+    import scipy.ndimage  # NOQA
+except ImportError:
+    pass
+
+
+def _generate_binary_structure(rank, connectivity):
+    if connectivity < 1:
+        connectivity = 1
+    if rank < 1:
+        return numpy.array(True, dtype=bool)
+    output = numpy.fabs(numpy.indices([3] * rank) - 1)
+    output = numpy.add.reduce(output, 0)
+    return output <= connectivity
+
+
+@testing.parameterize(*testing.product({
+    'ndim': [1, 2, 3, 4],
+    'size': [50, 100],
+    'density': [0.2, 0.3, 0.4],
+    'connectivity': [None, 2, 3],
+    'x_dtype': [bool, numpy.int8, numpy.int32, numpy.int64,
+                numpy.float32, numpy.float64],
+    'output': [None, numpy.int32, numpy.int64],
+    'o_type': [None, 'ndarray']
+}))
+@testing.gpu
+@testing.with_requires('scipy')
+class TestLabel(unittest.TestCase):
+
+    @testing.numpy_cupy_array_equal(scipy_name='scp')
+    def test_label(self, xp, scp):
+        size = int(pow(self.size, 1 / self.ndim))
+        x_shape = range(size, size + self.ndim)
+        x = xp.zeros(x_shape, dtype=self.x_dtype)
+        x[testing.shaped_random(x_shape, xp) < self.density] = 1
+        if self.connectivity is None:
+            structure = None
+        else:
+            structure = _generate_binary_structure(self.ndim,
+                                                   self.connectivity)
+        if self.o_type == 'ndarray' and self.output is not None:
+            output = xp.empty(x_shape, dtype=self.output)
+            num_features = scp.ndimage.label(x, structure=structure,
+                                             output=output)
+            return output
+        labels, num_features = scp.ndimage.label(x, structure=structure,
+                                                 output=self.output)
+        return labels
+
+
+@testing.gpu
+@testing.with_requires('scipy')
+class TestLabelSpecialCases(unittest.TestCase):
+
+    @testing.numpy_cupy_array_equal(scipy_name='scp')
+    def test_label_empty(self, xp, scp):
+        x = xp.empty(0)
+        labels, num_features = scp.ndimage.label(x)
+        return labels
+
+    @testing.numpy_cupy_array_equal(scipy_name='scp')
+    def test_label_0d_zero(self, xp, scp):
+        x = xp.zeros([])
+        labels, num_features = scp.ndimage.label(x)
+        return labels
+
+    @testing.numpy_cupy_array_equal(scipy_name='scp')
+    def test_label_0d_one(self, xp, scp):
+        x = xp.ones([])
+        labels, num_features = scp.ndimage.label(x)
+        return labels
+
+    @testing.numpy_cupy_array_equal(scipy_name='scp')
+    def test_label_swirl(self, xp, scp):
+        x = [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+             [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
+             [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1],
+             [1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1],
+             [1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1],
+             [1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1],
+             [1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1],
+             [1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 1],
+             [1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1],
+             [1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1],
+             [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+             [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
+        x = xp.array(x)
+        labels, num_features = scp.ndimage.label(x)
+        return labels