fix possible segfault in cucim.skimage.morphology.reconstruction

Signed-off-by: Gregory R. Lee <[email protected]>

Author: grlee77

python/cucim/src/cucim/skimage/morphology/grayreconstruct.py

@@ -70,7 +70,11 @@ def reconstruction(seed, mask, method="dilation", footprint=None, offset=None):
     Returns
     -------
     reconstructed : ndarray
-       The result of morphological reconstruction.
+       The result of morphological reconstruction. Note that scikit-image always
+       returns a floating-point image. cuCIM returns the same dtype as the input
+       except in the case of erosion, where it will have promoted any unsigned
+       integer dtype to a signed type (using the dtype returned by
+       ``cp.promote_types(seed.dtype, cp.int8)``).
 
     Examples
     --------
@@ -197,6 +201,12 @@ def reconstruction(seed, mask, method="dilation", footprint=None, offset=None):
     # CuPy Backend: modified to allow images_dtype based on input dtype
     #               instead of float64
     images_dtype = np.promote_types(seed.dtype, mask.dtype)
+    # For erosion, we need to negate the array, so ensure we use a signed type
+    # that can represent negative values without wraparound
+    if method == "erosion" and cp.issubdtype(images_dtype, cp.unsignedinteger):
+        # Promote unsigned types to signed types with sufficient range
+        images_dtype = cp.promote_types(images_dtype, cp.int8)
+
     images = cp.full(dims, pad_value, dtype=images_dtype)
     images[(0, *inside_slices)] = seed
     images[(1, *inside_slices)] = mask
@@ -259,5 +269,5 @@ def reconstruction(seed, mask, method="dilation", footprint=None, offset=None):
     value_rank = cp.asarray(value_rank[:image_stride])
 
     rec_img = value_map[value_rank]
-    rec_img.shape = dims[1:]
+    rec_img = rec_img.reshape(dims[1:])
     return rec_img[inside_slices]

python/cucim/src/cucim/skimage/morphology/tests/test_reconstruction.py

@@ -18,8 +18,10 @@
 import numpy as np
 import pytest
 from cupy.testing import assert_array_almost_equal
+from skimage import data
 from skimage.morphology import reconstruction as reconstruction_cpu
 
+from cucim.skimage.exposure import rescale_intensity
 from cucim.skimage.morphology import reconstruction
 
 
@@ -102,6 +104,35 @@ def test_invalid_seed():
         reconstruction(seed * 0.5, mask, method="erosion")
 
 
+@pytest.mark.parametrize(
+    "unsigned_dtype", [cp.uint8, cp.uint16, cp.uint32, cp.uint64]
+)
+def test_erosion_uint8_input(unsigned_dtype):
+    image = cp.asarray(data.moon())
+    # Rescale image intensity so that we can see dim features.
+    image = rescale_intensity(image, in_range=(50, 200))
+
+    seed = cp.copy(image)
+    seed[1:-1, 1:-1] = image.max()
+    mask = image
+
+    seed = seed.astype(unsigned_dtype)
+    filled = reconstruction(seed, mask, method="erosion")
+
+    expected_out_type = cp.promote_types(seed.dtype, cp.int8)
+    assert filled.dtype == expected_out_type
+    # promoted to signed dtype (or float in the case of cp.uint64)
+    assert filled.dtype.kind == "i" if unsigned_dtype != cp.uint64 else "f"
+    assert filled.dtype.itemsize >= seed.dtype.itemsize
+
+    # compare to scikit-image CPU result
+    filled_cpu = reconstruction_cpu(
+        cp.asnumpy(seed), cp.asnumpy(mask), method="erosion"
+    )
+    # filled_cpu will be np.float64, so convert to the type returned by cuCIM
+    cp.testing.assert_allclose(filled, filled_cpu.astype(filled.dtype))
+
+
 def test_invalid_footprint():
     seed = cp.ones((5, 5))
     mask = cp.ones((5, 5))