diff --git a/python/cucim/src/cucim/skimage/measure/_label.py b/python/cucim/src/cucim/skimage/measure/_label.py index 3c168fc96..58a0dd9e8 100644 --- a/python/cucim/src/cucim/skimage/measure/_label.py +++ b/python/cucim/src/cucim/skimage/measure/_label.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: 2009-2022 the scikit-image team -# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. # SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause import cupy as cp @@ -7,6 +7,12 @@ from ._label_kernels import _label +_INT32_LABEL_SIZE_LIMIT = 2**31 - 2 + + +def _label_dtype_for_size(size): + return cp.int64 if size >= _INT32_LABEL_SIZE_LIMIT else cp.int32 + def _get_structure(ndim, connectivity): if connectivity is None: @@ -17,7 +23,6 @@ def _get_structure(ndim, connectivity): return cpu_ndi.generate_binary_structure(ndim, connectivity) -# TODO: currently uses int32 for the labels. should add int64 option as well def label(label_image, background=None, return_num=False, connectivity=None): r"""Label connected regions of an integer array. @@ -76,9 +81,8 @@ def label(label_image, background=None, return_num=False, connectivity=None): Notes ----- - Currently the cucim implementation of this function always uses 32-bit - integers for the label array. This is done for performance. In the future - 64-bit integer support may also be added for better skimage compatibility. + The cucim implementation of this function uses 32-bit integers for the + label array unless the input size requires 64-bit labels. Examples -------- @@ -121,7 +125,11 @@ def label(label_image, background=None, return_num=False, connectivity=None): # same here for non-integer dtypes. label_image = label_image.astype(cp.intp) - labels = cp.empty(label_image.shape, order="C", dtype=cp.int32) + labels = cp.empty( + label_image.shape, + order="C", + dtype=_label_dtype_for_size(label_image.size), + ) num = _label(label_image, structure, labels, greyscale_mode=True) if return_num: diff --git a/python/cucim/src/cucim/skimage/measure/_label_kernels.py b/python/cucim/src/cucim/skimage/measure/_label_kernels.py index 24535f754..1fa3bdf0e 100644 --- a/python/cucim/src/cucim/skimage/measure/_label_kernels.py +++ b/python/cucim/src/cucim/skimage/measure/_label_kernels.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: 2009-2022 the scikit-image team -# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. # SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause """Kernels for scikit-image label. @@ -24,7 +24,7 @@ def _label(x, structure, y, greyscale_mode=False): 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) + count = cupy.zeros(2, dtype=y.dtype) _kernel_init()(x, y) if greyscale_mode: _kernel_connect(True)(x, y_shape, dirs, ndirs, x.ndim, y, size=y.size) @@ -32,7 +32,7 @@ def _label(x, structure, y, greyscale_mode=False): _kernel_connect(False)(y_shape, dirs, ndirs, x.ndim, y, size=y.size) _kernel_count()(y, count, size=y.size) maxlabel = int(count[0]) # synchronize - labels = cupy.empty(maxlabel, dtype=numpy.int32) + labels = cupy.empty(maxlabel, dtype=y.dtype) _kernel_labels()(y, count, labels, size=y.size) _kernel_finalize()(maxlabel, cupy.sort(labels), y, size=y.size) return maxlabel @@ -78,18 +78,20 @@ def _kernel_connect(greyscale_mode=False, int_t="int"): # binary mode -> all non-background voxels treated the same x_condition = "" - # Note: atomicCAS is implemented for int, unsigned short, unsigned int, and - # unsigned long long - code = f""" + using atomic_t = typename cupy::type_traits::conditional< + cupy::type_traits::is_same::value, + unsigned long long, + Y + >::type; if (y[i] < 0) continue; for (int dr = 0; dr < ndirs; dr++) {{ - {int_t} j = i; - {int_t} rest = j; - {int_t} stride = 1; - {int_t} k = 0; + Y j = i; + Y rest = j; + Y stride = 1; + Y k = 0; for (int dm = ndim-1; dm >= 0; dm--) {{ - int pos = rest % shape[dm] + dirs[dm + dr * ndim]; + Y pos = rest % shape[dm] + dirs[dm + dr * ndim]; if (pos < 0 || pos >= shape[dm]) {{ k = -1; break; @@ -106,12 +108,12 @@ def _kernel_connect(greyscale_mode=False, int_t="int"): while (k != y[k]) {{ k = y[k]; }} if (j == k) break; if (j < k) {{ - {int_t} old = atomicCAS( &y[k], (Y)k, (Y)j ); + Y old = atomicCAS((atomic_t*)&y[k], k, j); if (old == k) break; k = old; }} else {{ - {int_t} old = atomicCAS( &y[j], (Y)j, (Y)k ); + Y old = atomicCAS((atomic_t*)&y[j], j, k); if (old == j) break; j = old; }} @@ -130,10 +132,10 @@ def _kernel_connect(greyscale_mode=False, int_t="int"): def _kernel_count(): return cupy.ElementwiseKernel( "", - "raw Y y, raw int32 count", + "raw Y y, raw Y count", """ if (y[i] < 0) continue; - int j = i; + Y j = i; while (j != y[j]) { j = y[j]; } if (j != i) y[i] = j; else atomicAdd(&count[0], 1); @@ -145,10 +147,10 @@ def _kernel_count(): def _kernel_labels(): return cupy.ElementwiseKernel( "", - "raw Y y, raw int32 count, raw int32 labels", + "raw Y y, raw Y count, raw Y labels", """ if (y[i] != i) continue; - int j = atomicAdd(&count[1], 1); + Y j = atomicAdd(&count[1], 1); labels[j] = i; """, "cucim_skimage_measure_label_labels", @@ -157,17 +159,17 @@ def _kernel_labels(): def _kernel_finalize(): return cupy.ElementwiseKernel( - "int32 maxlabel", - "raw int32 labels, raw Y y", + "Y maxlabel", + "raw Y 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; + Y yi = y[i]; + Y j_min = 0; + Y j_max = maxlabel - 1; + Y j = (j_min + j_max) / 2; while (j_min < j_max) { if (yi == labels[j]) break; if (yi < labels[j]) j_max = j - 1; diff --git a/python/cucim/src/cucim/skimage/measure/tests/test_ccomp.py b/python/cucim/src/cucim/skimage/measure/tests/test_ccomp.py index bbb3270bf..0801dc0ee 100644 --- a/python/cucim/src/cucim/skimage/measure/tests/test_ccomp.py +++ b/python/cucim/src/cucim/skimage/measure/tests/test_ccomp.py @@ -1,20 +1,19 @@ # SPDX-FileCopyrightText: 2009-2022 the scikit-image team -# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. # SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause # Note: These test cases originated in skimage/morphology/tests/test_ccomp.py import cupy as cp - -# import numpy as np from cupy.testing import assert_array_equal from cucim.skimage.measure import label - -# import pytest - -# import cucim.skimage.measure._ccomp as ccomp - +from cucim.skimage.measure._label import ( + _INT32_LABEL_SIZE_LIMIT, + _get_structure, + _label_dtype_for_size, +) +from cucim.skimage.measure._label_kernels import _label BG = 0 # background value @@ -140,6 +139,31 @@ def test_return_num(self): assert_array_equal(label(x, background=-1, return_num=True)[1], 4) + def test_label_dtype_for_size(self): + assert ( + cp.dtype(_label_dtype_for_size(_INT32_LABEL_SIZE_LIMIT - 1)) + == cp.int32 + ) + assert ( + cp.dtype(_label_dtype_for_size(_INT32_LABEL_SIZE_LIMIT)) == cp.int64 + ) + + def test_label_uses_int32_for_small_image(self): + labels = label(cp.eye(3, dtype=cp.uint8)) + + assert labels.dtype == cp.int32 + + def test_label_kernels_support_int64_output(self): + x = cp.eye(3, dtype=cp.uint8) + structure = _get_structure(x.ndim, connectivity=1) + labels = cp.empty(x.shape, dtype=cp.int64) + + num = _label(x, structure, labels, greyscale_mode=True) + + assert labels.dtype == cp.int64 + assert num == 3 + assert_array_equal(labels, cp.diag(cp.arange(1, 4, dtype=cp.int64))) + class TestConnectedComponents3d: def setup_method(self):