feat(NII): split a label at a cumulative-volume fraction or at its center of mass

TPTBox/core/nii_wrapper.py

@@ -3007,6 +3007,126 @@ 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 _split_axis(self, axis: DIRECTIONS | int) -> tuple[int, bool]:
+        """Resolves ``axis`` to an array axis and whether that direction runs with increasing index."""
+        if isinstance(axis, int):
+            return axis, True
+        # get_axis falls back to the opposite letter, so recover which way the axis actually runs.
+        return self.get_axis(axis), axis in self.orientation
+
+    def _split_label_profile(self, labels: LABEL_REFERENCE, ax: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Returns the seg array, the mask of ``labels`` and its per-slice voxel count along ``ax``."""
+        arr = self.get_seg_array()
+        selected = arr > 0 if labels is None else np_isin(arr, labels)
+        counts = selected.sum(axis=tuple(i for i in range(arr.ndim) if i != ax))
+        return arr, selected, counts
+
+    def _split_label_at_index(self, arr: np.ndarray, selected: np.ndarray, ax: int, region: slice, offset: int, inplace: bool) -> Self:
+        """Adds ``offset`` to the selected voxels that lie inside ``region`` along axis ``ax``."""
+        part = np.zeros_like(selected)
+        part[(slice(None),) * ax + (region,)] = True
+        part &= selected
+        if offset != 0 and part.any():
+            arr = arr.astype(np.promote_types(arr.dtype, np.min_scalar_type(int(arr.max()) + offset)), copy=False)
+            arr[part] += offset
+        return self.set_array(arr, inplace=inplace)
+
+    def split_label_at_cumulative_volume_fraction(
+        self,
+        labels: LABEL_REFERENCE = None,
+        axis: DIRECTIONS | int = "P",
+        fraction: float = 0.5,
+        offset: int = 100,
+        inplace: bool = False,
+    ) -> Self:
+        """Splits labels along an axis at the plane that cuts off a given fraction of their volume.
+
+        Slices are walked inwards from the ``axis``-most end of the mask, summing voxels until
+        ``fraction`` of the total volume is covered; that leading slab gets ``offset`` added to its
+        label values, the rest is left untouched. With the default ``axis="P"`` and
+        ``fraction=1/3`` the posterior third (by volume, not by extent) is split off.
+
+        The direction is resolved against :attr:`orientation`, so no reorientation round-trip is
+        needed. Passing an ``int`` instead selects that array axis, pointing at high indices.
+
+        Args:
+            labels (LABEL_REFERENCE, optional): Label(s) to split. Defaults to None (every label).
+            axis (DIRECTIONS | int, optional): Anatomical direction the split-off part lies in,
+                or a raw array axis. Defaults to ``"P"``.
+            fraction (float, optional): Volume fraction that ends up in the split-off part, in
+                ``[0, 1]``. Defaults to 0.5.
+            offset (int, optional): Added to the labels of the split-off part. Defaults to 100.
+            inplace (bool, optional): If True, modifies this NII in place. Defaults to False.
+
+        Returns:
+            NII: The segmentation with the split-off part relabeled to ``label + offset``.
+
+        Examples:
+            >>> ethmoid.split_label_at_cumulative_volume_fraction([7, 8], "P", 1 / 3)  # doctest: +SKIP
+            # posterior third of labels 7 and 8 becomes 107 and 108
+        """
+        assert 0 <= fraction <= 1, f"fraction has to be in [0,1], got {fraction}"
+        ax, forward = self._split_axis(axis)
+        arr, selected, counts = self._split_label_profile(labels, ax)
+        total = int(counts.sum())
+        if total == 0:
+            return self if inplace else self.copy()
+        n = len(counts)
+        # walk inwards from the axis-most end until the slab holds fraction of the volume
+        order = counts[::-1] if forward else counts
+        target = fraction * total
+        taken = 0 if target <= 0 else min(int(np.searchsorted(np.cumsum(order), target)) + 1, n)
+        region = slice(n - taken, None) if forward else slice(None, taken)
+        return self._split_label_at_index(arr, selected, ax, region, offset, inplace)
+
+    def split_label_at_center_of_mass(
+        self,
+        labels: LABEL_REFERENCE = None,
+        axis: DIRECTIONS | int = "L",
+        offset: int = 100,
+        inplace: bool = False,
+    ) -> Self:
+        """Splits labels along an axis at their center of mass.
+
+        The center of mass of ``labels`` is projected onto ``axis`` and rounded to the nearest
+        slice; everything from there towards ``axis`` gets ``offset`` added to its label values.
+        With the default ``axis="L"`` this separates the left half of a structure from its right
+        half, and flipping the letter to ``"R"`` splits off the other side. The slice the plane
+        lands on always goes to the ``axis`` side, so the two halves overlap in that one slice
+        when the center of mass sits exactly on a voxel center.
+
+        The direction is resolved against :attr:`orientation` and the rounding is done in the
+        anatomical frame, so the split is the same no matter which orientation the image is
+        stored in and no reorientation round-trip is needed. Passing an ``int`` instead selects
+        that array axis, pointing at high indices.
+
+        Args:
+            labels (LABEL_REFERENCE, optional): Label(s) to split. Defaults to None (every label).
+            axis (DIRECTIONS | int, optional): Anatomical direction the split-off part lies in,
+                or a raw array axis. Defaults to ``"L"``.
+            offset (int, optional): Added to the labels of the split-off part. Defaults to 100.
+            inplace (bool, optional): If True, modifies this NII in place. Defaults to False.
+
+        Returns:
+            NII: The segmentation with the split-off part relabeled to ``label + offset``.
+
+        Examples:
+            >>> ethmoid.split_label_at_center_of_mass([7, 15], "L")  # doctest: +SKIP
+            # left half of labels 7 and 15 becomes 107 and 115
+        """
+        ax, forward = self._split_axis(axis)
+        arr, selected, counts = self._split_label_profile(labels, ax)
+        total = int(counts.sum())
+        if total == 0:
+            return self if inplace else self.copy()
+        n = len(counts)
+        com = float((counts * np.arange(n)).sum()) / total
+        # round in the anatomical frame, not the array frame, so the result does not depend on
+        # whether the axis happens to be stored forwards or backwards
+        c = int(np.clip(np.rint(com if forward else n - 1 - com), 0, n))
+        region = slice(c, None) if forward else slice(None, n - c)
+        return self._split_label_at_index(arr, selected, ax, region, offset, inplace)
+
     def filter_connected_components_by_bbox_chain(
         self,
         margin_mm: float = 0.0,

unit_tests/test_nii_extended.py

@@ -606,5 +606,93 @@ def test_matches_the_array_level_function(self):
         self.assertTrue(np.allclose(nii.label_interface_thickness(1, 2), expected, equal_nan=True))
 
 
+class Test_split_label_at(unittest.TestCase):
+    def _two_slabs(self) -> NII:
+        """Label 1 fills index 0-9 along axis 0, label 2 fills 10-19; 100 voxels per slice."""
+        arr = np.zeros((20, 10, 10), dtype=np.uint8)
+        arr[0:10] = 1
+        arr[10:20] = 2
+        return _make_nii(arr)
+
+    def test_fraction_splits_off_the_axis_most_part(self):
+        # RAS: axis 0 points to R, so the R-most quarter of the 2000 voxels is the last 5 slices
+        out = self._two_slabs().split_label_at_cumulative_volume_fraction(None, "R", 0.25).volumes()
+        self.assertEqual(dict(out), {1: 1000, 2: 500, 102: 500})
+
+    def test_fraction_is_by_volume_not_by_extent(self):
+        arr = np.zeros((20, 10, 10), dtype=np.uint8)
+        arr[0:16, 0:2] = 1  # thin: 20 voxels per slice
+        arr[16:20] = 1  # thick: 100 voxels per slice
+        # 70% of the 720 voxels is reached after the 4 thick slices plus 6 thin ones,
+        # i.e. 10 of 20 slices - taking 70% of the extent would give 14
+        out = _make_nii(arr).split_label_at_cumulative_volume_fraction(1, 0, 0.7).volumes()
+        self.assertEqual(out[101], 400 + 20 * 6)
+
+    def test_fraction_only_touches_the_requested_labels(self):
+        out = self._two_slabs().split_label_at_cumulative_volume_fraction(1, "R", 0.5).volumes()
+        self.assertEqual(dict(out), {1: 500, 101: 500, 2: 1000})
+
+    def test_fraction_bounds(self):
+        nii = self._two_slabs()
+        self.assertEqual(dict(nii.split_label_at_cumulative_volume_fraction(None, 0, 0.0).volumes()), {1: 1000, 2: 1000})
+        self.assertEqual(dict(nii.split_label_at_cumulative_volume_fraction(None, 0, 1.0).volumes()), {101: 1000, 102: 1000})
+        with self.assertRaises(AssertionError):
+            nii.split_label_at_cumulative_volume_fraction(None, 0, 1.5)
+
+    def test_center_of_mass_halves_a_label(self):
+        # label 2 spans 10..19, center of mass 14.5 -> plane at 14 (rint rounds .5 to even)
+        out = self._two_slabs().split_label_at_center_of_mass(2, 0).volumes()
+        self.assertEqual(dict(out), {1: 1000, 2: 400, 102: 600})
+
+    def test_center_of_mass_is_orientation_independent(self):
+        nii = self._two_slabs()
+        ref = nii.split_label_at_center_of_mass(None, "R")
+        for ori in [("L", "A", "S"), ("A", "R", "S"), ("I", "P", "L"), ("S", "L", "P")]:
+            got = nii.reorient(ori).split_label_at_center_of_mass(None, "R").reorient(nii.orientation)
+            self.assertTrue(np.array_equal(got.get_seg_array(), ref.get_seg_array()), ori)
+
+    def test_fraction_is_orientation_independent(self):
+        nii = self._two_slabs()
+        ref = nii.split_label_at_cumulative_volume_fraction(None, "P", 1 / 3)
+        for ori in [("L", "A", "S"), ("A", "R", "S"), ("I", "P", "L"), ("S", "L", "P")]:
+            got = nii.reorient(ori).split_label_at_cumulative_volume_fraction(None, "P", 1 / 3).reorient(nii.orientation)
+            self.assertTrue(np.array_equal(got.get_seg_array(), ref.get_seg_array()), ori)
+
+    def test_missing_label_is_a_no_op(self):
+        nii = self._two_slabs()
+        for out in (nii.split_label_at_center_of_mass(5, 0), nii.split_label_at_cumulative_volume_fraction(5, 0, 0.5)):
+            self.assertTrue(np.array_equal(out.get_seg_array(), nii.get_seg_array()))
+
+    def test_dtype_is_promoted_when_the_offset_overflows(self):
+        out = self._two_slabs().split_label_at_center_of_mass(None, 0, offset=1000)
+        self.assertEqual(dict(out.volumes()), {1: 1000, 1002: 1000})
+
+    def test_works_in_2d(self):
+        arr = np.zeros((10, 10), dtype=np.uint8)
+        arr[:, 0:6] = 3
+        out = _make_nii(arr).split_label_at_center_of_mass(3, 1, offset=10).volumes()
+        self.assertEqual(dict(out), {3: 20, 13: 40})
+
+    def test_inplace(self):
+        nii = self._two_slabs()
+        out = nii.split_label_at_center_of_mass(None, 0, inplace=True)
+        self.assertIs(out, nii)
+        self.assertIn(102, nii.volumes())
+
+    def test_ethmoid_style_pipeline(self):
+        """The motivating use case: posterior third, then left/right of the remaining front."""
+        arr = np.zeros((30, 25, 24), dtype=np.uint8)
+        arr[5:25, 5:20, 3:12] = 7
+        arr[5:25, 5:20, 12:21] = 8
+        seg = _make_nii(arr).reorient_(("P", "I", "R"))
+        out = seg.split_label_at_cumulative_volume_fraction([7, 8], "P", 1 / 3, 100)
+        out = out.split_label_at_center_of_mass([7], "R", 200).split_label_at_center_of_mass([8], "R", 200)
+        vols = out.volumes()
+        self.assertEqual(set(vols), {7, 8, 107, 108, 207, 208})
+        # the posterior third of each label, and the front split into two halves
+        self.assertEqual(vols[107] + vols[7] + vols[207], 20 * 15 * 9)
+        self.assertAlmostEqual(vols[107] / (20 * 15 * 9), 1 / 3, places=2)
+
+
 if __name__ == "__main__":
     unittest.main()