diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index e10adc6..cb46fbc 100755 --- a/TPTBox/core/np_utils.py +++ b/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 @@ -592,6 +594,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, diff --git a/unit_tests/test_nputils.py b/unit_tests/test_nputils.py index 439666d..db66f96 100755 --- a/unit_tests/test_nputils.py +++ b/unit_tests/test_nputils.py @@ -747,6 +747,144 @@ def test_max_count_component_drops_smaller_components(self): self.assertEqual(len(np_utils.np_label_interface_thickness(seg, 1, 2, max_count_component=2)), len(every)) +def _brute_voronoi(arr: np.ndarray, zoom=None) -> np.ndarray: + """Reference partition: for every voxel, scan every region and keep the closest, ties to zero.""" + labels = sorted(int(x) for x in np.unique(arr) if x != 0) + sampling = np.ones(arr.ndim) if zoom is None else np.asarray(zoom, dtype=float) + positions = {label: np.argwhere(arr == label) for label in labels} + out = np.zeros(arr.shape, dtype=np.int32) + for coord in np.argwhere(np.ones(arr.shape, dtype=bool)): + dists = {label: float((((pos - coord) * sampling) ** 2).sum(1).min()) for label, pos in positions.items()} + best = min(dists.values()) + closest = [label for label, d in dists.items() if abs(d - best) <= 1e-9] + value = 0 if len(closest) > 1 else closest[0] + out[tuple(coord)] = -value if arr[tuple(coord)] == 0 else value + return out + + +class Test_voronoi_labels(unittest.TestCase): + def test_signed_background_and_tie(self): + arr = np.zeros((5, 1, 1), dtype=np.uint8) + arr[0] = 1 + arr[4] = 2 + # the middle voxel is equidistant to both regions + self.assertEqual(np_utils.np_voronoi_labels(arr).ravel().tolist(), [1, -1, 0, -2, 2]) + + def test_matches_a_brute_force_partition(self): + rng = np.random.default_rng(0) + for shape in [(9, 9), (7, 7, 7)]: + for _ in range(4): + arr = np.zeros(shape, dtype=np.uint8) + for label in range(1, 5): + arr[tuple(rng.integers(0, s) for s in shape)] = label + with self.subTest(shape=shape, arr=arr): + got = np_utils.np_voronoi_labels(arr).astype(np.int32) + self.assertTrue(np.array_equal(got, _brute_voronoi(arr))) + + def test_matches_a_brute_force_partition_anisotropic(self): + rng = np.random.default_rng(1) + zoom = (2.0, 1.0, 0.5) + for _ in range(4): + arr = np.zeros((7, 7, 7), dtype=np.uint8) + for label in range(1, 4): + arr[tuple(rng.integers(0, 7, 3))] = label + with self.subTest(arr=arr): + got = np_utils.np_voronoi_labels(arr, zoom=zoom).astype(np.int32) + self.assertTrue(np.array_equal(got, _brute_voronoi(arr, zoom=zoom))) + + def test_unsigned_keeps_the_labels_positive(self): + arr = np.zeros((5, 1, 1), dtype=np.uint8) + arr[0] = 1 + arr[4] = 2 + out = np_utils.np_voronoi_labels(arr, signed_background=False) + self.assertEqual(out.ravel().tolist(), [1, 1, 0, 2, 2]) + + def test_without_tie_to_zero_nothing_stays_unassigned(self): + arr = np.zeros((5, 1, 1), dtype=np.uint8) + arr[0] = 1 + arr[4] = 2 + out = np_utils.np_voronoi_labels(arr, signed_background=False, tie_to_zero=False) + self.assertNotIn(0, out.ravel().tolist()) + + def test_max_distance_leaves_far_voxels_empty(self): + arr = np.zeros((11, 1, 1), dtype=np.uint8) + arr[0] = 1 + out = np_utils.np_voronoi_labels(arr, max_distance=3, signed_background=False) + self.assertEqual(out.ravel().tolist(), [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0]) + + def test_zoom_scales_max_distance(self): + arr = np.zeros((11, 1, 1), dtype=np.uint8) + arr[0] = 1 + out = np_utils.np_voronoi_labels(arr, max_distance=3, zoom=(2.0, 1.0, 1.0), signed_background=False) + self.assertEqual(int(np.count_nonzero(out)), 2) + + def test_label_ref_treats_the_rest_as_background(self): + arr = np.zeros((5, 1, 1), dtype=np.uint8) + arr[0] = 1 + arr[4] = 2 + out = np_utils.np_voronoi_labels(arr, label_ref=1) + self.assertEqual(out.ravel().tolist(), [1, -1, -1, -1, -1]) + + def test_empty_input(self): + arr = np.zeros((4, 4, 4), dtype=np.uint8) + self.assertEqual(int(np.count_nonzero(np_utils.np_voronoi_labels(arr))), 0) + + def test_single_label_fills_everything(self): + arr = np.zeros((4, 4, 4), dtype=np.uint8) + arr[0, 0, 0] = 7 + out = np_utils.np_voronoi_labels(arr) + self.assertEqual(sorted(set(out.ravel().tolist())), [-7, 7]) + + def test_dtype_holds_the_negated_labels(self): + arr = np.zeros((4, 4), dtype=np.uint8) + arr[0, 0] = 200 + out = np_utils.np_voronoi_labels(arr) + self.assertTrue(np.issubdtype(out.dtype, np.signedinteger)) + self.assertEqual(int(out.min()), -200) + + def test_cost_does_not_grow_with_the_label_count(self): + """One feature transform for the whole volume, not one per label.""" + rng = np.random.default_rng(2) + arr = np.zeros((40, 40, 40), dtype=np.uint16) + for label in range(1, 201): + arr[tuple(rng.integers(0, 40, 3))] = label + out = np_utils.np_voronoi_labels(arr, signed_background=False, tie_to_zero=False) + self.assertEqual(set(np_utils.np_unique_withoutzero(out)), set(np_utils.np_unique_withoutzero(arr))) + + +class Test_expand_labels(unittest.TestCase): + def test_grows_by_the_given_distance(self): + arr = np.zeros((11, 1, 1), dtype=np.uint8) + arr[5] = 3 + out = np_utils.np_expand_labels(arr, distance=2) + self.assertEqual(out.ravel().tolist(), [0, 0, 0, 3, 3, 3, 3, 3, 0, 0, 0]) + + def test_labels_do_not_eat_into_each_other(self): + arr = np.zeros((10, 1, 1), dtype=np.uint8) + arr[0] = 1 + arr[9] = 2 + out = np_utils.np_expand_labels(arr) + self.assertEqual(out.ravel().tolist(), [1, 1, 1, 1, 1, 2, 2, 2, 2, 2]) + + def test_unbounded_fills_the_volume(self): + arr = np.zeros((6, 6), dtype=np.uint8) + arr[0, 0] = 4 + self.assertTrue((np_utils.np_expand_labels(arr) == 4).all()) + + def test_agrees_with_dilate_msk_euclid(self): + rng = np.random.default_rng(3) + arr = np.zeros((20, 20, 20), dtype=np.uint8) + for label in (1, 2, 3): + p = rng.integers(2, 17, 3) + arr[p[0] : p[0] + 2, p[1] : p[1] + 2, p[2] : p[2] + 2] = label + expected = np_utils.np_dilate_msk_euclid(arr.copy(), n_pixel=3, use_crop=False) + got = np_utils.np_expand_labels(arr, distance=3) + # the two disagree only where a voxel is equidistant to two regions and the tie-break differs + differing = got != expected + self.assertLess(int(differing.sum()), 0.02 * differing.size) + self.assertTrue(np.array_equal(got != 0, expected != 0)) + + if __name__ == "__main__": unittest.main()