diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 0c0ffd0a..98379493 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -45,6 +45,7 @@ np_fill_holes, np_fill_holes_global_with_majority_voting, np_filter_connected_components, + np_filter_connected_components_by_bbox_chain, np_get_connected_components_center_of_mass, np_is_empty, np_isin, @@ -3005,6 +3006,55 @@ def relabel_by_position_(self, axis: DIRECTIONS = "I", offset: int = 0, verbose: """In-place version of :meth:`relabel_by_position`.""" return self.relabel_by_position(axis=axis, offset=offset, inplace=True, verbose=verbose) + def filter_connected_components_by_bbox_chain( + self, + margin_mm: float = 0.0, + extra_margin_mm: float = 0.0, + extra_margin_axis: DIRECTIONS | None = None, + connectivity: int = 3, + inplace: bool = False, + ) -> Self: + """Keeps only components whose bounding boxes chain onto the largest component. + + The mask is binarized and split into connected components. Starting from the largest one, + any component whose bounding box (grown by ``margin_mm``) overlaps the growing region on + every axis is kept, repeating until nothing new is added; everything else is removed. + Labels of the kept voxels are preserved. + + This is the "keep the spine, drop the unrelated blobs" filter: a structure broken into + several pieces along its length stays, while a component sitting off to the side goes. + + Margins are given in millimetres and converted per axis using :attr:`zoom`, so the region + grows by the same physical distance regardless of anisotropy. + + Args: + margin_mm (float, optional): Bounding-box margin in mm applied on every axis. + Defaults to 0.0. + extra_margin_mm (float, optional): Additional margin in mm along ``extra_margin_axis`` + only, to tolerate gaps along the structure's main direction. Defaults to 0.0. + extra_margin_axis (DIRECTIONS | None, optional): Anatomical direction the extra margin + applies to (e.g. ``"I"`` for a spine). Required when ``extra_margin_mm`` is set. + Defaults to None. + connectivity (int, optional): Connectivity used to find the components. Defaults to 3. + inplace (bool, optional): If True, modifies this NII in place. Defaults to False. + + Returns: + NII: The filtered segmentation. + """ + assert extra_margin_mm == 0 or extra_margin_axis is not None, "extra_margin_mm needs extra_margin_axis" + zoom = self.zoom + margin = [margin_mm / z for z in zoom] + axis = self.get_axis(extra_margin_axis) if extra_margin_axis is not None else None + extra = extra_margin_mm / zoom[axis] if axis is not None else 0.0 + arr = np_filter_connected_components_by_bbox_chain( + self.get_seg_array(), + margin=margin, + extra_margin=extra, + extra_margin_axis=axis, + connectivity=connectivity, + ) + return self.set_array(arr, inplace=inplace) + diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index ecadbb72..1408dc1b 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -315,6 +315,94 @@ def np_bounding_boxes(arr: UINTARRAY) -> dict[int, tuple[slice, slice, slice]]: return {idx: v for idx, v in enumerate(stats["bounding_boxes"]) if idx != 0 and vc[idx] > 0} +def np_slices_overlap(slice1: slice, slice2: slice) -> bool: + """Checks whether two ranges given as slices overlap or touch. + + Borders count as overlapping, so ``slice(0, 5)`` and ``slice(5, 9)`` overlap. + + Args: + slice1 (slice): First range; only ``start`` and ``stop`` are used. + slice2 (slice): Second range; only ``start`` and ``stop`` are used. + + Returns: + bool: True if the two ranges intersect or touch at a border. + """ + return slice1.start <= slice2.stop and slice2.start <= slice1.stop + + +def np_filter_connected_components_by_bbox_chain( + arr: np.ndarray, + margin: Sequence[float] | float = 0.0, + extra_margin: float = 0.0, + extra_margin_axis: int | None = None, + connectivity: int = 3, +) -> np.ndarray: + """Keeps only connected components whose bounding boxes chain onto the largest component. + + Starting from the largest connected component, any other component whose (margin-grown) + bounding box overlaps the growing region on *every* axis is incorporated, and the process + repeats until nothing new is added. Everything else is dropped. This keeps a structure that + is fragmented into several pieces along its length while discarding unrelated blobs + elsewhere in the volume. + + Args: + arr (np.ndarray): Input array. Treated as binary -- every non-zero voxel is foreground. + margin (Sequence[float] | float, optional): Bounding-box margin in voxels, either one + value for all axes or one per axis. Defaults to 0.0. + extra_margin (float, optional): Additional margin in voxels applied only along + ``extra_margin_axis``, to tolerate gaps along the structure's main direction. + Defaults to 0.0. + extra_margin_axis (int | None, optional): Axis that ``extra_margin`` applies to. + Required when ``extra_margin`` is non-zero. Defaults to None. + connectivity (int, optional): Connectivity used to find the components. Defaults to 3. + + Returns: + np.ndarray: A copy of ``arr`` with every non-incorporated component zeroed out. + """ + assert extra_margin == 0 or extra_margin_axis is not None, "extra_margin needs extra_margin_axis" + ndim = arr.ndim + margins = np.broadcast_to(np.asarray(margin, dtype=float), (ndim,)) + + cc, n = np_connected_components(arr != 0, connectivity=connectivity) + if n <= 1: + return arr.copy() + + boxes = np_bounding_boxes(cc) + # Largest component first, so the chain starts from the anchor the caller expects. + volumes = np_volume(cc) + order = sorted(volumes, key=lambda k: volumes[k], reverse=True) + + def grow(box): + return tuple(slice(floor(sl.start - margins[ax]), ceil(sl.stop + margins[ax])) for ax, sl in enumerate(box)) + + def widen(box): + if extra_margin_axis is None or extra_margin == 0: + return box + return tuple( + slice(floor(sl.start - extra_margin), ceil(sl.stop + extra_margin)) if ax == extra_margin_axis else sl + for ax, sl in enumerate(box) + ) + + anchor_label = order[0] + incorporated = [anchor_label] + region = [grow(boxes[anchor_label])] + changed = True + while changed: + changed = False + for k in [l for l in order if l not in incorporated]: + candidate = grow(boxes[k]) + # The already-incorporated box gets the extra axial margin, so a gap along the + # structure's main direction does not break the chain. + if any(all(np_slices_overlap(w, c) for w, c in zip(widen(box), candidate)) for box in region): + region.append(candidate) + incorporated.append(k) + changed = True + + out = arr.copy() + out[~np_isin(cc, incorporated)] = 0 + return out + + def np_contacts(arr: UINTARRAY, connectivity: int) -> dict[tuple[int, int], int]: """Calculates the contacting labels and the amount of touching voxels based on connectivity. diff --git a/seg_all.ipynb b/tutorials/seg_all.ipynb similarity index 100% rename from seg_all.ipynb rename to tutorials/seg_all.ipynb diff --git a/unit_tests/test_nii_extended.py b/unit_tests/test_nii_extended.py index ad28b7f3..b9e1b8d0 100644 --- a/unit_tests/test_nii_extended.py +++ b/unit_tests/test_nii_extended.py @@ -516,5 +516,75 @@ def test_labels_are_preserved_in_count(self): self.assertEqual(np.count_nonzero(out.get_seg_array()), np.count_nonzero(nii.get_seg_array())) +class Test_bbox_chain_filter(unittest.TestCase): + def _spine_with_outlier(self, zoom=(1.0, 1.0, 1.0)): + """Three blobs stacked along axis 2 with 6-voxel gaps, plus one unrelated blob far away.""" + arr = np.zeros((40, 40, 60), dtype=np.uint16) + arr[18:22, 18:22, 5:12] = 1 + arr[18:22, 18:22, 18:25] = 2 + arr[18:22, 18:22, 31:45] = 3 # largest -> the chain anchor + arr[2:6, 2:6, 50:56] = 4 # unrelated + return _make_nii(arr, zoom=zoom) + + def test_zero_margin_keeps_only_the_largest_component(self): + nii = self._spine_with_outlier() + self.assertEqual(sorted(nii.filter_connected_components_by_bbox_chain(margin_mm=0.0).unique()), [3]) + + def test_margin_chains_across_gaps_but_excludes_the_outlier(self): + nii = self._spine_with_outlier() + self.assertEqual(sorted(nii.filter_connected_components_by_bbox_chain(margin_mm=4.0).unique()), [1, 2, 3]) + + def test_extra_margin_applies_only_to_its_axis(self): + nii = self._spine_with_outlier() + out = nii.filter_connected_components_by_bbox_chain(margin_mm=0.0, extra_margin_mm=8.0, extra_margin_axis="S") + self.assertEqual(sorted(out.unique()), [1, 2, 3]) + + def test_large_margin_reaches_everything(self): + nii = self._spine_with_outlier() + self.assertEqual(sorted(nii.filter_connected_components_by_bbox_chain(margin_mm=40.0).unique()), [1, 2, 3, 4]) + + def test_margin_is_physical_not_voxel(self): + """At 4 mm spacing a 4 mm margin is one voxel, so the 6-voxel gap must not close.""" + nii = self._spine_with_outlier(zoom=(1.0, 1.0, 4.0)) + self.assertEqual(sorted(nii.filter_connected_components_by_bbox_chain(margin_mm=4.0).unique()), [3]) + + def test_original_labels_are_preserved(self): + nii = self._spine_with_outlier() + out = nii.filter_connected_components_by_bbox_chain(margin_mm=4.0) + kept = out.get_seg_array() + original = nii.get_seg_array() + self.assertTrue(np.array_equal(kept[kept != 0], original[kept != 0])) + + def test_single_component_is_untouched(self): + arr = np.zeros((10, 10, 10), dtype=np.uint16) + arr[2:6, 2:6, 2:6] = 9 + nii = _make_nii(arr) + out = nii.filter_connected_components_by_bbox_chain(margin_mm=1.0) + self.assertTrue(np.array_equal(out.get_seg_array(), arr)) + + def test_extra_margin_without_axis_is_rejected(self): + nii = self._spine_with_outlier() + with self.assertRaises(AssertionError): + nii.filter_connected_components_by_bbox_chain(extra_margin_mm=5.0) + + +class Test_slices_overlap(unittest.TestCase): + def test_overlap_cases(self): + from TPTBox.core.np_utils import np_slices_overlap + + self.assertTrue(np_slices_overlap(slice(0, 5), slice(5, 9))) # touching borders count + self.assertTrue(np_slices_overlap(slice(2, 4), slice(0, 9))) # nested + self.assertTrue(np_slices_overlap(slice(0, 9), slice(2, 4))) # nested, other way + self.assertTrue(np_slices_overlap(slice(0, 6), slice(4, 9))) # partial + self.assertFalse(np_slices_overlap(slice(0, 5), slice(7, 9))) + self.assertFalse(np_slices_overlap(slice(7, 9), slice(0, 5))) + + def test_is_symmetric(self): + from TPTBox.core.np_utils import np_slices_overlap + + for a, b in [(slice(0, 5), slice(5, 9)), (slice(0, 5), slice(7, 9)), (slice(1, 8), slice(3, 4))]: + self.assertEqual(np_slices_overlap(a, b), np_slices_overlap(b, a)) + + if __name__ == "__main__": unittest.main()