From 43c49a452b8ede8c40d883eb9b3ede931b968e7d Mon Sep 17 00:00:00 2001 From: iback Date: Mon, 31 Aug 2026 06:54:13 +0000 Subject: [PATCH] feat(NII): add split_label_at_cumulative_volume_fraction and split_label_at_center_of_mass Two reusable label-splitting kernels lifted out of paraside's split_ethmoid, which hard-coded the ethmoid label ids (7/8/15/16) and offsets (+100/+200) around otherwise generic logic. split_label_at_cumulative_volume_fraction walks slices inwards from the axis-most end of a mask, summing voxels until a given fraction of the total volume is covered, and adds an offset to the labels of that slab. Note this is a volume fraction, not an extent fraction, so a structure that is thin at one end is not cut in the middle of its bounding box. split_label_at_center_of_mass cuts at the rounded center of mass along an axis and offsets the labels on the axis side of the plane. Changes from the source implementation: - the axis is an anatomical direction resolved against `orientation`, so the reorient round-trip the original did (and the reorient back at the end) is gone, and either method can also take a raw array axis as an int - the center-of-mass rounding happens in the anatomical frame, so the split is identical whichever orientation the image is stored in; rounding in the array frame shifted it by a slice when the center of mass landed on a voxel center - the offset is applied via a mask instead of `arr[slab] += offset` followed by `arr[arr == offset] = 0`, which silently breaks for any label equal to the offset, and the dtype is promoted when the offset would overflow it - fraction 0 now yields an empty slab; the original had no such case and its loop left `cutoff_slice` as None when the threshold was never reached, which then crashed on an undefined variable - labels not present in the mask are a no-op instead of a KeyError A test composes the two into the original ethmoid pipeline; the result is voxel-identical to split_ethmoid. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01GFZf44LHBB6N5jaf6YFzhx --- TPTBox/core/nii_wrapper.py | 120 ++++++++++++++++++++++++++++++++ unit_tests/test_nii_extended.py | 88 +++++++++++++++++++++++ 2 files changed, 208 insertions(+) diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index bb3f8ec..8388bcb 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/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, diff --git a/unit_tests/test_nii_extended.py b/unit_tests/test_nii_extended.py index 827531b..2dbf8f5 100644 --- a/unit_tests/test_nii_extended.py +++ b/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()