From d6b1a47bfb484fe7c2da0758ec2600014ebbc86d Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 27 Aug 2026 14:29:02 +0000 Subject: [PATCH] feat(nii): add relabel_by_position to renumber instances along an axis Renumbers instance labels 1, 2, 3, ... in the order their centers of mass appear when travelling along an anatomical direction. With the default axis="I" the most superior instance becomes label 1 and numbering runs downwards. Three sibling repos each carried their own version of this (spineps label_instance_top_to_bottom, nako-segmentation label_top2bot and label_top2bot_advanced), all of which reorient to PIR, sort by axis 1, remap and reorient back. This version resolves the direction against the image's existing orientation instead, so no reorientation round-trip is needed and the result is identical whatever orientation the caller happens to be in. Co-Authored-By: Claude Opus 5 --- TPTBox/core/nii_wrapper.py | 43 +++++++++++++++++++++++++++ unit_tests/test_nii_extended.py | 51 +++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 833761c0..94e5072a 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2932,6 +2932,49 @@ def center_of_masses(self) -> dict[int, COORDINATE]: """Returns a dict stating the center of mass for each present label (not including zero!).""" return np_center_of_mass(self.get_seg_array()) + def relabel_by_position( + self, + axis: DIRECTIONS = "I", + offset: int = 0, + inplace: bool = False, + verbose: logging = False, + ) -> Self: + """Relabels instances consecutively by their position along an anatomical axis. + + Labels are renumbered ``1 + offset``, ``2 + offset``, ... in the order their centers of + mass appear when travelling along ``axis``. With the default ``axis="I"`` the most + superior instance becomes label ``1 + offset`` and numbering runs downwards. + + This works in whatever orientation the image already has -- the direction is resolved + against :attr:`orientation`, so no reorientation round-trip is needed. + + Args: + axis (DIRECTIONS, optional): Anatomical direction the numbering advances along. + Defaults to ``"I"`` (superior to inferior). + offset (int, optional): Added to every new label, so numbering starts at + ``1 + offset``. Defaults to 0. + inplace (bool, optional): If True, modifies this NII in place. Defaults to False. + verbose (logging, optional): Passed through to :meth:`map_labels`. Defaults to False. + + Returns: + NII: The relabeled segmentation. + + Examples: + >>> vert.relabel_by_position("I") # doctest: +SKIP + # topmost vertebra -> 1, next one down -> 2, ... + """ + ax = self.get_axis(axis) + # get_axis falls back to the opposite letter, so recover which way the axis actually runs. + forward = axis in self.orientation + coms = np_center_of_mass(self.get_seg_array()) + ordered = sorted(coms.items(), key=lambda kv: kv[1][ax], reverse=not forward) + label_map = {int(label): idx + 1 + offset for idx, (label, _) in enumerate(ordered)} + return self.map_labels(label_map, verbose=verbose, inplace=inplace) + + def relabel_by_position_(self, axis: DIRECTIONS = "I", offset: int = 0, verbose: logging = False) -> Self: + """In-place version of :meth:`relabel_by_position`.""" + return self.relabel_by_position(axis=axis, offset=offset, inplace=True, verbose=verbose) + diff --git a/unit_tests/test_nii_extended.py b/unit_tests/test_nii_extended.py index b5bf4c26..ad28b7f3 100644 --- a/unit_tests/test_nii_extended.py +++ b/unit_tests/test_nii_extended.py @@ -465,5 +465,56 @@ def test_returns_array_when_not_seg(self): self.assertEqual(result.shape, arr.shape) +class Test_relabel_by_position(unittest.TestCase): + def _stacked(self, labels=(7, 3, 5)): + """Three blobs stacked along axis 2, listed from low index to high index.""" + arr = np.zeros((10, 10, 24), dtype=np.uint16) + for i, label in enumerate(labels): + arr[4:6, 4:6, 2 + 8 * i : 5 + 8 * i] = label + return _make_nii(arr) + + def test_numbering_runs_along_the_requested_direction(self): + nii = self._stacked() + self.assertNotIn("I", nii.orientation) # axis 2 is 'S', so the direction is flipped + coms = nii.relabel_by_position("I").center_of_masses() + self.assertGreater(coms[1][2], coms[2][2]) + self.assertGreater(coms[2][2], coms[3][2]) + + def test_opposite_direction_reverses_numbering(self): + nii = self._stacked() + coms = nii.relabel_by_position("S").center_of_masses() + self.assertLess(coms[1][2], coms[2][2]) + self.assertLess(coms[2][2], coms[3][2]) + + def test_offset_shifts_the_label_range(self): + nii = self._stacked() + self.assertEqual(sorted(nii.relabel_by_position("I", offset=100).unique()), [101, 102, 103]) + + def test_result_does_not_depend_on_input_orientation(self): + nii = self._stacked() + expected = nii.relabel_by_position("I") + for ori in (("P", "I", "R"), ("L", "A", "S"), ("A", "S", "L")): + rotated = nii.reorient(ori).relabel_by_position("I").reorient(nii.orientation) + self.assertTrue(np.array_equal(rotated.get_seg_array(), expected.get_seg_array()), f"differs for {ori}") + + def test_inplace_variant_matches(self): + nii = self._stacked() + expected = nii.relabel_by_position("I") + clone = nii.copy() + clone.relabel_by_position_("I") + self.assertTrue(np.array_equal(clone.get_seg_array(), expected.get_seg_array())) + + def test_single_label_is_renumbered_to_one(self): + arr = np.zeros((8, 8, 8), dtype=np.uint16) + arr[2:5, 2:5, 2:5] = 42 + self.assertEqual(sorted(_make_nii(arr).relabel_by_position("I").unique()), [1]) + + def test_labels_are_preserved_in_count(self): + nii = self._stacked(labels=(11, 4, 9)) + out = nii.relabel_by_position("I") + self.assertEqual(len(out.unique()), len(nii.unique())) + self.assertEqual(np.count_nonzero(out.get_seg_array()), np.count_nonzero(nii.get_seg_array())) + + if __name__ == "__main__": unittest.main()