fix: remove the padding again in predict_instance_mask(pad_size=n)

spineps/phase_instance.py

@@ -137,7 +137,8 @@ def predict_instance_mask(
         seg_nii (NII): Subregion (semantic) segmentation mask used as input.
         model (SegmentationModel): Instance model producing the per-vertebra-body cutout predictions.
         debug_data (dict): Dictionary for collecting intermediate results across the pipeline.
-        pad_size (int, optional): Edge padding added before processing and removed afterwards. Defaults to 0.
+        pad_size (int, optional): Edge padding (in voxels of the model resolution) added around the mask
+            before the cutouts are taken and removed again from the result. Defaults to 0.
         proc_inst_fill_3d_holes (bool, optional): Whether to fill 3D holes in the final vertebra mask. Defaults to True.
         proc_detect_and_solve_merged_corpi (bool, optional): Whether to detect and split merged vertebral bodies. Defaults to True.
         proc_corpus_clean (bool, optional): Whether to clean small corpus connected-component artifacts. Defaults to True.
@@ -170,15 +171,16 @@ def predict_instance_mask(
         seg_nii_rdy = seg_nii.reorient(verbose=logger)
         debug_put(debug_data, "inst_uncropped_Subreg_nii_a_PIR", seg_nii_rdy.copy)
 
-        # Padding?
-        if pad_size > 0:
-            seg_nii_rdy = seg_nii_rdy.apply_pad(pad_size)
-
         zms = seg_nii_rdy.zoom
         logger.print("zms", zms, verbose=verbose)
         expected_zms = model.calc_recommended_resampling_zoom(seg_nii_rdy.zoom)
         if not seg_nii_rdy.assert_affine(zoom=expected_zms, raise_error=False):
             seg_nii_rdy.rescale_(expected_zms, verbose=logger)  # in PIR
+
+        # Padding? After the rescale, so that removing it at the end is the exact inverse: padding first
+        # meant `pad_size` input voxels became a different number of voxels in the rescaled grid.
+        if pad_size > 0:
+            seg_nii_rdy = seg_nii_rdy.apply_pad(pad_size)
         seg_nii_uncropped = seg_nii_rdy.copy()
         logger.print(
             "Vertebra seg_nii_uncropped", seg_nii_uncropped.zoom, seg_nii_uncropped.orientation, seg_nii_uncropped.shape, verbose=verbose
@@ -262,9 +264,11 @@ def predict_instance_mask(
         whole_vert_nii_uncropped = seg_nii_uncropped.set_array(uncropped_vert_mask)
         debug_put(debug_data, "inst_uncropped_vert_arr_a", whole_vert_nii_uncropped.copy)
 
-        # Uncrop again
+        # Remove the padding again. apply_pad() is not in place and returns a new NII, so the result has
+        # to be rebound -- dropping it left the padding on the returned mask, which is then neither the
+        # shape nor the field of view the caller handed in.
         if pad_size > 0:
-            whole_vert_nii_uncropped.apply_pad(-pad_size)
+            whole_vert_nii_uncropped = whole_vert_nii_uncropped.apply_pad(-pad_size, verbose=logger)
 
     return whole_vert_nii_uncropped, ErrCode.OK
 

unit_tests/test_instance_pad_size.py

@@ -0,0 +1,97 @@
+# Call 'python -m unittest' on this folder
+"""``predict_instance_mask(pad_size=n)`` must return a mask in the field of view it was handed.
+
+The padding added before the cutouts were taken was never removed: `whole_vert_nii_uncropped.apply_pad(-pad_size)`
+is not an in-place call, and its result was dropped. So any `pad_size > 0` returned a mask that is
+`2 * pad_size` voxels larger per axis than the input -- while the docstring promised the padding was
+"removed afterwards". Padding also happened *before* the rescale to the model resolution, so even a
+rebound `apply_pad(-pad_size)` would have removed the wrong number of voxels on a mask whose zoom differs
+from the model's.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+from TPTBox import NII, No_Logger
+from TPTBox.tests.test_utils import get_test_mri
+from typing_extensions import Self
+
+from spineps.phase_instance import predict_instance_mask
+from spineps.seg_enums import ErrCode, OutputType
+from spineps.seg_model import Segmentation_Inference_Config, SegmentationModel
+
+logger = No_Logger()
+
+
+class _ThirdsInstanceModel(SegmentationModel):
+    """Deterministic stand-in: splits each cutout into three bands along the inferior axis."""
+
+    def __init__(self, zoom: tuple[float, float, float] = (1.5, 1.5, 1.5)) -> None:
+        self.logger = No_Logger()
+        config = Segmentation_Inference_Config(
+            logger=self.logger,
+            modality=["SEG"],
+            acquisition="sag",
+            log_name="ThirdsInstanceModel",
+            modeltype="unet",
+            model_expected_orientation=("P", "I", "R"),
+            available_folds=1,
+            inference_augmentation=False,
+            resolution_range=list(zoom),
+            default_step_size=0.5,
+            labels={1: 1, 2: 2, 3: 3},
+            expected_inputs=["seg"],
+            cutout_size=(24, 24, 24),
+        )
+        super().__init__(__file__, config, default_verbose=False, default_allow_tqdm=False)
+
+    def load(self, folds: tuple[str, ...] | None = None) -> Self:  # noqa: ARG002
+        self.predictor = object()
+        return self
+
+    def run(self, input_nii: list[NII], verbose: bool = False) -> dict[OutputType, NII | None]:  # noqa: ARG002
+        nii = input_nii[0]
+        arr = nii.get_seg_array()
+        out = np.zeros_like(arr)
+        height = arr.shape[1]
+        for band, lo in enumerate((0, height // 3, 2 * height // 3)):
+            hi = height if band == 2 else (band + 1) * height // 3
+            band_slice = out[:, lo:hi, :]
+            band_slice[arr[:, lo:hi, :] != 0] = band + 1
+        return {OutputType.seg: nii.set_array(out), OutputType.softmax_logits: None}
+
+
+def _run(pad_size: int, zoom: tuple[float, float, float] = (1.5, 1.5, 1.5)) -> NII:
+    _mri, subreg, _vert, _label = get_test_mri()
+    model = _ThirdsInstanceModel(zoom).load()
+    vert_nii, errcode = predict_instance_mask(subreg.copy(), model, debug_data={}, pad_size=pad_size, verbose=False)
+    assert errcode == ErrCode.OK, errcode
+    assert vert_nii is not None
+    return vert_nii
+
+
+class Test_Instance_Pad_Size(unittest.TestCase):
+    def test_padding_is_removed_again(self):
+        unpadded = _run(pad_size=0)
+        padded = _run(pad_size=4)
+        self.assertEqual(padded.shape, unpadded.shape)
+        self.assertTrue(padded.assert_affine(other=unpadded), "pad_size changed the field of view")
+
+    def test_padding_does_not_change_the_result(self):
+        unpadded = _run(pad_size=0)
+        padded = _run(pad_size=4)
+        np.testing.assert_array_equal(padded.get_seg_array(), unpadded.get_seg_array())
+
+    def test_padding_is_removed_when_the_model_rescales(self):
+        # The fixture is at 1.5mm; a model at 3.0mm halves the grid, so `pad_size` input voxels and
+        # `pad_size` model voxels are different distances. The output must still match the unpadded run.
+        unpadded = _run(pad_size=0, zoom=(3.0, 3.0, 3.0))
+        padded = _run(pad_size=3, zoom=(3.0, 3.0, 3.0))
+        self.assertEqual(padded.shape, unpadded.shape)
+        self.assertTrue(padded.assert_affine(other=unpadded))
+
+
+if __name__ == "__main__":
+    unittest.main()