diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index f27341f..0c0ffd0 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2962,6 +2962,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 b5bf4c2..ad28b7f 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()