From c1adf524e17cfa46bbd4619a1d7cf142b3956de1 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 10 Nov 2026 09:03:55 +0100 Subject: [PATCH 1/6] fix(vision_augmentation): keep gpu anisotropy image on the mask grid GpuPatchAugmentation downsampled with F.interpolate(mode='nearest'), which takes the first voxel of each block, so after upsampling the image sat (factor - 1) / 2 voxels away from the unchanged mask (0.5 to 1.5 voxels for factors 2 to 4). 'nearest-exact' samples the block centres: the shift is 0 for odd factors and 0.5 voxel for even ones, as in TorchIO's RandomAnisotropy. Models trained with anisotropy (the VS config uses p=0.25) saw this as label noise on about a quarter of the patches; no retraining is needed for it alone. --- fastMONAI/vision_augmentation.py | 6 ++++-- nbs/03_vision_augment.ipynb | 17 ++++++++++++++--- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/fastMONAI/vision_augmentation.py b/fastMONAI/vision_augmentation.py index 206ead2..ddb25e9 100644 --- a/fastMONAI/vision_augmentation.py +++ b/fastMONAI/vision_augmentation.py @@ -1261,8 +1261,10 @@ def _apply_anisotropy(self, img, mask): down_size = list(orig_size) down_size[axis] = max(1, round(orig_size[axis] / factor)) - # Downsample with nearest, upsample with trilinear - down = F.interpolate(sample, size=down_size, mode='nearest') + # Downsample with nearest at the sample centres, so the image moves at + # most half a voxel against the mask (as in TorchIO), then upsample with + # trilinear + down = F.interpolate(sample, size=down_size, mode='nearest-exact') up = F.interpolate(down, size=orig_size, mode='trilinear', align_corners=False) img[i:i+1] = up diff --git a/nbs/03_vision_augment.ipynb b/nbs/03_vision_augment.ipynb index 010cad1..3a0ab7a 100644 --- a/nbs/03_vision_augment.ipynb +++ b/nbs/03_vision_augment.ipynb @@ -1583,8 +1583,10 @@ " down_size = list(orig_size)\n", " down_size[axis] = max(1, round(orig_size[axis] / factor))\n", "\n", - " # Downsample with nearest, upsample with trilinear\n", - " down = F.interpolate(sample, size=down_size, mode='nearest')\n", + " # Downsample with nearest at the sample centres, so the image moves at\n", + " # most half a voxel against the mask (as in TorchIO), then upsample with\n", + " # trilinear\n", + " down = F.interpolate(sample, size=down_size, mode='nearest-exact')\n", " up = F.interpolate(down, size=orig_size, mode='trilinear',\n", " align_corners=False)\n", " img[i:i+1] = up\n", @@ -2115,7 +2117,16 @@ "# Determinism: same seed -> identical output\n", "torch.manual_seed(123); _a, _ = _full(_img.clone(), _mask.clone())\n", "torch.manual_seed(123); _b, _ = _full(_img.clone(), _mask.clone())\n", - "test_eq(torch.equal(_a, _b), True)" + "test_eq(torch.equal(_a, _b), True)\n", + "\n", + "# Anisotropy moves the image at most half a voxel against the mask: a ramp along the\n", + "# downsampled axis comes back within half a voxel of where it was\n", + "_ramp = torch.arange(48.).view(1, 1, 1, 1, 48).repeat(1, 1, 4, 4, 1)\n", + "for _f in (2, 3, 4):\n", + " _an = GpuPatchAugmentation(anisotropy={'axes': (2,), 'downsampling': (_f, _f), 'p': 1.0})\n", + " _out, _ = _an(_ramp.clone(), None)\n", + " _shift = (_out - _ramp)[0, 0, 2, 2, 8:40].mean().item()\n", + " assert abs(_shift) <= 0.5, f\"factor {_f}: image shifted {_shift:.2f} voxels\"" ] }, { From fccfe7abdc9460695f193165663bb12caebdcb66 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 10 Nov 2026 09:04:04 +0100 Subject: [PATCH 2/6] fix(mlflow): find training runs in the fastmonai store fastMONAI points MLflow at its own store only on its first MLflow call. Code that called raw mlflow first in a fresh process searched /mlflow.db instead, so `prepare_model_bundle.py --run` failed with "Run ... not found" and tutorial 11f with "No MLflow experiment found", although training had logged the runs. Both now set fastMONAI's store first unless MLFLOW_TRACKING_URI is set. The commented MLflow option in 12c uses find_model_artifacts, which does the same. --- nbs/11f_tutorial_inference.ipynb | 3 +++ nbs/12c_tutorial_patch_inference.ipynb | 4 +--- .../deployment/pacs/prepare_model_bundle.py | 4 ++++ .../tests/deployment/test_prepare_model_bundle.py | 15 +++++++++++++++ 4 files changed, 23 insertions(+), 3 deletions(-) diff --git a/nbs/11f_tutorial_inference.ipynb b/nbs/11f_tutorial_inference.ipynb index 30fb3ad..ade03ad 100644 --- a/nbs/11f_tutorial_inference.ipynb +++ b/nbs/11f_tutorial_inference.ipynb @@ -113,6 +113,9 @@ "# MODEL_RUN_ID = \"your-mlflow-run-id\"\n", "\n", "if MODEL_RUN_ID is None:\n", + " # Search the store 11d logged to (fastMONAI's default unless MLFLOW_TRACKING_URI is set).\n", + " if not mlflow.is_tracking_uri_set():\n", + " set_mlflow_tracking_uri()\n", " experiment = mlflow.get_experiment_by_name(task)\n", " if experiment is None:\n", " raise RuntimeError(f\"No MLflow experiment found for {task!r}. Run tutorial 11d first.\")\n", diff --git a/nbs/12c_tutorial_patch_inference.ipynb b/nbs/12c_tutorial_patch_inference.ipynb index e46b76c..6d86101 100644 --- a/nbs/12c_tutorial_patch_inference.ipynb +++ b/nbs/12c_tutorial_patch_inference.ipynb @@ -200,10 +200,8 @@ "model_path = Path('models/final_model.safetensors')\n", "\n", "# Option 2: download the final artifact from MLflow (uncomment to use).\n", - "# import mlflow\n", "# run_id = \"your_run_id\" # the final_all_data run, from the MLflow UI\n", - "# model_path = Path(mlflow.artifacts.download_artifacts(\n", - "# run_id=run_id, artifact_path=\"model/final_model.safetensors\", dst_path=\"./\"))\n", + "# model_path = find_model_artifacts(run_ids={\"model\": run_id}, artifact_role=\"final\")[\"model\"]\n", "\n", "metadata = read_safetensors_metadata(model_path)\n", "if metadata['artifact_role'] != 'final':\n", diff --git a/research/vestibular_schwannoma/deployment/pacs/prepare_model_bundle.py b/research/vestibular_schwannoma/deployment/pacs/prepare_model_bundle.py index d07575b..4c3fe25 100644 --- a/research/vestibular_schwannoma/deployment/pacs/prepare_model_bundle.py +++ b/research/vestibular_schwannoma/deployment/pacs/prepare_model_bundle.py @@ -38,6 +38,7 @@ load_safetensors_model, make_output_spec, read_safetensors_metadata, + set_mlflow_tracking_uri, ) from deployment_hashing import ( bundle_sha256 as _bundle_sha256, @@ -97,6 +98,9 @@ def _declared_sources(args) -> list[MemberSource]: def _download_run_artifact(run_id: str, artifact_path: str) -> Path: import mlflow + # Use fastMONAI's store, where training logged, unless MLFLOW_TRACKING_URI is set. + if not mlflow.is_tracking_uri_set(): + set_mlflow_tracking_uri() return Path(mlflow.artifacts.download_artifacts( run_id=run_id, artifact_path=artifact_path )) diff --git a/research/vestibular_schwannoma/tests/deployment/test_prepare_model_bundle.py b/research/vestibular_schwannoma/tests/deployment/test_prepare_model_bundle.py index 1a7f441..06b9648 100644 --- a/research/vestibular_schwannoma/tests/deployment/test_prepare_model_bundle.py +++ b/research/vestibular_schwannoma/tests/deployment/test_prepare_model_bundle.py @@ -33,6 +33,21 @@ def parse_run(self, role, *extra): ] ) + def test_run_download_uses_fastmonai_store_unless_uri_is_configured(self): + import mlflow + + for configured in (False, True): + with self.subTest(configured=configured), \ + patch.object(mlflow, "is_tracking_uri_set", return_value=configured), \ + patch.object(bundle, "set_mlflow_tracking_uri") as set_uri, \ + patch.object(mlflow.artifacts, "download_artifacts", + return_value="/tmp/model.safetensors") as download: + path = bundle._download_run_artifact("run-id", "model/best_model.safetensors") + self.assertEqual(path, Path("/tmp/model.safetensors")) + self.assertEqual(set_uri.called, not configured) + download.assert_called_once_with( + run_id="run-id", artifact_path="model/best_model.safetensors") + def test_pacs_and_research_expect_the_same_output_contract(self): self.assertEqual( bundle.make_output_spec("multiclass_segmentation", classes=2), From 7da301b08fc6bf282ec4c8dd77fd018348835cf9 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 10 Nov 2026 09:04:12 +0100 Subject: [PATCH 3/6] fix(utils): fail a tracking run left open by an interrupted fit fastai skips after_fit when a fit raises or is interrupted (Ctrl-C), so the run the callback started stays active. Only the same callback instance cleaned it up at its next fit. A new callback, as when a notebook cell is run again, raised "Run ... already active" with auto_start=True and silently joined the stale run with auto_start=False. Runs the callback starts are now recorded until after_fit or mark_failed ends them, and the next fit of any ModelTrackingCallback ends such a run as FAILED before starting its own. A caller's own run is still joined. --- fastMONAI/_modidx.py | 1 + fastMONAI/utils.py | 33 +++++++++++++++++++++++------- nbs/07_utils.ipynb | 48 +++++++++++++++++++++++++++++++++++++------- 3 files changed, 68 insertions(+), 14 deletions(-) diff --git a/fastMONAI/_modidx.py b/fastMONAI/_modidx.py index f5010a6..4566ac4 100644 --- a/fastMONAI/_modidx.py +++ b/fastMONAI/_modidx.py @@ -148,6 +148,7 @@ 'fastMONAI.utils._extract_size_from_transforms': ( 'utils.html#_extract_size_from_transforms', 'fastMONAI/utils.py'), 'fastMONAI.utils._extract_standard_config': ('utils.html#_extract_standard_config', 'fastMONAI/utils.py'), + 'fastMONAI.utils._fail_open_callback_run': ('utils.html#_fail_open_callback_run', 'fastMONAI/utils.py'), 'fastMONAI.utils._get_x_column': ('utils.html#_get_x_column', 'fastMONAI/utils.py'), 'fastMONAI.utils._json_round_trip': ('utils.html#_json_round_trip', 'fastMONAI/utils.py'), 'fastMONAI.utils._monai_constructor': ('utils.html#_monai_constructor', 'fastMONAI/utils.py'), diff --git a/fastMONAI/utils.py b/fastMONAI/utils.py index 9f11c5c..0f1875c 100644 --- a/fastMONAI/utils.py +++ b/fastMONAI/utils.py @@ -950,6 +950,25 @@ def _read_preprocessing_manifest(path, dataset_version=None): # %% ../nbs/07_utils.ipynb #030876d2 +# IDs of runs started by a ModelTrackingCallback and not ended yet. A fit that raises +# or is interrupted before after_fit leaves its run active; the next fit ends it as +# FAILED instead of joining it. +_OPEN_CALLBACK_RUNS = set() + + +def _fail_open_callback_run(): + """End the active run as FAILED if a ModelTrackingCallback started it and never ended it. + + Returns True if a run was ended. + """ + active_run = mlflow.active_run() + if active_run is None or active_run.info.run_id not in _OPEN_CALLBACK_RUNS: + return False + _OPEN_CALLBACK_RUNS.discard(active_run.info.run_id) + mlflow.end_run(status='FAILED') + return True + + class ModelTrackingCallback(Callback): """ A FastAI callback for comprehensive MLflow experiment tracking. @@ -1247,15 +1266,12 @@ def before_fit(self) -> None: With ``auto_start=False`` the active run is used. If no run is active, the callback starts its own, so MLflow does not open an implicit run that stays active after fit. A run started by this callback is ended with status FAILED - if this setup raises, or at the next fit if an earlier fit raised before - after_fit. + if this setup raises. A run that a ModelTrackingCallback started but never + ended, because its fit raised or was interrupted before after_fit, is ended + with status FAILED at the next fit. """ active_run = mlflow.active_run() - if ( - self._auto_started and active_run is not None - and active_run.info.run_id == self.run_id - ): - self.mark_failed() + if _fail_open_callback_run(): active_run = None if self.auto_start or active_run is None: @@ -1265,6 +1281,7 @@ def before_fit(self) -> None: self.run_name = self.run_name or f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}" started_run = mlflow.start_run(run_name=self.run_name) self.run_id = started_run.info.run_id + _OPEN_CALLBACK_RUNS.add(self.run_id) self._auto_started = True self._owns_run = True self._run_failed = False @@ -1324,6 +1341,7 @@ def after_fit(self) -> None: if self._auto_started: mlflow.end_run() + _OPEN_CALLBACK_RUNS.discard(self.run_id) self._auto_started = False def mark_failed(self) -> None: @@ -1338,6 +1356,7 @@ def mark_failed(self) -> None: mlflow.tracking.MlflowClient().set_terminated( self.run_id, status='FAILED' ) + _OPEN_CALLBACK_RUNS.discard(self.run_id) self._run_failed = True def log_metrics(self, metrics: dict) -> None: diff --git a/nbs/07_utils.ipynb b/nbs/07_utils.ipynb index d5e0ff0..c47818b 100644 --- a/nbs/07_utils.ipynb +++ b/nbs/07_utils.ipynb @@ -1085,6 +1085,25 @@ "outputs": [], "source": [ "#| export\n", + "# IDs of runs started by a ModelTrackingCallback and not ended yet. A fit that raises\n", + "# or is interrupted before after_fit leaves its run active; the next fit ends it as\n", + "# FAILED instead of joining it.\n", + "_OPEN_CALLBACK_RUNS = set()\n", + "\n", + "\n", + "def _fail_open_callback_run():\n", + " \"\"\"End the active run as FAILED if a ModelTrackingCallback started it and never ended it.\n", + "\n", + " Returns True if a run was ended.\n", + " \"\"\"\n", + " active_run = mlflow.active_run()\n", + " if active_run is None or active_run.info.run_id not in _OPEN_CALLBACK_RUNS:\n", + " return False\n", + " _OPEN_CALLBACK_RUNS.discard(active_run.info.run_id)\n", + " mlflow.end_run(status='FAILED')\n", + " return True\n", + "\n", + "\n", "class ModelTrackingCallback(Callback):\n", " \"\"\"\n", " A FastAI callback for comprehensive MLflow experiment tracking.\n", @@ -1382,15 +1401,12 @@ " With ``auto_start=False`` the active run is used. If no run is active, the\n", " callback starts its own, so MLflow does not open an implicit run that stays\n", " active after fit. A run started by this callback is ended with status FAILED\n", - " if this setup raises, or at the next fit if an earlier fit raised before\n", - " after_fit.\n", + " if this setup raises. A run that a ModelTrackingCallback started but never\n", + " ended, because its fit raised or was interrupted before after_fit, is ended\n", + " with status FAILED at the next fit.\n", " \"\"\"\n", " active_run = mlflow.active_run()\n", - " if (\n", - " self._auto_started and active_run is not None\n", - " and active_run.info.run_id == self.run_id\n", - " ):\n", - " self.mark_failed()\n", + " if _fail_open_callback_run():\n", " active_run = None\n", "\n", " if self.auto_start or active_run is None:\n", @@ -1400,6 +1416,7 @@ " self.run_name = self.run_name or f\"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}\"\n", " started_run = mlflow.start_run(run_name=self.run_name)\n", " self.run_id = started_run.info.run_id\n", + " _OPEN_CALLBACK_RUNS.add(self.run_id)\n", " self._auto_started = True\n", " self._owns_run = True\n", " self._run_failed = False\n", @@ -1459,6 +1476,7 @@ " \n", " if self._auto_started:\n", " mlflow.end_run()\n", + " _OPEN_CALLBACK_RUNS.discard(self.run_id)\n", " self._auto_started = False\n", "\n", " def mark_failed(self) -> None:\n", @@ -1473,6 +1491,7 @@ " mlflow.tracking.MlflowClient().set_terminated(\n", " self.run_id, status='FAILED'\n", " )\n", + " _OPEN_CALLBACK_RUNS.discard(self.run_id)\n", " self._run_failed = True\n", "\n", " def log_metrics(self, metrics: dict) -> None:\n", @@ -2298,6 +2317,21 @@ " test_eq(mlflow.active_run().info.run_id, _cb.run_id)\n", " mlflow.end_run()\n", "\n", + " # The same when the fit was interrupted and a new callback runs the next fit,\n", + " # with either auto_start: the stale run is failed, not joined.\n", + " for _auto in (True, False):\n", + " _cb = _tracking_cb(auto_start=_auto, log_split=False)\n", + " _cb.before_fit()\n", + " _stale_run_id = _cb.run_id\n", + " _cb = _tracking_cb(auto_start=_auto, log_split=False)\n", + " _cb.before_fit()\n", + " test_eq(_client.get_run(_stale_run_id).info.status, 'FAILED')\n", + " test_eq(mlflow.active_run().info.run_id, _cb.run_id)\n", + " test_eq(_cb.run_id == _stale_run_id, False)\n", + " with _patch.object(_cb, '_save_model_artifacts'):\n", + " _cb.after_fit()\n", + " test_eq(_client.get_run(_cb.run_id).info.status, 'FINISHED')\n", + "\n", " # A caller-owned run stays active and is not marked failed.\n", " with mlflow.start_run() as _external:\n", " _cb = _tracking_cb()\n", From 36fae046c76f58d29ecf26c37505f9b4fe9abf3f Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 10 Nov 2026 09:04:20 +0100 Subject: [PATCH 4/6] fix(vision_inference): restore predictions onto the original image grid Predictions were resized to the original shape with tio.Resize, then reoriented and given the original affine. Resize only matches the shape: on many grids its output came out one slice short and was padded at one end, so the mask sat half a slice off, e.g. 0.6 mm on 120 slices of 1.2 mm, and grids that do not divide evenly were slightly stretched. The VS training grids happened to be unaffected, but clinical grids (1.2 mm x 120/160 slices, 0.6 and 0.8 mm) are not. inference, inference_from_config and PatchInferenceEngine now resample the prediction with tio.Resample onto the original image's own shape and affine, which also undoes the reordering. The prediction's edge voxels are repeated once first, so original edge slices beyond the last model voxel keep the edge value instead of 0. Shape-based label restore and probability maps use the same path. inference and inference_from_config no longer need org_size to use caller-supplied images; the argument is still accepted and ignored. Deployed VS masks and vote maps move by up to half a slice on affected grids once the PACS container is rebuilt with this version. --- fastMONAI/_modidx.py | 12 ++- fastMONAI/vision_inference.py | 82 +++++++++--------- fastMONAI/vision_patch.py | 23 +++-- nbs/06_vision_inference.ipynb | 153 +++++++++++++++++++++------------- nbs/10_vision_patch.ipynb | 82 +++++++++++------- 5 files changed, 198 insertions(+), 154 deletions(-) diff --git a/fastMONAI/_modidx.py b/fastMONAI/_modidx.py index 4566ac4..510c3e3 100644 --- a/fastMONAI/_modidx.py +++ b/fastMONAI/_modidx.py @@ -459,18 +459,16 @@ 'fastMONAI/vision_data.py'), 'fastMONAI.vision_data.show_batch': ('vision_data.html#show_batch', 'fastMONAI/vision_data.py'), 'fastMONAI.vision_data.show_results': ('vision_data.html#show_results', 'fastMONAI/vision_data.py')}, - 'fastMONAI.vision_inference': { 'fastMONAI.vision_inference._do_resize': ( 'vision_inference.html#_do_resize', - 'fastMONAI/vision_inference.py'), - 'fastMONAI.vision_inference._prediction_save_path': ( 'vision_inference.html#_prediction_save_path', + 'fastMONAI.vision_inference': { 'fastMONAI.vision_inference._prediction_save_path': ( 'vision_inference.html#_prediction_save_path', 'fastMONAI/vision_inference.py'), - 'fastMONAI.vision_inference._resize_labels': ( 'vision_inference.html#_resize_labels', - 'fastMONAI/vision_inference.py'), + 'fastMONAI.vision_inference._resample_labels': ( 'vision_inference.html#_resample_labels', + 'fastMONAI/vision_inference.py'), + 'fastMONAI.vision_inference._resample_to_grid': ( 'vision_inference.html#_resample_to_grid', + 'fastMONAI/vision_inference.py'), 'fastMONAI.vision_inference._restore_prediction': ( 'vision_inference.html#_restore_prediction', 'fastMONAI/vision_inference.py'), 'fastMONAI.vision_inference._signed_depth': ( 'vision_inference.html#_signed_depth', 'fastMONAI/vision_inference.py'), - 'fastMONAI.vision_inference._to_original_orientation': ( 'vision_inference.html#_to_original_orientation', - 'fastMONAI/vision_inference.py'), 'fastMONAI.vision_inference._update_uid': ( 'vision_inference.html#_update_uid', 'fastMONAI/vision_inference.py'), 'fastMONAI.vision_inference.compute_binary_tumor_volume': ( 'vision_inference.html#compute_binary_tumor_volume', diff --git a/fastMONAI/vision_inference.py b/fastMONAI/vision_inference.py index 30517f5..ec7f31d 100644 --- a/fastMONAI/vision_inference.py +++ b/fastMONAI/vision_inference.py @@ -13,7 +13,7 @@ import numpy as np from scipy.ndimage import distance_transform_edt, label from skimage.morphology import remove_small_objects -from torchio import Resize, Image, ScalarImage +from torchio import Pad, Resample, Image, ScalarImage from .vision_core import * from .vision_augmentation import do_pad_or_crop, transforms_from_specs from fastMONAI.utils import ( @@ -45,28 +45,22 @@ def save_series_pred(series_obj, save_dir, val='1234'): series_obj.write(save_dir, opts={'keep_uid': True}, formats=['dicom']) -# %% ../nbs/06_vision_inference.ipynb #e4a2771e-25aa-4669-bda2-abcbe29d0504 -def _to_original_orientation(input_img, org_orientation): - """Reorients the image to its original orientation.""" - - from SimpleITK import DICOMOrient, GetArrayFromImage - orientation_itk = DICOMOrient(input_img, org_orientation) - reoriented_array = GetArrayFromImage(orientation_itk).transpose() - - return reoriented_array[None] - # %% ../nbs/06_vision_inference.ipynb #75a1169f-7385-4c48-9a24-51994c80732c -def _do_resize(o, target_shape, image_interpolation='linear', - label_interpolation='nearest'): - """Resample images so the output shape matches the given target shape.""" - - resize = Resize( - target_shape, - image_interpolation=image_interpolation, - label_interpolation=label_interpolation +def _resample_to_grid(o, shape, affine, image_interpolation='linear', + label_interpolation='nearest'): + """Resample an image onto the voxel grid given by shape and affine. + + Predictions are restored onto the original image's own grid, so they line up with it + in any orientation. The input's edge voxels are repeated once first: a resampled grid + can end almost one of its own voxels short of the original image's edge voxels, and + those voxels take the edge value instead of 0. + """ + resample = Resample( + (tuple(int(n) for n in shape), np.asarray(affine, dtype=float)), + image_interpolation=image_interpolation, + label_interpolation=label_interpolation, ) - - return resize(o) + return resample(Pad(1, padding_mode='edge')(o)) # %% ../nbs/06_vision_inference.ipynb #shape-based-label-resize def _signed_depth(mask: np.ndarray, spacing, margin: int = 8) -> np.ndarray: @@ -91,11 +85,11 @@ def _signed_depth(mask: np.ndarray, spacing, margin: int = 8) -> np.ndarray: return depth -def _resize_labels(o, target_shape, label_interpolation='shape'): - """Resize a single-channel label image so its shape matches target_shape. +def _resample_labels(o, shape, affine, label_interpolation='shape'): + """Resample a single-channel label image onto the voxel grid given by shape and affine. 'shape' uses shape-based interpolation: the signed distance to the edge of the - foreground (all labels together) is resized with cubic B-spline interpolation and + foreground (all labels together) is resampled with cubic B-spline interpolation and decides which voxels are foreground. With several labels, the same is done for each label and every foreground voxel takes the label it lies deepest inside, so labels that touch keep touching. A slice that falls between two source slices gets an @@ -103,22 +97,22 @@ def _resize_labels(o, target_shape, label_interpolation='shape'): nearest source voxel. """ if label_interpolation == 'nearest': - return _do_resize(o, target_shape, image_interpolation='nearest') + return _resample_to_grid(o, shape, affine, image_interpolation='nearest') if label_interpolation != 'shape': raise ValueError( f"label_interpolation must be 'shape' or 'nearest', got {label_interpolation!r}") - if tuple(o.spatial_shape) == tuple(target_shape): + if tuple(o.spatial_shape) == tuple(shape) and np.allclose(o.affine, affine): return o if o.data.shape[0] != 1: raise ValueError(f"Expected a single-channel label image, got {o.data.shape[0]} channels") labels = o.data[0].numpy() values = [v for v in np.unique(labels) if v != 0] if not values: - return _do_resize(o, target_shape, image_interpolation='nearest') + return _resample_to_grid(o, shape, affine, image_interpolation='nearest') masks = [labels != 0] + ([labels == v for v in values] if len(values) > 1 else []) depth = np.stack([_signed_depth(m, o.spacing) for m in masks]) - resized = _do_resize(ScalarImage(tensor=torch.from_numpy(depth), affine=o.affine), - target_shape, image_interpolation='bspline') + resized = _resample_to_grid(ScalarImage(tensor=torch.from_numpy(depth), affine=o.affine), + shape, affine, image_interpolation='bspline') values = torch.as_tensor(values, dtype=torch.float32) label = values[resized.data[1:].argmax(0)] if len(values) > 1 else values[0] restored = torch.where(resized.data[0] > 0, label, 0.) @@ -217,17 +211,17 @@ def _prediction_save_path(save_path, fn, org_img): return Path(save_path) / ('pred_' + Path(source_fn).name) -def _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn=None): +def _restore_prediction(pred, org_img, input_img, label_interpolation, save_fn=None): """Map a predicted mask back to the original image space, then save or return it. - The mask is padded or cropped to the preprocessed shape, resized to the original shape, - reoriented to the original orientation and given the original affine. + The mask is padded or cropped to the preprocessed shape, then resampled onto the + original image's grid, which gives it the original shape, orientation and affine. """ pred_mask = do_pad_or_crop(pred.float(), input_img.shape[1:], padding_mode=0, mask_name=None) input_img.set_data(pred_mask) - input_img = _resize_labels(input_img, org_size, label_interpolation) - reoriented_array = _to_original_orientation(input_img.as_sitk(), ''.join(org_img.orientation)) - pred_img = ScalarImage(tensor=reoriented_array, affine=org_img.affine.copy()) + restored = _resample_labels(input_img, org_img.spatial_shape, org_img.affine, + label_interpolation) + pred_img = ScalarImage(tensor=restored.data.float(), affine=org_img.affine.copy()) if save_fn is None: return pred_img pred_img.save(save_fn) @@ -253,8 +247,8 @@ def inference_from_config(model, inference_config, fn: (str, Path) = '', apply_reorder = inference_config['apply_reorder'] target_spacing = inference_config['target_spacing'] - if org_img is None or input_img is None or org_size is None: - org_img, input_img, org_size = med_img_reader( + if org_img is None or input_img is None: + org_img, input_img, _ = med_img_reader( fn, apply_reorder, target_spacing, only_tensor=False ) else: @@ -271,7 +265,7 @@ def inference_from_config(model, inference_config, fn: (str, Path) = '', 'Decoded segmentation must have shape [D, H, W] or [1, D, H, W]; ' f'got {tuple(pred.shape)}.' ) - return _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn) + return _restore_prediction(pred, org_img, input_img, label_interpolation, save_fn) # %% ../nbs/06_vision_inference.ipynb #2ba94a55-2cd1-4fae-b8b4-e82051be0cdc def inference(learn_inf, apply_reorder, target_spacing, fn: (str, Path) = '', @@ -286,13 +280,13 @@ def inference(learn_inf, apply_reorder, target_spacing, fn: (str, Path) = '', apply_reorder: Whether to reorder the image to RAS+, as in training. target_spacing: Voxel spacing used in training, or None to keep the original spacing. fn: Image path, or a list of paths stacked as channels. Not needed when - `org_img`, `input_img` and `org_size` are given. + `org_img` and `input_img` are given. save_path: Directory to save the mask in as 'pred_'. If None, the mask is returned instead. org_img: Original image, as returned by `med_img_reader(..., only_tensor=False)`. input_img: Preprocessed image from the same call. It is cast to float32 before prediction. - org_size: Original image shape from the same call. + org_size: Not used; accepted so existing calls keep working. label_interpolation: 'shape' (default) restores the mask with shape-based interpolation; 'nearest' copies the nearest voxel instead. @@ -303,15 +297,15 @@ def inference(learn_inf, apply_reorder, target_spacing, fn: (str, Path) = '', learn_inf = unwrap_compiled_model(learn_inf) - if None in [org_img, input_img, org_size]: - org_img, input_img, org_size = med_img_reader(fn, apply_reorder, target_spacing, - only_tensor=False) + if org_img is None or input_img is None: + org_img, input_img, _ = med_img_reader(fn, apply_reorder, target_spacing, + only_tensor=False) else: org_img, input_img = copy(org_img), copy(input_img) save_fn = _prediction_save_path(save_path, fn, org_img) pred, *_ = learn_inf.predict(input_img.data.float()) - return _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn) + return _restore_prediction(pred, org_img, input_img, label_interpolation, save_fn) # %% ../nbs/06_vision_inference.ipynb #1730df1e-e6d7-4fa7-b40e-83f4dda26dc5 def compute_binary_tumor_volume(mask_data: Image): diff --git a/fastMONAI/vision_patch.py b/fastMONAI/vision_patch.py index 2505e01..d522e06 100644 --- a/fastMONAI/vision_patch.py +++ b/fastMONAI/vision_patch.py @@ -27,7 +27,7 @@ from fastai.learner import Learner from .vision_core import MedImage, MedMask, med_img_reader from .vision_plot import find_max_slice -from .vision_inference import _do_resize, _resize_labels, keep_largest +from .vision_inference import _resample_to_grid, _resample_labels, keep_largest from .dataset_info import MedDataset, suggest_patch_size from .vision_augmentation import transforms_to_specs, transforms_from_specs from .utils import unwrap_compiled_model @@ -1232,7 +1232,6 @@ class _PreparedSubject: subject: tio.Subject org_img: tio.Image input_img: tio.Image - org_size: tuple padding: tuple valid_slices: tuple grid_sampler: tio.GridSampler @@ -1355,8 +1354,8 @@ def _prepare_subject(self, img_path: Path | str) -> _PreparedSubject: Thread-safe: creates only new local objects and reads only immutable self config. """ - # Keep org_img and org_size for post-processing - org_img, input_img, org_size = med_img_reader( + # Keep org_img: predictions are restored onto its grid + org_img, input_img, _ = med_img_reader( img_path, apply_reorder=self.apply_reorder, target_spacing=self.target_spacing, only_tensor=False ) @@ -1398,7 +1397,7 @@ def _prepare_subject(self, img_path: Path | str) -> _PreparedSubject: return _PreparedSubject( subject=subject, org_img=org_img, input_img=input_img, - org_size=org_size, padding=padding, valid_slices=valid_slices, + padding=padding, valid_slices=valid_slices, grid_sampler=grid_sampler, aggregator=aggregator, patch_loader=patch_loader ) @@ -1439,7 +1438,7 @@ def _postprocess( prepared: _PreparedSubject, return_probabilities: bool = False ) -> tuple[torch.Tensor, np.ndarray]: - """Post-process the aggregated output (threshold/argmax, resize, reorient); returns (result, affine). + """Post-process the aggregated output (threshold/argmax, restore to the original grid); returns (result, affine). return_probabilities keeps the probability map and restores it with linear interpolation. Decoded labels are restored with self.label_interpolation. @@ -1478,15 +1477,13 @@ def _postprocess( else: pred_img = tio.LabelMap(tensor=result.float(), affine=prepared.input_img.affine) + # Resample onto the original image's grid, which also undoes the reordering + org_shape, org_affine = prepared.org_img.spatial_shape, prepared.org_img.affine if return_probabilities: - pred_img = _do_resize(pred_img, prepared.org_size, image_interpolation='linear') + pred_img = _resample_to_grid(pred_img, org_shape, org_affine, + image_interpolation='linear') else: - pred_img = _resize_labels(pred_img, prepared.org_size, self.label_interpolation) - - # Reorient to original orientation (if reorder was applied) - if self.apply_reorder: - target_orientation = ''.join(prepared.org_img.orientation) - pred_img = tio.ToOrientation(target_orientation)(pred_img) + pred_img = _resample_labels(pred_img, org_shape, org_affine, self.label_interpolation) result = pred_img.data.cpu() if not return_probabilities: diff --git a/nbs/06_vision_inference.ipynb b/nbs/06_vision_inference.ipynb index 5be13fa..a0e6efa 100644 --- a/nbs/06_vision_inference.ipynb +++ b/nbs/06_vision_inference.ipynb @@ -24,7 +24,7 @@ "import numpy as np\n", "from scipy.ndimage import distance_transform_edt, label\n", "from skimage.morphology import remove_small_objects\n", - "from torchio import Resize, Image, ScalarImage\n", + "from torchio import Pad, Resample, Image, ScalarImage\n", "from fastMONAI.vision_core import *\n", "from fastMONAI.vision_augmentation import do_pad_or_crop, transforms_from_specs\n", "from fastMONAI.utils import (\n", @@ -82,24 +82,6 @@ " series_obj.write(save_dir, opts={'keep_uid': True}, formats=['dicom'])" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "e4a2771e-25aa-4669-bda2-abcbe29d0504", - "metadata": {}, - "outputs": [], - "source": [ - "#| export\n", - "def _to_original_orientation(input_img, org_orientation):\n", - " \"\"\"Reorients the image to its original orientation.\"\"\"\n", - " \n", - " from SimpleITK import DICOMOrient, GetArrayFromImage\n", - " orientation_itk = DICOMOrient(input_img, org_orientation)\n", - " reoriented_array = GetArrayFromImage(orientation_itk).transpose()\n", - " \n", - " return reoriented_array[None]" - ] - }, { "cell_type": "code", "execution_count": null, @@ -108,17 +90,21 @@ "outputs": [], "source": [ "#| export\n", - "def _do_resize(o, target_shape, image_interpolation='linear', \n", - " label_interpolation='nearest'):\n", - " \"\"\"Resample images so the output shape matches the given target shape.\"\"\"\n", - "\n", - " resize = Resize(\n", - " target_shape, \n", - " image_interpolation=image_interpolation, \n", - " label_interpolation=label_interpolation\n", + "def _resample_to_grid(o, shape, affine, image_interpolation='linear',\n", + " label_interpolation='nearest'):\n", + " \"\"\"Resample an image onto the voxel grid given by shape and affine.\n", + "\n", + " Predictions are restored onto the original image's own grid, so they line up with it\n", + " in any orientation. The input's edge voxels are repeated once first: a resampled grid\n", + " can end almost one of its own voxels short of the original image's edge voxels, and\n", + " those voxels take the edge value instead of 0.\n", + " \"\"\"\n", + " resample = Resample(\n", + " (tuple(int(n) for n in shape), np.asarray(affine, dtype=float)),\n", + " image_interpolation=image_interpolation,\n", + " label_interpolation=label_interpolation,\n", " )\n", - " \n", - " return resize(o)" + " return resample(Pad(1, padding_mode='edge')(o))" ] }, { @@ -151,11 +137,11 @@ " return depth\n", "\n", "\n", - "def _resize_labels(o, target_shape, label_interpolation='shape'):\n", - " \"\"\"Resize a single-channel label image so its shape matches target_shape.\n", + "def _resample_labels(o, shape, affine, label_interpolation='shape'):\n", + " \"\"\"Resample a single-channel label image onto the voxel grid given by shape and affine.\n", "\n", " 'shape' uses shape-based interpolation: the signed distance to the edge of the\n", - " foreground (all labels together) is resized with cubic B-spline interpolation and\n", + " foreground (all labels together) is resampled with cubic B-spline interpolation and\n", " decides which voxels are foreground. With several labels, the same is done for each\n", " label and every foreground voxel takes the label it lies deepest inside, so labels\n", " that touch keep touching. A slice that falls between two source slices gets an\n", @@ -163,22 +149,22 @@ " nearest source voxel.\n", " \"\"\"\n", " if label_interpolation == 'nearest':\n", - " return _do_resize(o, target_shape, image_interpolation='nearest')\n", + " return _resample_to_grid(o, shape, affine, image_interpolation='nearest')\n", " if label_interpolation != 'shape':\n", " raise ValueError(\n", " f\"label_interpolation must be 'shape' or 'nearest', got {label_interpolation!r}\")\n", - " if tuple(o.spatial_shape) == tuple(target_shape):\n", + " if tuple(o.spatial_shape) == tuple(shape) and np.allclose(o.affine, affine):\n", " return o\n", " if o.data.shape[0] != 1:\n", " raise ValueError(f\"Expected a single-channel label image, got {o.data.shape[0]} channels\")\n", " labels = o.data[0].numpy()\n", " values = [v for v in np.unique(labels) if v != 0]\n", " if not values:\n", - " return _do_resize(o, target_shape, image_interpolation='nearest')\n", + " return _resample_to_grid(o, shape, affine, image_interpolation='nearest')\n", " masks = [labels != 0] + ([labels == v for v in values] if len(values) > 1 else [])\n", " depth = np.stack([_signed_depth(m, o.spacing) for m in masks])\n", - " resized = _do_resize(ScalarImage(tensor=torch.from_numpy(depth), affine=o.affine),\n", - " target_shape, image_interpolation='bspline')\n", + " resized = _resample_to_grid(ScalarImage(tensor=torch.from_numpy(depth), affine=o.affine),\n", + " shape, affine, image_interpolation='bspline')\n", " values = torch.as_tensor(values, dtype=torch.float32)\n", " label = values[resized.data[1:].argmax(0)] if len(values) > 1 else values[0]\n", " restored = torch.where(resized.data[0] > 0, label, 0.)\n", @@ -207,13 +193,24 @@ " return sum(m[..., k].any() and np.array_equal(m[..., k], m[..., k + 1])\n", " for k in range(m.shape[-1] - 1))\n", "\n", + "def _target_grid(img, shape):\n", + " \"A grid with this shape over the same extent and centre as img.\"\n", + " shape = np.array(shape)\n", + " spacing = np.array(img.spatial_shape) * np.array(img.spacing) / shape\n", + " direction = img.affine[:3, :3] / np.array(img.spacing)\n", + " centre = img.affine[:3, :3] @ ((np.array(img.spatial_shape) - 1) / 2) + img.affine[:3, 3]\n", + " affine = np.eye(4)\n", + " affine[:3, :3] = direction * spacing\n", + " affine[:3, 3] = centre - affine[:3, :3] @ ((shape - 1) / 2)\n", + " return tuple(int(n) for n in shape), affine\n", + "\n", "_src_spacing = (0.5, 0.5, 1.5)\n", "_src = tio.LabelMap(\n", " tensor=torch.from_numpy(_sphere((40, 40, 16), _src_spacing, (10., 10., 12.), 6.))[None].float(),\n", " affine=np.diag([*_src_spacing, 1.0])\n", ")\n", - "_shape = _resize_labels(_src, (40, 40, 24))\n", - "_nearest = _resize_labels(_src, (40, 40, 24), label_interpolation='nearest')\n", + "_shape = _resample_labels(_src, *_target_grid(_src, (40, 40, 24)))\n", + "_nearest = _resample_labels(_src, *_target_grid(_src, (40, 40, 24)), label_interpolation='nearest')\n", "test_eq(tuple(_shape.spatial_shape), (40, 40, 24))\n", "test_eq(type(_shape), tio.LabelMap)\n", "test_eq(_shape.data.dtype, _src.data.dtype)\n", @@ -224,19 +221,40 @@ "assert _identical_neighbours(_s) < _identical_neighbours(_n)\n", "assert abs(_s.sum() / _truth.sum() - 1) < 0.1\n", "\n", - "# Label values survive, empty masks stay empty, and an unchanged shape is returned as is.\n", + "# Label values survive, empty masks stay empty, and an unchanged grid is returned as is.\n", "_two = torch.zeros(1, 20, 20, 8)\n", "_two[0, 4:10, 4:16, 2:6] = 1\n", "_two[0, 10:16, 4:16, 2:6] = 2\n", "_two_img = tio.LabelMap(tensor=_two, affine=np.diag([*_src_spacing, 1.0]))\n", - "test_eq(torch.unique(_resize_labels(_two_img, (20, 20, 12)).data).tolist(), [0., 1., 2.])\n", + "test_eq(torch.unique(_resample_labels(_two_img, *_target_grid(_two_img, (20, 20, 12))).data).tolist(),\n", + " [0., 1., 2.])\n", "_empty = tio.LabelMap(tensor=torch.zeros(1, 20, 20, 8), affine=np.diag([*_src_spacing, 1.0]))\n", - "_empty_out = _resize_labels(_empty, (20, 20, 12))\n", + "_empty_out = _resample_labels(_empty, *_target_grid(_empty, (20, 20, 12)))\n", "test_eq(tuple(_empty_out.spatial_shape), (20, 20, 12))\n", "test_eq(int(_empty_out.data.sum()), 0)\n", - "assert _resize_labels(_src, (40, 40, 16)) is _src\n", - "test_fail(lambda: _resize_labels(_src, (40, 40, 24), label_interpolation='linear'),\n", - " contains='label_interpolation')\n" + "assert _resample_labels(_src, _src.spatial_shape, _src.affine) is _src\n", + "test_fail(lambda: _resample_labels(_src, *_target_grid(_src, (40, 40, 24)), label_interpolation='linear'),\n", + " contains='label_interpolation')\n", + "\n", + "# Resampling back onto the original grid puts every voxel where it was: a coordinate ramp\n", + "# on 120 slices of 1.2 mm, taken through 1.5 mm and back, keeps its values.\n", + "_ramp = tio.ScalarImage(tensor=torch.arange(120.).view(1, 1, 1, 120).repeat(1, 4, 4, 1),\n", + " affine=np.diag([1., 1., 1.2, 1.]))\n", + "_coarse = tio.Resample((1., 1., 1.5))(_ramp)\n", + "_back = _resample_to_grid(_coarse, _ramp.spatial_shape, _ramp.affine)\n", + "assert np.allclose(_back.affine, _ramp.affine)\n", + "_err = (_back.data - _ramp.data)[0, 2, 2, 10:110]\n", + "assert float(_err.abs().max()) < 0.05, float(_err.abs().max())\n", + "\n", + "# The whole original grid is covered, including edge slices beyond the last model slice\n", + "# (53 slices of 1 mm run at 3 mm), for masks and for probabilities.\n", + "_full = tio.LabelMap(tensor=torch.ones(1, 4, 4, 53), affine=np.eye(4))\n", + "_model = tio.Resample((1., 1., 3.))(_full)\n", + "test_eq(_model.spatial_shape[-1], 17)\n", + "for _mode in ('shape', 'nearest'):\n", + " test_eq(int(_resample_labels(_model, _full.spatial_shape, np.eye(4), _mode).data.sum()), 4 * 4 * 53)\n", + "_probs = tio.ScalarImage(tensor=torch.cat([0.3 * _model.data, 0.7 * _model.data]), affine=_model.affine)\n", + "assert torch.allclose(_resample_to_grid(_probs, (4, 4, 53), np.eye(4)).data.sum(0), torch.ones(4, 4, 53))\n" ] }, { @@ -248,15 +266,16 @@ "source": [ "#| hide\n", "# Labels touching along the resampled axis leave no background between them:\n", - "# the multi-label foreground covers the same voxels as the resized union of the labels.\n", + "# the multi-label foreground covers the same voxels as the resampled union of the labels.\n", "_slabs = torch.zeros(1, 20, 20, 12)\n", "_slabs[0, 4:16, 4:16, 1:5] = 1\n", "_slabs[0, 4:16, 4:16, 5:11] = 2\n", "_slabs_img = tio.LabelMap(tensor=_slabs, affine=np.diag([*_src_spacing, 1.0]))\n", "_union_img = tio.LabelMap(tensor=(_slabs > 0).float(), affine=_slabs_img.affine)\n", "for _n in (16, 17, 18, 19):\n", - " _out = _resize_labels(_slabs_img, (20, 20, _n)).data[0]\n", - " _union = _resize_labels(_union_img, (20, 20, _n)).data[0] > 0\n", + " _grid = _target_grid(_slabs_img, (20, 20, _n))\n", + " _out = _resample_labels(_slabs_img, *_grid).data[0]\n", + " _union = _resample_labels(_union_img, *_grid).data[0] > 0\n", " test_eq(int(((_out == 0) & _union).sum()), 0)\n", " test_eq(torch.unique(_out).tolist(), [0., 1., 2.])\n" ] @@ -384,17 +403,17 @@ " return Path(save_path) / ('pred_' + Path(source_fn).name)\n", "\n", "\n", - "def _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn=None):\n", + "def _restore_prediction(pred, org_img, input_img, label_interpolation, save_fn=None):\n", " \"\"\"Map a predicted mask back to the original image space, then save or return it.\n", "\n", - " The mask is padded or cropped to the preprocessed shape, resized to the original shape,\n", - " reoriented to the original orientation and given the original affine.\n", + " The mask is padded or cropped to the preprocessed shape, then resampled onto the\n", + " original image's grid, which gives it the original shape, orientation and affine.\n", " \"\"\"\n", " pred_mask = do_pad_or_crop(pred.float(), input_img.shape[1:], padding_mode=0, mask_name=None)\n", " input_img.set_data(pred_mask)\n", - " input_img = _resize_labels(input_img, org_size, label_interpolation)\n", - " reoriented_array = _to_original_orientation(input_img.as_sitk(), ''.join(org_img.orientation))\n", - " pred_img = ScalarImage(tensor=reoriented_array, affine=org_img.affine.copy())\n", + " restored = _resample_labels(input_img, org_img.spatial_shape, org_img.affine,\n", + " label_interpolation)\n", + " pred_img = ScalarImage(tensor=restored.data.float(), affine=org_img.affine.copy())\n", " if save_fn is None:\n", " return pred_img\n", " pred_img.save(save_fn)\n", @@ -427,8 +446,8 @@ "\n", " apply_reorder = inference_config['apply_reorder']\n", " target_spacing = inference_config['target_spacing']\n", - " if org_img is None or input_img is None or org_size is None:\n", - " org_img, input_img, org_size = med_img_reader(\n", + " if org_img is None or input_img is None:\n", + " org_img, input_img, _ = med_img_reader(\n", " fn, apply_reorder, target_spacing, only_tensor=False\n", " )\n", " else:\n", @@ -445,7 +464,7 @@ " 'Decoded segmentation must have shape [D, H, W] or [1, D, H, W]; '\n", " f'got {tuple(pred.shape)}.'\n", " )\n", - " return _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn)" + " return _restore_prediction(pred, org_img, input_img, label_interpolation, save_fn)" ] }, { @@ -562,7 +581,7 @@ "id": "2ba94a55-2cd1-4fae-b8b4-e82051be0cdc", "metadata": {}, "outputs": [], - "source": "#| export\ndef inference(learn_inf, apply_reorder, target_spacing, fn: (str, Path) = '',\n save_path: (str, Path) = None, org_img=None, input_img=None,\n org_size=None, label_interpolation='shape'):\n \"\"\"Segment an image with an exported learner and restore the mask to the original image space.\n\n Only segmentation learners are supported: the decoded prediction is treated as a mask.\n\n Args:\n learn_inf: Exported segmentation learner, e.g. from `load_learner`.\n apply_reorder: Whether to reorder the image to RAS+, as in training.\n target_spacing: Voxel spacing used in training, or None to keep the original spacing.\n fn: Image path, or a list of paths stacked as channels. Not needed when\n `org_img`, `input_img` and `org_size` are given.\n save_path: Directory to save the mask in as 'pred_'. If None, the mask\n is returned instead.\n org_img: Original image, as returned by `med_img_reader(..., only_tensor=False)`.\n input_img: Preprocessed image from the same call. It is cast to float32 before\n prediction.\n org_size: Original image shape from the same call.\n label_interpolation: 'shape' (default) restores the mask with shape-based\n interpolation; 'nearest' copies the nearest voxel instead.\n\n Returns:\n The saved file path if `save_path` is given, otherwise the mask as a `ScalarImage`\n with the original shape, orientation and affine.\n \"\"\"\n\n learn_inf = unwrap_compiled_model(learn_inf)\n\n if None in [org_img, input_img, org_size]:\n org_img, input_img, org_size = med_img_reader(fn, apply_reorder, target_spacing,\n only_tensor=False)\n else:\n org_img, input_img = copy(org_img), copy(input_img)\n save_fn = _prediction_save_path(save_path, fn, org_img)\n\n pred, *_ = learn_inf.predict(input_img.data.float())\n return _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn)" + "source": "#| export\ndef inference(learn_inf, apply_reorder, target_spacing, fn: (str, Path) = '',\n save_path: (str, Path) = None, org_img=None, input_img=None,\n org_size=None, label_interpolation='shape'):\n \"\"\"Segment an image with an exported learner and restore the mask to the original image space.\n\n Only segmentation learners are supported: the decoded prediction is treated as a mask.\n\n Args:\n learn_inf: Exported segmentation learner, e.g. from `load_learner`.\n apply_reorder: Whether to reorder the image to RAS+, as in training.\n target_spacing: Voxel spacing used in training, or None to keep the original spacing.\n fn: Image path, or a list of paths stacked as channels. Not needed when\n `org_img` and `input_img` are given.\n save_path: Directory to save the mask in as 'pred_'. If None, the mask\n is returned instead.\n org_img: Original image, as returned by `med_img_reader(..., only_tensor=False)`.\n input_img: Preprocessed image from the same call. It is cast to float32 before\n prediction.\n org_size: Not used; accepted so existing calls keep working.\n label_interpolation: 'shape' (default) restores the mask with shape-based\n interpolation; 'nearest' copies the nearest voxel instead.\n\n Returns:\n The saved file path if `save_path` is given, otherwise the mask as a `ScalarImage`\n with the original shape, orientation and affine.\n \"\"\"\n\n learn_inf = unwrap_compiled_model(learn_inf)\n\n if org_img is None or input_img is None:\n org_img, input_img, _ = med_img_reader(fn, apply_reorder, target_spacing,\n only_tensor=False)\n else:\n org_img, input_img = copy(org_img), copy(input_img)\n save_fn = _prediction_save_path(save_path, fn, org_img)\n\n pred, *_ = learn_inf.predict(input_img.data.float())\n return _restore_prediction(pred, org_img, input_img, label_interpolation, save_fn)" }, { "cell_type": "code", @@ -603,11 +622,27 @@ " test_eq(inference(_learner, False, None, fn=_fns, save_path=_tmp).name, 'pred_int16_1.nii.gz')\n", " _org, _inp, _size = med_img_reader(_fns[0], False, None, only_tensor=False)\n", " _saved = inference(_learner, False, None, save_path=_tmp, org_img=_org, input_img=_inp, org_size=_size)\n", + " test_eq(type(inference(_learner, False, None, org_img=_org, input_img=_inp)), tio.ScalarImage)\n", " test_eq(_saved.name, 'pred_int16_0.nii.gz')\n", " assert _saved.exists()\n", " _mem = tio.ScalarImage(tensor=_org.data, affine=_org.affine)\n", " test_fail(lambda: inference(_learner, False, None, save_path=_tmp, org_img=_mem,\n", - " input_img=_inp, org_size=_size), contains='fn is required')\n" + " input_img=_inp, org_size=_size), contains='fn is required')\n", + "\n", + " # Masks come back on the original grid: exactly after reordering an LPS image, and\n", + " # without a shift after resampling 120 slices of 1.2 mm through 1.5 mm.\n", + " _block = torch.zeros(1, 12, 10, 120)\n", + " _block[0, 3:9, 2:7, 40:70] = 100\n", + " _lps = np.diag([-0.9, -0.9, 1.2, 1.0])\n", + " _lps_fn = Path(_tmp) / 'lps.nii.gz'\n", + " tio.ScalarImage(tensor=_block, affine=_lps).save(_lps_fn)\n", + " _reordered = inference(_RecordingLearner(), True, None, fn=str(_lps_fn))\n", + " test_eq(torch.equal(_reordered.data > 0, _block > 0), True)\n", + " assert np.allclose(_reordered.affine, _lps)\n", + " _resampled = inference(_RecordingLearner(), True, [0.9, 0.9, 1.5], fn=str(_lps_fn))\n", + " test_eq(tuple(_resampled.shape), tuple(_block.shape))\n", + " _z = lambda m: m[0].nonzero()[:, 2].float().mean().item()\n", + " assert abs(_z(_resampled.data > 0) - _z(_block > 0)) < 0.2\n" ] }, { diff --git a/nbs/10_vision_patch.ipynb b/nbs/10_vision_patch.ipynb index cf1bd11..2935176 100644 --- a/nbs/10_vision_patch.ipynb +++ b/nbs/10_vision_patch.ipynb @@ -48,7 +48,7 @@ "from fastai.learner import Learner\n", "from fastMONAI.vision_core import MedImage, MedMask, med_img_reader\n", "from fastMONAI.vision_plot import find_max_slice\n", - "from fastMONAI.vision_inference import _do_resize, _resize_labels, keep_largest\n", + "from fastMONAI.vision_inference import _resample_to_grid, _resample_labels, keep_largest\n", "from fastMONAI.dataset_info import MedDataset, suggest_patch_size\n", "from fastMONAI.vision_augmentation import transforms_to_specs, transforms_from_specs\n", "from fastMONAI.utils import unwrap_compiled_model" @@ -1572,7 +1572,6 @@ " subject: tio.Subject\n", " org_img: tio.Image\n", " input_img: tio.Image\n", - " org_size: tuple\n", " padding: tuple\n", " valid_slices: tuple\n", " grid_sampler: tio.GridSampler\n", @@ -1695,8 +1694,8 @@ "\n", " Thread-safe: creates only new local objects and reads only immutable self config.\n", " \"\"\"\n", - " # Keep org_img and org_size for post-processing\n", - " org_img, input_img, org_size = med_img_reader(\n", + " # Keep org_img: predictions are restored onto its grid\n", + " org_img, input_img, _ = med_img_reader(\n", " img_path, apply_reorder=self.apply_reorder, target_spacing=self.target_spacing, only_tensor=False\n", " )\n", "\n", @@ -1738,7 +1737,7 @@ "\n", " return _PreparedSubject(\n", " subject=subject, org_img=org_img, input_img=input_img,\n", - " org_size=org_size, padding=padding, valid_slices=valid_slices,\n", + " padding=padding, valid_slices=valid_slices,\n", " grid_sampler=grid_sampler, aggregator=aggregator,\n", " patch_loader=patch_loader\n", " )\n", @@ -1779,7 +1778,7 @@ " prepared: _PreparedSubject,\n", " return_probabilities: bool = False\n", " ) -> tuple[torch.Tensor, np.ndarray]:\n", - " \"\"\"Post-process the aggregated output (threshold/argmax, resize, reorient); returns (result, affine).\n", + " \"\"\"Post-process the aggregated output (threshold/argmax, restore to the original grid); returns (result, affine).\n", "\n", " return_probabilities keeps the probability map and restores it with linear\n", " interpolation. Decoded labels are restored with self.label_interpolation.\n", @@ -1818,15 +1817,13 @@ " else:\n", " pred_img = tio.LabelMap(tensor=result.float(), affine=prepared.input_img.affine)\n", "\n", + " # Resample onto the original image's grid, which also undoes the reordering\n", + " org_shape, org_affine = prepared.org_img.spatial_shape, prepared.org_img.affine\n", " if return_probabilities:\n", - " pred_img = _do_resize(pred_img, prepared.org_size, image_interpolation='linear')\n", + " pred_img = _resample_to_grid(pred_img, org_shape, org_affine,\n", + " image_interpolation='linear')\n", " else:\n", - " pred_img = _resize_labels(pred_img, prepared.org_size, self.label_interpolation)\n", - "\n", - " # Reorient to original orientation (if reorder was applied)\n", - " if self.apply_reorder:\n", - " target_orientation = ''.join(prepared.org_img.orientation)\n", - " pred_img = tio.ToOrientation(target_orientation)(pred_img)\n", + " pred_img = _resample_labels(pred_img, org_shape, org_affine, self.label_interpolation)\n", "\n", " result = pred_img.data.cpu()\n", " if not return_probabilities:\n", @@ -2219,7 +2216,7 @@ " assert isinstance(prepared.grid_sampler, tio.GridSampler)\n", " assert isinstance(prepared.aggregator, tio.GridAggregator)\n", " assert isinstance(prepared.patch_loader, DataLoader)\n", - " assert prepared.org_size is not None\n", + " assert prepared.org_img is not None\n", "\n", " # Test 2: Decomposed path equals predict() output\n", " pred_decomposed, affine_decomposed = _engine._postprocess(\n", @@ -2615,20 +2612,18 @@ "_restore_affine = np.eye(4)\n", "_foreground = torch.tensor([[[0., 0.], [0., 0.]], [[1., 1.], [1., 1.]]])\n", "_restore_output = torch.stack((1 - _foreground, _foreground))\n", - "_restore_input = tio.ScalarImage(tensor=torch.zeros(1, 2, 2, 2), affine=_restore_affine)\n", "_restore_org = tio.ScalarImage(tensor=torch.zeros(1, 3, 3, 3), affine=_restore_affine)\n", - "_restore_prepared = _SimpleNamespace(\n", - " input_img=_restore_input, org_img=_restore_org, org_size=(3, 3, 3)\n", - ")\n", + "_restore_input = tio.Resample(1.5)(_restore_org)\n", + "_restore_prepared = _SimpleNamespace(input_img=_restore_input, org_img=_restore_org)\n", "_restore_engine = PatchInferenceEngine(\n", " torch.nn.Conv3d(1, 2, 1), PatchConfig(patch_size=[16, 16, 16]), apply_reorder=False\n", ")\n", "_linear_prob, _ = _restore_engine._postprocess(\n", " _restore_output, _restore_prepared, return_probabilities=True\n", ")\n", - "_nearest_prob = _do_resize(\n", - " tio.ScalarImage(tensor=_restore_output, affine=_restore_affine),\n", - " (3, 3, 3), image_interpolation='nearest'\n", + "_nearest_prob = _resample_to_grid(\n", + " tio.ScalarImage(tensor=_restore_output, affine=_restore_input.affine),\n", + " (3, 3, 3), _restore_affine, image_interpolation='nearest'\n", ").data\n", "test_eq(tuple(_linear_prob.shape), (2, 3, 3, 3))\n", "assert torch.allclose(_linear_prob[:, 1, 1, 1], torch.tensor([0.5, 0.5]))\n", @@ -2636,6 +2631,32 @@ "assert torch.allclose(_linear_prob.sum(0), torch.ones(3, 3, 3), atol=1e-6)\n", "assert 0 <= float(_linear_prob.min()) <= float(_linear_prob.max()) <= 1\n", "\n", + "# A mask predicted on the reordered, resampled grid comes back on the original grid: in\n", + "# place for an LPS image with 120 slices of 1.2 mm run at 1.5 mm.\n", + "import tempfile as _tempfile\n", + "_block = torch.zeros(1, 12, 10, 120)\n", + "_block[0, 3:9, 2:7, 40:70] = 1\n", + "_lps = np.diag([-0.9, -0.9, 1.2, 1.0])\n", + "_grid_engine = PatchInferenceEngine(\n", + " torch.nn.Conv3d(1, 2, 1), PatchConfig(patch_size=[16, 16, 16]),\n", + " apply_reorder=True, target_spacing=[1.0, 1.0, 1.5]\n", + ")\n", + "with _tempfile.TemporaryDirectory() as _tmp:\n", + " _lps_fn = Path(_tmp) / 'lps.nii.gz'\n", + " tio.ScalarImage(tensor=_block, affine=_lps).save(_lps_fn)\n", + " _grid_prepared = _grid_engine._prepare_subject(_lps_fn)\n", + " # The model output covers the padded grid; the mask is the image itself.\n", + " _fg = torch.zeros(_grid_prepared.subject['image'].spatial_shape)\n", + " _fg[_grid_prepared.valid_slices] = (_grid_prepared.input_img.data[0] > 0.5).float()\n", + " _grid_mask, _grid_affine = _grid_engine._postprocess(\n", + " torch.stack((1 - _fg, _fg)), _grid_prepared)\n", + "test_eq(tuple(_grid_mask.shape), tuple(_block.shape))\n", + "assert np.allclose(_grid_affine, _lps)\n", + "_z = lambda m: m[0].nonzero()[:, 2].float().mean().item()\n", + "assert abs(_z(_grid_mask > 0) - _z(_block > 0)) < 0.1\n", + "_a, _b = _grid_mask[0] > 0, _block[0] > 0\n", + "assert 2 * (_a & _b).sum() / (_a.sum() + _b.sum()) > 0.95\n", + "\n", "# Masks: a small ball on 6 slices restored onto 9, like 1.5 mm model slices onto a\n", "# 1.0 mm scan. The restored distance to the label edge is at least 0.2 everywhere, so no\n", "# voxel depends on rounding, and shape-based restoration differs from the nearest copy.\n", @@ -2644,15 +2665,16 @@ " sum((g - c) ** 2 for g, c in zip(_grid, (2.5, 4.5, 4.5))) <= 2.6 ** 2\n", ").float()\n", "_label_output = torch.stack((1 - _ball, _ball))\n", + "_label_org = tio.ScalarImage(tensor=torch.zeros(1, 9, 10, 10), affine=_restore_affine)\n", "_label_prepared = _SimpleNamespace(\n", - " input_img=tio.ScalarImage(tensor=torch.zeros(1, 6, 10, 10), affine=_restore_affine),\n", - " org_img=tio.ScalarImage(tensor=torch.zeros(1, 9, 10, 10), affine=_restore_affine),\n", - " org_size=(9, 10, 10)\n", + " input_img=tio.Resample((1.5, 1., 1.))(_label_org), org_img=_label_org\n", ")\n", - "_restore_labels = tio.LabelMap(tensor=_ball[None], affine=_restore_affine)\n", - "_shape_expected = _resize_labels(_restore_labels, (9, 10, 10)).data.long()\n", - "_nearest_expected = _do_resize(\n", - " _restore_labels, (9, 10, 10), image_interpolation='linear', label_interpolation='nearest'\n", + "test_eq(tuple(_label_prepared.input_img.spatial_shape), (6, 10, 10))\n", + "_restore_labels = tio.LabelMap(tensor=_ball[None], affine=_label_prepared.input_img.affine)\n", + "_shape_expected = _resample_labels(_restore_labels, (9, 10, 10), _restore_affine).data.long()\n", + "_nearest_expected = _resample_to_grid(\n", + " _restore_labels, (9, 10, 10), _restore_affine,\n", + " image_interpolation='linear', label_interpolation='nearest'\n", ").data.long()\n", "assert not torch.equal(_shape_expected, _nearest_expected)\n", "_label, _ = _restore_engine._postprocess(\n", @@ -2685,9 +2707,7 @@ "_pair_prob[3, 3, 4] = 0.9\n", "_pair_output = torch.stack((1 - _pair_prob, _pair_prob))\n", "_pair_input = tio.ScalarImage(tensor=torch.zeros(1, 5, 5, 5), affine=_restore_affine)\n", - "_pair_prepared = _SimpleNamespace(\n", - " input_img=_pair_input, org_img=_pair_input, org_size=(5, 5, 5)\n", - ")\n", + "_pair_prepared = _SimpleNamespace(input_img=_pair_input, org_img=_pair_input)\n", "_pair_engine = PatchInferenceEngine(\n", " torch.nn.Conv3d(1, 2, 1),\n", " PatchConfig(patch_size=[16, 16, 16], keep_largest_component=True),\n", From cdeea23fef076d22ec9bf2cece67329f1b4cf79e Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 10 Nov 2026 09:04:26 +0100 Subject: [PATCH 5/6] fix(vs-training): let folds from other machines pass the merge contract The training contract hashed the whole preprocessing manifest, which records absolute source paths and file times, the preprocessing worker count, which defaults to the CPU count, and the loss factory's module path, which is workflow.models in notebook 01 and vestibular_schwannoma.workflow.models in train_5fold.py. Folds from two machines, or from the notebook and the script, never merged. The contract now keeps the preprocessing settings and the label fingerprint from the manifest, drops preprocess_workers and records the loss factory by name. The workflow revision is now 2, so folds trained before this change do not merge with folds trained after it. The README says when folds from different machines merge. --- research/vestibular_schwannoma/README.md | 5 ++ .../tests/workflow/test_training.py | 52 +++++++++++++++++++ .../workflow/training.py | 22 ++++---- 3 files changed, 70 insertions(+), 9 deletions(-) diff --git a/research/vestibular_schwannoma/README.md b/research/vestibular_schwannoma/README.md index 2da925b..f8baf17 100644 --- a/research/vestibular_schwannoma/README.md +++ b/research/vestibular_schwannoma/README.md @@ -74,6 +74,11 @@ export MLFLOW_TRACKING_URI=http://mlflow.example:5000 Never merge run IDs from different tracking stores. +Folds trained on different machines, or in notebook 01 and with the script, can be merged +if every machine prepared the data with `scripts/prepare_data.py`, `verify` passed, and all +jobs use the same commit, conda environment and training settings, including +`--queue-workers` and `--queue-length`. + ## Merge parallel folds Each launcher updates `completed_run_ids.json` after every successful fold. A subset run is not diff --git a/research/vestibular_schwannoma/tests/workflow/test_training.py b/research/vestibular_schwannoma/tests/workflow/test_training.py index 499aaa8..27e00e8 100644 --- a/research/vestibular_schwannoma/tests/workflow/test_training.py +++ b/research/vestibular_schwannoma/tests/workflow/test_training.py @@ -1,5 +1,6 @@ import json import tempfile +import types import unittest from dataclasses import replace from contextlib import redirect_stderr, redirect_stdout @@ -557,9 +558,15 @@ def contract( encoding="utf-8", ) changed_preprocessing = contract((3,)) + preprocessing_manifest.write_text( + '{"dataset_version": "test", "contract": {"target_spacing": [1, 1, 1]}}', + encoding="utf-8", + ) + changed_preprocessing_settings = contract((3,)) self.assertEqual(first["sha256"], second["sha256"]) for changed in ( + changed_preprocessing_settings, changed_epoch, changed_patch, changed_assignment, @@ -577,6 +584,51 @@ def contract( model_config.loss_spec, ) + def test_training_contract_is_the_same_on_every_machine_and_entry_point(self): + train_df = pd.DataFrame({"case_id": ["a", "b"], "fold": [1, 2]}) + model_config = TRAINING_MODEL_CONFIGS["unet"] + loss = model_config.make_loss + # The same loss factory imported as workflow.models (notebook 01) instead of + # vestibular_schwannoma.workflow.models (train_5fold.py). + notebook_loss = types.FunctionType( + loss.__code__, loss.__globals__, loss.__name__, loss.__defaults__, + loss.__closure__, + ) + notebook_loss.__qualname__ = loss.__qualname__ + notebook_loss.__module__ = "workflow.models" + + def contract(directory, source, mtime_ns, workers, model): + manifest = Path(directory) / "preprocessing_manifest.json" + manifest.write_text(json.dumps({ + "contract": {"target_spacing": [0.4102, 0.4102, 1.5]}, + "dataset_version": "labels", + "preprocessing_cache_version": f"cache-{source}", + "sources": [{"image": {"path": source, "mtime_ns": mtime_ns}}], + }), encoding="utf-8") + return training._make_training_contracts( + {"unet": model}, + train_df, + ExperimentConfig( + model_keys=("unet",), + epochs=1, + compile_models=False, + preprocess_workers=workers, + ), + SimpleNamespace(patch_overlap=0.5), + manifest, + )["unet"] + + with tempfile.TemporaryDirectory() as directory: + here = contract(directory, "/home/a/nii_data/x.nii.gz", 1, 32, model_config) + there = contract( + directory, + "/home/b/data/x.nii.gz", + 2, + 16, + replace(model_config, make_loss=notebook_loss), + ) + self.assertEqual(here["sha256"], there["sha256"]) + def test_sweep_rejects_an_existing_results_root(self): experiment = ExperimentConfig( model_keys=("unet",), diff --git a/research/vestibular_schwannoma/workflow/training.py b/research/vestibular_schwannoma/workflow/training.py index 29ec362..a205ffa 100644 --- a/research/vestibular_schwannoma/workflow/training.py +++ b/research/vestibular_schwannoma/workflow/training.py @@ -40,7 +40,7 @@ ) -VS_TRAINING_WORKFLOW_REVISION = 1 +VS_TRAINING_WORKFLOW_REVISION = 2 @dataclass(frozen=True) @@ -125,17 +125,23 @@ def _make_training_contracts( manifest_path = Path(preprocessing_manifest).expanduser() if not manifest_path.is_file(): raise FileNotFoundError(f"Preprocessing manifest not found: {manifest_path}") + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + # The contract holds only what changes the trained model, so folds trained on + # different machines, or from the notebook and the script, can be merged. experiment_payload = asdict(experiment) - for field_name in ("model_keys", "folds", "continue_on_error"): + for field_name in ("model_keys", "folds", "continue_on_error", "preprocess_workers"): experiment_payload.pop(field_name) shared_payload = { "experiment": experiment_payload, "fold_assignments_sha256": _json_sha256(assignments), "patch_config": patch_config_to_dict(patch_config), - "preprocessing_manifest_sha256": hashlib.sha256( - manifest_path.read_bytes() - ).hexdigest(), + # Preprocessing settings and the label-content fingerprint. The rest of the + # manifest records source paths and file times, which differ between machines. + "preprocessing": { + "contract": manifest.get("contract"), + "dataset_version": manifest.get("dataset_version"), + }, "training_case_count": len(assignments), "training_workflow_revision": VS_TRAINING_WORKFLOW_REVISION, } @@ -145,10 +151,8 @@ def _make_training_contracts( contracts[model_key] = make_training_contract( { **shared_payload, - "loss_factory": ( - f"{model_config.make_loss.__module__}." - f"{model_config.make_loss.__qualname__}" - ), + # The module path depends on how the package was imported. + "loss_factory": model_config.make_loss.__qualname__, "loss_spec": model_config.loss_spec, "model_key": model_key, "model_spec": model_config.model_spec, From 370b8600ba72e38fe58afdb1ba72f88b4fa2a569 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 10 Nov 2026 09:04:27 +0100 Subject: [PATCH 6/6] build: bump version to 0.13.1 --- fastMONAI/__init__.py | 2 +- settings.ini | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/fastMONAI/__init__.py b/fastMONAI/__init__.py index f23a6b3..7e0dc0e 100644 --- a/fastMONAI/__init__.py +++ b/fastMONAI/__init__.py @@ -1 +1 @@ -__version__ = "0.13.0" +__version__ = "0.13.1" diff --git a/settings.ini b/settings.ini index ecf38f5..239bbad 100644 --- a/settings.ini +++ b/settings.ini @@ -5,7 +5,7 @@ ### Python Library ### lib_name = fastMONAI min_python = 3.10 -version = 0.13.0 +version = 0.13.1 ### OPTIONAL ### requirements = fastai>=2.8.6 monai>=1.5.2 torchio>=1.2.1 xlrd>=2.0.1 scikit-image>=0.23.2,<2 imagedata>=3.8.14 mlflow>=3.14 ipython huggingface-hub gdown plum-dispatch<2.10 safetensors>=0.7.0