feat(nii): add bounding-box-chain connected-component filter

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)
+
 
 
 

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.
 

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()