feat(np_utils): Voronoi / nearest-label expansion

TPTBox/core/nii_wrapper.py

@@ -58,6 +58,7 @@
     np_unique,
     np_unique_withoutzero,
     np_volume,
+    np_voronoi_labels,
 )
 from TPTBox.core.vert_constants import (
     AFFINE,
@@ -1808,6 +1809,59 @@ def dilate_msk_(self, n_pixel: int = 5, labels: LABEL_REFERENCE = None, connecti
         """In-place variant of `dilate_msk`."""
         return self.dilate_msk(n_pixel=n_pixel, labels=labels, connectivity=connectivity, mask=mask, inplace=True, verbose=verbose,use_crop=use_crop,ignore_direction=ignore_direction)
 
+    def voronoi_labels(self, labels: LABEL_REFERENCE = None, max_distance: float | None = None, use_mm: bool = True, signed_background: bool = True, tie_to_zero: bool = True, inplace=False, verbose: logging = True) -> Self:
+        """Assigns every voxel the label of the nearest labelled region (a Voronoi partition).
+
+        Args:
+            labels (LABEL_REFERENCE, optional): Labels to partition by. Everything else is treated as
+                background. Defaults to None (all labels present in the segmentation).
+            max_distance (float | None, optional): Leave background voxels further than this from any
+                region at 0. In mm when `use_mm` is True, else in voxels. Defaults to None (no limit).
+            use_mm (bool, optional): If True, distances are measured in mm using this image's voxel
+                spacing, so the partition is correct on anisotropic images. If False, voxels are
+                treated as isotropic. Defaults to True.
+            signed_background (bool, optional): If True, voxels filled in from the background get the
+                negated label of their region, so the input mask stays recoverable as `out > 0`. The
+                returned NII then has a signed dtype. Defaults to True.
+            tie_to_zero (bool, optional): If True, voxels equidistant to two different regions are set
+                to 0 instead of being assigned to one of them. Defaults to True.
+            inplace (bool, optional): Whether to modify this NII in place or return a new object.
+                Defaults to False.
+            verbose (bool, optional): Whether to print a message. Defaults to True.
+
+        Returns:
+            NII: The partition, signed when `signed_background` is set.
+
+        Notes:
+            Runtime is independent of the number of labels (one Euclidean feature transform for the
+            whole volume), unlike `dilate_msk`.
+            Set `signed_background=False, tie_to_zero=False` and a `max_distance` for a plain,
+            unsigned expansion of the labels (`np_expand_labels`).
+
+        Examples:
+            >>> nii.voronoi_labels(max_distance=10)  # doctest: +SKIP
+            # every voxel within 1 cm of a labelled region takes its label
+        """
+        assert self.seg
+        log.print("voronoi labels", end="\r", verbose=verbose)
+        out = np_voronoi_labels(
+            self.get_seg_array(),
+            label_ref=labels,
+            max_distance=max_distance,
+            zoom=self.zoom if use_mm else None,
+            signed_background=signed_background,
+            tie_to_zero=tie_to_zero,
+        )
+        if not signed_background:
+            # keep the image dtype; the negated labels of a signed partition would not fit an unsigned one
+            out = out.astype(self.dtype)
+        log.print("Voronoi partition of", len(np_unique_withoutzero(np.abs(out))), "labels", verbose=verbose)
+        return self.set_array(out, inplace=inplace)
+
+    def voronoi_labels_(self, labels: LABEL_REFERENCE = None, max_distance: float | None = None, use_mm: bool = True, signed_background: bool = True, tie_to_zero: bool = True, verbose: logging = True) -> Self:
+        """In-place variant of `voronoi_labels`."""
+        return self.voronoi_labels(labels=labels, max_distance=max_distance, use_mm=use_mm, signed_background=signed_background, tie_to_zero=tie_to_zero, inplace=True, verbose=verbose)
+
 
     def fill_holes(self, labels: LABEL_REFERENCE = None, slice_wise_dim: int|str | None = None, verbose:logging=False, inplace=False,use_crop=True) -> Self:  # noqa: ARG002
         """Fills holes in segmentation.

TPTBox/core/np_utils.py

@@ -24,6 +24,8 @@
     distance_transform_edt,
     gaussian_filter,
     generate_binary_structure,
+    maximum_filter,
+    minimum_filter,
     sobel,
 )
 from scipy.spatial.distance import cdist
@@ -598,6 +600,145 @@ def np_dilate_msk_euclid(arr: np.ndarray, n_pixel: int = 3, use_crop=True, label
     return arr
 
 
+def _np_voronoi_ties(arr: UINTARRAY, indices: np.ndarray, nearest: np.ndarray, zoom: Sequence[float] | None) -> np.ndarray:
+    """Finds the voxels that are equidistant to two differently labelled regions.
+
+    ``distance_transform_edt`` picks one nearest source per voxel and never reports that the choice
+    was arbitrary, so the ties have to be recovered afterwards. Only voxels next to a change in the
+    nearest-label assignment can sit on a boundary between two regions, so the search is restricted
+    to those and then each of their neighbours' nearest sources is tested for being at exactly the
+    same distance. Any neighbour's source is a real foreground voxel, so a hit is always a genuine
+    tie -- the test can miss one, but never invent one.
+
+    Args:
+        arr (UINTARRAY): The labelled input array.
+        indices (np.ndarray): Feature transform of ``arr``, shape ``(ndim, *arr.shape)``.
+        nearest (np.ndarray): ``arr`` gathered at ``indices``.
+        zoom (Sequence[float] | None): Voxel spacing used for the distances.
+
+    Returns:
+        np.ndarray: Boolean mask of the tied voxels.
+    """
+    ties = np.zeros(arr.shape, dtype=bool)
+    boundary = maximum_filter(nearest, size=3) != minimum_filter(nearest, size=3)
+    pts = np.argwhere(boundary)
+    if len(pts) == 0:
+        return ties
+    sampling = np.ones(arr.ndim) if zoom is None else np.asarray(zoom, dtype=float)
+    limit = np.asarray(arr.shape) - 1
+    key = (slice(None), *pts.T)
+    dist2 = (((indices[key].T - pts) * sampling) ** 2).sum(1)
+    label = nearest[tuple(pts.T)]
+    tied = np.zeros(len(pts), dtype=bool)
+    for offset in itertools.product((-1, 0, 1), repeat=arr.ndim):
+        if not any(offset):
+            continue
+        neighbour = np.clip(pts + offset, 0, limit)
+        source = indices[(slice(None), *neighbour.T)]
+        cand2 = (((source.T - pts) * sampling) ** 2).sum(1)
+        tied |= (arr[tuple(source)] != label) & (np.abs(cand2 - dist2) <= 1e-9 * np.maximum(dist2, 1.0))
+    ties[tuple(pts[tied].T)] = True
+    return ties
+
+
+def np_voronoi_labels(
+    arr: UINTARRAY,
+    label_ref: LABEL_REFERENCE = None,
+    max_distance: float | None = None,
+    zoom: Sequence[float] | None = None,
+    signed_background: bool = True,
+    tie_to_zero: bool = True,
+) -> INTARRAY:
+    """Assigns every voxel the label of the nearest labelled region (a Voronoi partition).
+
+    Computed with a single Euclidean feature transform of the background: the nearest background
+    voxel of ``arr == 0`` is exactly the nearest labelled voxel of ``arr``, so one pass yields the
+    assignment for the whole volume, independent of how many labels there are.
+
+    By default the partition is signed: voxels that were already labelled keep their positive
+    label, voxels filled in from the background carry the *negative* label of the region they were
+    assigned to, so the original mask can still be recovered from the result. Voxels that are
+    exactly equidistant to two different regions are set to 0 rather than being broken arbitrarily.
+
+    For a bounded expansion that overwrites the input in place and crops to the mask first, see
+    :func:`np_dilate_msk_euclid`; this function is the unbounded, spacing-aware variant.
+
+    Args:
+        arr (UINTARRAY): Labelled input array with background 0.
+        label_ref (int | list[int] | None, optional): Labels to partition by. Everything else is
+            treated as background. Defaults to None (all labels found in ``arr``).
+        max_distance (float | None, optional): Leave background voxels further than this from any
+            region at 0. In the units of ``zoom``, i.e. voxels when ``zoom`` is None.
+            Defaults to None (no limit).
+        zoom (Sequence[float] | None, optional): Voxel spacing, so distances are measured in mm on
+            anisotropic images. Defaults to None (isotropic voxels).
+        signed_background (bool, optional): If True, background voxels get the negated label of
+            their region. Defaults to True.
+        tie_to_zero (bool, optional): If True, voxels equidistant to two different regions are set
+            to 0 instead of being assigned to one of them. Defaults to True.
+
+    Returns:
+        INTARRAY: The partition, signed when ``signed_background`` is set.
+
+    Examples:
+        >>> arr = np.zeros((5, 1, 1), dtype=np.uint8)
+        >>> arr[0] = 1
+        >>> arr[4] = 2
+        >>> np_voronoi_labels(arr).ravel().tolist()
+        [1, -1, 0, -2, 2]
+    """
+    assert 2 <= arr.ndim <= 3, f"expected 2D or 3D, but got {arr.ndim}"
+    assert np.min(arr) >= 0, f"expected non-negative labels, got {np.min(arr)}"
+    assert zoom is None or len(zoom) == arr.ndim, f"zoom has to have {arr.ndim} entries, got {zoom}"
+    if label_ref is not None:
+        arr = np.where(np_isin(arr, label_ref), arr, 0)
+    background = arr == 0
+    if background.all():
+        return np.zeros(arr.shape, dtype=np.int8 if signed_background else arr.dtype)
+
+    want_distance = max_distance is not None
+    result = distance_transform_edt(background, sampling=zoom, return_distances=want_distance, return_indices=True)
+    distances, indices = result if want_distance else (None, result)
+    nearest = arr[tuple(indices)]
+
+    out = nearest.astype(np.min_scalar_type(-int(arr.max()) - 1) if signed_background else arr.dtype, copy=True)
+    if signed_background:
+        out[background] = -out[background]
+    if want_distance:
+        out[background & (distances > max_distance)] = 0
+    if tie_to_zero:
+        out[_np_voronoi_ties(arr, indices, nearest, zoom)] = 0
+    return out
+
+
+def np_expand_labels(
+    arr: UINTARRAY,
+    distance: float | None = None,
+    label_ref: LABEL_REFERENCE = None,
+    zoom: Sequence[float] | None = None,
+) -> UINTARRAY:
+    """Grows every label into the surrounding background, up to ``distance``.
+
+    Each background voxel within reach takes the label of the nearest region; labelled voxels are
+    left alone, so labels never eat into each other. This is :func:`np_voronoi_labels` with the
+    plain, unsigned semantics -- see there for the signed partition and the handling of voxels that
+    are equidistant to two regions.
+
+    Args:
+        arr (UINTARRAY): Labelled input array with background 0.
+        distance (float | None, optional): How far to grow, in the units of ``zoom``, i.e. voxels
+            when ``zoom`` is None. Defaults to None (fill the whole volume).
+        label_ref (int | list[int] | None, optional): Labels to grow. Everything else is treated as
+            background. Defaults to None (all labels found in ``arr``).
+        zoom (Sequence[float] | None, optional): Voxel spacing, so ``distance`` is in mm on
+            anisotropic images. Defaults to None (isotropic voxels).
+
+    Returns:
+        UINTARRAY: A copy of ``arr`` with the labels grown.
+    """
+    return np_voronoi_labels(arr, label_ref=label_ref, max_distance=distance, zoom=zoom, signed_background=False, tie_to_zero=False)
+
+
 def np_dilate_msk(
     arr: np.ndarray,
     label_ref: LABEL_REFERENCE = None,

unit_tests/test_nii_extended.py

@@ -694,5 +694,62 @@ def test_ethmoid_style_pipeline(self):
         self.assertAlmostEqual(vols[107] / (20 * 15 * 9), 1 / 3, places=2)
 
 
+class Test_voronoi_labels_nii(unittest.TestCase):
+    @staticmethod
+    def _two_points(zoom=(1.0, 1.0, 1.0)) -> NII:
+        arr = np.zeros((5, 1, 1), dtype=np.uint8)
+        arr[0] = 1
+        arr[4] = 2
+        return _make_nii(arr, zoom=zoom)
+
+    def test_signed_partition_and_tie(self):
+        out = self._two_points().voronoi_labels(verbose=False)
+        self.assertEqual(out.get_seg_array().ravel().tolist(), [1, -1, 0, -2, 2])
+        self.assertTrue(np.issubdtype(out.dtype, np.signedinteger))
+
+    def test_input_mask_is_recoverable(self):
+        nii = self._two_points()
+        out = nii.voronoi_labels(verbose=False)
+        np.testing.assert_array_equal((out.get_seg_array() > 0) * out.get_seg_array(), nii.get_seg_array())
+
+    def test_unsigned_keeps_the_image_dtype(self):
+        nii = self._two_points()
+        out = nii.voronoi_labels(signed_background=False, tie_to_zero=False, verbose=False)
+        self.assertEqual(out.dtype, nii.dtype)
+        self.assertEqual(out.get_seg_array().ravel().tolist(), [1, 1, 1, 2, 2])
+
+    def test_use_mm_scales_max_distance_with_zoom(self):
+        arr = np.zeros((11, 1, 1), dtype=np.uint8)
+        arr[0] = 1
+        nii = _make_nii(arr, zoom=(2.0, 1.0, 1.0))
+        in_mm = nii.voronoi_labels(max_distance=3, signed_background=False, verbose=False)
+        in_vox = nii.voronoi_labels(max_distance=3, use_mm=False, signed_background=False, verbose=False)
+        self.assertEqual(int(np.count_nonzero(in_mm.get_seg_array())), 2)
+        self.assertEqual(int(np.count_nonzero(in_vox.get_seg_array())), 4)
+
+    def test_labels_treats_the_rest_as_background(self):
+        out = self._two_points().voronoi_labels(labels=1, verbose=False)
+        self.assertEqual(out.get_seg_array().ravel().tolist(), [1, -1, -1, -1, -1])
+
+    def test_grid_is_preserved(self):
+        nii = self._two_points(zoom=(2.0, 1.0, 0.5))
+        out = nii.voronoi_labels(verbose=False)
+        self.assertTrue(out.assert_affine(nii, raise_error=False))
+        self.assertEqual(out.shape, nii.shape)
+
+    def test_inplace_variant(self):
+        nii = self._two_points()
+        out = nii.voronoi_labels_(verbose=False)
+        self.assertIs(out, nii)
+        self.assertEqual(nii.get_seg_array().ravel().tolist(), [1, -1, 0, -2, 2])
+
+    def test_matches_the_array_level_function(self):
+        from TPTBox.core.np_utils import np_voronoi_labels
+
+        nii = self._two_points(zoom=(1.5, 1.0, 1.0))
+        expected = np_voronoi_labels(nii.get_seg_array(), zoom=nii.zoom)
+        np.testing.assert_array_equal(nii.voronoi_labels(verbose=False).get_seg_array(), expected)
+
+
 if __name__ == "__main__":
     unittest.main()