feat(nii): add relabel_by_position to renumber instances along an axis

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

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