From 63c049c197a7ff2e7e98e44d7807495a7aff990e Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 27 Aug 2026 14:38:16 +0000 Subject: [PATCH] feat(np_utils): add np_split_connected_component with erosion and mincut backends Splits a mask that is one connected component but should be two - two merged vertebral bodies, for instance - into two spatially separate parts. Unifies two independent implementations behind one API: - "erosion" (from spineps get_separating_components): erode until the component breaks apart, then re-grow both halves. Fast and dependency-free. - "mincut" (from nako-segmentation split_cc): erode only until two seeds appear, then find the minimal separating surface by min-cut/max-flow over the voxel adjacency graph, with edge capacities from the voxel spacing. Needs networkx, slower, but the cut is minimal-area rather than incidental. Also adds np_connected_component_contact_map, which dilates two disjoint parts until they touch and returns the 0/1/2/3 contact map - the reusable second half of the erosion algorithm. The two backends originally disagreed on what they returned: the erosion one recovered only the eroded cores while the mincut one partitioned the whole volume. A full_partition flag (default True) now grows the erosion cores back over the removed voxels so both agree; pass False for the raw seeds, which is what you want when fitting a separating plane to them. Two latent bugs in the mincut source are fixed, both now covered by tests: - add_2d_edges raised NameError on an undefined `geom.norm` - max_ignore=None raised NameError, since `max_errors` was only assigned inside the branch that the same argument disables networkx is imported lazily with a clear error message and declared only as a dev dependency, so the core package does not gain a runtime requirement. Co-Authored-By: Claude Opus 5 --- TPTBox/core/np_utils.py | 293 +++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + unit_tests/test_nputils.py | 111 ++++++++++++++ 3 files changed, 405 insertions(+) diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index c7200f48..174221a7 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -1056,6 +1056,299 @@ def np_filter_connected_components( return cc_out +def np_connected_component_contact_map(part_a: np.ndarray, part_b: np.ndarray, connectivity: int = 3) -> UINTARRAY: + """Dilates two disjoint components until they touch and returns their contact map. + + Both parts are grown one voxel at a time, alternating, until their dilations overlap. + + Args: + part_a (np.ndarray): Binary mask of the first component. + part_b (np.ndarray): Binary mask of the second component. + connectivity (int, optional): Connectivity used for the dilation. Defaults to 3. + + Returns: + UINTARRAY: Map with 0 = background, 1 = only ``part_a``'s dilation, 2 = only ``part_b``'s + dilation, 3 = the contact zone where the two dilations overlap. + """ + # Dilate copies: np_dilate_msk works in place and returns its input, so dilating the + # caller's arrays would grow them too and the "two separated components" would smear. + a_dil = np_dilate_msk(part_a.copy().astype(np.uint8), n_pixel=1, connectivity=connectivity) + b_dil = np_dilate_msk(part_b.copy().astype(np.uint8), n_pixel=1, connectivity=connectivity) + contact = (a_dil + (b_dil * 2)).astype(np.uint8) + while 3 not in np_volume(contact): + a_dil = np_dilate_msk(a_dil, n_pixel=1, connectivity=connectivity) + contact = (a_dil + (b_dil * 2)).astype(np.uint8) + if 3 in np_volume(contact): + break + b_dil = np_dilate_msk(b_dil, n_pixel=1, connectivity=connectivity) + contact = (a_dil + (b_dil * 2)).astype(np.uint8) + return contact + + +def _expand_seeds_to_mask(seeds: UINTARRAY, mask: np.ndarray) -> UINTARRAY: + """Assigns every voxel of ``mask`` the label of its nearest non-zero seed.""" + unassigned = (mask != 0) & (seeds == 0) + if not unassigned.any(): + return seeds + _, indices = distance_transform_edt(seeds == 0, return_indices=True) + out = seeds.copy() + out[unassigned] = seeds[tuple(indices[:, unassigned])] + return out + + +def _split_cc_erosion(arr: np.ndarray, connectivity: int, max_iter: int, full_partition: bool) -> UINTARRAY: + """Erosion backend of :func:`np_split_connected_component`.""" + check_connectivity = 3 + vol = arr.copy().astype(np.uint8) + vol_old = vol.copy() + iterations = 0 + while True: + # np_erode_msk mutates its input and returns the same object, so erode a copy. Without it + # `vol`, `vol_old` and `vol_erode` all alias one array after the first iteration and the + # "iteration before" needed by the wiped-out branch below is lost. + vol_erode = np_erode_msk(vol.copy(), n_pixel=1, connectivity=connectivity) + subreg_cc, subreg_cc_n = np_connected_components(vol_erode, connectivity=check_connectivity) + if subreg_cc_n > 1: + vol = subreg_cc + break + if subreg_cc_n == 0: + # Erosion wiped everything out; step back and grow the previous state into two parts. + vol_dilated = np_dilate_msk(vol.copy(), n_pixel=1, connectivity=connectivity, mask=vol.copy()) + vol[vol_old != 0] = 2 + vol[vol_dilated == 1] = 1 + volume = np_volume(vol) + if 1 not in volume or 2 not in volume: + raise ValueError(f"cannot split volume into two parts after {iterations} iterations, got regions {volume}.") + while volume[1] / (volume[1] + volume[2]) < 0.5: + vol_dilated = np_dilate_msk(vol_dilated, n_pixel=1, connectivity=connectivity, mask=vol.copy()) + vol[vol_dilated == 1] = 1 + volume = np_volume(vol) + if 1 not in volume or 2 not in volume: + raise ValueError("could not divide into two parts while re-growing the eroded volume.") + vol_1 = np_filter_connected_components(vol == 1, largest_k_components=1, connectivity=check_connectivity).astype(np.uint8) + vol_2 = np_filter_connected_components(vol == 2, largest_k_components=1, connectivity=check_connectivity).astype(np.uint8) + vol_2 *= 2 + vol[vol == 1] = vol_1[vol == 1] + vol[vol == 2] = vol_2[vol == 2] + break + vol_old = vol + vol = vol_erode + iterations += 1 + if iterations > max_iter: + raise ValueError(f"could not divide into two parts after max_iter={max_iter} erosions.") + + if len(np_volume(vol)) != 2: + vol = np_filter_connected_components(vol, largest_k_components=2, connectivity=check_connectivity, return_original_labels=False) + part_a = vol == 1 + part_b = vol == 2 + if part_a.sum() == 0 or part_b.sum() == 0: + raise ValueError("one of the two split parts is empty.") + out = (part_a.astype(np.uint8) + part_b.astype(np.uint8) * 2).astype(np.uint8) + # Erosion only ever recovers the two cores; grow them back over the voxels the erosion ate so + # the result partitions the input, matching what the mincut backend returns. + return _expand_seeds_to_mask(out, arr) if full_partition else out + + +def _split_cc_mincut( # noqa: C901 + arr: np.ndarray, + connectivity: int, + separator: np.ndarray | None, + structure, + min_volume: int | None, + max_cut: float | None, + max_ignore: int | None, + zoom: Sequence[float] | None, + add_diagonal_edges: bool, +) -> UINTARRAY: + """Min-cut/max-flow backend of :func:`np_split_connected_component`.""" + try: + import networkx as nx + except ImportError as e: # pragma: no cover - depends on the environment + raise ImportError("the 'mincut' method needs networkx; install it with `pip install networkx`.") from e + + cc_connectivity = 6 if connectivity == 1 else (18 if connectivity == 2 else 26) + vol = arr != 0 + _, n = _connected_components(vol, connectivity=cc_connectivity, return_N=True) + if n != 1: + raise ValueError(f"volume separates into {n} parts at connectivity={connectivity}; it must be a single component.") + + structures = [structure] if isinstance(structure, np.ndarray) else structure + vol_erode = vol + iterations = 0 + max_errors = 0 + cc_erode = None + while True: + struct_now = structures[iterations % len(structures)] if structures is not None else None + if separator is not None: + separator = _binary_dilation(separator, struct_now) + vol_erode = np.where(separator, 0, vol) + else: + vol_erode = _binary_erosion(vol_erode, struct_now) + cc_erode, n = _connected_components(vol_erode, connectivity=cc_connectivity, return_N=True) + iterations += 1 + if n > 1: + if max_ignore is not None: + values, counts = np.unique(cc_erode, return_counts=True) + max_errors = sum(c for v, c in zip(values, counts) if v > 0 and c <= max_ignore) + keep = [v for v, c in zip(values, counts) if v > 0 and c > max_ignore] + n = len(keep) + if n > 2: + break + if n == 2: + relabeled = np.zeros(cc_erode.shape, dtype=cc_erode.dtype) + relabeled[cc_erode == keep[0]] = 1 + relabeled[cc_erode == keep[1]] = 2 + cc_erode = relabeled + break + else: + break + if n == 0: + raise ValueError(f"cannot split volume into two parts after {iterations} iterations, erosion emptied it.") + if iterations > 100: + raise ValueError("could not split the volume into two parts within 100 erosions.") + if n > 2: + raise ValueError(f"erosion produced {n} components after {iterations} iterations, expected 2.") + + source_mask = cc_erode == 1 + sink_mask = cc_erode == 2 + boundary = vol ^ vol_erode + + for name, mask in (("source", source_mask), ("sink", sink_mask)): + if min_volume is not None and mask.sum() < min_volume: + raise ValueError(f"after erosion the {name} part has volume {mask.sum()}, below min_volume={min_volume}.") + + if zoom is None: + voxel_dim = np.ones(3) + capacity_end = 1000.0 + else: + voxel_dim = np.asarray(zoom, dtype=float) + capacity_end = float(np.prod(voxel_dim) / np.min(voxel_dim) * 1000) + + source_dil = _binary_dilation(source_mask, struct_now) + sink_dil = _binary_dilation(sink_mask, struct_now) + to_source = np.argwhere(source_dil & boundary) + to_sink = np.argwhere(sink_dil & boundary) + if len(to_source) == 0 or len(to_sink) == 0: + raise ValueError("no connection between the separated parts and the remaining voxels.") + + graph = nx.Graph() + + def add_edges(points, diff1, diff2, capacity): + graph.add_edges_from(zip(map(tuple, points + diff1), map(tuple, points + diff2)), capacity=capacity) + + for x, y, z in itertools.product([0, 1], repeat=3): + vec = np.array([x, y, z]) + xe, ye, ze = np.array(boundary.shape) - vec + overlap = boundary[x:, y:, z:] & boundary[:xe, :ye, :ze] + if x + y + z == 1: + add_edges(np.argwhere(overlap), [0, 0, 0], [x, y, z], float(np.prod(voxel_dim[vec == 0]))) + elif x + y + z == 2 and add_diagonal_edges: + # Diagonal in two dimensions, extruded along the third. + capacity = float(voxel_dim[vec == 0][0] * np.linalg.norm(voxel_dim[vec == 1])) + add_edges(np.argwhere(overlap), [0, 0, 0], [x, y, z], capacity) + if x == 1: + add_edges(np.argwhere(boundary[:xe, y:, z:] & boundary[x:, :ye, :ze]), [x, 0, 0], [0, y, z], capacity) + else: + add_edges(np.argwhere(boundary[x:, :ye, z:] & boundary[:xe, y:, :ze]), [0, y, 0], [x, 0, z], capacity) + + graph.add_edges_from([((x, y, z), "t") for x, y, z in to_sink], capacity=capacity_end) + graph.add_edges_from([("s", (x, y, z)) for x, y, z in to_source], capacity=capacity_end) + if not nx.has_path(graph, "s", "t"): + raise ValueError("no path exists between the two parts in the adjacency graph.") + + cut_value, (source_side, sink_side) = nx.minimum_cut(graph, "s", "t") + if max_cut is not None and cut_value > max_cut: + raise ValueError(f"cut size is {cut_value}, above the allowed max_cut={max_cut}.") + source_side = set(source_side) - {"s"} + sink_side = set(sink_side) - {"t"} + if len(source_side) == 0 or len(sink_side) == 0: + raise ValueError("one side of the cut is empty.") + + out = cc_erode.astype(np.uint8) + out[tuple(np.asarray(list(source_side)).reshape([-1, 3]).transpose())] = 1 + out[tuple(np.asarray(list(sink_side)).reshape([-1, 3]).transpose())] = 2 + lost = int(abs((out > 0).sum() - vol.sum())) + if lost > max_errors: + raise ValueError(f"lost {lost} voxels while splitting, but only {max_errors} are allowed.") + return out + + +def np_split_connected_component( + arr: np.ndarray, + method: str = "erosion", + connectivity: int = 3, + max_iter: int = 10, + full_partition: bool = True, + separator: np.ndarray | None = None, + structure: np.ndarray | Sequence[np.ndarray] | None = None, + min_volume: int | None = None, + max_cut: float | None = None, + max_ignore: int | None = 6, + zoom: Sequence[float] | None = None, + add_diagonal_edges: bool = False, +) -> UINTARRAY: + """Splits one connected component into two spatially separate parts. + + For a mask that is a single connected component but should be two things -- two merged + vertebral bodies, say -- this finds the separation. Two backends are available: + + * ``"erosion"``: erode until the component breaks apart, then re-grow both halves. Fast and + dependency-free, but the cut follows the erosion front rather than any optimality criterion. + * ``"mincut"``: erode only until two seeds appear, then find the minimal separating surface + between them by min-cut/max-flow over the voxel adjacency graph, with edge capacities taken + from ``zoom`` so anisotropic voxels are weighted correctly. Needs ``networkx`` and is + considerably slower, but the cut is minimal-area rather than incidental. + + Args: + arr (np.ndarray): Binary (or labeled) array holding exactly one connected component. + method (str, optional): ``"erosion"`` or ``"mincut"``. Defaults to ``"erosion"``. + connectivity (int, optional): Connectivity for the morphology and component labelling + (1 = faces, 2 = +edges, 3 = +corners). Defaults to 3. + max_iter (int, optional): Erosion backend only -- maximum erosion iterations. Defaults to 10. + full_partition (bool, optional): Erosion backend only -- grow the two eroded cores back over + the voxels the erosion removed, so the result covers all of ``arr``. Set False to get + just the eroded seeds, which is what you want when fitting a separating plane to them. + Defaults to True. + separator (np.ndarray | None, optional): Min-cut backend only -- if given, this mask is + dilated into the volume instead of eroding the volume itself. Defaults to None. + structure (np.ndarray | Sequence[np.ndarray] | None, optional): Min-cut backend only -- + structuring element(s) to erode with; a sequence is cycled through. Defaults to None. + min_volume (int | None, optional): Min-cut backend only -- reject the split if either seed + is smaller than this. Defaults to None. + max_cut (float | None, optional): Min-cut backend only -- reject the split if the cut + exceeds this cost. Defaults to None. + max_ignore (int | None, optional): Min-cut backend only -- components at or below this + size are treated as noise and their voxels are allowed to be lost. Defaults to 6. + zoom (Sequence[float] | None, optional): Min-cut backend only -- voxel spacing used to + weight the graph edges. Defaults to None (isotropic). + add_diagonal_edges (bool, optional): Min-cut backend only -- also connect diagonal + neighbours in the adjacency graph. Defaults to False. + + Returns: + UINTARRAY: Array with 0 = background, 1 = first part, 2 = second part. + + Raises: + ValueError: If the volume cannot be split into exactly two non-empty parts, or if a + sanity limit (``min_volume``, ``max_cut``, ``max_ignore``) is exceeded. + ImportError: If ``method="mincut"`` and ``networkx`` is not installed. + """ + if method == "erosion": + return _split_cc_erosion(arr, connectivity=connectivity, max_iter=max_iter, full_partition=full_partition) + if method == "mincut": + return _split_cc_mincut( + arr, + connectivity=connectivity, + separator=separator, + structure=structure, + min_volume=min_volume, + max_cut=max_cut, + max_ignore=max_ignore, + zoom=zoom, + add_diagonal_edges=add_diagonal_edges, + ) + raise ValueError(f"unknown method {method!r}, expected 'erosion' or 'mincut'.") + + def np_get_connected_components_center_of_mass( arr: UINTARRAY, label: int, connectivity: int = 3, sort_by_axis: int | None = None ) -> list[COORDINATE]: diff --git a/pyproject.toml b/pyproject.toml index 9fd2f542..be555724 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ pre-commit = "*" pyvista = "^0.43.2" coverage = ">=7.0.1" pytest-mock = "^3.6.0" +networkx = "*" # optional: np_split_connected_component(method="mincut") exceptiongroup = { version = "^1.2", python = "<3.11" } tomli = {version = "*", python = "<3.11" } diff --git a/unit_tests/test_nputils.py b/unit_tests/test_nputils.py index 4477540f..5a71570a 100755 --- a/unit_tests/test_nputils.py +++ b/unit_tests/test_nputils.py @@ -463,6 +463,117 @@ def test_np_binary_fill_holes_and_set_inter_labels_based_on_majority(self): self.assertTrue(np.array_equal(filled, expected)) +def _dumbbell(neck: int = 2) -> np.ndarray: + """Two cubes joined by a thin neck, forming a single connected component.""" + arr = np.zeros((40, 24, 24), dtype=np.uint8) + arr[4:16, 6:18, 6:18] = 1 + arr[24:36, 6:18, 6:18] = 1 + c = 12 + arr[16:24, c - neck : c + neck, c - neck : c + neck] = 1 + return arr + + +class Test_split_connected_component(unittest.TestCase): + def test_input_really_is_one_component(self): + _, n = np_utils.np_connected_components(_dumbbell(), connectivity=3) + self.assertEqual(n, 1) + + def test_erosion_splits_into_two_parts(self): + out = np_utils.np_split_connected_component(_dumbbell(), method="erosion") + self.assertEqual(sorted(int(v) for v in np.unique(out)), [0, 1, 2]) + self.assertGreater((out == 1).sum(), 0) + self.assertGreater((out == 2).sum(), 0) + + def test_erosion_full_partition_keeps_every_voxel(self): + arr = _dumbbell() + out = np_utils.np_split_connected_component(arr, method="erosion", full_partition=True) + self.assertEqual((out != 0).sum(), (arr != 0).sum()) + # and never paints outside the input + self.assertEqual(((out != 0) & (arr == 0)).sum(), 0) + + def test_erosion_without_full_partition_returns_only_the_cores(self): + arr = _dumbbell() + out = np_utils.np_split_connected_component(arr, method="erosion", full_partition=False) + self.assertLess((out != 0).sum(), (arr != 0).sum()) + + def test_the_two_parts_land_on_opposite_cubes(self): + arr = _dumbbell() + out = np_utils.np_split_connected_component(arr, method="erosion") + low = out[4:16][out[4:16] != 0] + high = out[24:36][out[24:36] != 0] + # each cube must be dominated by a single, and different, label + self.assertNotEqual(np.bincount(low).argmax(), np.bincount(high).argmax()) + + def test_unknown_method_is_rejected(self): + with self.assertRaises(ValueError): + np_utils.np_split_connected_component(_dumbbell(), method="bogus") + + def test_solid_block_that_cannot_split_raises(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[5:15, 5:15, 5:15] = 1 + with self.assertRaises(ValueError): + np_utils.np_split_connected_component(arr, method="erosion", max_iter=2) + + +class Test_split_connected_component_mincut(unittest.TestCase): + def setUp(self): + try: + import networkx as nx # noqa: F401 + except ImportError: + self.skipTest("networkx not installed") + + def test_mincut_partitions_the_whole_volume(self): + arr = _dumbbell() + out = np_utils.np_split_connected_component(arr, method="mincut", connectivity=1) + self.assertEqual(sorted(int(v) for v in np.unique(out)), [0, 1, 2]) + self.assertEqual((out != 0).sum(), (arr != 0).sum()) + + def test_rejects_input_that_is_already_two_components(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[2:6, 2:6, 2:6] = 1 + arr[14:18, 14:18, 14:18] = 1 + with self.assertRaises(ValueError): + np_utils.np_split_connected_component(arr, method="mincut", connectivity=1) + + def test_min_volume_guard(self): + with self.assertRaises(ValueError): + np_utils.np_split_connected_component(_dumbbell(), method="mincut", connectivity=1, min_volume=10**6) + + def test_max_cut_guard(self): + with self.assertRaises(ValueError): + np_utils.np_split_connected_component(_dumbbell(), method="mincut", connectivity=1, max_cut=0.5) + + def test_anisotropic_zoom_is_accepted(self): + out = np_utils.np_split_connected_component(_dumbbell(), method="mincut", connectivity=1, zoom=(1.0, 1.0, 3.0)) + self.assertEqual(sorted(int(v) for v in np.unique(out)), [0, 1, 2]) + + def test_diagonal_edges_branch_runs(self): + """The upstream version raised NameError here (undefined ``geom.norm``).""" + out = np_utils.np_split_connected_component(_dumbbell(), method="mincut", connectivity=1, add_diagonal_edges=True) + self.assertEqual(sorted(int(v) for v in np.unique(out)), [0, 1, 2]) + + def test_max_ignore_none_runs(self): + """The upstream version raised NameError here (``max_errors`` never assigned).""" + out = np_utils.np_split_connected_component(_dumbbell(), method="mincut", connectivity=1, max_ignore=None) + self.assertEqual(sorted(int(v) for v in np.unique(out)), [0, 1, 2]) + + +class Test_connected_component_contact_map(unittest.TestCase): + def test_contact_map_labels_and_contact_zone(self): + out = np_utils.np_split_connected_component(_dumbbell(), method="erosion", full_partition=False) + contact = np_utils.np_connected_component_contact_map(out == 1, out == 2) + self.assertTrue({int(v) for v in np.unique(contact)} <= {0, 1, 2, 3}) + self.assertGreater((contact == 3).sum(), 0) + + def test_inputs_are_not_mutated(self): + out = np_utils.np_split_connected_component(_dumbbell(), method="erosion", full_partition=False) + a, b = out == 1, out == 2 + a_before, b_before = a.copy(), b.copy() + np_utils.np_connected_component_contact_map(a, b) + self.assertTrue(np.array_equal(a, a_before)) + self.assertTrue(np.array_equal(b, b_before)) + + if __name__ == "__main__": unittest.main()