Skip to content

Commit 304ab1d

Browse files
Hendrik-codeclaude
andcommitted
feat(NII): add voronoi_labels wrapping np_voronoi_labels
Distances default to mm using the image zoom, since the partition itself changes with anisotropic spacing; use_mm=False falls back to voxels. The signed dtype chosen by np_voronoi_labels is kept, as the negated labels would wrap on an unsigned image. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1 parent 004dea7 commit 304ab1d

2 files changed

Lines changed: 111 additions & 0 deletions

File tree

‎TPTBox/core/nii_wrapper.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
np_unique,
5959
np_unique_withoutzero,
6060
np_volume,
61+
np_voronoi_labels,
6162
)
6263
from TPTBox.core.vert_constants import (
6364
AFFINE,
@@ -1808,6 +1809,59 @@ def dilate_msk_(self, n_pixel: int = 5, labels: LABEL_REFERENCE = None, connecti
18081809
"""In-place variant of `dilate_msk`."""
18091810
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)
18101811

1812+
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:
1813+
"""Assigns every voxel the label of the nearest labelled region (a Voronoi partition).
1814+
1815+
Args:
1816+
labels (LABEL_REFERENCE, optional): Labels to partition by. Everything else is treated as
1817+
background. Defaults to None (all labels present in the segmentation).
1818+
max_distance (float | None, optional): Leave background voxels further than this from any
1819+
region at 0. In mm when `use_mm` is True, else in voxels. Defaults to None (no limit).
1820+
use_mm (bool, optional): If True, distances are measured in mm using this image's voxel
1821+
spacing, so the partition is correct on anisotropic images. If False, voxels are
1822+
treated as isotropic. Defaults to True.
1823+
signed_background (bool, optional): If True, voxels filled in from the background get the
1824+
negated label of their region, so the input mask stays recoverable as `out > 0`. The
1825+
returned NII then has a signed dtype. Defaults to True.
1826+
tie_to_zero (bool, optional): If True, voxels equidistant to two different regions are set
1827+
to 0 instead of being assigned to one of them. Defaults to True.
1828+
inplace (bool, optional): Whether to modify this NII in place or return a new object.
1829+
Defaults to False.
1830+
verbose (bool, optional): Whether to print a message. Defaults to True.
1831+
1832+
Returns:
1833+
NII: The partition, signed when `signed_background` is set.
1834+
1835+
Notes:
1836+
Runtime is independent of the number of labels (one Euclidean feature transform for the
1837+
whole volume), unlike `dilate_msk`.
1838+
Set `signed_background=False, tie_to_zero=False` and a `max_distance` for a plain,
1839+
unsigned expansion of the labels (`np_expand_labels`).
1840+
1841+
Examples:
1842+
>>> nii.voronoi_labels(max_distance=10) # doctest: +SKIP
1843+
# every voxel within 1 cm of a labelled region takes its label
1844+
"""
1845+
assert self.seg
1846+
log.print("voronoi labels", end="\r", verbose=verbose)
1847+
out = np_voronoi_labels(
1848+
self.get_seg_array(),
1849+
label_ref=labels,
1850+
max_distance=max_distance,
1851+
zoom=self.zoom if use_mm else None,
1852+
signed_background=signed_background,
1853+
tie_to_zero=tie_to_zero,
1854+
)
1855+
if not signed_background:
1856+
# keep the image dtype; the negated labels of a signed partition would not fit an unsigned one
1857+
out = out.astype(self.dtype)
1858+
log.print("Voronoi partition of", len(np_unique_withoutzero(np.abs(out))), "labels", verbose=verbose)
1859+
return self.set_array(out, inplace=inplace)
1860+
1861+
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:
1862+
"""In-place variant of `voronoi_labels`."""
1863+
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)
1864+
18111865

18121866
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
18131867
"""Fills holes in segmentation.

‎unit_tests/test_nii_extended.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -694,5 +694,62 @@ def test_ethmoid_style_pipeline(self):
694694
self.assertAlmostEqual(vols[107] / (20 * 15 * 9), 1 / 3, places=2)
695695

696696

697+
class Test_voronoi_labels_nii(unittest.TestCase):
698+
@staticmethod
699+
def _two_points(zoom=(1.0, 1.0, 1.0)) -> NII:
700+
arr = np.zeros((5, 1, 1), dtype=np.uint8)
701+
arr[0] = 1
702+
arr[4] = 2
703+
return _make_nii(arr, zoom=zoom)
704+
705+
def test_signed_partition_and_tie(self):
706+
out = self._two_points().voronoi_labels(verbose=False)
707+
self.assertEqual(out.get_seg_array().ravel().tolist(), [1, -1, 0, -2, 2])
708+
self.assertTrue(np.issubdtype(out.dtype, np.signedinteger))
709+
710+
def test_input_mask_is_recoverable(self):
711+
nii = self._two_points()
712+
out = nii.voronoi_labels(verbose=False)
713+
np.testing.assert_array_equal((out.get_seg_array() > 0) * out.get_seg_array(), nii.get_seg_array())
714+
715+
def test_unsigned_keeps_the_image_dtype(self):
716+
nii = self._two_points()
717+
out = nii.voronoi_labels(signed_background=False, tie_to_zero=False, verbose=False)
718+
self.assertEqual(out.dtype, nii.dtype)
719+
self.assertEqual(out.get_seg_array().ravel().tolist(), [1, 1, 1, 2, 2])
720+
721+
def test_use_mm_scales_max_distance_with_zoom(self):
722+
arr = np.zeros((11, 1, 1), dtype=np.uint8)
723+
arr[0] = 1
724+
nii = _make_nii(arr, zoom=(2.0, 1.0, 1.0))
725+
in_mm = nii.voronoi_labels(max_distance=3, signed_background=False, verbose=False)
726+
in_vox = nii.voronoi_labels(max_distance=3, use_mm=False, signed_background=False, verbose=False)
727+
self.assertEqual(int(np.count_nonzero(in_mm.get_seg_array())), 2)
728+
self.assertEqual(int(np.count_nonzero(in_vox.get_seg_array())), 4)
729+
730+
def test_labels_treats_the_rest_as_background(self):
731+
out = self._two_points().voronoi_labels(labels=1, verbose=False)
732+
self.assertEqual(out.get_seg_array().ravel().tolist(), [1, -1, -1, -1, -1])
733+
734+
def test_grid_is_preserved(self):
735+
nii = self._two_points(zoom=(2.0, 1.0, 0.5))
736+
out = nii.voronoi_labels(verbose=False)
737+
self.assertTrue(out.assert_affine(nii, raise_error=False))
738+
self.assertEqual(out.shape, nii.shape)
739+
740+
def test_inplace_variant(self):
741+
nii = self._two_points()
742+
out = nii.voronoi_labels_(verbose=False)
743+
self.assertIs(out, nii)
744+
self.assertEqual(nii.get_seg_array().ravel().tolist(), [1, -1, 0, -2, 2])
745+
746+
def test_matches_the_array_level_function(self):
747+
from TPTBox.core.np_utils import np_voronoi_labels
748+
749+
nii = self._two_points(zoom=(1.5, 1.0, 1.0))
750+
expected = np_voronoi_labels(nii.get_seg_array(), zoom=nii.zoom)
751+
np.testing.assert_array_equal(nii.voronoi_labels(verbose=False).get_seg_array(), expected)
752+
753+
697754
if __name__ == "__main__":
698755
unittest.main()

0 commit comments

Comments
 (0)