diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index e24b1de..0fa4366 100755 --- a/spineps/phase_instance.py +++ b/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 diff --git a/unit_tests/test_instance_pad_size.py b/unit_tests/test_instance_pad_size.py new file mode 100644 index 0000000..46877d3 --- /dev/null +++ b/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()