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/fastMONAI/_modidx.py b/fastMONAI/_modidx.py index f5010a6..510c3e3 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'), @@ -458,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/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/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/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/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\"" ] }, { 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/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", 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", 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/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/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), 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, 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