From 4120375a8385cd4a9697412a229ddaee8d859184 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:25:31 +0100 Subject: [PATCH 01/11] fix(vision_inference): keep touching labels together in shape-based restore Shape-based label restore decided foreground per label, so voxels where two labels meet could fall back to background on the resampled axis. Foreground now comes from the signed depth of all labels together and the per-label depth only picks the label. Binary masks are bit-identical to before, so the VS deployment output does not change. Also in vision_core, vision_data, vision_plot and vision_inference: - repr/print of MedImage and MedMask no longer raises - MedImageDataLoaders.from_df passes batch_tfms on and applies path/folder as a file name prefix like fastai - inference entry points cast the input to float32, as in training - inference() can save predictions for a list of files or without fn - show_med_img accepts numpy arrays - the duplicated restore code of inference/inference_from_config is one helper; dead affine fallbacks and local imports removed - vision_all re-exports all of vision_plot - docstrings, page descriptions, hidden test cells --- fastMONAI/vision_all.py | 2 +- fastMONAI/vision_core.py | 28 +++--- fastMONAI/vision_data.py | 25 ++++- fastMONAI/vision_inference.py | 159 +++++++++++++++--------------- fastMONAI/vision_plot.py | 28 +++--- nbs/00_vision_plot.ipynb | 46 ++++++--- nbs/01_vision_core.ipynb | 42 +++++--- nbs/02_vision_data.ipynb | 67 ++++++++++--- nbs/06_vision_inference.ipynb | 176 +++++++++++++++++++++++++--------- 9 files changed, 387 insertions(+), 186 deletions(-) diff --git a/fastMONAI/vision_all.py b/fastMONAI/vision_all.py index 9c78fcc..29309ad 100644 --- a/fastMONAI/vision_all.py +++ b/fastMONAI/vision_all.py @@ -1,5 +1,5 @@ #imports all of the functions and classes from the fastMONAI library -from .vision_plot import show_segmentation_comparison +from .vision_plot import * from .vision_core import * from .vision_data import * from .vision_augmentation import * diff --git a/fastMONAI/vision_core.py b/fastMONAI/vision_core.py index 5187766..86506f8 100644 --- a/fastMONAI/vision_core.py +++ b/fastMONAI/vision_core.py @@ -1,3 +1,5 @@ +"""Load medical images and represent them as `MedImage` and `MedMask` tensors.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_vision_core.ipynb. # %% auto #0 @@ -17,7 +19,7 @@ def _preprocess(obj, apply_reorder, target_spacing): Args: obj: The object to preprocess. apply_reorder: Whether to reorder the object. - target_spacing: Whether to resample the object. + target_spacing: Voxel spacing to resample to. None keeps the original spacing. Returns: The preprocessed object and its original size. @@ -45,7 +47,7 @@ def _load_and_preprocess(file_path, apply_reorder, target_spacing, dtype): Args: file_path: Image file path. apply_reorder: Whether to reorder data for canonical (RAS+) orientation. - target_spacing: Whether to resample image to different voxel sizes and dimensions. + target_spacing: Voxel spacing to resample to. None keeps the original spacing. dtype: Desired datatype for output. Returns: @@ -62,9 +64,10 @@ def _multi_channel(image_paths: L | list, apply_reorder: bool, target_spacing: l Load and preprocess multisequence data. Args: - image_paths: List of image paths (e.g., T1, T2, T1CE, DWI). + image_paths: List of image paths (e.g., T1, T2, T1CE, DWI). The first channel of + each file is stacked; the original image and size come from the last file. apply_reorder: Whether to reorder data for canonical (RAS+) orientation. - target_spacing: Whether to resample image to different voxel sizes and dimensions. + target_spacing: Voxel spacing to resample to. None keeps the original spacing. only_tensor: Whether to return only image tensor. dtype: Desired datatype for output. @@ -89,11 +92,13 @@ def med_img_reader(file_path: str | Path | L | list, apply_reorder: bool = False """Loads and preprocesses a medical image. Args: - file_path: Path to the image. Can be a string, Path object or a list. + file_path: Path to the image, or a list of paths (e.g. T1, T2) that are + stacked as channels of one image, using the first channel of each file. + For a list, the original image and size come from the last file. apply_reorder: Whether to reorder the data to be closest to canonical (RAS+) orientation. Defaults to False. - target_spacing: Whether to resample image to different voxel sizes and - image dimensions. Defaults to None. + target_spacing: Voxel spacing to resample to. Defaults to None, which keeps + the original spacing. only_tensor: Whether to return only image tensor. Defaults to True. dtype: Datatype for the return value. Defaults to torch.Tensor. @@ -146,8 +151,9 @@ def create(cls, fn: (Path, str, L, list, torch.Tensor), **kwargs) -> torch.Tenso If `fn` is a torch.Tensor, it's cast to MedBase object. Args: - fn : (Path, str, torch.Tensor) - Image path or a 4D torch.Tensor. + fn : (Path, str, list, torch.Tensor) + Image path, a list of image paths stacked as channels + (see `med_img_reader`), or a 4D torch.Tensor. kwargs : dict Additional parameters for the medical image reader. @@ -240,7 +246,7 @@ def show(self, ctx=None, channel: int = 0, slice_index: int = None, def __repr__(self) -> str: """Returns the string representation of the MedBase instance.""" - return f'{self.__class__.__name__} mode={self.mode} size={"x".join([str(d) for d in self.size])}' + return f'{self.__class__.__name__}(shape={tuple(self.shape)}, dtype={self.dtype})' # %% ../nbs/01_vision_core.ipynb #450b02cf class MedImage(MedBase): @@ -249,5 +255,5 @@ class MedImage(MedBase): # %% ../nbs/01_vision_core.ipynb #c9a702da class MedMask(MedBase): - """Subclass of MedBase that represents an mask object.""" + """Subclass of MedBase that represents a mask object.""" _show_args = {'alpha':0.5, 'cmap':'tab20'} diff --git a/fastMONAI/vision_data.py b/fastMONAI/vision_data.py index cfb2edc..f5a83e5 100644 --- a/fastMONAI/vision_data.py +++ b/fastMONAI/vision_data.py @@ -1,3 +1,5 @@ +"""Build fastai `DataLoaders` for medical images and show batches, results and top losses.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_vision_data.ipynb. # %% auto #0 @@ -87,9 +89,20 @@ class MedImageDataLoaders(DataLoaders): @delegates(DataLoaders.from_dblock) def from_df(cls, df, valid_pct=0.2, seed=None, fn_col=0, folder=None, suff='', label_col=1, label_delim=None, y_block=None, valid_col=None, - item_tfms=None, batch_tfms=None, apply_reorder=False, target_spacing=None, **kwargs): - """Create from DataFrame.""" + item_tfms=None, batch_tfms=None, apply_reorder=False, target_spacing=None, + path=None, **kwargs): + """Create from DataFrame. + + File names in `fn_col` are used as given. If `path` or `folder` is set, they are + read relative to `path/folder`, as in fastai's `ImageDataLoaders.from_df`. The + prefix is applied to single file names only, so give full paths for multi-channel + rows that list several files. + """ + pref = '' + if path is not None or folder is not None: + pref = f"{Path(path or '.') / (folder or '')}{os.path.sep}" + if y_block is None: is_multi = (is_listy(label_col) and len(label_col) > 1) or label_delim is not None y_block = MultiCategoryBlock if is_multi else CategoryBlock @@ -99,15 +112,16 @@ def from_df(cls, df, valid_pct=0.2, seed=None, fn_col=0, folder=None, suff='', dblock = MedDataBlock( blocks=(ImageBlock(cls=MedImage), y_block), - get_x=ColReader(fn_col, suff=suff), + get_x=ColReader(fn_col, pref=pref, suff=suff), get_y=ColReader(label_col, label_delim=label_delim), splitter=splitter, item_tfms=item_tfms, + batch_tfms=batch_tfms, apply_reorder=apply_reorder, target_spacing=target_spacing ) - dls = cls.from_dblock(dblock, df, **kwargs) + dls = cls.from_dblock(dblock, df, path=path or '.', **kwargs) # Retain column identities; actual membership remains in train_ds/valid_ds. dls._img_col = fn_col dls._mask_col = label_col @@ -268,7 +282,8 @@ def plot_top_losses(x: MedImage, y: TensorMultiCategory, samples, outs, raws, losses, nrows: int | None = None, ncols: int | None = None, figsize=None, channel: int = 0, slice_index: int | None = None, anatomical_plane: int = 0, **kwargs): - # TODO: not tested yet + """Show images in top_losses for multi-label classification, with a table of + their targets, predictions, probabilities and losses.""" axs = get_grid(len(samples), nrows=nrows, ncols=ncols, figsize=figsize) for i, (ax, s) in enumerate(zip(axs, samples)): diff --git a/fastMONAI/vision_inference.py b/fastMONAI/vision_inference.py index 4528e6f..30517f5 100644 --- a/fastMONAI/vision_inference.py +++ b/fastMONAI/vision_inference.py @@ -1,3 +1,5 @@ +"""Run trained segmentation models on new images and restore the masks to the original image space.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_vision_inference.ipynb. # %% auto #0 @@ -68,15 +70,13 @@ def _do_resize(o, target_shape, image_interpolation='linear', # %% ../nbs/06_vision_inference.ipynb #shape-based-label-resize def _signed_depth(mask: np.ndarray, spacing, margin: int = 8) -> np.ndarray: - """Signed distance to the mask edge in mm: positive inside, negative outside. + """Signed distance to the edge of a non-empty mask in mm: positive inside, negative outside. Computed on the mask's bounding box plus margin voxels, which is wide enough for cubic interpolation to match a full-volume distance map. Voxels outside the box get a value below every value inside it, so they stay outside. """ outside = -4.0 * float(max(spacing)) - if not mask.any(): - return np.full(mask.shape, outside, dtype=np.float32) if mask.all(): return np.full(mask.shape, -outside, dtype=np.float32) coords = np.argwhere(mask) @@ -94,11 +94,13 @@ def _signed_depth(mask: np.ndarray, spacing, margin: int = 8) -> np.ndarray: def _resize_labels(o, target_shape, label_interpolation='shape'): """Resize a single-channel label image so its shape matches target_shape. - 'shape' uses shape-based interpolation: the signed distance to each label's edge is - resized with cubic B-spline interpolation and every voxel takes the label it lies - deepest inside. A slice that falls between two source slices then gets an outline - between theirs instead of a copy of the nearest one. 'nearest' copies the nearest - source voxel. + '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 + 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 + outline between theirs instead of a copy of the nearest one. 'nearest' copies the + nearest source voxel. """ if label_interpolation == 'nearest': return _do_resize(o, target_shape, image_interpolation='nearest') @@ -113,11 +115,13 @@ def _resize_labels(o, target_shape, label_interpolation='shape'): values = [v for v in np.unique(labels) if v != 0] if not values: return _do_resize(o, target_shape, image_interpolation='nearest') - depth = np.stack([_signed_depth(labels == v, o.spacing) for v in values]) + 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') - best, idx = resized.data.max(0) - restored = torch.where(best > 0, torch.as_tensor(values, dtype=torch.float32)[idx], 0.) + 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.) return type(o)(tensor=restored.unsqueeze(0).to(o.data.dtype), affine=resized.affine) @@ -151,7 +155,10 @@ def load_model_resources(model_path, device='cpu'): # %% ../nbs/06_vision_inference.ipynb #model-only-predict def predict_from_config(model, inference_config, input): - """Apply serialized item transforms, execute one input, and decode its output contract.""" + """Apply serialized item transforms, execute one input, and decode its output contract. + + The input is cast to float32 first, as images are during training. + """ if not isinstance(inference_config, dict) or inference_config.get('workflow') != 'standard': raise ValueError("predict_from_config requires a workflow='standard' inference config.") output_spec = inference_config.get('output') @@ -170,7 +177,7 @@ def predict_from_config(model, inference_config, input): f'got {tuple(tensor.shape)}.' ) - transformed = MedImage.create(tensor) + transformed = MedImage.create(tensor.float()) for transform in transforms_from_specs(inference_config.get('item_tfms') or []): transformed = transform(transformed, split_idx=1) @@ -194,6 +201,39 @@ def predict_from_config(model, inference_config, input): decoded = decoded[0] return decoded +# %% ../nbs/06_vision_inference.ipynb #restore-prediction-helpers +def _prediction_save_path(save_path, fn, org_img): + """Path to save a prediction to: save_path/'pred_', or None without save_path. + + For a list of files, the last one names the prediction, as it is the one that gives the + original image. Without `fn`, the file `org_img` was read from is used. + """ + if not save_path: + return None + source_fn = fn if fn is None or isinstance(fn, (str, Path)) else (fn[-1] if len(fn) else None) + source_fn = source_fn or getattr(org_img, 'path', None) + if not source_fn: + raise ValueError('fn is required when save_path is provided and org_img has no file path.') + return Path(save_path) / ('pred_' + Path(source_fn).name) + + +def _restore_prediction(pred, org_img, input_img, org_size, 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. + """ + 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()) + if save_fn is None: + return pred_img + pred_img.save(save_fn) + return save_fn + + # %% ../nbs/06_vision_inference.ipynb #model-only-spatial-inference def inference_from_config(model, inference_config, fn: (str, Path) = '', save_path: (str, Path) = None, org_img=None, @@ -219,6 +259,7 @@ def inference_from_config(model, inference_config, fn: (str, Path) = '', ) else: org_img, input_img = copy(org_img), copy(input_img) + save_fn = _prediction_save_path(save_path, fn, org_img) pred = predict_from_config(model, inference_config, input_img) if not isinstance(pred, torch.Tensor): @@ -230,83 +271,47 @@ 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)}.' ) - - 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) - ) - - from torchio import ScalarImage - - if hasattr(org_img, 'affine') and org_img.affine is not None: - affine_matrix = org_img.affine.copy() - elif MedBase.affine_matrix is not None: - affine_matrix = MedBase.affine_matrix.copy() - else: - affine_matrix = np.eye(4, dtype=np.float64) - - pred_img = ScalarImage(tensor=reoriented_array, affine=affine_matrix) - if save_path: - source_fn = fn[-1] if isinstance(fn, (list, tuple)) else fn - if not source_fn: - raise ValueError('fn is required when save_path is provided.') - save_fn = Path(save_path) / ('pred_' + Path(source_fn).name) - pred_img.save(save_fn) - return save_fn - return pred_img + return _restore_prediction(pred, org_img, input_img, org_size, 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) = '', save_path: (str, Path) = None, org_img=None, input_img=None, org_size=None, label_interpolation='shape'): - """Predict on new data using exported model. + """Segment an image with an exported learner and restore the mask to the original image space. - The mask is restored with shape-based interpolation; pass - label_interpolation='nearest' to copy the nearest voxel instead. + Only segmentation learners are supported: the decoded prediction is treated as a mask. + + Args: + learn_inf: Exported segmentation learner, e.g. from `load_learner`. + 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. + 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. + label_interpolation: 'shape' (default) restores the mask with shape-based + interpolation; 'nearest' copies the nearest voxel instead. + + Returns: + The saved file path if `save_path` is given, otherwise the mask as a `ScalarImage` + with the original shape, orientation and affine. """ 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, + 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) - else: - org_img, input_img = copy(org_img), copy(input_img) - - pred, *_ = learn_inf.predict(input_img.data) - - 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)) - - # Spatial properties (spacing, origin, direction) are automatically derived from affine - from torchio import ScalarImage - import numpy as np - - if hasattr(org_img, 'affine') and org_img.affine is not None: - affine_matrix = org_img.affine.copy() - elif MedBase.affine_matrix is not None: - affine_matrix = MedBase.affine_matrix.copy() else: - affine_matrix = np.eye(4, dtype=np.float64) - - pred_img = ScalarImage(tensor=reoriented_array, affine=affine_matrix) + org_img, input_img = copy(org_img), copy(input_img) + save_fn = _prediction_save_path(save_path, fn, org_img) - if save_path: - save_fn = Path(save_path)/('pred_' + Path(fn).parts[-1]) - pred_img.save(save_fn) - return save_fn - - return pred_img + pred, *_ = learn_inf.predict(input_img.data.float()) + return _restore_prediction(pred, org_img, input_img, org_size, 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_plot.py b/fastMONAI/vision_plot.py index 75e50cc..71ca4c6 100644 --- a/fastMONAI/vision_plot.py +++ b/fastMONAI/vision_plot.py @@ -1,3 +1,5 @@ +"""Plot slices of 3D medical images and segmentation masks.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_vision_plot.ipynb. # %% auto #0 @@ -56,13 +58,13 @@ def validate_anatomical_plane(anatomical_plane): raise ValueError("anatomical_plane should be either 0, 1, or 2.") # %% ../nbs/00_vision_plot.ipynb #7955e15b-7580-4219-838e-93ff094e146a -@delegates(plt.Axes.imshow, keep=True, but=['shape', 'imlim']) +@delegates(plt.Axes.imshow, keep=True) def show_med_img(im, ctx, channel: int, slice_index: int, anatomical_plane: int, voxel_size: (int, list), ax=None, figsize=None, title=None, **kwargs): """Show an image on `ax`. This is a modified code from the fastai function `show_image`. Args: - im: The input image. + im: The input image: a MedImage, MedMask, tensor or numpy array of shape [C, H, W, D]. ctx: The context. channel: Channel of the image. slice_index: Index of the 2D slice. @@ -79,15 +81,19 @@ def show_med_img(im, ctx, channel: int, slice_index: int, anatomical_plane: int, validate_anatomical_plane(anatomical_plane) - if hasattrs(im, ('data', 'cpu', 'permute')): - im = im.data.cpu() - im, aspect = _get_slice( - im, - channel=channel, - anatomical_plane=anatomical_plane, - voxel_size=voxel_size, - slice_index=slice_index - ) + if isinstance(im, np.ndarray): + im = torch.as_tensor(np.ascontiguousarray(im)) + if not hasattrs(im, ('data', 'cpu', 'permute')): + raise TypeError(f'Expected a tensor or numpy array, got {type(im).__name__}.') + + im = im.data.cpu() + im, aspect = _get_slice( + im, + channel=channel, + anatomical_plane=anatomical_plane, + voxel_size=voxel_size, + slice_index=slice_index + ) ax = ax if ax is not None else ctx diff --git a/nbs/00_vision_plot.ipynb b/nbs/00_vision_plot.ipynb index 4b23f40..1016320 100644 --- a/nbs/00_vision_plot.ipynb +++ b/nbs/00_vision_plot.ipynb @@ -24,7 +24,7 @@ "metadata": {}, "source": [ "# Vision plot\n", - ">" + "> Plot slices of 3D medical images and segmentation masks." ] }, { @@ -98,13 +98,13 @@ "outputs": [], "source": [ "#| export\n", - "@delegates(plt.Axes.imshow, keep=True, but=['shape', 'imlim'])\n", + "@delegates(plt.Axes.imshow, keep=True)\n", "def show_med_img(im, ctx, channel: int, slice_index: int, anatomical_plane: int,\n", " voxel_size: (int, list), ax=None, figsize=None, title=None, **kwargs):\n", " \"\"\"Show an image on `ax`. This is a modified code from the fastai function `show_image`.\n", "\n", " Args:\n", - " im: The input image.\n", + " im: The input image: a MedImage, MedMask, tensor or numpy array of shape [C, H, W, D].\n", " ctx: The context.\n", " channel: Channel of the image.\n", " slice_index: Index of the 2D slice.\n", @@ -121,15 +121,19 @@ " \n", " validate_anatomical_plane(anatomical_plane)\n", " \n", - " if hasattrs(im, ('data', 'cpu', 'permute')):\n", - " im = im.data.cpu()\n", - " im, aspect = _get_slice(\n", - " im, \n", - " channel=channel, \n", - " anatomical_plane=anatomical_plane,\n", - " voxel_size=voxel_size, \n", - " slice_index=slice_index\n", - " )\n", + " if isinstance(im, np.ndarray):\n", + " im = torch.as_tensor(np.ascontiguousarray(im))\n", + " if not hasattrs(im, ('data', 'cpu', 'permute')):\n", + " raise TypeError(f'Expected a tensor or numpy array, got {type(im).__name__}.')\n", + "\n", + " im = im.data.cpu()\n", + " im, aspect = _get_slice(\n", + " im, \n", + " channel=channel, \n", + " anatomical_plane=anatomical_plane,\n", + " voxel_size=voxel_size, \n", + " slice_index=slice_index\n", + " )\n", "\n", " ax = ax if ax is not None else ctx \n", "\n", @@ -146,6 +150,24 @@ " return ax" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "show-med-img-numpy-test", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# A numpy array is shown the same way as the equivalent tensor.\n", + "from fastcore.test import test_eq\n", + "_arr = np.zeros((1, 6, 8, 4), dtype=np.float32)\n", + "_arr[0, 2:4, 3:6, 1:3] = 1\n", + "_ax_np = show_med_img(_arr, ctx=None, channel=0, slice_index=None, anatomical_plane=2, voxel_size=1)\n", + "_ax_t = show_med_img(torch.from_numpy(_arr), ctx=None, channel=0, slice_index=None, anatomical_plane=2, voxel_size=1)\n", + "test_eq(np.asarray(_ax_np.images[0].get_array()), np.asarray(_ax_t.images[0].get_array()))\n", + "plt.close('all')\n" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/nbs/01_vision_core.ipynb b/nbs/01_vision_core.ipynb index 9f02394..aa9ed5e 100644 --- a/nbs/01_vision_core.ipynb +++ b/nbs/01_vision_core.ipynb @@ -41,7 +41,7 @@ "metadata": {}, "source": [ "# Vision core\n", - ">" + "> Load medical images and represent them as `MedImage` and `MedMask` tensors." ] }, { @@ -67,7 +67,7 @@ " Args:\n", " obj: The object to preprocess.\n", " apply_reorder: Whether to reorder the object.\n", - " target_spacing: Whether to resample the object.\n", + " target_spacing: Voxel spacing to resample to. None keeps the original spacing.\n", "\n", " Returns:\n", " The preprocessed object and its original size.\n", @@ -103,7 +103,7 @@ " Args:\n", " file_path: Image file path.\n", " apply_reorder: Whether to reorder data for canonical (RAS+) orientation.\n", - " target_spacing: Whether to resample image to different voxel sizes and dimensions.\n", + " target_spacing: Voxel spacing to resample to. None keeps the original spacing.\n", " dtype: Desired datatype for output.\n", "\n", " Returns:\n", @@ -128,9 +128,10 @@ " Load and preprocess multisequence data.\n", "\n", " Args:\n", - " image_paths: List of image paths (e.g., T1, T2, T1CE, DWI).\n", + " image_paths: List of image paths (e.g., T1, T2, T1CE, DWI). The first channel of\n", + " each file is stacked; the original image and size come from the last file.\n", " apply_reorder: Whether to reorder data for canonical (RAS+) orientation.\n", - " target_spacing: Whether to resample image to different voxel sizes and dimensions.\n", + " target_spacing: Voxel spacing to resample to. None keeps the original spacing.\n", " only_tensor: Whether to return only image tensor.\n", " dtype: Desired datatype for output.\n", "\n", @@ -163,11 +164,13 @@ " \"\"\"Loads and preprocesses a medical image.\n", "\n", " Args:\n", - " file_path: Path to the image. Can be a string, Path object or a list.\n", + " file_path: Path to the image, or a list of paths (e.g. T1, T2) that are\n", + " stacked as channels of one image, using the first channel of each file.\n", + " For a list, the original image and size come from the last file.\n", " apply_reorder: Whether to reorder the data to be closest to canonical \n", " (RAS+) orientation. Defaults to False.\n", - " target_spacing: Whether to resample image to different voxel sizes and \n", - " image dimensions. Defaults to None.\n", + " target_spacing: Voxel spacing to resample to. Defaults to None, which keeps\n", + " the original spacing.\n", " only_tensor: Whether to return only image tensor. Defaults to True.\n", " dtype: Datatype for the return value. Defaults to torch.Tensor.\n", "\n", @@ -194,7 +197,7 @@ "id": "971e94f0", "metadata": {}, "source": [ - "## Basic types -" + "## Basic types" ] }, { @@ -244,8 +247,9 @@ " If `fn` is a torch.Tensor, it's cast to MedBase object.\n", "\n", " Args:\n", - " fn : (Path, str, torch.Tensor)\n", - " Image path or a 4D torch.Tensor.\n", + " fn : (Path, str, list, torch.Tensor)\n", + " Image path, a list of image paths stacked as channels\n", + " (see `med_img_reader`), or a 4D torch.Tensor.\n", " kwargs : dict\n", " Additional parameters for the medical image reader.\n", "\n", @@ -338,7 +342,7 @@ "\n", " def __repr__(self) -> str:\n", " \"\"\"Returns the string representation of the MedBase instance.\"\"\"\n", - " return f'{self.__class__.__name__} mode={self.mode} size={\"x\".join([str(d) for d in self.size])}'" + " return f'{self.__class__.__name__}(shape={tuple(self.shape)}, dtype={self.dtype})'" ] }, { @@ -363,10 +367,22 @@ "source": [ "#| export\n", "class MedMask(MedBase):\n", - " \"\"\"Subclass of MedBase that represents an mask object.\"\"\"\n", + " \"\"\"Subclass of MedBase that represents a mask object.\"\"\"\n", " _show_args = {'alpha':0.5, 'cmap':'tab20'}" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "medbase-repr-test", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "test_eq(repr(MedImage(torch.zeros(1, 4, 5, 6))), 'MedImage(shape=(1, 4, 5, 6), dtype=torch.float32)')\n", + "test_eq(str(MedMask(torch.zeros(1, 2, 2, 2, dtype=torch.uint8))), 'MedMask(shape=(1, 2, 2, 2), dtype=torch.uint8)')\n" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/nbs/02_vision_data.ipynb b/nbs/02_vision_data.ipynb index 319155a..4027b20 100644 --- a/nbs/02_vision_data.ipynb +++ b/nbs/02_vision_data.ipynb @@ -35,7 +35,7 @@ "metadata": {}, "source": [ "# Vision data\n", - ">" + "> Build fastai `DataLoaders` for medical images and show batches, results and top losses." ] }, { @@ -108,7 +108,7 @@ "id": "28008ac9", "metadata": {}, "source": [ - "## `DataBlock` for vision -" + "## `DataBlock` for vision" ] }, { @@ -145,7 +145,7 @@ "id": "0879e254", "metadata": {}, "source": [ - "## `DataLoaders` for vision -" + "## `DataLoaders` for vision" ] }, { @@ -163,9 +163,20 @@ " @delegates(DataLoaders.from_dblock)\n", " def from_df(cls, df, valid_pct=0.2, seed=None, fn_col=0, folder=None, suff='',\n", " label_col=1, label_delim=None, y_block=None, valid_col=None,\n", - " item_tfms=None, batch_tfms=None, apply_reorder=False, target_spacing=None, **kwargs):\n", - " \"\"\"Create from DataFrame.\"\"\"\n", + " item_tfms=None, batch_tfms=None, apply_reorder=False, target_spacing=None,\n", + " path=None, **kwargs):\n", + " \"\"\"Create from DataFrame.\n", + "\n", + " File names in `fn_col` are used as given. If `path` or `folder` is set, they are\n", + " read relative to `path/folder`, as in fastai's `ImageDataLoaders.from_df`. The\n", + " prefix is applied to single file names only, so give full paths for multi-channel\n", + " rows that list several files.\n", + " \"\"\"\n", " \n", + " pref = ''\n", + " if path is not None or folder is not None:\n", + " pref = f\"{Path(path or '.') / (folder or '')}{os.path.sep}\"\n", + "\n", " if y_block is None:\n", " is_multi = (is_listy(label_col) and len(label_col) > 1) or label_delim is not None\n", " y_block = MultiCategoryBlock if is_multi else CategoryBlock\n", @@ -175,27 +186,59 @@ "\n", " dblock = MedDataBlock(\n", " blocks=(ImageBlock(cls=MedImage), y_block),\n", - " get_x=ColReader(fn_col, suff=suff),\n", + " get_x=ColReader(fn_col, pref=pref, suff=suff),\n", " get_y=ColReader(label_col, label_delim=label_delim),\n", " splitter=splitter,\n", " item_tfms=item_tfms,\n", + " batch_tfms=batch_tfms,\n", " apply_reorder=apply_reorder,\n", " target_spacing=target_spacing\n", " )\n", "\n", - " dls = cls.from_dblock(dblock, df, **kwargs)\n", + " dls = cls.from_dblock(dblock, df, path=path or '.', **kwargs)\n", " # Retain column identities; actual membership remains in train_ds/valid_ds.\n", " dls._img_col = fn_col\n", " dls._mask_col = label_col\n", " return dls" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "from-df-tests", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# batch_tfms reach the dataloaders, and file names are read relative to folder.\n", + "import tempfile\n", + "import torchio as tio\n", + "from fastcore.test import test_eq\n", + "\n", + "class _MarkBatch(Transform):\n", + " order = 99\n", + " def encodes(self, x: MedImage): return x\n", + "\n", + "with tempfile.TemporaryDirectory() as _tmp:\n", + " _names = [f'img{_i}.nii.gz' for _i in range(5)]\n", + " for _name in _names:\n", + " tio.ScalarImage(tensor=torch.rand(1, 4, 4, 4), affine=np.eye(4)).save(Path(_tmp) / _name)\n", + " _df = pd.DataFrame({'fn': _names, 'path': [str(Path(_tmp) / _n) for _n in _names],\n", + " 'label': ['a', 'b', 'a', 'b', 'a']})\n", + " _dls = MedImageDataLoaders.from_df(_df, fn_col='fn', folder=_tmp, label_col='label',\n", + " batch_tfms=[_MarkBatch()], bs=2, num_workers=0)\n", + " assert any(isinstance(_t, _MarkBatch) for _t in _dls.train.after_batch.fs)\n", + " test_eq(tuple(_dls.one_batch()[0].shape), (2, 1, 4, 4, 4))\n", + " _dls = MedImageDataLoaders.from_df(_df, fn_col='path', label_col='label', bs=2, num_workers=0)\n", + " test_eq(tuple(_dls.one_batch()[0].shape), (2, 1, 4, 4, 4))\n" + ] + }, { "cell_type": "markdown", "id": "963ff9fe", "metadata": {}, "source": [ - "## Show methods - " + "## Show methods" ] }, { @@ -203,7 +246,7 @@ "id": "fb2b67d8", "metadata": {}, "source": [ - "### Show batch of images - " + "### Show batch of images" ] }, { @@ -227,7 +270,7 @@ "id": "f913e202", "metadata": {}, "source": [ - "### Show results - " + "### Show results" ] }, { @@ -259,7 +302,7 @@ "id": "6eb00e73", "metadata": {}, "source": [ - "### Show top losses - " + "### Show top losses" ] }, { @@ -276,7 +319,7 @@ "id": "f7346f94", "metadata": {}, "outputs": [], - "source": "#| export\n@dispatch\ndef plot_top_losses(x: MedImage, y: TensorMultiCategory, samples, outs, raws, \n losses, nrows: int | None = None, ncols: int | None = None, figsize=None, \n channel: int = 0, slice_index: int | None = None, \n anatomical_plane: int = 0, **kwargs):\n # TODO: not tested yet\n axs = get_grid(len(samples), nrows=nrows, ncols=ncols, figsize=figsize)\n\n for i, (ax, s) in enumerate(zip(axs, samples)):\n s[0].show(ctx=ax, title=f'Image {i}', channel=channel, \n slice_index=slice_index, anatomical_plane=anatomical_plane, **kwargs)\n\n rows = get_empty_df(len(samples))\n outs = L(s[1:] + o + (TitledStr(r), TitledFloat(l.item())) \n for s, o, r, l in zip(samples, outs, raws, losses))\n\n for i, l in enumerate([\"target\", \"predicted\", \"probabilities\", \"loss\"]):\n rows = [b.show(ctx=r, label=l, channel=channel, slice_index=slice_index, \n anatomical_plane=anatomical_plane, **kwargs) \n for b, r in zip(outs.itemgot(i), rows)]\n\n display_df(pd.DataFrame(rows))" + "source": "#| export\n@dispatch\ndef plot_top_losses(x: MedImage, y: TensorMultiCategory, samples, outs, raws, \n losses, nrows: int | None = None, ncols: int | None = None, figsize=None, \n channel: int = 0, slice_index: int | None = None, \n anatomical_plane: int = 0, **kwargs):\n \"\"\"Show images in top_losses for multi-label classification, with a table of\n their targets, predictions, probabilities and losses.\"\"\"\n axs = get_grid(len(samples), nrows=nrows, ncols=ncols, figsize=figsize)\n\n for i, (ax, s) in enumerate(zip(axs, samples)):\n s[0].show(ctx=ax, title=f'Image {i}', channel=channel, \n slice_index=slice_index, anatomical_plane=anatomical_plane, **kwargs)\n\n rows = get_empty_df(len(samples))\n outs = L(s[1:] + o + (TitledStr(r), TitledFloat(l.item())) \n for s, o, r, l in zip(samples, outs, raws, losses))\n\n for i, l in enumerate([\"target\", \"predicted\", \"probabilities\", \"loss\"]):\n rows = [b.show(ctx=r, label=l, channel=channel, slice_index=slice_index, \n anatomical_plane=anatomical_plane, **kwargs) \n for b, r in zip(outs.itemgot(i), rows)]\n\n display_df(pd.DataFrame(rows))" } ], "metadata": { diff --git a/nbs/06_vision_inference.ipynb b/nbs/06_vision_inference.ipynb index 51dfb77..5be13fa 100644 --- a/nbs/06_vision_inference.ipynb +++ b/nbs/06_vision_inference.ipynb @@ -43,7 +43,7 @@ "metadata": {}, "source": [ "# Vision inference\n", - ">" + "> Run trained segmentation models on new images and restore the masks to the original image space." ] }, { @@ -130,15 +130,13 @@ "source": [ "#| export\n", "def _signed_depth(mask: np.ndarray, spacing, margin: int = 8) -> np.ndarray:\n", - " \"\"\"Signed distance to the mask edge in mm: positive inside, negative outside.\n", + " \"\"\"Signed distance to the edge of a non-empty mask in mm: positive inside, negative outside.\n", "\n", " Computed on the mask's bounding box plus margin voxels, which is wide enough for cubic\n", " interpolation to match a full-volume distance map. Voxels outside the box get a value\n", " below every value inside it, so they stay outside.\n", " \"\"\"\n", " outside = -4.0 * float(max(spacing))\n", - " if not mask.any():\n", - " return np.full(mask.shape, outside, dtype=np.float32)\n", " if mask.all():\n", " return np.full(mask.shape, -outside, dtype=np.float32)\n", " coords = np.argwhere(mask)\n", @@ -156,11 +154,13 @@ "def _resize_labels(o, target_shape, label_interpolation='shape'):\n", " \"\"\"Resize a single-channel label image so its shape matches target_shape.\n", "\n", - " 'shape' uses shape-based interpolation: the signed distance to each label's edge is\n", - " resized with cubic B-spline interpolation and every voxel takes the label it lies\n", - " deepest inside. A slice that falls between two source slices then gets an outline\n", - " between theirs instead of a copy of the nearest one. 'nearest' copies the nearest\n", - " source voxel.\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", + " 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", + " outline between theirs instead of a copy of the nearest one. 'nearest' copies the\n", + " nearest source voxel.\n", " \"\"\"\n", " if label_interpolation == 'nearest':\n", " return _do_resize(o, target_shape, image_interpolation='nearest')\n", @@ -175,11 +175,13 @@ " 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", - " depth = np.stack([_signed_depth(labels == v, o.spacing) for v in values])\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", - " best, idx = resized.data.max(0)\n", - " restored = torch.where(best > 0, torch.as_tensor(values, dtype=torch.float32)[idx], 0.)\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", " return type(o)(tensor=restored.unsqueeze(0).to(o.data.dtype), affine=resized.affine)\n" ] }, @@ -237,6 +239,28 @@ " contains='label_interpolation')\n" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "touching-labels-resize-tests", + "metadata": {}, + "outputs": [], + "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", + "_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", + " test_eq(int(((_out == 0) & _union).sum()), 0)\n", + " test_eq(torch.unique(_out).tolist(), [0., 1., 2.])\n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -290,7 +314,10 @@ "source": [ "#| export\n", "def predict_from_config(model, inference_config, input):\n", - " \"\"\"Apply serialized item transforms, execute one input, and decode its output contract.\"\"\"\n", + " \"\"\"Apply serialized item transforms, execute one input, and decode its output contract.\n", + "\n", + " The input is cast to float32 first, as images are during training.\n", + " \"\"\"\n", " if not isinstance(inference_config, dict) or inference_config.get('workflow') != 'standard':\n", " raise ValueError(\"predict_from_config requires a workflow='standard' inference config.\")\n", " output_spec = inference_config.get('output')\n", @@ -309,7 +336,7 @@ " f'got {tuple(tensor.shape)}.'\n", " )\n", "\n", - " transformed = MedImage.create(tensor)\n", + " transformed = MedImage.create(tensor.float())\n", " for transform in transforms_from_specs(inference_config.get('item_tfms') or []):\n", " transformed = transform(transformed, split_idx=1)\n", "\n", @@ -334,6 +361,46 @@ " return decoded" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "restore-prediction-helpers", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "def _prediction_save_path(save_path, fn, org_img):\n", + " \"\"\"Path to save a prediction to: save_path/'pred_', or None without save_path.\n", + "\n", + " For a list of files, the last one names the prediction, as it is the one that gives the\n", + " original image. Without `fn`, the file `org_img` was read from is used.\n", + " \"\"\"\n", + " if not save_path:\n", + " return None\n", + " source_fn = fn if fn is None or isinstance(fn, (str, Path)) else (fn[-1] if len(fn) else None)\n", + " source_fn = source_fn or getattr(org_img, 'path', None)\n", + " if not source_fn:\n", + " raise ValueError('fn is required when save_path is provided and org_img has no file path.')\n", + " 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", + " \"\"\"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", + " \"\"\"\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", + " if save_fn is None:\n", + " return pred_img\n", + " pred_img.save(save_fn)\n", + " return save_fn\n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -366,6 +433,7 @@ " )\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 = predict_from_config(model, inference_config, input_img)\n", " if not isinstance(pred, torch.Tensor):\n", @@ -377,34 +445,7 @@ " 'Decoded segmentation must have shape [D, H, W] or [1, D, H, W]; '\n", " f'got {tuple(pred.shape)}.'\n", " )\n", - "\n", - " pred_mask = do_pad_or_crop(\n", - " pred.float(), input_img.shape[1:], padding_mode=0, mask_name=None\n", - " )\n", - " input_img.set_data(pred_mask)\n", - " input_img = _resize_labels(input_img, org_size, label_interpolation)\n", - " reoriented_array = _to_original_orientation(\n", - " input_img.as_sitk(), ''.join(org_img.orientation)\n", - " )\n", - "\n", - " from torchio import ScalarImage\n", - "\n", - " if hasattr(org_img, 'affine') and org_img.affine is not None:\n", - " affine_matrix = org_img.affine.copy()\n", - " elif MedBase.affine_matrix is not None:\n", - " affine_matrix = MedBase.affine_matrix.copy()\n", - " else:\n", - " affine_matrix = np.eye(4, dtype=np.float64)\n", - "\n", - " pred_img = ScalarImage(tensor=reoriented_array, affine=affine_matrix)\n", - " if save_path:\n", - " source_fn = fn[-1] if isinstance(fn, (list, tuple)) else fn\n", - " if not source_fn:\n", - " raise ValueError('fn is required when save_path is provided.')\n", - " save_fn = Path(save_path) / ('pred_' + Path(source_fn).name)\n", - " pred_img.save(save_fn)\n", - " return save_fn\n", - " return pred_img" + " return _restore_prediction(pred, org_img, input_img, org_size, label_interpolation, save_fn)" ] }, { @@ -414,6 +455,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "from fastcore.test import test_eq, test_fail\n", "from fastMONAI.utils import make_output_spec, standard_inference_config\n", "\n", @@ -520,7 +562,53 @@ "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 \"\"\"Predict on new data using exported model.\n\n The mask is restored with shape-based interpolation; pass\n label_interpolation='nearest' to copy the nearest voxel instead.\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 \n pred, *_ = learn_inf.predict(input_img.data)\n \n pred_mask = do_pad_or_crop(pred.float(), input_img.shape[1:], padding_mode=0, \n mask_name=None)\n input_img.set_data(pred_mask)\n \n input_img = _resize_labels(input_img, org_size, label_interpolation)\n \n reoriented_array = _to_original_orientation(input_img.as_sitk(), \n ('').join(org_img.orientation))\n \n # Spatial properties (spacing, origin, direction) are automatically derived from affine\n from torchio import ScalarImage\n import numpy as np\n \n if hasattr(org_img, 'affine') and org_img.affine is not None:\n affine_matrix = org_img.affine.copy()\n elif MedBase.affine_matrix is not None:\n affine_matrix = MedBase.affine_matrix.copy()\n else:\n affine_matrix = np.eye(4, dtype=np.float64)\n \n pred_img = ScalarImage(tensor=reoriented_array, affine=affine_matrix)\n\n if save_path:\n save_fn = Path(save_path)/('pred_' + Path(fn).parts[-1])\n pred_img.save(save_fn)\n return save_fn\n \n return pred_img" + "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)" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "inference-input-and-save-tests", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# Inputs are cast to float32, and predictions are saved for a list of files or for a\n", + "# preloaded image without fn.\n", + "class _RecordingLearner:\n", + " def predict(self, x):\n", + " self.dtype = x.dtype\n", + " return ((x[:1] > 50).long(),)\n", + "\n", + "class _Conv(torch.nn.Module):\n", + " def __init__(self):\n", + " super().__init__()\n", + " self.conv = torch.nn.Conv3d(1, 1, 1)\n", + " def forward(self, x): return self.conv(x)\n", + "\n", + "with tempfile.TemporaryDirectory() as _tmp:\n", + " _fns = []\n", + " for _i in range(2):\n", + " _fn = Path(_tmp) / f'int16_{_i}.nii.gz'\n", + " tio.ScalarImage(tensor=torch.randint(0, 100, (1, 6, 6, 6)).short(), affine=np.eye(4)).save(_fn)\n", + " _fns.append(str(_fn))\n", + " _learner = _RecordingLearner()\n", + " test_eq(type(inference(_learner, False, None, fn=_fns[0])), tio.ScalarImage)\n", + " test_eq(_learner.dtype, torch.float32)\n", + " _conv_config = standard_inference_config(\n", + " apply_reorder=False, target_spacing=None, item_tfms=[],\n", + " output_spec=make_output_spec('binary_segmentation', threshold=0.5),\n", + " )\n", + " test_eq(tuple(inference_from_config(_Conv(), _conv_config, fn=_fns[0]).shape), (1, 6, 6, 6))\n", + "\n", + " 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(_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" + ] }, { "cell_type": "code", @@ -560,7 +648,7 @@ "id": "02ee71aa", "metadata": {}, "outputs": [], - "source": "from fastcore.test import test_eq, test_fail\n\n# large 4x4x4 (64 vox) + small 2x2x2 (8 vox) disjoint cubes\n_m = np.zeros((10, 10, 10), dtype=np.uint8)\n_m[1:5, 1:5, 1:5] = 1\n_m[7:9, 7:9, 7:9] = 1\n\n# cutoff 64*0.2=12.8 removes the small cube, keeps the large\n_out = refine_binary_pred_mask(_m, remove_size=64, percentage=0.2)\ntest_eq(float(_out.sum()), 64.)\ntest_eq(float(_out[7:9, 7:9, 7:9].sum()), 0.)\n\n# cutoff 8*0.2=1.6 keeps both cubes\n_out2 = refine_binary_pred_mask(_m, remove_size=8, percentage=0.2)\ntest_eq(float(_out2.sum()), 72.)\n\n# empty mask returns zeros\n_empty = np.zeros((4, 4, 4), dtype=np.uint8)\ntest_eq(float(refine_binary_pred_mask(_empty, remove_size=8).sum()), 0.)\n\n# remove_size is required and must be positive\ntest_fail(lambda: refine_binary_pred_mask(_m))\ntest_fail(lambda: refine_binary_pred_mask(_m, remove_size=None))\ntest_fail(lambda: refine_binary_pred_mask(_m, remove_size=0))" + "source": "#| hide\nfrom fastcore.test import test_eq, test_fail\n\n# large 4x4x4 (64 vox) + small 2x2x2 (8 vox) disjoint cubes\n_m = np.zeros((10, 10, 10), dtype=np.uint8)\n_m[1:5, 1:5, 1:5] = 1\n_m[7:9, 7:9, 7:9] = 1\n\n# cutoff 64*0.2=12.8 removes the small cube, keeps the large\n_out = refine_binary_pred_mask(_m, remove_size=64, percentage=0.2)\ntest_eq(float(_out.sum()), 64.)\ntest_eq(float(_out[7:9, 7:9, 7:9].sum()), 0.)\n\n# cutoff 8*0.2=1.6 keeps both cubes\n_out2 = refine_binary_pred_mask(_m, remove_size=8, percentage=0.2)\ntest_eq(float(_out2.sum()), 72.)\n\n# empty mask returns zeros\n_empty = np.zeros((4, 4, 4), dtype=np.uint8)\ntest_eq(float(refine_binary_pred_mask(_empty, remove_size=8).sum()), 0.)\n\n# remove_size is required and must be positive\ntest_fail(lambda: refine_binary_pred_mask(_m))\ntest_fail(lambda: refine_binary_pred_mask(_m, remove_size=None))\ntest_fail(lambda: refine_binary_pred_mask(_m, remove_size=0))" }, { "cell_type": "code", From 29bd973b0247aa16ea45bde04ec91b7f88aab475 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:25:31 +0100 Subject: [PATCH 02/11] fix(vision_augmentation)!: apply GPU affine parameters on the matching axis GpuPatchAugmentation built the affine_grid theta in (x, y, z) = (W, H, D) order from (D, H, W) vectors, so per-axis degrees and translation acted on the reversed array axis, and scales were inverted (0.5 enlarged content). Index i now acts on array axis i and scales follow TorchIO's meaning. Noise std is drawn per sample from (0, std) like TorchIO, and a blur sigma of 0 leaves the sample unchanged instead of producing NaN. BREAKING CHANGE: GPU augmentation with per-axis degrees, translation or scales behaves differently, so retraining with the same config (e.g. the VS config degrees=(5, 5, 30)) sees the intended policy, not the old one. Already trained models are unaffected. Also: corrected list of GPU vs CPU differences, documented accepted config forms and PadOrCrop mask_name, OneOf accepts fastMONAI wrappers, removed the unused normalization registry alias and FocalLoss import, loss docstrings, hidden test cells. --- fastMONAI/vision_augmentation.py | 162 ++++++++++++------ fastMONAI/vision_loss.py | 8 +- nbs/03_vision_augment.ipynb | 256 +++++++++++++++++++++++------ nbs/04_vision_loss_functions.ipynb | 6 +- 4 files changed, 322 insertions(+), 110 deletions(-) diff --git a/fastMONAI/vision_augmentation.py b/fastMONAI/vision_augmentation.py index 1d4e72d..206ead2 100644 --- a/fastMONAI/vision_augmentation.py +++ b/fastMONAI/vision_augmentation.py @@ -1,3 +1,5 @@ +"""TorchIO and MONAI transform wrappers for 3D medical images, serializable inference transforms and GPU-batched patch augmentation.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_vision_augment.ipynb. # %% auto #0 @@ -23,7 +25,7 @@ class CustomDictTransform(ItemTransform): the image and the target (if it's a mask). """ - split_idx = 0 # Only perform transformations on training data. Use TTA() for transformations on validation data. + split_idx = 0 # Training only def __init__(self, aug): """Constructs CustomDictTransform object. @@ -76,7 +78,16 @@ def do_pad_or_crop(o, target_shape, padding_mode, mask_name, dtype=torch.Tensor) # %% ../nbs/03_vision_augment.ipynb #9c75bb84 class PadOrCrop(DisplayedTransform): - """Resize image using TorchIO CropOrPad.""" + """Resize image using TorchIO CropOrPad. + + Args: + size: Target shape, an int or three ints. + padding_mode: Padding value or mode passed to tio.CropOrPad. + mask_name: Ignored by this transform, which crops or pads each image and + mask on its own around the volume centre. Still accepted and saved in + to_spec() so older saved specs load. Only the raw tio_transform + honours it, when called on a Subject that has that key. + """ order = 0 @@ -302,8 +313,6 @@ def encodes(self, o: MedMask): 'RescaleIntensity': RescaleIntensity, 'NormalizeIntensity': NormalizeIntensity, } -# Backward-compatible private alias for code that inspected the old registry. -_NORM_TFM_REGISTRY = _TRANSFORM_SPEC_REGISTRY def _validate_transform_spec_json(spec): @@ -921,9 +930,20 @@ def __init__(self, axes='LRAPIS', p=0.5): # %% ../nbs/03_vision_augment.ipynb #ddd7b99b class OneOf(CustomDictTransform): - """Apply only one of the given transforms using TorchIO `OneOf`.""" + """Apply only one of the given transforms using TorchIO `OneOf`. + + transform_dict maps transforms to probabilities, or is a sequence of + transforms picked with equal probability. Transforms can be TorchIO + transforms or fastMONAI wrappers, which are unwrapped through their + tio_transform property. + """ def __init__(self, transform_dict, p=1): + unwrap = lambda t: getattr(t, 'tio_transform', t) + if isinstance(transform_dict, dict): + transform_dict = {unwrap(t): w for t, w in transform_dict.items()} + else: + transform_dict = [unwrap(t) for t in transform_dict] super().__init__(tio.OneOf(transform_dict, p=p)) # %% ../nbs/03_vision_augment.ipynb #lqet5pabzy @@ -1023,17 +1043,37 @@ class GpuPatchAugmentation: same parameters to both image and mask. Intensity transforms skip the mask. Each transform is controlled by a parameter dict with at minimum a 'p' key - for per-sample probability. Pass None to disable a transform. + for per-sample probability. Pass None to disable a transform. As in TorchIO, + index i of a per-axis value refers to array axis i of the patch (D, H, W), + and lengths are in voxels. Args: - affine: dict with keys 'scales', 'degrees', 'translation', - 'default_pad_value', 'p'. None to disable. - anisotropy: dict with keys 'axes', 'downsampling', 'p'. None to disable. - flip: dict with keys 'axes', 'p'. None to disable. - gamma: dict with keys 'log_gamma', 'p'. None to disable. - intensity_scale: dict with keys 'scale_range', 'p'. None to disable. - noise: dict with keys 'std', 'p'. None to disable. - blur: dict with keys 'std', 'p'. None to disable. + affine: dict with keys + 'scales': (lo, hi). Each axis gets a factor from U(lo, hi); values + below 1 shrink the content, as in TorchIO. A single number is not + accepted. + 'degrees': d or per-axis maxima (d0, d1, d2). The rotation about + axis i is drawn from U(-d_i, d_i). + 'translation': t or per-axis maxima (t0, t1, t2), in voxels. The + shift along axis i is drawn from U(-t_i, t_i). + 'default_pad_value': value for voxels moved in from outside the + patch (default 0). + 'p': probability. + TorchIO's range forms for degrees and translation, (a, b) or six + values, are not supported. None to disable. + anisotropy: dict with keys 'axes' (one is picked at random), + 'downsampling' as (lo, hi), 'p'. None to disable. + flip: dict with keys 'axes' and 'p'; each axis is flipped + independently with probability p. None to disable. + gamma: dict with keys 'log_gamma' as (lo, hi), 'p'. None to disable. + intensity_scale: dict with keys 'scale_range' as (lo, hi), 'p'. + None to disable. + noise: dict with keys 'std', 'p'. Zero-mean noise whose std is drawn + from U(0, s) for a number s, or from U(lo, hi) for (lo, hi). + None to disable. + blur: dict with keys 'std', 'p'. The sigma in voxels is drawn from + U(0, s) for a number s, or from U(lo, hi) for (lo, hi); a sample + whose sigma is 0 is left unchanged. None to disable. Example:: @@ -1052,14 +1092,27 @@ class GpuPatchAugmentation: Ada (full pipeline, 128^3): this batched GPU path ~3 ms/batch (B=2) vs MONAI per-sample-on-GPU ~77 ms (~25x slower) vs TorchIO CPU ~287 ms (~100x slower); the gap is structural and grows with batch size, so the - GPU path is kept. It is an approximate reimplementation (voxel-space - affine, per-axis flip, image-only anisotropy, sign-preserving gamma), - not bit-identical to the CPU/TorchIO path. - The voxel-space affine's geometric error scales with the patch's physical-extent - anisotropy (patch_size x target_spacing), not voxel spacing alone: it is sub-voxel - when the physical extent is near-isotropic (the common case, including cubic patches - on isotropic data) and deviates materially only for strongly anisotropic physical - extents (e.g. a cubic patch on anisotropic spacing). + GPU path is kept. + + Differences from the CPU/TorchIO path (suggest_patch_augmentations): + Each parameter means the same thing in both paths. The CPU patch path + hands patches to TorchIO with an identity affine, so TorchIO works in + voxels there too. The GPU path is not bit-identical: + + - Rotations act in affine_grid's normalised [-1, 1] patch coordinates. + A rotation between two axes of different patch length is stretched by + their length ratio, so it is rigid in voxel space only between axes of + equal length (e.g. a 256 x 256 plane). It is close to rigid in + physical space when patch_size x target_spacing is near isotropic. + TorchIO rotates rigidly in voxel space. + - Blur uses one isotropic sigma per sample, a kernel cut at 3 sigma and + replicate padding. TorchIO draws a sigma per axis and uses scipy's + gaussian_filter (cut at 4 sigma, reflect padding). + - Anisotropy follows the same recipe as TorchIO (nearest down, linear + up, image only), but resamples with F.interpolate on the voxel grid, + so sample positions and the downsampled size can differ slightly. + - Gamma uses one exponent per sample, while TorchIO draws one per + channel. This only matters for multi-channel input. """ def __init__(self, affine=None, anisotropy=None, flip=None, @@ -1141,13 +1194,6 @@ def _apply_affine(self, img, mask): ], dim=1) # [n, 3] angles_rad = angles_deg * (math.pi / 180.0) - # Build rotation matrices [n, 3, 3] - R = _build_rotation_matrix_3d(angles_rad) - - # Scale matrix: S @ R -> [n, 3, 3] - S = torch.diag_embed(scales) # [n, 3, 3] - SR = S @ R # [n, 3, 3] - # Random translation (voxels -> normalized [-1, 1] coords) translation = cfg['translation'] if not isinstance(translation, (list, tuple)): @@ -1160,8 +1206,17 @@ def _apply_affine(self, img, mask): for i in range(3) ], dim=1) # [n, 3] - # Assemble theta for active samples - theta[idx, :3, :3] = SR + # The vectors above are indexed by array axis (D, H, W), while theta + # uses affine_grid's (x, y, z) = (W, H, D) order, so reverse them. + angles_rad = angles_rad.flip(1) + scales = scales.flip(1) + t_norm = t_norm.flip(1) + + # theta maps output to input coordinates, so 1/scales makes a scale + # below 1 shrink the content (as in TorchIO). Rotate, then scale along + # the output axes, matching TorchIO's composition. + R = _build_rotation_matrix_3d(angles_rad) # [n, 3, 3] + theta[idx, :3, :3] = R @ torch.diag_embed(1.0 / scales) theta[idx, :3, 3] = t_norm # Apply grid_sample @@ -1183,11 +1238,9 @@ def _apply_affine(self, img, mask): def _apply_anisotropy(self, img, mask): """Per-sample anisotropy simulation via F.interpolate. - Downsample along a random axis with nearest interpolation, - then upsample back with trilinear. Approximates TorchIO's - RandomAnisotropy and degrades the image only (TorchIO also - resamples the label map). - Only affects img, not mask (anisotropy is intensity degradation). + Downsample along a random axis with nearest interpolation, then + upsample back with trilinear. Only the image is changed, as in + TorchIO's RandomAnisotropy with scalars_only=True (the default). """ cfg = self.anisotropy B = img.shape[0] @@ -1276,7 +1329,11 @@ def _apply_intensity_scale(self, img): return img def _apply_noise(self, img): - """Batched additive Gaussian noise with per-sample random std.""" + """Batched zero-mean Gaussian noise with a per-sample std. + + A number s draws the std from U(0, s), as TorchIO RandomNoise does; + a pair (lo, hi) draws it from U(lo, hi). + """ cfg = self.noise B = img.shape[0] device = img.device @@ -1285,9 +1342,9 @@ def _apply_noise(self, img): std_val = cfg['std'] if isinstance(std_val, (list, tuple)): std_lo, std_hi = std_val - per_std = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi) else: - per_std = torch.full((B,), std_val, device=device, dtype=dtype) + std_lo, std_hi = 0.0, std_val + per_std = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi) # Zero std for inactive samples active = torch.rand(B, device=device) < cfg['p'] @@ -1298,7 +1355,10 @@ def _apply_noise(self, img): return img def _apply_blur(self, img): - """Batched separable 3D Gaussian blur via F.conv3d with groups trick.""" + """Batched separable 3D Gaussian blur via F.conv3d with groups trick. + + One isotropic sigma (in voxels) per sample. + """ cfg = self.blur B, C, D, H, W = img.shape device = img.device @@ -1310,9 +1370,9 @@ def _apply_blur(self, img): else: std_lo, std_hi = 0.0, std_val - # Per-sample sigma + # Per-sample sigma; samples drawn with sigma <= 0 are left unblurred sigma = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi) - active = torch.rand(B, device=device) < cfg['p'] + active = (torch.rand(B, device=device) < cfg['p']) & (sigma > 0) if not active.any(): return img @@ -1324,7 +1384,7 @@ def _apply_blur(self, img): # Build per-sample 1D Gaussian kernels [B, kernel_size] x = torch.arange(-kernel_radius, kernel_radius + 1, device=device, dtype=dtype) - # Avoid division by zero for sigma=0 + # Avoid division by zero for inactive samples safe_sigma = torch.where(active, sigma, torch.ones_like(sigma)) kernels = torch.exp(-x.unsqueeze(0)**2 / (2 * safe_sigma.unsqueeze(1)**2)) kernels = kernels / kernels.sum(dim=1, keepdim=True) @@ -1378,15 +1438,13 @@ def gpu_patch_augmentations(patch_size, target_spacing, noise_p=0.1, blur_p=0.2, flip_p=0.5): """Create GpuPatchAugmentation with nnU-Net-inspired defaults. - Approximate GPU reimplementation of suggest_patch_augmentations: it shares - the same parameter logic (via _compute_patch_aug_params) and probabilities, - but is NOT bit-identical: the affine is voxel-space (not spacing-aware), - flip is per-axis independent, anisotropy degrades the image only, and gamma - is sign-preserving. Runs batched on GPU (~25x faster than the per-sample - CPU/TorchIO path); see GpuPatchAugmentation for the benchmark and rationale. - The voxel-space affine's geometric error scales with physical-extent anisotropy - (patch_size x target_spacing), so it is sub-voxel for near-isotropic extents and only - matters for strongly anisotropic ones (e.g. a cubic patch on anisotropic spacing). + GPU counterpart of suggest_patch_augmentations: it shares the same parameter + logic (via _compute_patch_aug_params) and probabilities, and each value means + the same thing in both paths. It runs batched on GPU, much faster than the + per-sample CPU/TorchIO path, but is not bit-identical. See + GpuPatchAugmentation for the benchmark and the remaining differences + (rotation on non-cubic patches, isotropic blur sigma, anisotropy resampling + grid, per-sample gamma). Args: patch_size: List/tuple of 3 ints, patch dimensions. diff --git a/fastMONAI/vision_loss.py b/fastMONAI/vision_loss.py index 0d26721..01b691d 100644 --- a/fastMONAI/vision_loss.py +++ b/fastMONAI/vision_loss.py @@ -1,3 +1,5 @@ +"""Segmentation loss wrappers, a focal Tversky loss and a DynUNet deep-supervision adapter.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_vision_loss_functions.ipynb. # %% auto #0 @@ -7,7 +9,7 @@ import torch from .vision_core import * from .vision_data import pred_to_binary_mask, batch_pred_to_multiclass_mask -from monai.losses import TverskyLoss, FocalLoss +from monai.losses import TverskyLoss from torch.nn.modules.loss import _Loss # %% ../nbs/04_vision_loss_functions.ipynb #e0c0c220-aaeb-46c6-8d18-f72cd9da0555 @@ -55,6 +57,7 @@ def decodes(self, x) -> torch.Tensor: class TverskyFocalLoss(_Loss): """Focal Tversky loss: ``(1 - TI)**gamma`` (TI = Tversky index; see ``monai.losses.TverskyLoss``). ``gamma=1`` reduces to the plain Tversky loss. + gamma is applied after TverskyLoss's mean over batch and classes, i.e. ``mean(1 - TI)**gamma``. """ def __init__( @@ -91,7 +94,8 @@ def __init__( def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: """ Args: - input: the shape should be [B, C, W, H, D]. The input should be the original logits. + input: the shape should be [B, C, W, H, D]. No activation is applied by default, + so pass probabilities, or set sigmoid=True or softmax=True to pass raw logits. target: the shape should be[B, C, W, H, D]. Raises: diff --git a/nbs/03_vision_augment.ipynb b/nbs/03_vision_augment.ipynb index 77af42b..010cad1 100644 --- a/nbs/03_vision_augment.ipynb +++ b/nbs/03_vision_augment.ipynb @@ -44,7 +44,7 @@ "metadata": {}, "source": [ "# Data augmentation\n", - ">" + "> TorchIO and MONAI transform wrappers for 3D medical images, serializable inference transforms and GPU-batched patch augmentation." ] }, { @@ -68,7 +68,7 @@ " the image and the target (if it's a mask).\n", " \"\"\"\n", "\n", - " split_idx = 0 # Only perform transformations on training data. Use TTA() for transformations on validation data.\n", + " split_idx = 0 # Training only\n", "\n", " def __init__(self, aug):\n", " \"\"\"Constructs CustomDictTransform object.\n", @@ -145,7 +145,16 @@ "source": [ "#| export \n", "class PadOrCrop(DisplayedTransform):\n", - " \"\"\"Resize image using TorchIO CropOrPad.\"\"\"\n", + " \"\"\"Resize image using TorchIO CropOrPad.\n", + "\n", + " Args:\n", + " size: Target shape, an int or three ints.\n", + " padding_mode: Padding value or mode passed to tio.CropOrPad.\n", + " mask_name: Ignored by this transform, which crops or pads each image and\n", + " mask on its own around the volume centre. Still accepted and saved in\n", + " to_spec() so older saved specs load. Only the raw tio_transform\n", + " honours it, when called on a Subject that has that key.\n", + " \"\"\"\n", " \n", " order = 0\n", "\n", @@ -411,8 +420,6 @@ " 'RescaleIntensity': RescaleIntensity,\n", " 'NormalizeIntensity': NormalizeIntensity,\n", "}\n", - "# Backward-compatible private alias for code that inspected the old registry.\n", - "_NORM_TFM_REGISTRY = _TRANSFORM_SPEC_REGISTRY\n", "\n", "\n", "def _validate_transform_spec_json(spec):\n", @@ -1182,12 +1189,46 @@ "source": [ "#| export\n", "class OneOf(CustomDictTransform):\n", - " \"\"\"Apply only one of the given transforms using TorchIO `OneOf`.\"\"\"\n", + " \"\"\"Apply only one of the given transforms using TorchIO `OneOf`.\n", + "\n", + " transform_dict maps transforms to probabilities, or is a sequence of\n", + " transforms picked with equal probability. Transforms can be TorchIO\n", + " transforms or fastMONAI wrappers, which are unwrapped through their\n", + " tio_transform property.\n", + " \"\"\"\n", "\n", " def __init__(self, transform_dict, p=1):\n", + " unwrap = lambda t: getattr(t, 'tio_transform', t)\n", + " if isinstance(transform_dict, dict):\n", + " transform_dict = {unwrap(t): w for t, w in transform_dict.items()}\n", + " else:\n", + " transform_dict = [unwrap(t) for t in transform_dict]\n", " super().__init__(tio.OneOf(transform_dict, p=p))" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "oneof-wrapper-tests", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# OneOf accepts fastMONAI wrappers as well as raw TorchIO transforms\n", + "from fastcore.test import test_eq\n", + "\n", + "_x = MedImage(torch.randn(1, 8, 8, 8))\n", + "_m = MedMask(torch.zeros(1, 8, 8, 8)); _m[0, 2:5, 2:5, 2:5] = 1\n", + "_one_of = OneOf({RandomAffine(degrees=10, p=1): 0.5, RandomFlip(p=1): 0.5})\n", + "test_eq([type(t) for t in _one_of.tio_transform.transforms_dict],\n", + " [tio.RandomAffine, tio.RandomFlip])\n", + "_xo, _mo = _one_of.encodes((_x, _m))\n", + "test_eq((type(_xo), type(_mo)), (MedImage, MedMask))\n", + "_xo, _ = OneOf({tio.RandomAffine(p=1): 0.5, tio.RandomFlip(): 0.5}).encodes((_x, _m))\n", + "test_eq(type(_xo), MedImage)\n", + "test_eq(type(list(OneOf([RandomGamma(p=1)]).tio_transform.transforms_dict)[0]), tio.RandomGamma)\n" + ] + }, { "cell_type": "markdown", "id": "3sz0qzsue2e", @@ -1324,17 +1365,37 @@ " same parameters to both image and mask. Intensity transforms skip the mask.\n", "\n", " Each transform is controlled by a parameter dict with at minimum a 'p' key\n", - " for per-sample probability. Pass None to disable a transform.\n", + " for per-sample probability. Pass None to disable a transform. As in TorchIO,\n", + " index i of a per-axis value refers to array axis i of the patch (D, H, W),\n", + " and lengths are in voxels.\n", "\n", " Args:\n", - " affine: dict with keys 'scales', 'degrees', 'translation',\n", - " 'default_pad_value', 'p'. None to disable.\n", - " anisotropy: dict with keys 'axes', 'downsampling', 'p'. None to disable.\n", - " flip: dict with keys 'axes', 'p'. None to disable.\n", - " gamma: dict with keys 'log_gamma', 'p'. None to disable.\n", - " intensity_scale: dict with keys 'scale_range', 'p'. None to disable.\n", - " noise: dict with keys 'std', 'p'. None to disable.\n", - " blur: dict with keys 'std', 'p'. None to disable.\n", + " affine: dict with keys\n", + " 'scales': (lo, hi). Each axis gets a factor from U(lo, hi); values\n", + " below 1 shrink the content, as in TorchIO. A single number is not\n", + " accepted.\n", + " 'degrees': d or per-axis maxima (d0, d1, d2). The rotation about\n", + " axis i is drawn from U(-d_i, d_i).\n", + " 'translation': t or per-axis maxima (t0, t1, t2), in voxels. The\n", + " shift along axis i is drawn from U(-t_i, t_i).\n", + " 'default_pad_value': value for voxels moved in from outside the\n", + " patch (default 0).\n", + " 'p': probability.\n", + " TorchIO's range forms for degrees and translation, (a, b) or six\n", + " values, are not supported. None to disable.\n", + " anisotropy: dict with keys 'axes' (one is picked at random),\n", + " 'downsampling' as (lo, hi), 'p'. None to disable.\n", + " flip: dict with keys 'axes' and 'p'; each axis is flipped\n", + " independently with probability p. None to disable.\n", + " gamma: dict with keys 'log_gamma' as (lo, hi), 'p'. None to disable.\n", + " intensity_scale: dict with keys 'scale_range' as (lo, hi), 'p'.\n", + " None to disable.\n", + " noise: dict with keys 'std', 'p'. Zero-mean noise whose std is drawn\n", + " from U(0, s) for a number s, or from U(lo, hi) for (lo, hi).\n", + " None to disable.\n", + " blur: dict with keys 'std', 'p'. The sigma in voxels is drawn from\n", + " U(0, s) for a number s, or from U(lo, hi) for (lo, hi); a sample\n", + " whose sigma is 0 is left unchanged. None to disable.\n", "\n", " Example::\n", "\n", @@ -1353,14 +1414,27 @@ " Ada (full pipeline, 128^3): this batched GPU path ~3 ms/batch (B=2) vs\n", " MONAI per-sample-on-GPU ~77 ms (~25x slower) vs TorchIO CPU ~287 ms\n", " (~100x slower); the gap is structural and grows with batch size, so the\n", - " GPU path is kept. It is an approximate reimplementation (voxel-space\n", - " affine, per-axis flip, image-only anisotropy, sign-preserving gamma),\n", - " not bit-identical to the CPU/TorchIO path.\n", - " The voxel-space affine's geometric error scales with the patch's physical-extent\n", - " anisotropy (patch_size x target_spacing), not voxel spacing alone: it is sub-voxel\n", - " when the physical extent is near-isotropic (the common case, including cubic patches\n", - " on isotropic data) and deviates materially only for strongly anisotropic physical\n", - " extents (e.g. a cubic patch on anisotropic spacing).\n", + " GPU path is kept.\n", + "\n", + " Differences from the CPU/TorchIO path (suggest_patch_augmentations):\n", + " Each parameter means the same thing in both paths. The CPU patch path\n", + " hands patches to TorchIO with an identity affine, so TorchIO works in\n", + " voxels there too. The GPU path is not bit-identical:\n", + "\n", + " - Rotations act in affine_grid's normalised [-1, 1] patch coordinates.\n", + " A rotation between two axes of different patch length is stretched by\n", + " their length ratio, so it is rigid in voxel space only between axes of\n", + " equal length (e.g. a 256 x 256 plane). It is close to rigid in\n", + " physical space when patch_size x target_spacing is near isotropic.\n", + " TorchIO rotates rigidly in voxel space.\n", + " - Blur uses one isotropic sigma per sample, a kernel cut at 3 sigma and\n", + " replicate padding. TorchIO draws a sigma per axis and uses scipy's\n", + " gaussian_filter (cut at 4 sigma, reflect padding).\n", + " - Anisotropy follows the same recipe as TorchIO (nearest down, linear\n", + " up, image only), but resamples with F.interpolate on the voxel grid,\n", + " so sample positions and the downsampled size can differ slightly.\n", + " - Gamma uses one exponent per sample, while TorchIO draws one per\n", + " channel. This only matters for multi-channel input.\n", " \"\"\"\n", "\n", " def __init__(self, affine=None, anisotropy=None, flip=None,\n", @@ -1442,13 +1516,6 @@ " ], dim=1) # [n, 3]\n", " angles_rad = angles_deg * (math.pi / 180.0)\n", "\n", - " # Build rotation matrices [n, 3, 3]\n", - " R = _build_rotation_matrix_3d(angles_rad)\n", - "\n", - " # Scale matrix: S @ R -> [n, 3, 3]\n", - " S = torch.diag_embed(scales) # [n, 3, 3]\n", - " SR = S @ R # [n, 3, 3]\n", - "\n", " # Random translation (voxels -> normalized [-1, 1] coords)\n", " translation = cfg['translation']\n", " if not isinstance(translation, (list, tuple)):\n", @@ -1461,8 +1528,17 @@ " for i in range(3)\n", " ], dim=1) # [n, 3]\n", "\n", - " # Assemble theta for active samples\n", - " theta[idx, :3, :3] = SR\n", + " # The vectors above are indexed by array axis (D, H, W), while theta\n", + " # uses affine_grid's (x, y, z) = (W, H, D) order, so reverse them.\n", + " angles_rad = angles_rad.flip(1)\n", + " scales = scales.flip(1)\n", + " t_norm = t_norm.flip(1)\n", + "\n", + " # theta maps output to input coordinates, so 1/scales makes a scale\n", + " # below 1 shrink the content (as in TorchIO). Rotate, then scale along\n", + " # the output axes, matching TorchIO's composition.\n", + " R = _build_rotation_matrix_3d(angles_rad) # [n, 3, 3]\n", + " theta[idx, :3, :3] = R @ torch.diag_embed(1.0 / scales)\n", " theta[idx, :3, 3] = t_norm\n", "\n", " # Apply grid_sample\n", @@ -1484,11 +1560,9 @@ " def _apply_anisotropy(self, img, mask):\n", " \"\"\"Per-sample anisotropy simulation via F.interpolate.\n", "\n", - " Downsample along a random axis with nearest interpolation,\n", - " then upsample back with trilinear. Approximates TorchIO's\n", - " RandomAnisotropy and degrades the image only (TorchIO also\n", - " resamples the label map).\n", - " Only affects img, not mask (anisotropy is intensity degradation).\n", + " Downsample along a random axis with nearest interpolation, then\n", + " upsample back with trilinear. Only the image is changed, as in\n", + " TorchIO's RandomAnisotropy with scalars_only=True (the default).\n", " \"\"\"\n", " cfg = self.anisotropy\n", " B = img.shape[0]\n", @@ -1577,7 +1651,11 @@ " return img\n", "\n", " def _apply_noise(self, img):\n", - " \"\"\"Batched additive Gaussian noise with per-sample random std.\"\"\"\n", + " \"\"\"Batched zero-mean Gaussian noise with a per-sample std.\n", + "\n", + " A number s draws the std from U(0, s), as TorchIO RandomNoise does;\n", + " a pair (lo, hi) draws it from U(lo, hi).\n", + " \"\"\"\n", " cfg = self.noise\n", " B = img.shape[0]\n", " device = img.device\n", @@ -1586,9 +1664,9 @@ " std_val = cfg['std']\n", " if isinstance(std_val, (list, tuple)):\n", " std_lo, std_hi = std_val\n", - " per_std = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi)\n", " else:\n", - " per_std = torch.full((B,), std_val, device=device, dtype=dtype)\n", + " std_lo, std_hi = 0.0, std_val\n", + " per_std = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi)\n", "\n", " # Zero std for inactive samples\n", " active = torch.rand(B, device=device) < cfg['p']\n", @@ -1599,7 +1677,10 @@ " return img\n", "\n", " def _apply_blur(self, img):\n", - " \"\"\"Batched separable 3D Gaussian blur via F.conv3d with groups trick.\"\"\"\n", + " \"\"\"Batched separable 3D Gaussian blur via F.conv3d with groups trick.\n", + "\n", + " One isotropic sigma (in voxels) per sample.\n", + " \"\"\"\n", " cfg = self.blur\n", " B, C, D, H, W = img.shape\n", " device = img.device\n", @@ -1611,9 +1692,9 @@ " else:\n", " std_lo, std_hi = 0.0, std_val\n", "\n", - " # Per-sample sigma\n", + " # Per-sample sigma; samples drawn with sigma <= 0 are left unblurred\n", " sigma = torch.empty(B, device=device, dtype=dtype).uniform_(std_lo, std_hi)\n", - " active = torch.rand(B, device=device) < cfg['p']\n", + " active = (torch.rand(B, device=device) < cfg['p']) & (sigma > 0)\n", " if not active.any():\n", " return img\n", "\n", @@ -1625,7 +1706,7 @@ " # Build per-sample 1D Gaussian kernels [B, kernel_size]\n", " x = torch.arange(-kernel_radius, kernel_radius + 1,\n", " device=device, dtype=dtype)\n", - " # Avoid division by zero for sigma=0\n", + " # Avoid division by zero for inactive samples\n", " safe_sigma = torch.where(active, sigma, torch.ones_like(sigma))\n", " kernels = torch.exp(-x.unsqueeze(0)**2 / (2 * safe_sigma.unsqueeze(1)**2))\n", " kernels = kernels / kernels.sum(dim=1, keepdim=True)\n", @@ -1687,15 +1768,13 @@ " noise_p=0.1, blur_p=0.2, flip_p=0.5):\n", " \"\"\"Create GpuPatchAugmentation with nnU-Net-inspired defaults.\n", "\n", - " Approximate GPU reimplementation of suggest_patch_augmentations: it shares\n", - " the same parameter logic (via _compute_patch_aug_params) and probabilities,\n", - " but is NOT bit-identical: the affine is voxel-space (not spacing-aware),\n", - " flip is per-axis independent, anisotropy degrades the image only, and gamma\n", - " is sign-preserving. Runs batched on GPU (~25x faster than the per-sample\n", - " CPU/TorchIO path); see GpuPatchAugmentation for the benchmark and rationale.\n", - " The voxel-space affine's geometric error scales with physical-extent anisotropy\n", - " (patch_size x target_spacing), so it is sub-voxel for near-isotropic extents and only\n", - " matters for strongly anisotropic ones (e.g. a cubic patch on anisotropic spacing).\n", + " GPU counterpart of suggest_patch_augmentations: it shares the same parameter\n", + " logic (via _compute_patch_aug_params) and probabilities, and each value means\n", + " the same thing in both paths. It runs batched on GPU, much faster than the\n", + " per-sample CPU/TorchIO path, but is not bit-identical. See\n", + " GpuPatchAugmentation for the benchmark and the remaining differences\n", + " (rotation on non-cubic patches, isotropic blur sigma, anisotropy resampling\n", + " grid, per-sample gamma).\n", "\n", " Args:\n", " patch_size: List/tuple of 3 ints, patch dimensions.\n", @@ -1821,6 +1900,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "from fastcore.test import test_eq, test_fail\n", "\n", "# Isotropic case\n", @@ -1873,6 +1953,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Tests for _compute_patch_aug_params\n", "params = _compute_patch_aug_params([128, 128, 128], [1.0, 1.0, 1.0])\n", "test_eq(params['degrees'], (30, 30, 30))\n", @@ -1953,6 +2034,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Tests for gpu_patch_augmentations factory\n", "gpu_aug = gpu_patch_augmentations([128, 128, 128], [1.0, 1.0, 1.0])\n", "test_eq(type(gpu_aug), GpuPatchAugmentation)\n", @@ -2036,6 +2118,72 @@ "test_eq(torch.equal(_a, _b), True)" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "gpu-aug-parity-tests", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# GPU parameters follow TorchIO: index i acts on array axis i, scales < 1 shrink,\n", + "# noise std is drawn from U(0, s) and a zero blur sigma is a no-op (CPU-only)\n", + "torch.manual_seed(0)\n", + "_c = torch.arange(24).float() - 11.5\n", + "_grid = torch.meshgrid(_c, _c, _c, indexing='ij')\n", + "\n", + "def _rod(axis, r=3.0):\n", + " \"\"\"Round rod along `axis` through the centre of a 24^3 volume.\"\"\"\n", + " a, b = [_grid[k] for k in range(3) if k != axis]\n", + " return (a**2 + b**2 <= r**2).float()[None, None]\n", + "\n", + "def _outside_rod(out, axis):\n", + " return ((out * (1 - _rod(axis, r=4.5))).sum() / out.sum()).item()\n", + "\n", + "# degrees[i] rotates about axis i: the rod along axis i stays put, the others tilt\n", + "for i in range(3):\n", + " _deg = [0, 0, 0]; _deg[i] = 45\n", + " _rot = GpuPatchAugmentation(affine={'scales': (1, 1), 'degrees': tuple(_deg),\n", + " 'translation': 0, 'p': 1.0})\n", + " for k in range(3):\n", + " _moved = [_outside_rod(_rot(_rod(k), None)[0], k) for _ in range(4)]\n", + " if k == i: assert max(_moved) < 0.01, (i, k, _moved)\n", + " else: assert max(_moved) > 0.1, (i, k, _moved)\n", + "\n", + "# translation[i] moves the centre of mass along axis i only\n", + "_cube = ((_grid[0].abs() < 4) & (_grid[1].abs() < 4) & (_grid[2].abs() < 4)).float()[None, None]\n", + "def _com(v): return torch.stack([(g * v[0, 0]).sum() / v.sum() for g in _grid])\n", + "for i in range(3):\n", + " _tr = [0, 0, 0]; _tr[i] = 6\n", + " _shift = GpuPatchAugmentation(affine={'scales': (1, 1), 'degrees': 0,\n", + " 'translation': tuple(_tr), 'p': 1.0})\n", + " _d = torch.stack([_com(_shift(_cube, None)[0]).abs() for _ in range(4)])\n", + " assert _d[:, i].max() > 1, _d\n", + " assert _d[:, [k for k in range(3) if k != i]].max() < 0.01, _d\n", + "\n", + "# scales: same voxel count as TorchIO (0.5 shrinks an 8^3 cube to about 4^3)\n", + "_c8 = torch.zeros(1, 1, 32, 32, 32); _c8[0, 0, 12:20, 12:20, 12:20] = 1\n", + "for _s in (0.5, 1.5):\n", + " _gpu = GpuPatchAugmentation(affine={'scales': (_s, _s), 'degrees': 0, 'translation': 0, 'p': 1.0})\n", + " _n_gpu = int((_gpu(_c8, None)[0] > 0.5).sum())\n", + " _n_tio = int((tio.RandomAffine(scales=(_s, _s), degrees=0, translation=0, default_pad_value=0.)(\n", + " tio.ScalarImage(tensor=_c8[0])).data > 0.5).sum())\n", + " assert abs(_n_gpu - _n_tio) <= 0.1 * _n_tio, (_s, _n_gpu, _n_tio)\n", + " assert (_n_gpu < 512) == (_s < 1)\n", + "\n", + "# noise: a number s draws the per-sample std from U(0, s); (lo, hi) from U(lo, hi)\n", + "_z = torch.zeros(64, 1, 8, 8, 8)\n", + "_stds = GpuPatchAugmentation(noise={'std': 0.1, 'p': 1.0})(_z, None)[0].flatten(1).std(1)\n", + "assert _stds.min() < 0.03 and _stds.max() < 0.115 and 0.03 < _stds.mean() < 0.07, _stds\n", + "_stds = GpuPatchAugmentation(noise={'std': (0.05, 0.1), 'p': 1.0})(_z, None)[0].flatten(1).std(1)\n", + "assert _stds.min() > 0.04 and _stds.max() < 0.115, _stds\n", + "\n", + "# blur: sigma 0 leaves the sample unchanged (no NaN)\n", + "_img = torch.randn(2, 1, 8, 8, 8)\n", + "for _std in (0.0, (0.0, 0.0)):\n", + " test_eq(torch.equal(GpuPatchAugmentation(blur={'std': _std, 'p': 1.0})(_img, None)[0], _img), True)\n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -2043,6 +2191,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Test .tio_transform property\n", "# CustomDictTransform-based wrappers\n", "test_eq(type(RandomAffine(degrees=10).tio_transform), tio.RandomAffine)\n", @@ -2074,6 +2223,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Test RandomCutout (ItemTransform - expects tuple input)\n", "import numpy as np\n", "\n", diff --git a/nbs/04_vision_loss_functions.ipynb b/nbs/04_vision_loss_functions.ipynb index 8a1b519..e35476c 100644 --- a/nbs/04_vision_loss_functions.ipynb +++ b/nbs/04_vision_loss_functions.ipynb @@ -21,7 +21,7 @@ "import torch\n", "from fastMONAI.vision_core import *\n", "from fastMONAI.vision_data import pred_to_binary_mask, batch_pred_to_multiclass_mask\n", - "from monai.losses import TverskyLoss, FocalLoss\n", + "from monai.losses import TverskyLoss\n", "from torch.nn.modules.loss import _Loss" ] }, @@ -31,7 +31,7 @@ "metadata": {}, "source": [ "# Custom loss functions\n", - ">" + "> Segmentation loss wrappers, a focal Tversky loss and a DynUNet deep-supervision adapter." ] }, { @@ -48,7 +48,7 @@ "id": "5052c7bc-3d9a-4e34-8b64-bceaf2fdc7b6", "metadata": {}, "outputs": [], - "source": "#| export\nclass TverskyFocalLoss(_Loss):\n \"\"\"Focal Tversky loss: ``(1 - TI)**gamma`` (TI = Tversky index; see ``monai.losses.TverskyLoss``).\n ``gamma=1`` reduces to the plain Tversky loss.\n \"\"\"\n\n def __init__(\n self,\n include_background: bool = True,\n to_onehot_y: bool = False,\n sigmoid: bool = False,\n softmax: bool = False,\n gamma: float = 2,\n alpha: float = 0.5, \n beta: float = 0.99):\n \"\"\"\n Args:\n include_background: if to calculate loss for the background class.\n to_onehot_y: whether to convert `y` into one-hot format.\n sigmoid: if True, apply a sigmoid function to the prediction.\n softmax: if True, apply a softmax function to the prediction.\n gamma: focal exponent in ``(1 - TI)**gamma``.\n alpha: the weight of false positive in Tversky loss calculation.\n beta: the weight of false negative in Tversky loss calculation.\n \"\"\"\n \n super().__init__()\n self.tversky = TverskyLoss(\n to_onehot_y=to_onehot_y, \n include_background=include_background, \n sigmoid=sigmoid, \n softmax=softmax, \n alpha=alpha, \n beta=beta\n )\n self.gamma = gamma\n\n def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n input: the shape should be [B, C, W, H, D]. The input should be the original logits.\n target: the shape should be[B, C, W, H, D].\n\n Raises:\n ValueError: When number of dimensions for input and target are different.\n \"\"\"\n if len(input.shape) != len(target.shape):\n raise ValueError(\"The number of dimensions for input and target should be the same.\")\n\n # (1 - TI)**gamma; MONAI's TverskyLoss already returns tversky_loss = 1 - TI\n tversky_loss = self.tversky(input, target)\n total_loss: torch.Tensor = tversky_loss ** self.gamma\n\n return total_loss" + "source": "#| export\nclass TverskyFocalLoss(_Loss):\n \"\"\"Focal Tversky loss: ``(1 - TI)**gamma`` (TI = Tversky index; see ``monai.losses.TverskyLoss``).\n ``gamma=1`` reduces to the plain Tversky loss.\n gamma is applied after TverskyLoss's mean over batch and classes, i.e. ``mean(1 - TI)**gamma``.\n \"\"\"\n\n def __init__(\n self,\n include_background: bool = True,\n to_onehot_y: bool = False,\n sigmoid: bool = False,\n softmax: bool = False,\n gamma: float = 2,\n alpha: float = 0.5, \n beta: float = 0.99):\n \"\"\"\n Args:\n include_background: if to calculate loss for the background class.\n to_onehot_y: whether to convert `y` into one-hot format.\n sigmoid: if True, apply a sigmoid function to the prediction.\n softmax: if True, apply a softmax function to the prediction.\n gamma: focal exponent in ``(1 - TI)**gamma``.\n alpha: the weight of false positive in Tversky loss calculation.\n beta: the weight of false negative in Tversky loss calculation.\n \"\"\"\n \n super().__init__()\n self.tversky = TverskyLoss(\n to_onehot_y=to_onehot_y, \n include_background=include_background, \n sigmoid=sigmoid, \n softmax=softmax, \n alpha=alpha, \n beta=beta\n )\n self.gamma = gamma\n\n def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n input: the shape should be [B, C, W, H, D]. No activation is applied by default,\n so pass probabilities, or set sigmoid=True or softmax=True to pass raw logits.\n target: the shape should be[B, C, W, H, D].\n\n Raises:\n ValueError: When number of dimensions for input and target are different.\n \"\"\"\n if len(input.shape) != len(target.shape):\n raise ValueError(\"The number of dimensions for input and target should be the same.\")\n\n # (1 - TI)**gamma; MONAI's TverskyLoss already returns tversky_loss = 1 - TI\n tversky_loss = self.tversky(input, target)\n total_loss: torch.Tensor = tversky_loss ** self.gamma\n\n return total_loss" }, { "cell_type": "markdown", From d5a3d52a9e5a2148a3a10cec9eae48a31e7374fd Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:25:31 +0100 Subject: [PATCH 03/11] fix(vision_metrics)!: score one-channel output and reject non-binary masks AccumulatedDice took an argmax over channels, so a one-channel sigmoid model always scored 0. One channel now means sigmoid >= 0.5, and any other channel count that does not match n_classes raises. calculate_all_metrics raises on masks that are not binary instead of returning wrong numbers. preprocess_dataset now validates its own cache when inputs are symlinks whose names differ from their targets, and a metadata cache hit returns the caller's path. BREAKING CHANGE: removed the deprecated calculate_haus and MedDataset.calculate_target_size, and the always-True `binary` parameter of calculate_all_metrics and evaluate_segmentations. Also: https IXI URLs, simpler nodule dataframe builder, docstrings, page descriptions, dropped a citation that argues against grid-based metrics. --- fastMONAI/dataset_info.py | 79 ++++++++-------------- fastMONAI/external_data.py | 25 +++---- fastMONAI/vision_metrics.py | 114 +++++++++++++++++-------------- nbs/05_vision_metrics.ipynb | 128 ++++++++++++++++++++++++++--------- nbs/08_dataset_info.ipynb | 129 +++++++++++++++++++++--------------- nbs/09_external_data.ipynb | 35 +++------- 6 files changed, 281 insertions(+), 229 deletions(-) diff --git a/fastMONAI/dataset_info.py b/fastMONAI/dataset_info.py index 69a8604..a6fecb2 100644 --- a/fastMONAI/dataset_info.py +++ b/fastMONAI/dataset_info.py @@ -1,3 +1,5 @@ +"""Inspect dataset metadata, suggest preprocessing parameters and build a reusable preprocessing cache.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/08_dataset_info.ipynb. # %% auto #0 @@ -40,14 +42,17 @@ def __init__(self, dataframe=None, image_col:str=None, mask_col:str="mask_path", """Constructs MedDataset object. Args: - dataframe: DataFrame containing image paths. + dataframe: DataFrame containing file paths. Only the files in mask_col + are analysed; image_col is used only by visualize_cases(). image_col: Column name for image paths (used for visualization). - mask_col: Column name for mask/label paths when using dataframe mode. + mask_col: Column whose files are analysed in dataframe mode. path: Directory path containing images. img_list: List of image file paths to analyze. postfix: File postfix filter when using path mode. apply_reorder: Whether to reorder images to RAS+ orientation. - dtype: MedImage for images or MedMask for segmentation masks. + dtype: MedImage for images or MedMask for segmentation masks. It is not + inferred from mask_col, so pass dtype=MedMask in dataframe mode to read + the masks as label maps and get per-label volumes. max_workers: Number of parallel workers for processing. use_cache: Enable metadata caching. When True (default) and cache_path is not provided, auto-generates a cache path in ~/.cache/fastmonai/. @@ -170,8 +175,8 @@ def get_suggestion(self, include_patch_size: bool = False): data-derived. Access dataset.apply_reorder directly if needed. Args: - include_patch_size: If True, includes suggested patch_size for - patch-based training. Requires vision_patch module. + include_patch_size: If True, also includes the patch_size from + suggest_patch_size() for patch-based training. Returns: dict: {'target_spacing': [voxel_0, voxel_1, voxel_2]} @@ -260,7 +265,7 @@ def _process_with_cache(self, file_list): and entry.get('size') == stat.st_size and entry.get('apply_reorder') == self.apply_reorder and entry.get('dtype') == self.dtype.__name__): - cached_results.append((i, entry['info'])) + cached_results.append((i, {**entry['info'], 'path': fn})) continue except OSError: pass @@ -310,39 +315,6 @@ def _process_with_cache(self, file_list): return all_results - def calculate_target_size(self, target_spacing: list = None) -> list: - """Calculate the target image size for the dataset. - - .. deprecated:: - Use `get_size_statistics(target_spacing)['max']` instead for consistency - with other size statistics methods. - - Args: - target_spacing: If provided, calculates size after resampling to this spacing. - If None, returns original dimensions. - - Returns: - list: [dim_0, dim_1, dim_2] largest dimensions in dataset. - """ - warnings.warn( - "calculate_target_size() is deprecated. " - "Use get_size_statistics(target_spacing)['max'] instead.", - DeprecationWarning, - stacklevel=2 - ) - if target_spacing is not None: - org_voxels = self.df[["voxel_0", "voxel_1", 'voxel_2']].values - org_dims = self.df[["dim_0", "dim_1", 'dim_2']].values - - ratio = org_voxels/target_spacing - new_dims = (org_dims * ratio).T - # Use floor() to match TorchIO's Resample dimension calculation - dims = [float(np.floor(new_dims[0].max())), float(np.floor(new_dims[1].max())), float(np.floor(new_dims[2].max()))] - else: - dims = [float(self.df.dim_0.max()), float(self.df.dim_1.max()), float(self.df.dim_2.max())] - - return dims - def get_size_statistics(self, target_spacing: list = None) -> dict: """Calculate comprehensive size statistics for the dataset. @@ -488,7 +460,7 @@ def suggest_patch_size( 1. Use min(median, min_volume) per axis for safety 2. Round down to nearest multiple of divisor (16 for UNet compatibility) 3. Clamp to [min_patch_size, max_patch_size] bounds - 4. Validate: error if min_patch_size exceeds smallest volume + 4. Validate: error if min_patch_size or divisor exceeds the smallest volume Args: dataset: MedDataset instance with analyzed images. @@ -544,8 +516,9 @@ def round_to_divisor(val, div): if p > min_dim: raise ValueError( f"Cannot suggest safe patch_size for dimension {i}: " - f"smallest volume is {min_dim:.0f} but min_patch_size={min_patch_size[i]}. " - f"Exclude small volumes or reduce min_patch_size." + f"smallest volume is {min_dim:.0f} but the patch would be {p} " + f"(min_patch_size={min_patch_size[i]}, divisor={divisor}). " + f"Exclude small volumes, or reduce min_patch_size or divisor." ) # Warn when suggestion was reduced to fit smallest volume @@ -651,7 +624,13 @@ def _safe_cache_relative_path(value): return relative -def _read_valid_preprocessing_manifest(cache_dir, expected_version, contract, sources): +def _read_valid_preprocessing_manifest(cache_dir, expected_version, contract, sources, + expected_paths): + """Return the manifest if it matches and every expected output exists, else None. + + expected_paths holds the cache-relative output paths ('images/', 'masks/'), + named after the input paths as given (not their symlink targets). + """ manifest_path = cache_dir / 'preprocessing_manifest.json' try: manifest = json.loads(manifest_path.read_text(encoding='utf-8')) @@ -666,14 +645,6 @@ def _read_valid_preprocessing_manifest(cache_dir, expected_version, contract, so if not isinstance(outputs, list) or len(outputs) != len(sources): return None - expected_paths = { - str(Path('images') / Path(source['image']['path']).name) - for source in sources - } - expected_paths.update( - str(Path('masks') / Path(source['mask']['path']).name) - for source in sources if source['mask'] is not None - ) manifest_paths = [] for output in outputs: for role in ('image', 'mask'): @@ -761,6 +732,10 @@ def preprocess_dataset(df, img_col, mask_col=None, output_dir='preprocessed', duplicates = {name for name in mask_names if mask_names.count(name) > 1} raise ValueError(f'Duplicate mask file names: {duplicates}') + expected_outputs = {str(Path('images') / name) for name in img_names} + if mask_col is not None: + expected_outputs.update(str(Path('masks') / name) for name in mask_names) + contract = _preprocessing_contract( target_spacing, apply_reorder, transforms, mask_col is not None ) @@ -787,7 +762,7 @@ def preprocess_dataset(df, img_col, mask_col=None, output_dir='preprocessed', mask_dir = cache_dir / 'masks' if mask_col is not None else None existing = _read_valid_preprocessing_manifest( - cache_dir, cache_version, contract, sources + cache_dir, cache_version, contract, sources, expected_outputs ) if cache_dir.is_dir() else None if existing is not None and skip_existing: df[f'{img_col}_preprocessed'] = [ diff --git a/fastMONAI/external_data.py b/fastMONAI/external_data.py index 69e2ef9..5f60825 100644 --- a/fastMONAI/external_data.py +++ b/fastMONAI/external_data.py @@ -1,3 +1,5 @@ +"""Download and prepare the public datasets and example data used in the tutorials.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/09_external_data.ipynb. # %% auto #0 @@ -19,8 +21,8 @@ class MURLs(): """A class with external medical dataset URLs.""" - IXI_DATA = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar' - IXI_DEMOGRAPHIC_INFORMATION = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls' + IXI_DATA = 'https://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar' + IXI_DEMOGRAPHIC_INFORMATION = 'https://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls' CHENGWEN_CHU_SPINE_DATA = 'https://drive.google.com/uc?id=1rbm9-KKAexpNm2mC9FsSbfnS8VJaF3Kn&confirm=t' EXAMPLE_SPINE_DATA = 'https://drive.google.com/uc?id=1K1J2wmWxYUPY0NbylcV4Otv88EqCL4Ys' MEDMNIST_DICT = {'OrganMNIST3D': 'https://zenodo.org/record/6496656/files/organmnist3d.npz?download=1', @@ -216,25 +218,16 @@ def _process_medmnist_img(path, idx_arr): img.save(fn) return str(fn) -# %% ../nbs/09_external_data.ipynb #b0199679-707b-445d-8467-3d342246322a -def _df_sort_and_add_columns(df, label_list, is_val): - """Sort the dataframe based on img_idx and add labels and if it is validation data column.""" - - df = df.sort_values(by='img_idx').reset_index(drop=True) - df['labels'], df['is_val'] = label_list, is_val - df = df.drop('img_idx', axis=1) - - return df - # %% ../nbs/09_external_data.ipynb #d82157b8-ab69-4a38-9323-af4b58c6b54a def _create_nodule_df(pool, output_dir, imgs, labels, is_val=False): """Create dataframe for MedMNIST data.""" + # pool.map keeps the input order, so row i holds image i img_path_list = pool.map(partial(_process_medmnist_img, output_dir), enumerate(imgs)) - img_idx = [float(Path(fn).parts[-1].split('_')[0]) for fn in img_path_list] - - df = pd.DataFrame(list(zip(img_path_list, img_idx)), columns=['img_path','img_idx']) - return _df_sort_and_add_columns(df, labels, is_val) + df = pd.DataFrame({'img_path': img_path_list}) + df['labels'], df['is_val'] = labels, is_val + + return df # %% ../nbs/09_external_data.ipynb #04095111-a38f-4d00-a0fb-2dbfac263eb6 def _dataset_available_locally(path: Path): diff --git a/fastMONAI/vision_metrics.py b/fastMONAI/vision_metrics.py index 5e55645..8fe8ff6 100644 --- a/fastMONAI/vision_metrics.py +++ b/fastMONAI/vision_metrics.py @@ -1,12 +1,14 @@ +"""Segmentation metrics for training and validation, plus a spacing-aware panel for benchmark evaluation.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_vision_metrics.ipynb. # %% auto #0 -__all__ = ['calculate_dsc', 'calculate_hd95', 'calculate_haus', 'binary_dice_score', 'multi_dice_score', - 'calculate_confusion_metrics', 'binary_sensitivity', 'multi_sensitivity', 'binary_precision', - 'multi_precision', 'calculate_lesion_detection_rate', 'binary_lesion_detection_rate', - 'multi_lesion_detection_rate', 'calculate_signed_rve', 'binary_signed_rve', 'multi_signed_rve', - 'calculate_surface_metrics', 'calculate_all_metrics', 'evaluate_segmentations', 'AccumulatedDice', - 'AccumulatedMultiDice', 'EMACheckpoint'] +__all__ = ['calculate_dsc', 'calculate_hd95', 'binary_dice_score', 'multi_dice_score', 'calculate_confusion_metrics', + 'binary_sensitivity', 'multi_sensitivity', 'binary_precision', 'multi_precision', + 'calculate_lesion_detection_rate', 'binary_lesion_detection_rate', 'multi_lesion_detection_rate', + 'calculate_signed_rve', 'binary_signed_rve', 'multi_signed_rve', 'calculate_surface_metrics', + 'calculate_all_metrics', 'evaluate_segmentations', 'AccumulatedDice', 'AccumulatedMultiDice', + 'EMACheckpoint'] # %% ../nbs/05_vision_metrics.ipynb #8b6a83ac import warnings @@ -25,12 +27,14 @@ def calculate_dsc(pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: """Calculate Dice score using MONAI's compute_dice. Accepts tensors of various shapes and automatically reshapes to 5D format. + The channel dimension must be 1 (one binary mask). Multi-channel input such as + a one-hot mask is not supported; use `multi_dice_score` for multi-class output. Args: pred: Binary prediction tensor. Accepts: - [D, H, W] single 3D volume - - [C, D, H, W] single volume with channel - - [B, C, D, H, W] batched volumes + - [1, D, H, W] single volume with channel + - [B, 1, D, H, W] batched volumes targ: Binary target tensor (same shape options as pred). Returns: @@ -67,6 +71,11 @@ def calculate_hd95(pred: torch.Tensor, targ: torch.Tensor, spacing=None) -> torc together with ASSD and NSD (and empty-mask guards), use :func:`calculate_surface_metrics`. + .. note:: + HD95 is NaN when ``pred`` or ``targ`` is empty. A ``nanmean`` over cases therefore + leaves out every empty prediction (a missed lesion) instead of counting it as a + failure. :func:`calculate_surface_metrics` scores that case as ``inf``. + Returns: HD95 value(s): a single value for 3D/4D input, a tensor for a 5D batch. """ @@ -90,27 +99,14 @@ def calculate_hd95(pred: torch.Tensor, targ: torch.Tensor, spacing=None) -> torc warnings.filterwarnings("ignore", message=".*always_return_as_numpy.*") return torch.Tensor([compute_hausdorff_distance(p[None], t[None], percentile=95, spacing=spacing) for p, t in zip(pred, targ)]) - -def calculate_haus(pred: torch.Tensor, targ: torch.Tensor, spacing=None) -> torch.Tensor: - """Deprecated alias for :func:`calculate_hd95` (renamed after 0.8.x). - - Kept so existing calls keep working; forwards to :func:`calculate_hd95` and - will be removed in a future release. - """ - warnings.warn( - "calculate_haus() has been renamed to calculate_hd95() and will be removed " - "in a future release; update your calls to calculate_hd95().", - DeprecationWarning, stacklevel=2, - ) - return calculate_hd95(pred, targ, spacing=spacing) - # %% ../nbs/05_vision_metrics.ipynb #430d64f8-8cd9-4a88-ad20-5f73bebbf12f def binary_dice_score(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: """Calculates the mean Dice score for binary semantic segmentation tasks. Args: - act: Activation tensor with dimensions [B, C, W, H, D]. - targ: Target masks with dimensions [B, C, W, H, D]. + act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For + softmax output with two or more channels, use `multi_dice_score`. + targ: Target masks with dimensions [B, 1, W, H, D]. Returns: Mean Dice score (empty-target samples are ignored). @@ -167,8 +163,9 @@ def binary_sensitivity(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: that are correctly identified. Args: - act: Activation tensor [B, C, W, H, D]. - targ: Target masks [B, C, W, H, D]. + act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For + softmax output with two or more channels, use `multi_sensitivity`. + targ: Target masks [B, 1, W, H, D]. Returns: Mean sensitivity score. @@ -207,8 +204,9 @@ def binary_precision(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: that are actually correct. Args: - act: Activation tensor [B, C, W, H, D]. - targ: Target masks [B, C, W, H, D]. + act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For + softmax output with two or more channels, use `multi_precision`. + targ: Target masks [B, 1, W, H, D]. Returns: Mean precision score. @@ -299,8 +297,9 @@ def binary_lesion_detection_rate(act: torch.Tensor, targ: torch.Tensor, threshol """Calculate mean lesion detection rate for binary segmentation. Args: - act: Activation tensor [B, C, W, H, D]. - targ: Target masks [B, C, W, H, D]. + act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For + softmax output with two or more channels, use `multi_lesion_detection_rate`. + targ: Target masks [B, 1, W, H, D]. threshold: Minimum Dice score for a lesion to be considered detected. Default 0.0 means any overlap counts as detected. @@ -370,8 +369,9 @@ def binary_signed_rve(act: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: """Calculate mean signed RVE for binary segmentation. Args: - act: Activation tensor [B, C, W, H, D]. - targ: Target masks [B, C, W, H, D]. + act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For + softmax output with two or more channels, use `multi_signed_rve`. + targ: Target masks [B, 1, W, H, D]. Returns: Mean signed RVE. @@ -510,34 +510,34 @@ def calculate_surface_metrics(pred, targ, spacing_mm, "status": "ok", **base} # %% ../nbs/05_vision_metrics.ipynb #ed336bbf -def calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0, 2.0), - binary: bool = True) -> dict: +def calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0, 2.0)) -> dict: """Full per-case metric panel for one segmentation vs its ground truth. Orchestrates the individual ``calculate_*`` functions into one dict: Dice, sensitivity, precision, lesion detection rate, signed RVE, and the spacing-aware surface metrics (ASSD / HD95 / NSD in mm). For a whole dataset, use ``evaluate_segmentations``. + Only binary (single foreground class) segmentation is supported. + + Empty masks: ``dsc``, ``sensitivity``, ``ldr`` and ``rve`` are NaN when ``targ`` is + empty, and ``precision`` is NaN when ``pred`` is empty, so a ``nanmean`` over cases + leaves those cases out. The surface metrics use the guard values of + ``calculate_surface_metrics`` instead (both empty: 0 mm and NSD 1.0; one empty: inf and + NSD 0.0). Use ``surface_status`` to filter these cases. Args: - pred: predicted mask (binary / argmax label). torch.Tensor or np.ndarray, shaped - [W,H,D], [C,W,H,D] or [B,C,W,H,D] with singleton batch/channel dims. - targ: ground-truth mask, same shape convention, on the SAME voxel grid as ``pred``. + pred: binary predicted mask (bool, or values in {0, 1}). torch.Tensor or np.ndarray, + shaped [W,H,D], [C,W,H,D] or [B,C,W,H,D] with singleton batch/channel dims. + Other values (e.g. 0/255 or multi-label maps) raise ValueError. + targ: binary ground-truth mask, same conventions, on the SAME voxel grid as ``pred``. spacing_mm: voxel spacing in mm in array-axis (i,j,k) order matching ``targ``. Pass it explicitly (read from the GT file); a permuted spacing silently corrupts the surface distances. nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to ``calculate_surface_metrics``. - binary: single-foreground-class segmentation (the only mode supported). The confusion - metrics run on the single label channel -- MONAI keeps a size-1 channel when - ``include_background=False``, so a single-channel input is scored correctly. Returns: dict: ``dsc``, ``sensitivity``, ``precision``, ``ldr``, ``rve``, the ``nsd_tau*_mm`` values, ``assd_mm``, ``hd95_mm``, ``surface_status`` and ``spacing_mm``. """ - if not binary: - raise NotImplementedError("calculate_all_metrics supports binary (single foreground " - "class) segmentation only.") - pred_t = pred if isinstance(pred, torch.Tensor) else torch.as_tensor(np.asarray(pred)) targ_t = targ if isinstance(targ, torch.Tensor) else torch.as_tensor(np.asarray(targ)) while pred_t.ndim < 5: pred_t = pred_t.unsqueeze(0) # -> [B,C,W,H,D] @@ -545,6 +545,10 @@ def calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0 if pred_t.shape[-3:] != targ_t.shape[-3:]: raise ValueError("pred and targ must share a voxel grid; got spatial shapes " f"{tuple(pred_t.shape[-3:])} vs {tuple(targ_t.shape[-3:])}.") + for name, t in (("pred", pred_t), ("targ", targ_t)): + if t.dtype != torch.bool and not ((t == 0) | (t == 1)).all(): + raise ValueError(f"calculate_all_metrics expects binary masks (bool or values in " + f"{{0, 1}}); {name} contains {torch.unique(t)[:10].tolist()}.") pred_t, targ_t = pred_t.float(), targ_t.float() sm = calculate_surface_metrics(pred_t, targ_t, spacing_mm, nsd_tolerances_mm=nsd_tolerances_mm) @@ -562,7 +566,7 @@ def calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0 # %% ../nbs/05_vision_metrics.ipynb #2e526be3 def evaluate_segmentations(preds, gt_paths, case_ids=None, *, - nsd_tolerances_mm=(0.5, 1.0, 2.0), binary: bool = True): + nsd_tolerances_mm=(0.5, 1.0, 2.0)): """Score a batch of predictions against ground-truth masks into a per-case DataFrame. For each (prediction, ground-truth) pair, the GT mask AND its voxel spacing are read from a @@ -570,14 +574,13 @@ def evaluate_segmentations(preds, gt_paths, case_ids=None, *, separately-loaded file is a silent source of axis-permutation errors that corrupt the surface metrics. Predictions may be tensors (as returned by ``patch_inference``) or paths to saved masks, so predictions written with ``patch_inference(save_dir=...)`` can be re-scored without - re-running inference. + re-running inference. Masks must be binary (see ``calculate_all_metrics``). Args: preds: sequence of predicted masks -- torch.Tensor / np.ndarray, or paths to mask files. gt_paths: sequence of ground-truth mask file paths (one per prediction, same voxel grid). case_ids: optional row labels; default is each GT file's name stem. nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to the surface metrics. - binary: binary (single foreground class) segmentation only. Returns: pandas.DataFrame with one row per case: ``case_id`` plus the columns from @@ -596,7 +599,7 @@ def evaluate_segmentations(preds, gt_paths, case_ids=None, *, case_id = case_ids[i] if case_ids is not None else Path(str(gt_path)).name.split(".")[0] row = {"case_id": case_id} row.update(calculate_all_metrics(pred_data, gt.data, gt.spacing, - nsd_tolerances_mm=nsd_tolerances_mm, binary=binary)) + nsd_tolerances_mm=nsd_tolerances_mm)) rows.append(row) return pd.DataFrame(rows) @@ -609,6 +612,10 @@ class AccumulatedDice(Metric): batches, then computes Dice from the totals. This gives more weight to batches with more foreground voxels and is more statistically robust. + The model output may have ``n_classes`` channels (the class is the argmax) or, for + binary segmentation (``n_classes=2``), a single sigmoid channel where foreground is + ``sigmoid >= 0.5``. Any other channel count raises a ValueError. + Args: n_classes: Number of classes including background (default: 2 for binary). include_background: Whether to include background in metric (default: False). @@ -643,7 +650,16 @@ def accumulate(self, learn): pred = learn.pred # Model output [B, C, D, H, W] targ = learn.y # Target [B, 1, D, H, W] or [B, D, H, W] - pred_seg = pred.argmax(dim=1) # [B, D, H, W] + n_channels = pred.shape[1] + if n_channels == 1 and self.n_classes == 2: + pred_seg = (torch.sigmoid(pred[:, 0]) >= 0.5).long() # [B, D, H, W] + elif n_channels == self.n_classes: + pred_seg = pred.argmax(dim=1) # [B, D, H, W] + else: + raise ValueError( + f"{type(self).__name__}(n_classes={self.n_classes}) got a model output with " + f"{n_channels} channels; expected {self.n_classes} channels, or 1 sigmoid " + "channel for binary segmentation (n_classes=2).") if targ.ndim == pred_seg.ndim + 1: targ = targ.squeeze(1) diff --git a/nbs/05_vision_metrics.ipynb b/nbs/05_vision_metrics.ipynb index 2d39073..a1fc2f6 100644 --- a/nbs/05_vision_metrics.ipynb +++ b/nbs/05_vision_metrics.ipynb @@ -24,7 +24,7 @@ "metadata": {}, "source": [ "# Vision metrics\n", - ">" + "> Segmentation metrics for training and validation, plus a spacing-aware panel for benchmark evaluation." ] }, { @@ -39,12 +39,14 @@ " \"\"\"Calculate Dice score using MONAI's compute_dice.\n", "\n", " Accepts tensors of various shapes and automatically reshapes to 5D format.\n", + " The channel dimension must be 1 (one binary mask). Multi-channel input such as\n", + " a one-hot mask is not supported; use `multi_dice_score` for multi-class output.\n", "\n", " Args:\n", " pred: Binary prediction tensor. Accepts:\n", " - [D, H, W] single 3D volume\n", - " - [C, D, H, W] single volume with channel\n", - " - [B, C, D, H, W] batched volumes\n", + " - [1, D, H, W] single volume with channel\n", + " - [B, 1, D, H, W] batched volumes\n", " targ: Binary target tensor (same shape options as pred).\n", "\n", " Returns:\n", @@ -89,6 +91,11 @@ " together with ASSD and NSD (and empty-mask guards), use\n", " :func:`calculate_surface_metrics`.\n", "\n", + " .. note::\n", + " HD95 is NaN when ``pred`` or ``targ`` is empty. A ``nanmean`` over cases therefore\n", + " leaves out every empty prediction (a missed lesion) instead of counting it as a\n", + " failure. :func:`calculate_surface_metrics` scores that case as ``inf``.\n", + "\n", " Returns:\n", " HD95 value(s): a single value for 3D/4D input, a tensor for a 5D batch.\n", " \"\"\"\n", @@ -110,21 +117,7 @@ " # MONAI's get_mask_edges passes its own deprecated `always_return_as_numpy` arg\n", " # internally (MONAI 1.5); silence just that internal deprecation.\n", " warnings.filterwarnings(\"ignore\", message=\".*always_return_as_numpy.*\")\n", - " return torch.Tensor([compute_hausdorff_distance(p[None], t[None], percentile=95, spacing=spacing) for p, t in zip(pred, targ)])\n", - "\n", - "\n", - "def calculate_haus(pred: torch.Tensor, targ: torch.Tensor, spacing=None) -> torch.Tensor:\n", - " \"\"\"Deprecated alias for :func:`calculate_hd95` (renamed after 0.8.x).\n", - "\n", - " Kept so existing calls keep working; forwards to :func:`calculate_hd95` and\n", - " will be removed in a future release.\n", - " \"\"\"\n", - " warnings.warn(\n", - " \"calculate_haus() has been renamed to calculate_hd95() and will be removed \"\n", - " \"in a future release; update your calls to calculate_hd95().\",\n", - " DeprecationWarning, stacklevel=2,\n", - " )\n", - " return calculate_hd95(pred, targ, spacing=spacing)" + " return torch.Tensor([compute_hausdorff_distance(p[None], t[None], percentile=95, spacing=spacing) for p, t in zip(pred, targ)])" ] }, { @@ -139,8 +132,9 @@ " \"\"\"Calculates the mean Dice score for binary semantic segmentation tasks.\n", " \n", " Args:\n", - " act: Activation tensor with dimensions [B, C, W, H, D].\n", - " targ: Target masks with dimensions [B, C, W, H, D].\n", + " act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For\n", + " softmax output with two or more channels, use `multi_dice_score`.\n", + " targ: Target masks with dimensions [B, 1, W, H, D].\n", "\n", " Returns:\n", " Mean Dice score (empty-target samples are ignored).\n", @@ -229,8 +223,9 @@ " that are correctly identified.\n", " \n", " Args:\n", - " act: Activation tensor [B, C, W, H, D].\n", - " targ: Target masks [B, C, W, H, D].\n", + " act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For\n", + " softmax output with two or more channels, use `multi_sensitivity`.\n", + " targ: Target masks [B, 1, W, H, D].\n", " \n", " Returns:\n", " Mean sensitivity score.\n", @@ -285,8 +280,9 @@ " that are actually correct.\n", " \n", " Args:\n", - " act: Activation tensor [B, C, W, H, D].\n", - " targ: Target masks [B, C, W, H, D].\n", + " act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For\n", + " softmax output with two or more channels, use `multi_precision`.\n", + " targ: Target masks [B, 1, W, H, D].\n", " \n", " Returns:\n", " Mean precision score.\n", @@ -409,8 +405,9 @@ " \"\"\"Calculate mean lesion detection rate for binary segmentation.\n", "\n", " Args:\n", - " act: Activation tensor [B, C, W, H, D].\n", - " targ: Target masks [B, C, W, H, D].\n", + " act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For\n", + " softmax output with two or more channels, use `multi_lesion_detection_rate`.\n", + " targ: Target masks [B, 1, W, H, D].\n", " threshold: Minimum Dice score for a lesion to be considered detected.\n", " Default 0.0 means any overlap counts as detected.\n", "\n", @@ -512,8 +509,9 @@ " \"\"\"Calculate mean signed RVE for binary segmentation.\n", " \n", " Args:\n", - " act: Activation tensor [B, C, W, H, D].\n", - " targ: Target masks [B, C, W, H, D].\n", + " act: Single-channel (sigmoid) activation tensor [B, 1, W, H, D]. For\n", + " softmax output with two or more channels, use `multi_signed_rve`.\n", + " targ: Target masks [B, 1, W, H, D].\n", " \n", " Returns:\n", " Mean signed RVE.\n", @@ -560,7 +558,7 @@ "source": [ "## Surface Distance Benchmark Metrics (MONAI grid-based, mm)\n", "\n", - "Canonical, spacing-aware boundary metrics for final benchmark reporting (NOT the training loop): symmetric ASSD, HD95, and Normalized Surface Dice (NSD) at one or more tolerances, all in millimetres, computed with MONAI's grid-based surface metrics (`monai.metrics`). Unlike `calculate_hd95` (MONAI, voxel units by default), these report a single reference implementation with explicit voxel spacing and tolerance, for reproducible benchmarking aligned with `arxiv:2410.02630`." + "Canonical, spacing-aware boundary metrics for final benchmark reporting (NOT the training loop): symmetric ASSD, HD95, and Normalized Surface Dice (NSD) at one or more tolerances, all in millimetres, computed with MONAI's grid-based surface metrics (`monai.metrics`). Unlike `calculate_hd95` (MONAI, voxel units by default), these report a single reference implementation with explicit voxel spacing and tolerance, for reproducible benchmarking." ] }, { @@ -690,7 +688,7 @@ "id": "ed336bbf", "metadata": {}, "outputs": [], - "source": "#| export\ndef calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0, 2.0),\n binary: bool = True) -> dict:\n \"\"\"Full per-case metric panel for one segmentation vs its ground truth.\n\n Orchestrates the individual ``calculate_*`` functions into one dict: Dice, sensitivity,\n precision, lesion detection rate, signed RVE, and the spacing-aware surface metrics\n (ASSD / HD95 / NSD in mm). For a whole dataset, use ``evaluate_segmentations``.\n\n Args:\n pred: predicted mask (binary / argmax label). torch.Tensor or np.ndarray, shaped\n [W,H,D], [C,W,H,D] or [B,C,W,H,D] with singleton batch/channel dims.\n targ: ground-truth mask, same shape convention, on the SAME voxel grid as ``pred``.\n spacing_mm: voxel spacing in mm in array-axis (i,j,k) order matching ``targ``. Pass it\n explicitly (read from the GT file); a permuted spacing silently corrupts the\n surface distances.\n nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to ``calculate_surface_metrics``.\n binary: single-foreground-class segmentation (the only mode supported). The confusion\n metrics run on the single label channel -- MONAI keeps a size-1 channel when\n ``include_background=False``, so a single-channel input is scored correctly.\n\n Returns:\n dict: ``dsc``, ``sensitivity``, ``precision``, ``ldr``, ``rve``, the ``nsd_tau*_mm``\n values, ``assd_mm``, ``hd95_mm``, ``surface_status`` and ``spacing_mm``.\n \"\"\"\n if not binary:\n raise NotImplementedError(\"calculate_all_metrics supports binary (single foreground \"\n \"class) segmentation only.\")\n\n pred_t = pred if isinstance(pred, torch.Tensor) else torch.as_tensor(np.asarray(pred))\n targ_t = targ if isinstance(targ, torch.Tensor) else torch.as_tensor(np.asarray(targ))\n while pred_t.ndim < 5: pred_t = pred_t.unsqueeze(0) # -> [B,C,W,H,D]\n while targ_t.ndim < 5: targ_t = targ_t.unsqueeze(0)\n if pred_t.shape[-3:] != targ_t.shape[-3:]:\n raise ValueError(\"pred and targ must share a voxel grid; got spatial shapes \"\n f\"{tuple(pred_t.shape[-3:])} vs {tuple(targ_t.shape[-3:])}.\")\n pred_t, targ_t = pred_t.float(), targ_t.float()\n\n sm = calculate_surface_metrics(pred_t, targ_t, spacing_mm, nsd_tolerances_mm=nsd_tolerances_mm)\n row = {\n \"dsc\": calculate_dsc(pred_t, targ_t).nanmean().item(),\n \"sensitivity\": calculate_confusion_metrics(pred_t, targ_t, \"sensitivity\").nanmean().item(),\n \"precision\": calculate_confusion_metrics(pred_t, targ_t, \"precision\").nanmean().item(),\n \"ldr\": calculate_lesion_detection_rate(pred_t, targ_t).nanmean().item(),\n \"rve\": calculate_signed_rve(pred_t, targ_t).nanmean().item(),\n }\n row.update({k: sm[k] for k in sm if k.startswith(\"nsd_tau\")})\n row[\"assd_mm\"], row[\"hd95_mm\"] = sm[\"assd_mm\"], sm[\"hd95_mm\"]\n row[\"surface_status\"], row[\"spacing_mm\"] = sm[\"status\"], sm[\"spacing_mm\"]\n return row" + "source": "#| export\ndef calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0, 2.0)) -> dict:\n \"\"\"Full per-case metric panel for one segmentation vs its ground truth.\n\n Orchestrates the individual ``calculate_*`` functions into one dict: Dice, sensitivity,\n precision, lesion detection rate, signed RVE, and the spacing-aware surface metrics\n (ASSD / HD95 / NSD in mm). For a whole dataset, use ``evaluate_segmentations``.\n Only binary (single foreground class) segmentation is supported.\n\n Empty masks: ``dsc``, ``sensitivity``, ``ldr`` and ``rve`` are NaN when ``targ`` is\n empty, and ``precision`` is NaN when ``pred`` is empty, so a ``nanmean`` over cases\n leaves those cases out. The surface metrics use the guard values of\n ``calculate_surface_metrics`` instead (both empty: 0 mm and NSD 1.0; one empty: inf and\n NSD 0.0). Use ``surface_status`` to filter these cases.\n\n Args:\n pred: binary predicted mask (bool, or values in {0, 1}). torch.Tensor or np.ndarray,\n shaped [W,H,D], [C,W,H,D] or [B,C,W,H,D] with singleton batch/channel dims.\n Other values (e.g. 0/255 or multi-label maps) raise ValueError.\n targ: binary ground-truth mask, same conventions, on the SAME voxel grid as ``pred``.\n spacing_mm: voxel spacing in mm in array-axis (i,j,k) order matching ``targ``. Pass it\n explicitly (read from the GT file); a permuted spacing silently corrupts the\n surface distances.\n nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to ``calculate_surface_metrics``.\n\n Returns:\n dict: ``dsc``, ``sensitivity``, ``precision``, ``ldr``, ``rve``, the ``nsd_tau*_mm``\n values, ``assd_mm``, ``hd95_mm``, ``surface_status`` and ``spacing_mm``.\n \"\"\"\n pred_t = pred if isinstance(pred, torch.Tensor) else torch.as_tensor(np.asarray(pred))\n targ_t = targ if isinstance(targ, torch.Tensor) else torch.as_tensor(np.asarray(targ))\n while pred_t.ndim < 5: pred_t = pred_t.unsqueeze(0) # -> [B,C,W,H,D]\n while targ_t.ndim < 5: targ_t = targ_t.unsqueeze(0)\n if pred_t.shape[-3:] != targ_t.shape[-3:]:\n raise ValueError(\"pred and targ must share a voxel grid; got spatial shapes \"\n f\"{tuple(pred_t.shape[-3:])} vs {tuple(targ_t.shape[-3:])}.\")\n for name, t in ((\"pred\", pred_t), (\"targ\", targ_t)):\n if t.dtype != torch.bool and not ((t == 0) | (t == 1)).all():\n raise ValueError(f\"calculate_all_metrics expects binary masks (bool or values in \"\n f\"{{0, 1}}); {name} contains {torch.unique(t)[:10].tolist()}.\")\n pred_t, targ_t = pred_t.float(), targ_t.float()\n\n sm = calculate_surface_metrics(pred_t, targ_t, spacing_mm, nsd_tolerances_mm=nsd_tolerances_mm)\n row = {\n \"dsc\": calculate_dsc(pred_t, targ_t).nanmean().item(),\n \"sensitivity\": calculate_confusion_metrics(pred_t, targ_t, \"sensitivity\").nanmean().item(),\n \"precision\": calculate_confusion_metrics(pred_t, targ_t, \"precision\").nanmean().item(),\n \"ldr\": calculate_lesion_detection_rate(pred_t, targ_t).nanmean().item(),\n \"rve\": calculate_signed_rve(pred_t, targ_t).nanmean().item(),\n }\n row.update({k: sm[k] for k in sm if k.startswith(\"nsd_tau\")})\n row[\"assd_mm\"], row[\"hd95_mm\"] = sm[\"assd_mm\"], sm[\"hd95_mm\"]\n row[\"surface_status\"], row[\"spacing_mm\"] = sm[\"status\"], sm[\"spacing_mm\"]\n return row" }, { "cell_type": "code", @@ -698,7 +696,7 @@ "id": "2e526be3", "metadata": {}, "outputs": [], - "source": "#| export\ndef evaluate_segmentations(preds, gt_paths, case_ids=None, *,\n nsd_tolerances_mm=(0.5, 1.0, 2.0), binary: bool = True):\n \"\"\"Score a batch of predictions against ground-truth masks into a per-case DataFrame.\n\n For each (prediction, ground-truth) pair, the GT mask AND its voxel spacing are read from a\n SINGLE object, so the array and its spacing can never disagree -- reading spacing from a\n separately-loaded file is a silent source of axis-permutation errors that corrupt the surface\n metrics. Predictions may be tensors (as returned by ``patch_inference``) or paths to saved\n masks, so predictions written with ``patch_inference(save_dir=...)`` can be re-scored without\n re-running inference.\n\n Args:\n preds: sequence of predicted masks -- torch.Tensor / np.ndarray, or paths to mask files.\n gt_paths: sequence of ground-truth mask file paths (one per prediction, same voxel grid).\n case_ids: optional row labels; default is each GT file's name stem.\n nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to the surface metrics.\n binary: binary (single foreground class) segmentation only.\n\n Returns:\n pandas.DataFrame with one row per case: ``case_id`` plus the columns from\n ``calculate_all_metrics``.\n \"\"\"\n import torchio as tio # lazy: only this benchmark helper needs torchio I/O\n preds, gt_paths = list(preds), list(gt_paths)\n if len(preds) != len(gt_paths):\n raise ValueError(f\"preds and gt_paths must have equal length; got {len(preds)} and \"\n f\"{len(gt_paths)}.\")\n\n rows = []\n for i, (pred, gt_path) in enumerate(zip(preds, gt_paths)):\n gt = tio.LabelMap(gt_path) # single source: .data and .spacing agree\n pred_data = tio.LabelMap(pred).data if isinstance(pred, (str, Path)) else pred\n case_id = case_ids[i] if case_ids is not None else Path(str(gt_path)).name.split(\".\")[0]\n row = {\"case_id\": case_id}\n row.update(calculate_all_metrics(pred_data, gt.data, gt.spacing,\n nsd_tolerances_mm=nsd_tolerances_mm, binary=binary))\n rows.append(row)\n return pd.DataFrame(rows)" + "source": "#| export\ndef evaluate_segmentations(preds, gt_paths, case_ids=None, *,\n nsd_tolerances_mm=(0.5, 1.0, 2.0)):\n \"\"\"Score a batch of predictions against ground-truth masks into a per-case DataFrame.\n\n For each (prediction, ground-truth) pair, the GT mask AND its voxel spacing are read from a\n SINGLE object, so the array and its spacing can never disagree -- reading spacing from a\n separately-loaded file is a silent source of axis-permutation errors that corrupt the surface\n metrics. Predictions may be tensors (as returned by ``patch_inference``) or paths to saved\n masks, so predictions written with ``patch_inference(save_dir=...)`` can be re-scored without\n re-running inference. Masks must be binary (see ``calculate_all_metrics``).\n\n Args:\n preds: sequence of predicted masks -- torch.Tensor / np.ndarray, or paths to mask files.\n gt_paths: sequence of ground-truth mask file paths (one per prediction, same voxel grid).\n case_ids: optional row labels; default is each GT file's name stem.\n nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to the surface metrics.\n\n Returns:\n pandas.DataFrame with one row per case: ``case_id`` plus the columns from\n ``calculate_all_metrics``.\n \"\"\"\n import torchio as tio # lazy: only this benchmark helper needs torchio I/O\n preds, gt_paths = list(preds), list(gt_paths)\n if len(preds) != len(gt_paths):\n raise ValueError(f\"preds and gt_paths must have equal length; got {len(preds)} and \"\n f\"{len(gt_paths)}.\")\n\n rows = []\n for i, (pred, gt_path) in enumerate(zip(preds, gt_paths)):\n gt = tio.LabelMap(gt_path) # single source: .data and .spacing agree\n pred_data = tio.LabelMap(pred).data if isinstance(pred, (str, Path)) else pred\n case_id = case_ids[i] if case_ids is not None else Path(str(gt_path)).name.split(\".\")[0]\n row = {\"case_id\": case_id}\n row.update(calculate_all_metrics(pred_data, gt.data, gt.spacing,\n nsd_tolerances_mm=nsd_tolerances_mm))\n rows.append(row)\n return pd.DataFrame(rows)" }, { "cell_type": "code", @@ -706,7 +704,30 @@ "id": "2fcc148b", "metadata": {}, "outputs": [], - "source": "#| hide\nimport tempfile\nfrom pathlib import Path as _Path\nimport torchio as tio\nfrom fastcore.test import test_eq, test_close, test_fail\n\n# calculate_all_metrics: a perfect prediction scores DSC ~1 and ASSD 0.\n_cube = np.zeros((16, 16, 16), np.float32); _cube[4:12, 4:12, 4:12] = 1.0\n_m = calculate_all_metrics(_cube, _cube, spacing_mm=(1.0, 1.0, 1.0))\ntest_close(_m[\"dsc\"], 1.0, eps=1e-4)\ntest_close(_m[\"assd_mm\"], 0.0, eps=1e-6)\ntest_eq(_m[\"sensitivity\"] > 0.99, True)\ntest_eq(_m[\"precision\"] > 0.99, True)\n\n# guards: multiclass not supported; mismatched grids raise loudly.\ntest_fail(lambda: calculate_all_metrics(_cube, _cube, (1, 1, 1), binary=False))\ntest_fail(lambda: calculate_all_metrics(_cube, np.zeros((8, 8, 8), np.float32), (1, 1, 1)))\n\n# evaluate_segmentations: pred as tensor AND as path; spacing read from the GT object.\nwith tempfile.TemporaryDirectory() as _d:\n _gt = _Path(_d) / \"case_seg.nii.gz\"\n tio.LabelMap(tensor=torch.from_numpy(_cube)[None], affine=np.diag([1., 1., 2., 1.])).save(_gt)\n _df = evaluate_segmentations([torch.from_numpy(_cube)[None], str(_gt)], [str(_gt), str(_gt)])\n test_eq(len(_df), 2)\n test_eq(\"case_id\" in _df.columns, True)\n test_close(_df[\"dsc\"].iloc[0], 1.0, eps=1e-4)\n test_eq(tuple(_df[\"spacing_mm\"].iloc[0]), (1.0, 1.0, 2.0)) # anisotropic spacing from GT" + "source": "#| hide\nimport tempfile\nfrom pathlib import Path as _Path\nimport torchio as tio\nfrom fastcore.test import test_eq, test_close, test_fail\n\n# calculate_all_metrics: a perfect prediction scores DSC ~1 and ASSD 0.\n_cube = np.zeros((16, 16, 16), np.float32); _cube[4:12, 4:12, 4:12] = 1.0\n_m = calculate_all_metrics(_cube, _cube, spacing_mm=(1.0, 1.0, 1.0))\ntest_close(_m[\"dsc\"], 1.0, eps=1e-4)\ntest_close(_m[\"assd_mm\"], 0.0, eps=1e-6)\ntest_eq(_m[\"sensitivity\"] > 0.99, True)\ntest_eq(_m[\"precision\"] > 0.99, True)\n\n# guard: mismatched grids raise loudly.\ntest_fail(lambda: calculate_all_metrics(_cube, np.zeros((8, 8, 8), np.float32), (1, 1, 1)))\n\n# evaluate_segmentations: pred as tensor AND as path; spacing read from the GT object.\nwith tempfile.TemporaryDirectory() as _d:\n _gt = _Path(_d) / \"case_seg.nii.gz\"\n tio.LabelMap(tensor=torch.from_numpy(_cube)[None], affine=np.diag([1., 1., 2., 1.])).save(_gt)\n _df = evaluate_segmentations([torch.from_numpy(_cube)[None], str(_gt)], [str(_gt), str(_gt)])\n test_eq(len(_df), 2)\n test_eq(\"case_id\" in _df.columns, True)\n test_close(_df[\"dsc\"].iloc[0], 1.0, eps=1e-4)\n test_eq(tuple(_df[\"spacing_mm\"].iloc[0]), (1.0, 1.0, 2.0)) # anisotropic spacing from GT" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "metricpanelbinaryinputtest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# calculate_all_metrics and evaluate_segmentations reject non-binary masks.\n", + "test_fail(lambda: calculate_all_metrics(_cube * 255, _cube * 255, (1, 1, 1)), contains=\"binary masks\")\n", + "test_fail(lambda: calculate_all_metrics(_cube, _cube * 255, (1, 1, 1)), contains=\"targ\")\n", + "_multi = _cube.copy(); _multi[4:8] = 2.0 # multi-label map\n", + "test_fail(lambda: calculate_all_metrics(_multi, _cube, (1, 1, 1)), contains=\"pred\")\n", + "_mb = calculate_all_metrics(_cube.astype(bool), torch.from_numpy(_cube).to(torch.uint8), (1, 1, 1))\n", + "test_close(_mb[\"dsc\"], 1.0, eps=1e-4) # bool and uint8 0/1 are fine\n", + "\n", + "with tempfile.TemporaryDirectory() as _d:\n", + " _gt255 = _Path(_d) / \"case255_seg.nii.gz\"\n", + " tio.LabelMap(tensor=torch.from_numpy(_cube * 255).to(torch.uint8)[None]).save(_gt255)\n", + " test_fail(lambda: evaluate_segmentations([torch.from_numpy(_cube)[None]], [str(_gt255)]),\n", + " contains=\"binary masks\")" + ] }, { "cell_type": "markdown", @@ -734,6 +755,10 @@ " batches, then computes Dice from the totals. This gives more weight to\n", " batches with more foreground voxels and is more statistically robust.\n", "\n", + " The model output may have ``n_classes`` channels (the class is the argmax) or, for\n", + " binary segmentation (``n_classes=2``), a single sigmoid channel where foreground is\n", + " ``sigmoid >= 0.5``. Any other channel count raises a ValueError.\n", + "\n", " Args:\n", " n_classes: Number of classes including background (default: 2 for binary).\n", " include_background: Whether to include background in metric (default: False).\n", @@ -768,7 +793,16 @@ " pred = learn.pred # Model output [B, C, D, H, W]\n", " targ = learn.y # Target [B, 1, D, H, W] or [B, D, H, W]\n", "\n", - " pred_seg = pred.argmax(dim=1) # [B, D, H, W]\n", + " n_channels = pred.shape[1]\n", + " if n_channels == 1 and self.n_classes == 2:\n", + " pred_seg = (torch.sigmoid(pred[:, 0]) >= 0.5).long() # [B, D, H, W]\n", + " elif n_channels == self.n_classes:\n", + " pred_seg = pred.argmax(dim=1) # [B, D, H, W]\n", + " else:\n", + " raise ValueError(\n", + " f\"{type(self).__name__}(n_classes={self.n_classes}) got a model output with \"\n", + " f\"{n_channels} channels; expected {self.n_classes} channels, or 1 sigmoid \"\n", + " \"channel for binary segmentation (n_classes=2).\")\n", "\n", " if targ.ndim == pred_seg.ndim + 1:\n", " targ = targ.squeeze(1)\n", @@ -1039,6 +1073,34 @@ "print(\"AccumulatedMultiDice test passed!\")" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "accumulateddicesigmoidtest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# AccumulatedDice also scores a single sigmoid channel and rejects other channel counts.\n", + "from unittest.mock import MagicMock\n", + "from fastcore.test import test_close, test_fail\n", + "\n", + "_learn = MagicMock()\n", + "_targ1 = torch.zeros((1, 1, 8, 8, 8)); _targ1[0, 0, :4] = 1\n", + "_logits1 = torch.full((1, 1, 8, 8, 8), -10.0); _logits1[0, 0, :4] = 10.0 # perfect 1-channel logits\n", + "_learn.pred, _learn.y = _logits1, _targ1\n", + "\n", + "_m = AccumulatedDice(n_classes=2); _m.reset(); _m.accumulate(_learn)\n", + "test_close(_m.value, 1.0, eps=1e-6)\n", + "_mm = AccumulatedMultiDice(n_classes=2); _mm.reset(); _mm.accumulate(_learn)\n", + "test_close(float(_mm.value[0]), 1.0, eps=1e-6)\n", + "\n", + "_m3 = AccumulatedDice(n_classes=3); _m3.reset()\n", + "test_fail(lambda: _m3.accumulate(_learn), contains=\"channels\") # 1 channel needs n_classes=2\n", + "_learn.pred = torch.zeros((1, 3, 8, 8, 8))\n", + "test_fail(lambda: _m.accumulate(_learn), contains=\"channels\") # 3 channels vs n_classes=2" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/nbs/08_dataset_info.ipynb b/nbs/08_dataset_info.ipynb index 8a611b5..37f052d 100644 --- a/nbs/08_dataset_info.ipynb +++ b/nbs/08_dataset_info.ipynb @@ -73,14 +73,17 @@ " \"\"\"Constructs MedDataset object.\n", "\n", " Args:\n", - " dataframe: DataFrame containing image paths.\n", + " dataframe: DataFrame containing file paths. Only the files in mask_col\n", + " are analysed; image_col is used only by visualize_cases().\n", " image_col: Column name for image paths (used for visualization).\n", - " mask_col: Column name for mask/label paths when using dataframe mode.\n", + " mask_col: Column whose files are analysed in dataframe mode.\n", " path: Directory path containing images.\n", " img_list: List of image file paths to analyze.\n", " postfix: File postfix filter when using path mode.\n", " apply_reorder: Whether to reorder images to RAS+ orientation.\n", - " dtype: MedImage for images or MedMask for segmentation masks.\n", + " dtype: MedImage for images or MedMask for segmentation masks. It is not\n", + " inferred from mask_col, so pass dtype=MedMask in dataframe mode to read\n", + " the masks as label maps and get per-label volumes.\n", " max_workers: Number of parallel workers for processing.\n", " use_cache: Enable metadata caching. When True (default) and cache_path\n", " is not provided, auto-generates a cache path in ~/.cache/fastmonai/.\n", @@ -203,8 +206,8 @@ " data-derived. Access dataset.apply_reorder directly if needed.\n", "\n", " Args:\n", - " include_patch_size: If True, includes suggested patch_size for\n", - " patch-based training. Requires vision_patch module.\n", + " include_patch_size: If True, also includes the patch_size from\n", + " suggest_patch_size() for patch-based training.\n", "\n", " Returns:\n", " dict: {'target_spacing': [voxel_0, voxel_1, voxel_2]}\n", @@ -293,7 +296,7 @@ " and entry.get('size') == stat.st_size\n", " and entry.get('apply_reorder') == self.apply_reorder\n", " and entry.get('dtype') == self.dtype.__name__):\n", - " cached_results.append((i, entry['info']))\n", + " cached_results.append((i, {**entry['info'], 'path': fn}))\n", " continue\n", " except OSError:\n", " pass\n", @@ -343,39 +346,6 @@ "\n", " return all_results\n", "\n", - " def calculate_target_size(self, target_spacing: list = None) -> list:\n", - " \"\"\"Calculate the target image size for the dataset.\n", - "\n", - " .. deprecated::\n", - " Use `get_size_statistics(target_spacing)['max']` instead for consistency\n", - " with other size statistics methods.\n", - "\n", - " Args:\n", - " target_spacing: If provided, calculates size after resampling to this spacing.\n", - " If None, returns original dimensions.\n", - "\n", - " Returns:\n", - " list: [dim_0, dim_1, dim_2] largest dimensions in dataset.\n", - " \"\"\"\n", - " warnings.warn(\n", - " \"calculate_target_size() is deprecated. \"\n", - " \"Use get_size_statistics(target_spacing)['max'] instead.\",\n", - " DeprecationWarning,\n", - " stacklevel=2\n", - " )\n", - " if target_spacing is not None:\n", - " org_voxels = self.df[[\"voxel_0\", \"voxel_1\", 'voxel_2']].values\n", - " org_dims = self.df[[\"dim_0\", \"dim_1\", 'dim_2']].values\n", - "\n", - " ratio = org_voxels/target_spacing\n", - " new_dims = (org_dims * ratio).T\n", - " # Use floor() to match TorchIO's Resample dimension calculation\n", - " dims = [float(np.floor(new_dims[0].max())), float(np.floor(new_dims[1].max())), float(np.floor(new_dims[2].max()))]\n", - " else:\n", - " dims = [float(self.df.dim_0.max()), float(self.df.dim_1.max()), float(self.df.dim_2.max())]\n", - "\n", - " return dims\n", - "\n", " def get_size_statistics(self, target_spacing: list = None) -> dict:\n", " \"\"\"Calculate comprehensive size statistics for the dataset.\n", "\n", @@ -529,7 +499,7 @@ " 1. Use min(median, min_volume) per axis for safety\n", " 2. Round down to nearest multiple of divisor (16 for UNet compatibility)\n", " 3. Clamp to [min_patch_size, max_patch_size] bounds\n", - " 4. Validate: error if min_patch_size exceeds smallest volume\n", + " 4. Validate: error if min_patch_size or divisor exceeds the smallest volume\n", "\n", " Args:\n", " dataset: MedDataset instance with analyzed images.\n", @@ -585,8 +555,9 @@ " if p > min_dim:\n", " raise ValueError(\n", " f\"Cannot suggest safe patch_size for dimension {i}: \"\n", - " f\"smallest volume is {min_dim:.0f} but min_patch_size={min_patch_size[i]}. \"\n", - " f\"Exclude small volumes or reduce min_patch_size.\"\n", + " f\"smallest volume is {min_dim:.0f} but the patch would be {p} \"\n", + " f\"(min_patch_size={min_patch_size[i]}, divisor={divisor}). \"\n", + " f\"Exclude small volumes, or reduce min_patch_size or divisor.\"\n", " )\n", "\n", " # Warn when suggestion was reduced to fit smallest volume\n", @@ -700,7 +671,13 @@ " return relative\n", "\n", "\n", - "def _read_valid_preprocessing_manifest(cache_dir, expected_version, contract, sources):\n", + "def _read_valid_preprocessing_manifest(cache_dir, expected_version, contract, sources,\n", + " expected_paths):\n", + " \"\"\"Return the manifest if it matches and every expected output exists, else None.\n", + "\n", + " expected_paths holds the cache-relative output paths ('images/', 'masks/'),\n", + " named after the input paths as given (not their symlink targets).\n", + " \"\"\"\n", " manifest_path = cache_dir / 'preprocessing_manifest.json'\n", " try:\n", " manifest = json.loads(manifest_path.read_text(encoding='utf-8'))\n", @@ -715,14 +692,6 @@ " if not isinstance(outputs, list) or len(outputs) != len(sources):\n", " return None\n", "\n", - " expected_paths = {\n", - " str(Path('images') / Path(source['image']['path']).name)\n", - " for source in sources\n", - " }\n", - " expected_paths.update(\n", - " str(Path('masks') / Path(source['mask']['path']).name)\n", - " for source in sources if source['mask'] is not None\n", - " )\n", " manifest_paths = []\n", " for output in outputs:\n", " for role in ('image', 'mask'):\n", @@ -810,6 +779,10 @@ " duplicates = {name for name in mask_names if mask_names.count(name) > 1}\n", " raise ValueError(f'Duplicate mask file names: {duplicates}')\n", "\n", + " expected_outputs = {str(Path('images') / name) for name in img_names}\n", + " if mask_col is not None:\n", + " expected_outputs.update(str(Path('masks') / name) for name in mask_names)\n", + "\n", " contract = _preprocessing_contract(\n", " target_spacing, apply_reorder, transforms, mask_col is not None\n", " )\n", @@ -836,7 +809,7 @@ " mask_dir = cache_dir / 'masks' if mask_col is not None else None\n", "\n", " existing = _read_valid_preprocessing_manifest(\n", - " cache_dir, cache_version, contract, sources\n", + " cache_dir, cache_version, contract, sources, expected_outputs\n", " ) if cache_dir.is_dir() else None\n", " if existing is not None and skip_existing:\n", " df[f'{img_col}_preprocessed'] = [\n", @@ -983,6 +956,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "import tempfile, shutil\n", "from fastcore.test import test_eq, test_fail\n", "\n", @@ -1118,13 +1092,62 @@ "shutil.rmtree(_tmp)\n" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "datasetcacheandpatchsizetest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "import tempfile, shutil\n", + "from fastcore.test import test_eq, test_fail\n", + "\n", + "# Symlinked inputs whose names differ from their targets reuse the published cache.\n", + "_tmp2 = Path(tempfile.mkdtemp())\n", + "(_tmp2 / 'raw').mkdir(); (_tmp2 / 'links').mkdir()\n", + "for i in range(2):\n", + " tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)).save(_tmp2 / 'raw' / f'img_raw_{i}.nii.gz')\n", + " tio.LabelMap(tensor=torch.randint(0, 2, (1, 10, 10, 10))).save(_tmp2 / 'raw' / f'seg_raw_{i}.nii.gz')\n", + " os.symlink(_tmp2 / 'raw' / f'img_raw_{i}.nii.gz', _tmp2 / 'links' / f'case_{i}_0000.nii.gz')\n", + " os.symlink(_tmp2 / 'raw' / f'seg_raw_{i}.nii.gz', _tmp2 / 'links' / f'case_{i}.nii.gz')\n", + "_df_links = pd.DataFrame({\n", + " 'img': [str(_tmp2 / 'links' / f'case_{i}_0000.nii.gz') for i in range(2)],\n", + " 'mask': [str(_tmp2 / 'links' / f'case_{i}.nii.gz') for i in range(2)],\n", + "})\n", + "_r1 = preprocess_dataset(_df_links.copy(), 'img', 'mask', output_dir=_tmp2 / 'cache', apply_reorder=False)\n", + "_df_again = _df_links.copy()\n", + "_r2 = preprocess_dataset(_df_again, 'img', 'mask', output_dir=_tmp2 / 'cache', apply_reorder=False)\n", + "test_eq(_r2.reused, True)\n", + "test_eq(_r2.cache_version, _r1.cache_version)\n", + "test_eq([Path(p).name for p in _df_again.img_preprocessed], ['case_0_0000.nii.gz', 'case_1_0000.nii.gz'])\n", + "test_eq(all(Path(p).is_file() for p in _df_again.mask_preprocessed), True)\n", + "\n", + "# A metadata cache hit reports the caller's own path spelling.\n", + "_seg0 = _tmp2 / 'raw' / 'seg_raw_0.nii.gz'\n", + "_meta_cache = _tmp2 / 'meta_cache.pkl'\n", + "MedDataset(img_list=[f'{_tmp2}/raw/./seg_raw_0.nii.gz'], dtype=MedMask,\n", + " apply_reorder=False, cache_path=_meta_cache)\n", + "_ds_hit = MedDataset(img_list=[str(_seg0)], dtype=MedMask, apply_reorder=False, cache_path=_meta_cache)\n", + "test_eq(_ds_hit.df.path.tolist(), [str(_seg0)])\n", + "\n", + "# suggest_patch_size names the divisor when the smallest volume is below it.\n", + "_ds_small = MedDataset.__new__(MedDataset)\n", + "_ds_small.df = pd.DataFrame({'dim_0': [64, 12], 'dim_1': [64, 64], 'dim_2': [64, 64],\n", + " 'voxel_0': [1.0] * 2, 'voxel_1': [1.0] * 2, 'voxel_2': [1.0] * 2})\n", + "test_fail(lambda: suggest_patch_size(_ds_small, target_spacing=[1, 1, 1], min_patch_size=[8, 8, 8]),\n", + " contains='divisor=16')\n", + "\n", + "shutil.rmtree(_tmp2)" + ] + }, { "cell_type": "markdown", "id": "74812108-f3eb-4a8d-9f2d-b93132619008", "metadata": {}, "source": [ "# Dataset information\n", - ">" + "> Inspect dataset metadata, suggest preprocessing parameters and build a reusable preprocessing cache." ] }, { diff --git a/nbs/09_external_data.ipynb b/nbs/09_external_data.ipynb index 24e37b7..13736d7 100644 --- a/nbs/09_external_data.ipynb +++ b/nbs/09_external_data.ipynb @@ -35,7 +35,7 @@ "metadata": {}, "source": [ "# External data\n", - ">" + "> Download and prepare the public datasets and example data used in the tutorials." ] }, { @@ -49,8 +49,8 @@ "class MURLs():\n", " \"\"\"A class with external medical dataset URLs.\"\"\"\n", "\n", - " IXI_DATA = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar'\n", - " IXI_DEMOGRAPHIC_INFORMATION = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls'\n", + " IXI_DATA = 'https://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-T1.tar'\n", + " IXI_DEMOGRAPHIC_INFORMATION = 'https://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI.xls'\n", " CHENGWEN_CHU_SPINE_DATA = 'https://drive.google.com/uc?id=1rbm9-KKAexpNm2mC9FsSbfnS8VJaF3Kn&confirm=t'\n", " EXAMPLE_SPINE_DATA = 'https://drive.google.com/uc?id=1K1J2wmWxYUPY0NbylcV4Otv88EqCL4Ys'\n", " MEDMNIST_DICT = {'OrganMNIST3D': 'https://zenodo.org/record/6496656/files/organmnist3d.npz?download=1',\t\n", @@ -302,24 +302,6 @@ " return str(fn)" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "b0199679-707b-445d-8467-3d342246322a", - "metadata": {}, - "outputs": [], - "source": [ - "#| export \n", - "def _df_sort_and_add_columns(df, label_list, is_val):\n", - " \"\"\"Sort the dataframe based on img_idx and add labels and if it is validation data column.\"\"\"\n", - " \n", - " df = df.sort_values(by='img_idx').reset_index(drop=True)\n", - " df['labels'], df['is_val'] = label_list, is_val \n", - " df = df.drop('img_idx', axis=1)\n", - " \n", - " return df " - ] - }, { "cell_type": "code", "execution_count": null, @@ -331,11 +313,12 @@ "def _create_nodule_df(pool, output_dir, imgs, labels, is_val=False): \n", " \"\"\"Create dataframe for MedMNIST data.\"\"\"\n", " \n", + " # pool.map keeps the input order, so row i holds image i\n", " img_path_list = pool.map(partial(_process_medmnist_img, output_dir), enumerate(imgs))\n", - " img_idx = [float(Path(fn).parts[-1].split('_')[0]) for fn in img_path_list]\n", - " \n", - " df = pd.DataFrame(list(zip(img_path_list, img_idx)), columns=['img_path','img_idx']) \n", - " return _df_sort_and_add_columns(df, labels, is_val)" + " df = pd.DataFrame({'img_path': img_path_list})\n", + " df['labels'], df['is_val'] = labels, is_val\n", + "\n", + " return df" ] }, { @@ -422,7 +405,7 @@ "id": "ed7e2a84-23f5-4779-bea4-facd9124fb83", "metadata": {}, "source": [ - "## Endometrical cancer " + "## Endometrial cancer" ] }, { From 4f102ae0fd4a7d2d9d1efab64866d78c553c5b64 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:25:47 +0100 Subject: [PATCH 04/11] fix(utils)!: stop tutorials 11d/11e failing at fit and leaking runs MedDataBlock on a dataframe with 'image'/'label' columns has no image-path column for split logging, so ModelTrackingCallback raised in before_fit and left the MLflow run active. The split identity now falls back to the column get_x reads. Runs the callback starts are ended (as FAILED on errors), also when auto_start=False finds no active run; a user's own run is never ended. Also: - log_metrics skips NaN and inf; MLflow stores inf as max float - a binary patch contract must use the same threshold for patch inference and output decoding (checked when created, float32-tolerant) - the MLflow UI runs with this Python's mlflow, keeps ownership of a server it started and shows the server's stderr when it fails to start - docstrings, section headings, hidden test cells, dead code removed BREAKING CHANGE: removed start_ui(auto_open=...) and MLflowUIManager.check_mlflow_installed. --- fastMONAI/utils.py | 302 ++++++++++++++++------- nbs/07_utils.ipynb | 597 ++++++++++++++++++++++++++++++++++++--------- 2 files changed, 696 insertions(+), 203 deletions(-) diff --git a/fastMONAI/utils.py b/fastMONAI/utils.py index 5080f89..9f11c5c 100644 --- a/fastMONAI/utils.py +++ b/fastMONAI/utils.py @@ -1,3 +1,5 @@ +"""Inference contracts, Safetensors model artifacts, MLflow experiment tracking and a notebook MLflow UI manager.""" + # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_utils.ipynb. # %% auto #0 @@ -21,13 +23,14 @@ from datetime import datetime from fastai.callback.core import Callback from fastai.learner import Learner -from fastcore.foundation import L +from fastai.data.transforms import ColReader from collections.abc import Callable, Mapping from typing import Any # Default MLflow backend store: a sqlite DB at the package/repo root, resolved from -# __file__ so it stays stable across os.chdir. NOT set at import time, which hijacked -# the user's MLflow config; set lazily instead (see _ensure_fastmonai_tracking_uri). +# __file__ so it stays stable across os.chdir. Importing this module never changes the +# tracking URI: _ensure_fastmonai_tracking_uri applies this default lazily, and only +# when the user has not configured MLflow. try: _mlruns_db = Path(__file__).resolve().parent.parent / 'mlruns.db' except NameError: @@ -131,7 +134,7 @@ def load_variables(config_fn: str | Path): try: with open(config_fn, 'r') as f: return json.load(f) - except (UnicodeDecodeError, json.JSONDecodeError, ValueError): + except ValueError: # also covers json.JSONDecodeError and UnicodeDecodeError if not str(config_fn).endswith('.pkl'): raise ValueError( f"{config_fn} is not valid JSON. Re-export it, or use a trusted legacy " @@ -248,14 +251,35 @@ def _validate_output_spec(spec: Mapping) -> dict: return expected +def _check_patch_threshold(config: Mapping) -> dict: + """Require a binary patch contract to use one threshold for inference and decoding. + + Checked when a contract is created, not when a saved one is read. + """ + import math + if config.get('workflow') == 'patch' and config['output']['kind'] == 'binary_segmentation': + threshold = config['patch_config']['binary_threshold'] + # Tolerance covers thresholds stored as float32. + if not math.isclose(threshold, config['output']['threshold'], rel_tol=1e-6): + raise ValueError( + f"PatchConfig binary_threshold ({threshold}) must equal the " + f"binary_segmentation output threshold ({config['output']['threshold']})" + ) + return config + + def patch_inference_config(config, output_spec: Mapping) -> dict: - """Build the complete versioned inference contract for a patch workflow.""" - return { + """Build the complete versioned inference contract for a patch workflow. + + For ``binary_segmentation`` outputs, ``config.binary_threshold`` must equal the + output threshold, so patch inference and output decoding agree. + """ + return _check_patch_threshold(_validate_inference_config({ 'config_schema': '1', 'workflow': 'patch', 'patch_config': patch_config_to_dict(config, inference_only=True), 'output': _validate_output_spec(output_spec), - } + })) def standard_inference_config(apply_reorder: bool, target_spacing, item_tfms: list, @@ -361,6 +385,7 @@ def strip_compile_prefix(state_dict, prefix: str = '_orig_mod.'): return state_dict +# %% ../nbs/07_utils.ipynb #safetensorsapi MODEL_ARTIFACT_SCHEMA = '1' _MODEL_CONSTRUCTORS = {} _MODEL_WRAPPERS = {} @@ -815,8 +840,34 @@ def _resolve_column(frame, value): return value if value in frame.columns else None +def _prefer_raw_column(frame, column): + """Map a preprocessed-cache column to its raw source column when both exist.""" + if isinstance(column, str) and column.endswith('_preprocessed'): + raw_column = column[:-len('_preprocessed')] + if raw_column in frame.columns: + return raw_column + return column + + +def _get_x_column(dls, frame): + """Column the DataLoaders read x from, when get_x is a single-column ColReader.""" + try: + x_pipeline = dls.train_ds.tfms[0] + except (AttributeError, IndexError, KeyError, TypeError): + return None + readers = [tfm for tfm in getattr(x_pipeline, 'fs', []) if isinstance(tfm, ColReader)] + if len(readers) != 1 or len(readers[0].cols) != 1: + return None + return _resolve_column(frame, readers[0].cols[0]) + + def _split_identity_column(dls, frame, sample_id_col=None): - """Choose an identity, preferring the raw image path over a cache path.""" + """Choose a stable sample identity column, preferring raw image paths over cache paths. + + Order: ``sample_id_col``, the DataLoaders' image column (``_img_col``), a + ``sample_id`` column, a single ``path``/``*_path`` column, then the column that + a ``ColReader`` get_x reads (e.g. ``'image'`` in a MedDataBlock). + """ if sample_id_col is not None: column = _resolve_column(frame, sample_id_col) if column is None: @@ -825,11 +876,7 @@ def _split_identity_column(dls, frame, sample_id_col=None): ) return column - column = _resolve_column(frame, getattr(dls, '_img_col', None)) - if isinstance(column, str) and column.endswith('_preprocessed'): - raw_column = column[:-len('_preprocessed')] - if raw_column in frame.columns: - column = raw_column + column = _prefer_raw_column(frame, _resolve_column(frame, getattr(dls, '_img_col', None))) if column is not None: return column if 'sample_id' in frame.columns: @@ -842,6 +889,9 @@ def _split_identity_column(dls, frame, sample_id_col=None): ] if len(candidates) == 1: return candidates[0] + column = _prefer_raw_column(frame, _get_x_column(dls, frame)) + if column is not None: + return column raise ValueError( 'Could not choose a stable sample identity for split_version. ' 'Pass sample_id_col to create_mlflow_callback().' @@ -946,8 +996,17 @@ def __init__( apply_reorder: Whether reordering augmentation is applied experiment_name: MLflow experiment name (used with auto_start) run_name: MLflow run name (auto-generated if None) - auto_start: If True, auto-starts/stops MLflow run + auto_start: If True, starts a run in before_fit and ends it in after_fit. + If False, logs to the active run, or starts and ends its own run when + none is active. patch_config: Patch configuration dict for logging (from MedPatchDataLoaders) + model_spec: Required model specification from make_model_spec(). It is + stored with the Safetensors artifacts so the model can be rebuilt. + output_spec: Output decoding contract from make_output_spec(). Patch + workflows only: combined with patch_config into the inference contract. + inference_config: Complete inference contract, e.g. from + standard_inference_config(). Required for standard workflows. Pass + either this or output_spec, not both. extra_params: Additional parameters to log extra_tags: MLflow tags to set on the run dataset_version: Optional annotation/dataset fingerprint for tracking. @@ -982,7 +1041,8 @@ def __init__( else: if output_spec is not None: raise ValueError("Pass output inside inference_config, not both arguments") - self.inference_config = _validate_inference_config(inference_config) + self.inference_config = _check_patch_threshold( + _validate_inference_config(inference_config)) self.extra_params = extra_params or {} self.extra_tags = extra_tags or {} self.preprocessing_manifest = None @@ -1077,7 +1137,9 @@ def _extract_epoch_metrics(self) -> dict[str, float]: """Extract metrics from the current epoch.""" recorder = self.learn.recorder - # Skip 'epoch' and 'time' + # Recorder.metric_names is [epoch, train_loss, valid_loss, *metrics, time]. Skip + # epoch and train_loss (added below); 'time' drops out of the zip because Recorder + # appends it to the log only after this callback's after_epoch has run. metric_names = recorder.metric_names[2:] raw_metric_values = recorder.log[2:] @@ -1107,7 +1169,6 @@ def _extract_epoch_metrics(self) -> dict[str, float]: def _save_model_artifacts(self, temp_dir: Path) -> None: """Save weights-only .pth and Safetensors model artifacts.""" compiled_model = self.learn.model - original_cbs = self.learn.cbs.copy() self.learn.model = unwrap_compiled_model(compiled_model) def log_safetensors(role): @@ -1150,7 +1211,6 @@ def log_safetensors(role): log_safetensors('best') finally: self.learn.model = compiled_model - self.learn.cbs = original_cbs def _log_split_df(self) -> None: """Log and hash actual train/validation membership selected by fastai.""" @@ -1182,18 +1242,45 @@ def _log_split_df(self) -> None: ) def before_fit(self) -> None: - """Log hyperparameters before training starts. Auto-start run if configured.""" - if self.auto_start: - if self.experiment_name: - mlflow.set_experiment(self.experiment_name) - self.run_name = self.run_name or f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + """Start or join an MLflow run and log hyperparameters before training. + + 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. + """ + 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() + active_run = None + + if self.auto_start or active_run is None: + if self.auto_start: + if self.experiment_name: + mlflow.set_experiment(self.experiment_name) + 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 self._auto_started = True self._owns_run = True - elif mlflow.active_run() is not None: - self.run_id = mlflow.active_run().info.run_id - + self._run_failed = False + else: + self.run_id = active_run.info.run_id + self._auto_started = False + self._owns_run = False + + try: + self._log_run_setup() + except Exception: + self.mark_failed() + raise + + def _log_run_setup(self) -> None: + """Log tags, parameters, the preprocessing manifest and the split to the run.""" if self.extra_tags: mlflow.set_tags(self.extra_tags) @@ -1219,16 +1306,20 @@ def after_epoch(self) -> None: mlflow.log_metrics(metrics, step=self.learn.epoch) def after_fit(self) -> None: - """Log model artifacts after training completion.""" + """Log model artifacts, then end the run if this callback started it. + + If saving the artifacts raises, a run started by this callback is ended + with status FAILED. + """ print("\nTraining finished. Logging model artifacts to MLflow...") - - with tempfile.TemporaryDirectory() as temp_dir: - temp_path = Path(temp_dir) - - self._save_model_artifacts(temp_path) - - if self.run_id is None and mlflow.active_run() is not None: - self.run_id = mlflow.active_run().info.run_id + + try: + with tempfile.TemporaryDirectory() as temp_dir: + self._save_model_artifacts(Path(temp_dir)) + except Exception: + self.mark_failed() + raise + print(f"MLflow run completed. Run ID: {self.run_id}") if self._auto_started: @@ -1254,7 +1345,8 @@ def log_metrics(self, metrics: dict) -> None: Reopens the run to log metrics, then closes it. Useful for post-training evaluation (e.g., validation set patch inference results). - NaN values are silently skipped (MLflow rejects them). + Values that are not finite are skipped: MLflow accepts NaN but stores inf + as the largest float, which would distort plots and comparisons. Args: metrics: Dictionary of metric name -> value pairs to log. @@ -1263,8 +1355,8 @@ def log_metrics(self, metrics: dict) -> None: raise RuntimeError("No MLflow run found. Train the model first.") import math - valid = {k: v for k, v in metrics.items() - if not (isinstance(v, float) and math.isnan(v))} + values = {k: float(v) for k, v in metrics.items()} + valid = {k: v for k, v in values.items() if math.isfinite(v)} if valid: with mlflow.start_run(run_id=self.run_id): @@ -1409,19 +1501,18 @@ def create_mlflow_callback( ModelTrackingCallback ready to use with learn.fit() Example: - >>> # Instead of this (6 manual params): - >>> # mlflow_callback = ModelTrackingCallback( - >>> # model_name=f"{task}_{model._get_name()}", - >>> # loss_function=loss_func.loss_func._get_name(), - >>> # item_tfms=item_tfms, - >>> # size=size, - >>> # target_spacing=target_spacing, - >>> # apply_reorder=True, - >>> # ) - >>> # with mlflow.start_run(run_name="training"): - >>> # learn.fit_one_cycle(30, lr, cbs=[mlflow_callback]) - >>> - >>> # Patch workflow: architecture and decoding remain explicit. + >>> # Standard workflow: pass the complete inference contract. + >>> model_spec = make_model_spec("monai.unet_inner", unet_kwargs) + >>> inference_config = standard_inference_config( + ... apply_reorder=True, target_spacing=target_spacing, + ... item_tfms=inference_transforms_to_specs(item_tfms), + ... output_spec=make_output_spec("binary_segmentation", threshold=0.5)) + >>> callback = create_mlflow_callback( + ... learn, experiment_name="Task02_Heart", + ... model_spec=model_spec, inference_config=inference_config) + >>> learn.fit_flat_cos(200, lr, cbs=[callback, save_best]) + >>> + >>> # Patch workflow: patch settings come from the DataLoaders; pass output_spec. >>> model_spec = make_model_spec("monai.unet", unet_kwargs) >>> output_spec = make_output_spec("multiclass_segmentation", classes=2) >>> callback = create_mlflow_callback( @@ -1535,22 +1626,19 @@ def find_model_artifacts(*, run_ids, artifact_role, expected_members=None): # %% ../nbs/07_utils.ipynb #a3d2d585 import subprocess -import threading import time import socket -import os import sys import atexit import signal import ctypes -from IPython.display import display, HTML -import shutil +import html # Ask the OS to SIGTERM the `mlflow ui` child the instant this (parent) process dies; # this covers a hard-killed kernel (SIGKILL/OOM), which atexit cannot. Linux-only; # a no-op elsewhere. libc is loaded once here, outside any fork, so the preexec_fn hook # below stays minimal (preexec_fn is documented as unsafe in multithreaded programs, and -# we launch mlflow ui from a daemon thread, so keeping the hook tiny is the accepted pattern). +# a Jupyter kernel runs several threads, so keeping the hook tiny is the accepted pattern). _libc = None if sys.platform.startswith('linux'): try: @@ -1572,17 +1660,18 @@ class MLflowUIManager: exits, gracefully (via an atexit handler) or on a hard kill (via Linux PR_SET_PDEATHSIG, see ``_die_with_parent``), so closing/restarting the notebook never leaves an orphaned server holding the port. Call ``stop()`` - to shut it down sooner. An externally-started UI is reused, not killed. + to shut it down sooner. An externally-started UI is reused, not killed, and + calling ``start_ui()`` again while its own server runs just shows the link. """ def __init__(self): self.process = None - self.thread = None self.port = 5001 self.host = '0.0.0.0' self.backend_store_uri = _resolve_backend_store_uri() self._owns_process = False self._reusing_external = False + self._log_path = None # Reap the mlflow ui subprocess on clean interpreter shutdown (graceful kernel # shutdown/restart runs atexit). PR_SET_PDEATHSIG covers hard kills. atexit.register(self._kill_process) @@ -1625,20 +1714,42 @@ def find_available_port(self, start_port=5001): return port return None - def check_mlflow_installed(self): - """Check if MLflow is installed.""" - return shutil.which('mlflow') is not None - - def start_ui(self, auto_open=True, quiet=False): - """Start MLflow UI with better error handling and user feedback.""" + def _log_tail(self, max_chars=2000): + """End of the server's stderr log, or '' when there is none.""" + try: + return Path(self._log_path).read_text(errors='replace')[-max_chars:].strip() + except (OSError, TypeError): + return '' + + def _remove_log(self): + """Delete the server's stderr log file, if any.""" + if self._log_path is not None: + try: + os.remove(self._log_path) + except OSError: + pass + self._log_path = None + + def start_ui(self, quiet=False): + """Start the MLflow UI, or show its link if it is already running. + + The server runs with this interpreter (``python -m mlflow ui``), so it uses the + same MLflow install as the notebook. Its stderr goes to a temporary log file, + and the end of that log is shown if the server fails to start. + """ + from IPython.display import display, HTML # Refresh in case the tracking URI was configured after __init__ self.backend_store_uri = _resolve_backend_store_uri() - - if not self.check_mlflow_installed(): - if not quiet: - display(HTML('
MLflow not installed. Run: pip install mlflow
')) - return False - + + if self._owns_process and self.process is not None and self.process.poll() is None: + display(HTML(f''' + + MLflow UI (Port {self.port}) + + ''')) + return True + existing_port = self._find_running_mlflow() if existing_port is not None: self.port = existing_port @@ -1675,30 +1786,34 @@ def start_ui(self, auto_open=True, quiet=False): return False self.port = available_port - - def run_mlflow(): - try: + + # stderr goes to a file rather than a pipe: nobody drains a pipe while the + # server runs, so a full pipe buffer would block it. + self._remove_log() + log_file = tempfile.NamedTemporaryFile( + mode='wb', prefix='mlflow_ui_', suffix='.log', delete=False) + self._log_path = log_file.name + try: + with log_file: self.process = subprocess.Popen([ - 'mlflow', 'ui', + sys.executable, '-m', 'mlflow', 'ui', '--host', self.host, '--port', str(self.port), '--backend-store-uri', self.backend_store_uri - ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ], stdout=subprocess.DEVNULL, stderr=log_file, preexec_fn=_die_with_parent if _libc is not None else None) - self.process.wait() - except Exception as e: - if not quiet: - display(HTML(f'
Error: {str(e)}
')) - - self.thread = threading.Thread(target=run_mlflow, daemon=True) - self.thread.start() - + except Exception as e: + self._remove_log() + if not quiet: + display(HTML(f'
Error: {html.escape(str(e))}
')) + return False + self._owns_process = True + self._reusing_external = False + max_wait = 10 for i in range(max_wait): time.sleep(1) if self.is_mlflow_running(): - self._owns_process = True - self._reusing_external = False if quiet: display(HTML(f''' ''')) return True - + if self.process.poll() is not None: + break + + exited = self.process.poll() is not None + log_tail = self._log_tail() + if exited: + self.process, self._owns_process = None, False + self._remove_log() if not quiet: - display(HTML('
Failed to start MLflow UI
')) + reason = 'the server exited' if exited else f'no response after {max_wait} s' + details = (f'
{html.escape(log_tail)}
' + if log_tail else '') + display(HTML(f'
Failed to start MLflow UI ({reason})
{details}')) return False def _kill_process(self): @@ -1747,9 +1872,11 @@ def _kill_process(self): proc.wait(timeout=5) except Exception: pass + self._remove_log() def stop(self): """Stop the MLflow UI server (terminate -> wait -> kill).""" + from IPython.display import display, HTML if self.process and self._owns_process: self._kill_process() display(HTML(''' @@ -1773,6 +1900,7 @@ def stop(self): def status(self): """Check MLflow UI status.""" + from IPython.display import display, HTML if self.is_mlflow_running(): display(HTML(f'''
diff --git a/nbs/07_utils.ipynb b/nbs/07_utils.ipynb index db0ad12..d5e0ff0 100644 --- a/nbs/07_utils.ipynb +++ b/nbs/07_utils.ipynb @@ -30,13 +30,14 @@ "from datetime import datetime\n", "from fastai.callback.core import Callback\n", "from fastai.learner import Learner\n", - "from fastcore.foundation import L\n", + "from fastai.data.transforms import ColReader\n", "from collections.abc import Callable, Mapping\n", "from typing import Any\n", "\n", "# Default MLflow backend store: a sqlite DB at the package/repo root, resolved from\n", - "# __file__ so it stays stable across os.chdir. NOT set at import time, which hijacked\n", - "# the user's MLflow config; set lazily instead (see _ensure_fastmonai_tracking_uri).\n", + "# __file__ so it stays stable across os.chdir. Importing this module never changes the\n", + "# tracking URI: _ensure_fastmonai_tracking_uri applies this default lazily, and only\n", + "# when the user has not configured MLflow.\n", "try:\n", " _mlruns_db = Path(__file__).resolve().parent.parent / 'mlruns.db'\n", "except NameError:\n", @@ -57,7 +58,7 @@ "metadata": {}, "source": [ "# Utils\n", - ">" + "> Inference contracts, Safetensors model artifacts, MLflow experiment tracking and a notebook MLflow UI manager." ] }, { @@ -74,7 +75,7 @@ "id": "c2db5512-171c-4dfd-a26e-561b773a6069", "metadata": {}, "outputs": [], - "source": "#| export\ndef load_variables(config_fn: str | Path):\n \"\"\"Load stored inference variables.\n\n Tries JSON first. Falls back to legacy pickle ONLY for files named ``.pkl``\n (so a non-JSON payload disguised as ``.json`` is refused rather than\n unpickled, since pickle can execute arbitrary code). Existing ``.pkl``\n artifacts still load.\n\n Args:\n config_fn: File path to load.\n\n Returns:\n The deserialized ``[size, apply_reorder, target_spacing]`` list.\n \"\"\"\n try:\n with open(config_fn, 'r') as f:\n return json.load(f)\n except (UnicodeDecodeError, json.JSONDecodeError, ValueError):\n if not str(config_fn).endswith('.pkl'):\n raise ValueError(\n f\"{config_fn} is not valid JSON. Re-export it, or use a trusted legacy \"\n f\"'.pkl' file (pickle can execute arbitrary code on load).\"\n )\n import warnings\n warnings.warn(f\"Loading legacy pickle '{config_fn}'; pickle can execute arbitrary \"\n f\"code -- only load files you trust.\")\n with open(config_fn, 'rb') as f:\n return pickle.load(f)" + "source": "#| export\ndef load_variables(config_fn: str | Path):\n \"\"\"Load stored inference variables.\n\n Tries JSON first. Falls back to legacy pickle ONLY for files named ``.pkl``\n (so a non-JSON payload disguised as ``.json`` is refused rather than\n unpickled, since pickle can execute arbitrary code). Existing ``.pkl``\n artifacts still load.\n\n Args:\n config_fn: File path to load.\n\n Returns:\n The deserialized ``[size, apply_reorder, target_spacing]`` list.\n \"\"\"\n try:\n with open(config_fn, 'r') as f:\n return json.load(f)\n except ValueError: # also covers json.JSONDecodeError and UnicodeDecodeError\n if not str(config_fn).endswith('.pkl'):\n raise ValueError(\n f\"{config_fn} is not valid JSON. Re-export it, or use a trusted legacy \"\n f\"'.pkl' file (pickle can execute arbitrary code on load).\"\n )\n import warnings\n warnings.warn(f\"Loading legacy pickle '{config_fn}'; pickle can execute arbitrary \"\n f\"code -- only load files you trust.\")\n with open(config_fn, 'rb') as f:\n return pickle.load(f)" }, { "cell_type": "markdown", @@ -196,14 +197,35 @@ " return expected\n", "\n", "\n", + "def _check_patch_threshold(config: Mapping) -> dict:\n", + " \"\"\"Require a binary patch contract to use one threshold for inference and decoding.\n", + "\n", + " Checked when a contract is created, not when a saved one is read.\n", + " \"\"\"\n", + " import math\n", + " if config.get('workflow') == 'patch' and config['output']['kind'] == 'binary_segmentation':\n", + " threshold = config['patch_config']['binary_threshold']\n", + " # Tolerance covers thresholds stored as float32.\n", + " if not math.isclose(threshold, config['output']['threshold'], rel_tol=1e-6):\n", + " raise ValueError(\n", + " f\"PatchConfig binary_threshold ({threshold}) must equal the \"\n", + " f\"binary_segmentation output threshold ({config['output']['threshold']})\"\n", + " )\n", + " return config\n", + "\n", + "\n", "def patch_inference_config(config, output_spec: Mapping) -> dict:\n", - " \"\"\"Build the complete versioned inference contract for a patch workflow.\"\"\"\n", - " return {\n", + " \"\"\"Build the complete versioned inference contract for a patch workflow.\n", + "\n", + " For ``binary_segmentation`` outputs, ``config.binary_threshold`` must equal the\n", + " output threshold, so patch inference and output decoding agree.\n", + " \"\"\"\n", + " return _check_patch_threshold(_validate_inference_config({\n", " 'config_schema': '1',\n", " 'workflow': 'patch',\n", " 'patch_config': patch_config_to_dict(config, inference_only=True),\n", " 'output': _validate_output_spec(output_spec),\n", - " }\n", + " }))\n", "\n", "\n", "def standard_inference_config(apply_reorder: bool, target_spacing, item_tfms: list,\n", @@ -279,6 +301,14 @@ "print(\"standard config JSON/pickle round-trip + security tests passed!\")\n" ] }, + { + "cell_type": "markdown", + "id": "colabgpumd", + "metadata": {}, + "source": [ + "## Colab GPU check" + ] + }, { "cell_type": "code", "execution_count": null, @@ -364,9 +394,75 @@ " from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present\n", "\n", " consume_prefix_in_state_dict_if_present(state_dict, prefix)\n", - " return state_dict\n", + " return state_dict\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "compilehelperstest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# torch.compile needs a C++ toolchain and is slow in CI, so an nn.Module holding the\n", + "# inner module as `_orig_mod` stands in for OptimizedModule: it reproduces both the\n", + "# attribute chain and the `_orig_mod.` state-dict prefix that these helpers handle.\n", + "import torch.nn as _nn\n", + "from fastcore.test import test_eq, test_fail\n", + "from fastai.test_utils import synth_learner\n", + "\n", + "\n", + "class _OrigModWrapper(_nn.Module):\n", + " def __init__(self, mod):\n", + " super().__init__()\n", + " self._orig_mod = mod\n", + "\n", + "\n", + "_bare = _nn.Sequential(_nn.Linear(3, 3))\n", + "\n", + "test_eq(unwrap_compiled_model(_bare) is _bare, True)\n", + "test_eq(unwrap_compiled_model(_OrigModWrapper(_bare)) is _bare, True)\n", + "test_eq(unwrap_compiled_model(_OrigModWrapper(_OrigModWrapper(_bare))) is _bare, True)\n", "\n", + "# a Learner is unwrapped in place and handed back as a Learner\n", + "_learn = synth_learner(n_trn=2, n_val=2)\n", + "_inner = _learn.model\n", + "_learn.model = _OrigModWrapper(_inner)\n", + "test_eq(unwrap_compiled_model(_learn) is _learn, True)\n", + "test_eq(_learn.model is _inner, True)\n", "\n", + "# the wrapper prefixes every key, and a bare module cannot load that until stripped\n", + "_sd = _OrigModWrapper(_bare).state_dict()\n", + "test_eq(all(k.startswith('_orig_mod.') for k in _sd), True)\n", + "test_fail(lambda: _bare.load_state_dict(_sd))\n", + "_bare.load_state_dict(strip_compile_prefix(_sd))\n", + "test_eq(sorted(_sd), sorted(_bare.state_dict()))\n", + "\n", + "# an already-clean state dict passes through untouched\n", + "_clean = _bare.state_dict()\n", + "test_eq(strip_compile_prefix(_clean) is _clean, True)\n", + "test_eq(sorted(_clean), sorted(_bare.state_dict()))" + ] + }, + { + "cell_type": "markdown", + "id": "safetensorsmd", + "metadata": {}, + "source": [ + "## Safetensors model artifacts\n", + "\n", + "Models are saved as Safetensors weights with metadata that holds the model specification and the inference contract. Loading rebuilds the model only through registered, trusted constructors and wrappers." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "safetensorsapi", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", "MODEL_ARTIFACT_SCHEMA = '1'\n", "_MODEL_CONSTRUCTORS = {}\n", "_MODEL_WRAPPERS = {}\n", @@ -637,49 +733,44 @@ { "cell_type": "code", "execution_count": null, - "id": "compilehelperstest", + "id": "patchthresholdtest", "metadata": {}, "outputs": [], "source": [ "#| hide\n", - "# torch.compile needs a C++ toolchain and is slow in CI, so an nn.Module holding the\n", - "# inner module as `_orig_mod` stands in for OptimizedModule: it reproduces both the\n", - "# attribute chain and the `_orig_mod.` state-dict prefix that these helpers handle.\n", - "import torch.nn as _nn\n", + "# A binary patch contract must use one threshold for patch inference and output decoding.\n", "from fastcore.test import test_eq, test_fail\n", - "from fastai.test_utils import synth_learner\n", - "\n", - "\n", - "class _OrigModWrapper(_nn.Module):\n", - " def __init__(self, mod):\n", - " super().__init__()\n", - " self._orig_mod = mod\n", - "\n", - "\n", - "_bare = _nn.Sequential(_nn.Linear(3, 3))\n", - "\n", - "test_eq(unwrap_compiled_model(_bare) is _bare, True)\n", - "test_eq(unwrap_compiled_model(_OrigModWrapper(_bare)) is _bare, True)\n", - "test_eq(unwrap_compiled_model(_OrigModWrapper(_OrigModWrapper(_bare))) is _bare, True)\n", - "\n", - "# a Learner is unwrapped in place and handed back as a Learner\n", - "_learn = synth_learner(n_trn=2, n_val=2)\n", - "_inner = _learn.model\n", - "_learn.model = _OrigModWrapper(_inner)\n", - "test_eq(unwrap_compiled_model(_learn) is _learn, True)\n", - "test_eq(_learn.model is _inner, True)\n", + "from fastMONAI.vision_patch import PatchConfig as _PatchConfig\n", "\n", - "# the wrapper prefixes every key, and a bare module cannot load that until stripped\n", - "_sd = _OrigModWrapper(_bare).state_dict()\n", - "test_eq(all(k.startswith('_orig_mod.') for k in _sd), True)\n", - "test_fail(lambda: _bare.load_state_dict(_sd))\n", - "_bare.load_state_dict(strip_compile_prefix(_sd))\n", - "test_eq(sorted(_sd), sorted(_bare.state_dict()))\n", + "_binary_out = make_output_spec('binary_segmentation', threshold=0.5)\n", + "_binary_infer = patch_inference_config(_PatchConfig(patch_size=[16, 16, 16]), _binary_out)\n", + "test_eq(_binary_infer['patch_config']['binary_threshold'], _binary_infer['output']['threshold'])\n", + "test_fail(lambda: patch_inference_config(\n", + " _PatchConfig(patch_size=[16, 16, 16], binary_threshold=0.4), _binary_out),\n", + " contains='binary_threshold')\n", + "_mismatch = {**_binary_infer,\n", + " 'patch_config': {**_binary_infer['patch_config'], 'binary_threshold': 0.4}}\n", + "test_fail(lambda: _check_patch_threshold(_mismatch), contains='binary_threshold')\n", + "# Saved contracts are still readable, and a float32 0.3 (0.30000001192092896)\n", + "# matches an output threshold of 0.3.\n", + "test_eq(_validate_inference_config(_mismatch)['patch_config']['binary_threshold'], 0.4)\n", + "patch_inference_config(\n", + " _PatchConfig(patch_size=[16, 16, 16], binary_threshold=0.30000001192092896),\n", + " make_output_spec('binary_segmentation', threshold=0.3))\n", + "\n", + "# Multiclass outputs decode with argmax, so binary_threshold is not compared.\n", + "patch_inference_config(_PatchConfig(patch_size=[16, 16, 16], binary_threshold=0.4),\n", + " make_output_spec('multiclass_segmentation', classes=2))\n" + ] + }, + { + "cell_type": "markdown", + "id": "trackingmd", + "metadata": {}, + "source": [ + "## MLflow experiment tracking\n", "\n", - "# an already-clean state dict passes through untouched\n", - "_clean = _bare.state_dict()\n", - "test_eq(strip_compile_prefix(_clean) is _clean, True)\n", - "test_eq(sorted(_clean), sorted(_bare.state_dict()))" + "`create_mlflow_callback` builds a `ModelTrackingCallback` from a Learner. The callback logs hyperparameters, per-epoch metrics, the actual train/validation split and Safetensors model artifacts to MLflow." ] }, { @@ -877,8 +968,34 @@ " return value if value in frame.columns else None\n", "\n", "\n", + "def _prefer_raw_column(frame, column):\n", + " \"\"\"Map a preprocessed-cache column to its raw source column when both exist.\"\"\"\n", + " if isinstance(column, str) and column.endswith('_preprocessed'):\n", + " raw_column = column[:-len('_preprocessed')]\n", + " if raw_column in frame.columns:\n", + " return raw_column\n", + " return column\n", + "\n", + "\n", + "def _get_x_column(dls, frame):\n", + " \"\"\"Column the DataLoaders read x from, when get_x is a single-column ColReader.\"\"\"\n", + " try:\n", + " x_pipeline = dls.train_ds.tfms[0]\n", + " except (AttributeError, IndexError, KeyError, TypeError):\n", + " return None\n", + " readers = [tfm for tfm in getattr(x_pipeline, 'fs', []) if isinstance(tfm, ColReader)]\n", + " if len(readers) != 1 or len(readers[0].cols) != 1:\n", + " return None\n", + " return _resolve_column(frame, readers[0].cols[0])\n", + "\n", + "\n", "def _split_identity_column(dls, frame, sample_id_col=None):\n", - " \"\"\"Choose an identity, preferring the raw image path over a cache path.\"\"\"\n", + " \"\"\"Choose a stable sample identity column, preferring raw image paths over cache paths.\n", + "\n", + " Order: ``sample_id_col``, the DataLoaders' image column (``_img_col``), a\n", + " ``sample_id`` column, a single ``path``/``*_path`` column, then the column that\n", + " a ``ColReader`` get_x reads (e.g. ``'image'`` in a MedDataBlock).\n", + " \"\"\"\n", " if sample_id_col is not None:\n", " column = _resolve_column(frame, sample_id_col)\n", " if column is None:\n", @@ -887,11 +1004,7 @@ " )\n", " return column\n", "\n", - " column = _resolve_column(frame, getattr(dls, '_img_col', None))\n", - " if isinstance(column, str) and column.endswith('_preprocessed'):\n", - " raw_column = column[:-len('_preprocessed')]\n", - " if raw_column in frame.columns:\n", - " column = raw_column\n", + " column = _prefer_raw_column(frame, _resolve_column(frame, getattr(dls, '_img_col', None)))\n", " if column is not None:\n", " return column\n", " if 'sample_id' in frame.columns:\n", @@ -904,6 +1017,9 @@ " ]\n", " if len(candidates) == 1:\n", " return candidates[0]\n", + " column = _prefer_raw_column(frame, _get_x_column(dls, frame))\n", + " if column is not None:\n", + " return column\n", " raise ValueError(\n", " 'Could not choose a stable sample identity for split_version. '\n", " 'Pass sample_id_col to create_mlflow_callback().'\n", @@ -1015,8 +1131,17 @@ " apply_reorder: Whether reordering augmentation is applied\n", " experiment_name: MLflow experiment name (used with auto_start)\n", " run_name: MLflow run name (auto-generated if None)\n", - " auto_start: If True, auto-starts/stops MLflow run\n", + " auto_start: If True, starts a run in before_fit and ends it in after_fit.\n", + " If False, logs to the active run, or starts and ends its own run when\n", + " none is active.\n", " patch_config: Patch configuration dict for logging (from MedPatchDataLoaders)\n", + " model_spec: Required model specification from make_model_spec(). It is\n", + " stored with the Safetensors artifacts so the model can be rebuilt.\n", + " output_spec: Output decoding contract from make_output_spec(). Patch\n", + " workflows only: combined with patch_config into the inference contract.\n", + " inference_config: Complete inference contract, e.g. from\n", + " standard_inference_config(). Required for standard workflows. Pass\n", + " either this or output_spec, not both.\n", " extra_params: Additional parameters to log\n", " extra_tags: MLflow tags to set on the run\n", " dataset_version: Optional annotation/dataset fingerprint for tracking.\n", @@ -1051,7 +1176,8 @@ " else:\n", " if output_spec is not None:\n", " raise ValueError(\"Pass output inside inference_config, not both arguments\")\n", - " self.inference_config = _validate_inference_config(inference_config)\n", + " self.inference_config = _check_patch_threshold(\n", + " _validate_inference_config(inference_config))\n", " self.extra_params = extra_params or {}\n", " self.extra_tags = extra_tags or {}\n", " self.preprocessing_manifest = None\n", @@ -1146,7 +1272,9 @@ " \"\"\"Extract metrics from the current epoch.\"\"\"\n", " recorder = self.learn.recorder\n", " \n", - " # Skip 'epoch' and 'time'\n", + " # Recorder.metric_names is [epoch, train_loss, valid_loss, *metrics, time]. Skip\n", + " # epoch and train_loss (added below); 'time' drops out of the zip because Recorder\n", + " # appends it to the log only after this callback's after_epoch has run.\n", " metric_names = recorder.metric_names[2:]\n", " raw_metric_values = recorder.log[2:]\n", " \n", @@ -1176,7 +1304,6 @@ " def _save_model_artifacts(self, temp_dir: Path) -> None:\n", " \"\"\"Save weights-only .pth and Safetensors model artifacts.\"\"\"\n", " compiled_model = self.learn.model\n", - " original_cbs = self.learn.cbs.copy()\n", " self.learn.model = unwrap_compiled_model(compiled_model)\n", "\n", " def log_safetensors(role):\n", @@ -1219,7 +1346,6 @@ " log_safetensors('best')\n", " finally:\n", " self.learn.model = compiled_model\n", - " self.learn.cbs = original_cbs\n", "\n", " def _log_split_df(self) -> None:\n", " \"\"\"Log and hash actual train/validation membership selected by fastai.\"\"\"\n", @@ -1251,18 +1377,45 @@ " )\n", "\n", " def before_fit(self) -> None:\n", - " \"\"\"Log hyperparameters before training starts. Auto-start run if configured.\"\"\"\n", - " if self.auto_start:\n", - " if self.experiment_name:\n", - " mlflow.set_experiment(self.experiment_name)\n", - " self.run_name = self.run_name or f\"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}\"\n", + " \"\"\"Start or join an MLflow run and log hyperparameters before training.\n", + "\n", + " 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", + " \"\"\"\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", + " active_run = None\n", + "\n", + " if self.auto_start or active_run is None:\n", + " if self.auto_start:\n", + " if self.experiment_name:\n", + " mlflow.set_experiment(self.experiment_name)\n", + " 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", " self._auto_started = True\n", " self._owns_run = True\n", - " elif mlflow.active_run() is not None:\n", - " self.run_id = mlflow.active_run().info.run_id\n", - " \n", + " self._run_failed = False\n", + " else:\n", + " self.run_id = active_run.info.run_id\n", + " self._auto_started = False\n", + " self._owns_run = False\n", + "\n", + " try:\n", + " self._log_run_setup()\n", + " except Exception:\n", + " self.mark_failed()\n", + " raise\n", + "\n", + " def _log_run_setup(self) -> None:\n", + " \"\"\"Log tags, parameters, the preprocessing manifest and the split to the run.\"\"\"\n", " if self.extra_tags:\n", " mlflow.set_tags(self.extra_tags)\n", "\n", @@ -1288,16 +1441,20 @@ " mlflow.log_metrics(metrics, step=self.learn.epoch)\n", " \n", " def after_fit(self) -> None:\n", - " \"\"\"Log model artifacts after training completion.\"\"\"\n", + " \"\"\"Log model artifacts, then end the run if this callback started it.\n", + "\n", + " If saving the artifacts raises, a run started by this callback is ended\n", + " with status FAILED.\n", + " \"\"\"\n", " print(\"\\nTraining finished. Logging model artifacts to MLflow...\")\n", - " \n", - " with tempfile.TemporaryDirectory() as temp_dir:\n", - " temp_path = Path(temp_dir)\n", - " \n", - " self._save_model_artifacts(temp_path)\n", - " \n", - " if self.run_id is None and mlflow.active_run() is not None:\n", - " self.run_id = mlflow.active_run().info.run_id\n", + "\n", + " try:\n", + " with tempfile.TemporaryDirectory() as temp_dir:\n", + " self._save_model_artifacts(Path(temp_dir))\n", + " except Exception:\n", + " self.mark_failed()\n", + " raise\n", + "\n", " print(f\"MLflow run completed. Run ID: {self.run_id}\")\n", " \n", " if self._auto_started:\n", @@ -1323,7 +1480,8 @@ "\n", " Reopens the run to log metrics, then closes it. Useful for\n", " post-training evaluation (e.g., validation set patch inference results).\n", - " NaN values are silently skipped (MLflow rejects them).\n", + " Values that are not finite are skipped: MLflow accepts NaN but stores inf\n", + " as the largest float, which would distort plots and comparisons.\n", "\n", " Args:\n", " metrics: Dictionary of metric name -> value pairs to log.\n", @@ -1332,8 +1490,8 @@ " raise RuntimeError(\"No MLflow run found. Train the model first.\")\n", "\n", " import math\n", - " valid = {k: v for k, v in metrics.items()\n", - " if not (isinstance(v, float) and math.isnan(v))}\n", + " values = {k: float(v) for k, v in metrics.items()}\n", + " valid = {k: v for k, v in values.items() if math.isfinite(v)}\n", "\n", " if valid:\n", " with mlflow.start_run(run_id=self.run_id):\n", @@ -1486,19 +1644,18 @@ " ModelTrackingCallback ready to use with learn.fit()\n", " \n", " Example:\n", - " >>> # Instead of this (6 manual params):\n", - " >>> # mlflow_callback = ModelTrackingCallback(\n", - " >>> # model_name=f\"{task}_{model._get_name()}\",\n", - " >>> # loss_function=loss_func.loss_func._get_name(),\n", - " >>> # item_tfms=item_tfms,\n", - " >>> # size=size,\n", - " >>> # target_spacing=target_spacing,\n", - " >>> # apply_reorder=True,\n", - " >>> # )\n", - " >>> # with mlflow.start_run(run_name=\"training\"):\n", - " >>> # learn.fit_one_cycle(30, lr, cbs=[mlflow_callback])\n", - " >>> \n", - " >>> # Patch workflow: architecture and decoding remain explicit.\n", + " >>> # Standard workflow: pass the complete inference contract.\n", + " >>> model_spec = make_model_spec(\"monai.unet_inner\", unet_kwargs)\n", + " >>> inference_config = standard_inference_config(\n", + " ... apply_reorder=True, target_spacing=target_spacing,\n", + " ... item_tfms=inference_transforms_to_specs(item_tfms),\n", + " ... output_spec=make_output_spec(\"binary_segmentation\", threshold=0.5))\n", + " >>> callback = create_mlflow_callback(\n", + " ... learn, experiment_name=\"Task02_Heart\",\n", + " ... model_spec=model_spec, inference_config=inference_config)\n", + " >>> learn.fit_flat_cos(200, lr, cbs=[callback, save_best])\n", + " >>>\n", + " >>> # Patch workflow: patch settings come from the DataLoaders; pass output_spec.\n", " >>> model_spec = make_model_spec(\"monai.unet\", unet_kwargs)\n", " >>> output_spec = make_output_spec(\"multiclass_segmentation\", classes=2)\n", " >>> callback = create_mlflow_callback(\n", @@ -1705,9 +1862,10 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Test auto-extraction helper functions\n", "from fastcore.test import test_eq, test_fail\n", - "from dataclasses import dataclass\n", + "from fastMONAI.vision_patch import PatchConfig as _PatchConfig\n", "\n", "# Test _detect_patch_workflow\n", "class MockStandardDls:\n", @@ -1716,27 +1874,9 @@ "mock_std = MockStandardDls()\n", "test_eq(_detect_patch_workflow(mock_std), False)\n", "\n", - "@dataclass\n", - "class MockPatchConfig:\n", - " patch_size: list = None\n", - " patch_overlap: float = 0.5\n", - " samples_per_volume: int = 8\n", - " sampler_type: str = 'uniform'\n", - " label_probabilities: dict = None\n", - " queue_length: int = 300\n", - " aggregation_mode: str = 'hann'\n", - " padding_mode: int = 0\n", - " keep_largest_component: bool = False\n", - " apply_reorder: bool = True\n", - " target_spacing: list = None\n", - " \n", - " def __post_init__(self):\n", - " if self.patch_size is None:\n", - " self.patch_size = [96, 96, 96]\n", - "\n", "class MockPatchDls:\n", " bs = 4\n", - " patch_config = MockPatchConfig()\n", + " patch_config = _PatchConfig()\n", "mock_patch = MockPatchDls()\n", "test_eq(_detect_patch_workflow(mock_patch), True)\n", "\n", @@ -1772,13 +1912,13 @@ "metadata": {}, "outputs": [], "source": [ - "#| hide\n", "#| hide\n", "# Artifact policy: unwrapped model, weights-only .pth plus Safetensors, no Learner pickle.\n", "import tempfile as _tf\n", "from pathlib import Path as _P\n", "from unittest.mock import patch as _patch\n", "from fastai.callback.tracker import TrackerCallback\n", + "from fastcore.foundation import L\n", "from fastcore.test import test_eq\n", "import torch.nn as _nn\n", "\n", @@ -1838,28 +1978,31 @@ "test_eq(_stub.model is _wrapped, True)\n", "test_eq(_cb in _stub.cbs, True)\n", "\n", - "#| hide\n", "# Safetensors contract: inference-only patch config, shared storage, strict rebuild/load, metadata.\n", "import tempfile as _tf\n", "from pathlib import Path as _P\n", "import torch as _torch\n", "import torch.nn as _nn\n", "from fastcore.test import test_eq\n", + "import dataclasses as _dataclasses\n", "from fastMONAI.vision_patch import PatchConfig as _PatchConfig\n", "\n", "\n", - "_patch = _PatchConfig(\n", + "_patch_cfg = _PatchConfig(\n", " patch_size=[16, 16, 16], patch_overlap=0.5, aggregation_mode='hann',\n", " apply_reorder=False, target_spacing=[1., 2., 3.], sampler_type='label',\n", " label_probabilities={0: 0.2, 1: 0.8}, samples_per_volume=2,\n", " queue_length=4, queue_num_workers=1, padding_mode='minimum',\n", " keep_largest_component=True, binary_threshold=0.4, preprocessed=True,\n", " normalization=[])\n", - "_patch_dict = patch_config_to_dict(_patch)\n", + "_patch_dict = patch_config_to_dict(_patch_cfg)\n", "test_eq(set(_patch_dict), set(_PATCH_CONFIG_DEFAULTS))\n", + "# The fallback defaults cover every PatchConfig field and match its default values.\n", + "test_eq({_f.name for _f in _dataclasses.fields(_PatchConfig)}, set(_PATCH_CONFIG_DEFAULTS))\n", + "test_eq(patch_config_to_dict(_PatchConfig()), _PATCH_CONFIG_DEFAULTS)\n", "_out_spec = make_output_spec('multiclass_segmentation', classes=2)\n", - "_infer = patch_inference_config(_patch, _out_spec)\n", - "_patch_inference_dict = patch_config_to_dict(_patch, inference_only=True)\n", + "_infer = patch_inference_config(_patch_cfg, _out_spec)\n", + "_patch_inference_dict = patch_config_to_dict(_patch_cfg, inference_only=True)\n", "test_eq(set(_patch_inference_dict), set(_PATCH_INFERENCE_FIELDS))\n", "test_eq(_infer['patch_config'], _patch_inference_dict)\n", "_training_only = {\n", @@ -1997,6 +2140,45 @@ " )\n" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "getxsplitidentitytest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# A MedDataBlock-style split on an image/label DataFrame is identified by the get_x column.\n", + "from types import SimpleNamespace as _SimpleNamespace\n", + "from fastai.data.block import DataBlock as _DataBlock, TransformBlock as _TransformBlock\n", + "from fastai.data.transforms import IndexSplitter as _IndexSplitter\n", + "from fastcore.test import test_eq, test_fail\n", + "\n", + "_image_label = pd.DataFrame({\n", + " 'image': ['img/a.nii', 'img/b.nii', 'img/c.nii'],\n", + " 'label': ['lbl/a.nii', 'lbl/b.nii', 'lbl/c.nii'],\n", + "})\n", + "\n", + "\n", + "def _image_label_dls(get_x):\n", + " _dsets = _DataBlock(\n", + " blocks=(_TransformBlock, _TransformBlock), get_x=get_x, get_y=ColReader('label'),\n", + " splitter=_IndexSplitter([1]),\n", + " ).datasets(_image_label)\n", + " return _SimpleNamespace(train_ds=_dsets.train, valid_ds=_dsets.valid)\n", + "\n", + "\n", + "_version, _audit, _identity = _split_version(_image_label_dls(ColReader('image')))\n", + "test_eq(_identity, 'image')\n", + "test_eq(_audit.set_index('sample_id').is_valid.to_dict(),\n", + " {'img/a.nii': False, 'img/b.nii': True, 'img/c.nii': False})\n", + "test_eq(_version, _split_version(_image_label_dls(ColReader('image')), 'image')[0])\n", + "\n", + "# Without a ColReader get_x there is no stable identity, and the error says how to fix it.\n", + "test_fail(lambda: _split_version(_image_label_dls(lambda row: row['image'])),\n", + " contains='Pass sample_id_col')\n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -2005,13 +2187,147 @@ "outputs": [], "source": "#| hide\nfrom types import SimpleNamespace as _SimpleNamespace\nfrom unittest.mock import MagicMock as _MagicMock, patch as _patch\nfrom fastcore.test import test_eq as _test_eq\n\n# An owned active run is failed once, even if failure handling is repeated.\n_owned_active = object.__new__(ModelTrackingCallback)\n_owned_active._owns_run = True\n_owned_active._auto_started = True\n_owned_active._run_failed = False\n_owned_active.run_id = 'owned-run'\n_active_run = _SimpleNamespace(info=_SimpleNamespace(run_id='owned-run'))\nwith (\n _patch.object(mlflow, 'active_run', return_value=_active_run),\n _patch.object(mlflow, 'end_run') as _end_run,\n):\n _owned_active.mark_failed()\n _owned_active.mark_failed()\n_end_run.assert_called_once_with(status='FAILED')\n_test_eq(_owned_active._auto_started, False)\n_test_eq(_owned_active._run_failed, True)\n\n# An owned run that is already closed is marked failed through the client.\n_owned_closed = object.__new__(ModelTrackingCallback)\n_owned_closed._owns_run = True\n_owned_closed._auto_started = False\n_owned_closed._run_failed = False\n_owned_closed.run_id = 'owned-run'\n_client = _MagicMock()\nwith (\n _patch.object(mlflow, 'active_run', return_value=None),\n _patch.object(mlflow.tracking, 'MlflowClient', return_value=_client),\n _patch.object(mlflow, 'end_run') as _end_run,\n):\n _owned_closed.mark_failed()\n_client.set_terminated.assert_called_once_with('owned-run', status='FAILED')\n_end_run.assert_not_called()\n_test_eq(_owned_closed._run_failed, True)\n\n# A caller-owned run is never modified by the callback.\n_external = object.__new__(ModelTrackingCallback)\n_external._owns_run = False\n_external._auto_started = False\n_external._run_failed = False\n_external.run_id = 'external-run'\nwith (\n _patch.object(mlflow, 'active_run') as _active_run,\n _patch.object(mlflow, 'end_run') as _end_run,\n _patch.object(mlflow.tracking, 'MlflowClient') as _client_factory,\n):\n _external.mark_failed()\n_active_run.assert_not_called()\n_end_run.assert_not_called()\n_client_factory.assert_not_called()\n_test_eq(_external._run_failed, False)\n" }, + { + "cell_type": "code", + "execution_count": null, + "id": "logmetricsfinitetest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# log_metrics skips values that are not finite, including numpy and tensor scalars.\n", + "import numpy as _np\n", + "from unittest.mock import patch as _patch\n", + "from fastcore.test import test_eq\n", + "\n", + "_metrics_cb = object.__new__(ModelTrackingCallback)\n", + "_metrics_cb.run_id = 'run-1'\n", + "with _patch.object(mlflow, 'start_run') as _start_run, \\\n", + " _patch.object(mlflow, 'log_metrics') as _log_metrics:\n", + " _metrics_cb.log_metrics({\n", + " 'dice': 0.9, 'hd95': float('inf'), 'nsd': float('nan'),\n", + " 'np_ok': _np.float32(0.5), 'np_inf': _np.float64('-inf'),\n", + " 'torch_ok': torch.tensor(2.0), 'torch_nan': torch.tensor(float('nan')),\n", + " })\n", + "_start_run.assert_called_once_with(run_id='run-1')\n", + "test_eq(_log_metrics.call_args.args[0], {'dice': 0.9, 'np_ok': 0.5, 'torch_ok': 2.0})\n", + "\n", + "with _patch.object(mlflow, 'start_run') as _start_run:\n", + " _metrics_cb.log_metrics({'hd95': float('inf')})\n", + "_start_run.assert_not_called()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "runlifecycletest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# A run the callback starts never stays active after an error in before_fit, and\n", + "# auto_start=False without an active run gets an owned run instead of an implicit one.\n", + "# Uses a throwaway sqlite store.\n", + "import tempfile as _tempfile\n", + "import torch.nn as _nn\n", + "from types import SimpleNamespace as _SimpleNamespace\n", + "from unittest.mock import patch as _patch\n", + "from fastcore.test import test_eq, test_fail\n", + "\n", + "register_model_constructor(\n", + " 'test.sequential', lambda: _nn.Sequential(_nn.Linear(3, 3)), replace=True)\n", + "\n", + "\n", + "def _tracking_cb(**kwargs):\n", + " _cb = ModelTrackingCallback(\n", + " model_name='m', loss_function='l', item_tfms=[], size=[8, 8, 8],\n", + " target_spacing=[1., 1., 1.], apply_reorder=True,\n", + " model_spec=make_model_spec('test.sequential', {}),\n", + " inference_config=standard_inference_config(\n", + " True, [1., 1., 1.], [], make_output_spec('regression', output_shape=[3])),\n", + " **kwargs)\n", + " # These DataLoaders expose no split, so logging the split raises in before_fit.\n", + " _cb.learn = _SimpleNamespace(n_epoch=1, lr=1e-3, opt_func=torch.optim.SGD,\n", + " dls=_SimpleNamespace(bs=2))\n", + " return _cb\n", + "\n", + "\n", + "_saved_set = mlflow.is_tracking_uri_set()\n", + "_saved_uri = mlflow.get_tracking_uri() if _saved_set else None\n", + "_saved_flag = _FASTMONAI_TRACKING_URI_SET\n", + "with _tempfile.TemporaryDirectory() as _d:\n", + " try:\n", + " mlflow.set_tracking_uri(f'sqlite:///{_d}/mlflow.db')\n", + " _client = mlflow.tracking.MlflowClient()\n", + " _no_split = 'train/validation membership'\n", + "\n", + " # auto_start=True: the run started in before_fit is ended as FAILED.\n", + " _cb = _tracking_cb(auto_start=True)\n", + " test_fail(_cb.before_fit, contains=_no_split)\n", + " test_eq(mlflow.active_run(), None)\n", + " test_eq(_client.get_run(_cb.run_id).info.status, 'FAILED')\n", + "\n", + " # auto_start=False with no active run: the callback starts, owns and ends its run.\n", + " _cb = _tracking_cb(log_split=False)\n", + " _cb.before_fit()\n", + " test_eq(mlflow.active_run().info.run_id, _cb.run_id)\n", + " with _patch.object(_cb, '_save_model_artifacts'):\n", + " _cb.after_fit()\n", + " test_eq(mlflow.active_run(), None)\n", + " test_eq(_client.get_run(_cb.run_id).info.status, 'FINISHED')\n", + "\n", + " _cb = _tracking_cb()\n", + " test_fail(_cb.before_fit, contains=_no_split)\n", + " test_eq(mlflow.active_run(), None)\n", + " test_eq(_client.get_run(_cb.run_id).info.status, 'FAILED')\n", + "\n", + " # A failure while saving artifacts ends the owned run as FAILED.\n", + " _cb = _tracking_cb(log_split=False)\n", + " _cb.before_fit()\n", + " with _patch.object(_cb, '_save_model_artifacts', side_effect=OSError('disk full')):\n", + " test_fail(_cb.after_fit, contains='disk full')\n", + " test_eq(mlflow.active_run(), None)\n", + " test_eq(_client.get_run(_cb.run_id).info.status, 'FAILED')\n", + "\n", + " # A fit that raised during training left its run open; the next fit ends it.\n", + " _cb = _tracking_cb(auto_start=True, log_split=False)\n", + " _cb.before_fit()\n", + " _stale_run_id = _cb.run_id\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", + " mlflow.end_run()\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", + " test_fail(_cb.before_fit, contains=_no_split)\n", + " test_eq(mlflow.active_run().info.run_id, _external.info.run_id)\n", + " test_eq(_client.get_run(_external.info.run_id).info.status, 'FINISHED')\n", + " finally:\n", + " while mlflow.active_run() is not None:\n", + " mlflow.end_run()\n", + " mlflow.set_tracking_uri(_saved_uri if _saved_set else None)\n", + " _FASTMONAI_TRACKING_URI_SET = _saved_flag\n" + ] + }, + { + "cell_type": "markdown", + "id": "mlflowuimd", + "metadata": {}, + "source": [ + "## MLflow UI\n", + "\n", + "`MLflowUIManager` starts a local `mlflow ui` server from a notebook and stops it with the kernel." + ] + }, { "cell_type": "code", "execution_count": null, "id": "a3d2d585", "metadata": {}, "outputs": [], - "source": "#| export\n\nimport subprocess\nimport threading\nimport time\nimport socket\nimport os\nimport sys\nimport atexit\nimport signal\nimport ctypes\nfrom IPython.display import display, HTML\nimport shutil\n\n# Ask the OS to SIGTERM the `mlflow ui` child the instant this (parent) process dies;\n# this covers a hard-killed kernel (SIGKILL/OOM), which atexit cannot. Linux-only;\n# a no-op elsewhere. libc is loaded once here, outside any fork, so the preexec_fn hook\n# below stays minimal (preexec_fn is documented as unsafe in multithreaded programs, and\n# we launch mlflow ui from a daemon thread, so keeping the hook tiny is the accepted pattern).\n_libc = None\nif sys.platform.startswith('linux'):\n try:\n _libc = ctypes.CDLL('libc.so.6', use_errno=True)\n except OSError:\n _libc = None\n\n\ndef _die_with_parent():\n \"\"\"preexec hook: runs in the child after fork(), before exec(). Sets Linux PR_SET_PDEATHSIG.\"\"\"\n if _libc is not None:\n _libc.prctl(1, signal.SIGTERM) # PR_SET_PDEATHSIG = 1\n\n\nclass MLflowUIManager:\n \"\"\"Launch and manage a local ``mlflow ui`` server from a notebook.\n\n The UI's lifetime is tied to the kernel: it is reaped when the interpreter\n exits, gracefully (via an atexit handler) or on a hard kill (via Linux\n PR_SET_PDEATHSIG, see ``_die_with_parent``), so closing/restarting the\n notebook never leaves an orphaned server holding the port. Call ``stop()``\n to shut it down sooner. An externally-started UI is reused, not killed.\n \"\"\"\n\n def __init__(self):\n self.process = None\n self.thread = None\n self.port = 5001\n self.host = '0.0.0.0'\n self.backend_store_uri = _resolve_backend_store_uri()\n self._owns_process = False\n self._reusing_external = False\n # Reap the mlflow ui subprocess on clean interpreter shutdown (graceful kernel\n # shutdown/restart runs atexit). PR_SET_PDEATHSIG covers hard kills.\n atexit.register(self._kill_process)\n \n def is_port_available(self, port):\n \"\"\"Check if a port is available.\"\"\"\n with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n try:\n s.bind(('localhost', port))\n return True\n except OSError:\n return False\n \n def is_mlflow_running(self):\n \"\"\"Check if MLflow UI is actually responding.\"\"\"\n import requests\n try:\n response = requests.get(f'http://localhost:{self.port}', timeout=2)\n return response.status_code == 200\n except (requests.RequestException, ConnectionError, OSError):\n return False\n \n def _find_running_mlflow(self):\n \"\"\"Find an existing MLflow UI in our port range. Returns port or None.\"\"\"\n import requests\n for port in range(self.port, self.port + 10):\n if not self.is_port_available(port):\n try:\n response = requests.get(f'http://localhost:{port}', timeout=2)\n if response.status_code == 200 and 'mlflow' in response.text.lower():\n return port\n except (requests.RequestException, ConnectionError, OSError):\n continue\n return None\n \n def find_available_port(self, start_port=5001):\n \"\"\"Find an available port starting from start_port.\"\"\"\n for port in range(start_port, start_port + 10):\n if self.is_port_available(port):\n return port\n return None\n \n def check_mlflow_installed(self):\n \"\"\"Check if MLflow is installed.\"\"\"\n return shutil.which('mlflow') is not None\n \n def start_ui(self, auto_open=True, quiet=False):\n \"\"\"Start MLflow UI with better error handling and user feedback.\"\"\"\n # Refresh in case the tracking URI was configured after __init__\n self.backend_store_uri = _resolve_backend_store_uri()\n \n if not self.check_mlflow_installed():\n if not quiet:\n display(HTML('
MLflow not installed. Run: pip install mlflow
'))\n return False\n \n existing_port = self._find_running_mlflow()\n if existing_port is not None:\n self.port = existing_port\n self._owns_process = False\n self._reusing_external = True\n if not quiet:\n display(HTML(f'''\n \n '''))\n else:\n display(HTML(f'''\n \n MLflow UI (Port {self.port}, reused)\n \n '''))\n return True\n \n available_port = self.find_available_port(self.port)\n if available_port is None:\n if not quiet:\n display(HTML('
No available ports found (5001-5010)
'))\n return False\n \n self.port = available_port\n \n def run_mlflow():\n try:\n self.process = subprocess.Popen([\n 'mlflow', 'ui', \n '--host', self.host,\n '--port', str(self.port),\n '--backend-store-uri', self.backend_store_uri\n ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n preexec_fn=_die_with_parent if _libc is not None else None)\n self.process.wait()\n except Exception as e:\n if not quiet:\n display(HTML(f'
Error: {str(e)}
'))\n \n self.thread = threading.Thread(target=run_mlflow, daemon=True)\n self.thread.start()\n \n max_wait = 10\n for i in range(max_wait):\n time.sleep(1)\n if self.is_mlflow_running():\n self._owns_process = True\n self._reusing_external = False\n if quiet:\n display(HTML(f'''\n \n MLflow UI (Port {self.port})\n \n '''))\n else:\n display(HTML(f'''\n
\n
\n MLflow UI is running successfully!\n
\n \n Open MLflow UI\n \n
\n
URL: http://localhost:{self.port}
\n
\n
\n '''))\n return True\n \n if not quiet:\n display(HTML('
Failed to start MLflow UI
'))\n return False\n\n def _kill_process(self):\n \"\"\"Terminate the mlflow ui subprocess: terminate -> wait(timeout) -> kill.\n\n Quiet (no IPython display) so it is safe to run during interpreter shutdown,\n and idempotent. Used by both ``stop()`` and the atexit handler. Only acts on a\n process this manager started (``_owns_process``); never an externally-started UI.\n \"\"\"\n proc = self.process\n if proc is None or not self._owns_process:\n return\n self.process, self._owns_process = None, False\n try:\n proc.terminate()\n try:\n proc.wait(timeout=5)\n except subprocess.TimeoutExpired:\n proc.kill()\n proc.wait(timeout=5)\n except Exception:\n pass\n\n def stop(self):\n \"\"\"Stop the MLflow UI server (terminate -> wait -> kill).\"\"\"\n if self.process and self._owns_process:\n self._kill_process()\n display(HTML('''\n
\n MLflow UI stopped\n
\n '''))\n elif self._reusing_external:\n self._reusing_external = False\n display(HTML('''\n
\n MLflow UI was started externally, not stopping it\n
\n '''))\n else:\n display(HTML('''\n
\n MLflow UI is not currently running\n
\n '''))\n \n def status(self):\n \"\"\"Check MLflow UI status.\"\"\"\n if self.is_mlflow_running():\n display(HTML(f'''\n
\n
MLflow UI is running
\n \n http://localhost:{self.port}\n \n
\n '''))\n else:\n display(HTML('''\n
\n
MLflow UI is not running
\n
\n Run mlflow_ui.start_ui() to start it.\n
\n
\n '''))" + "source": "#| export\n\nimport subprocess\nimport time\nimport socket\nimport sys\nimport atexit\nimport signal\nimport ctypes\nimport html\n\n# Ask the OS to SIGTERM the `mlflow ui` child the instant this (parent) process dies;\n# this covers a hard-killed kernel (SIGKILL/OOM), which atexit cannot. Linux-only;\n# a no-op elsewhere. libc is loaded once here, outside any fork, so the preexec_fn hook\n# below stays minimal (preexec_fn is documented as unsafe in multithreaded programs, and\n# a Jupyter kernel runs several threads, so keeping the hook tiny is the accepted pattern).\n_libc = None\nif sys.platform.startswith('linux'):\n try:\n _libc = ctypes.CDLL('libc.so.6', use_errno=True)\n except OSError:\n _libc = None\n\n\ndef _die_with_parent():\n \"\"\"preexec hook: runs in the child after fork(), before exec(). Sets Linux PR_SET_PDEATHSIG.\"\"\"\n if _libc is not None:\n _libc.prctl(1, signal.SIGTERM) # PR_SET_PDEATHSIG = 1\n\n\nclass MLflowUIManager:\n \"\"\"Launch and manage a local ``mlflow ui`` server from a notebook.\n\n The UI's lifetime is tied to the kernel: it is reaped when the interpreter\n exits, gracefully (via an atexit handler) or on a hard kill (via Linux\n PR_SET_PDEATHSIG, see ``_die_with_parent``), so closing/restarting the\n notebook never leaves an orphaned server holding the port. Call ``stop()``\n to shut it down sooner. An externally-started UI is reused, not killed, and\n calling ``start_ui()`` again while its own server runs just shows the link.\n \"\"\"\n\n def __init__(self):\n self.process = None\n self.port = 5001\n self.host = '0.0.0.0'\n self.backend_store_uri = _resolve_backend_store_uri()\n self._owns_process = False\n self._reusing_external = False\n self._log_path = None\n # Reap the mlflow ui subprocess on clean interpreter shutdown (graceful kernel\n # shutdown/restart runs atexit). PR_SET_PDEATHSIG covers hard kills.\n atexit.register(self._kill_process)\n \n def is_port_available(self, port):\n \"\"\"Check if a port is available.\"\"\"\n with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n try:\n s.bind(('localhost', port))\n return True\n except OSError:\n return False\n \n def is_mlflow_running(self):\n \"\"\"Check if MLflow UI is actually responding.\"\"\"\n import requests\n try:\n response = requests.get(f'http://localhost:{self.port}', timeout=2)\n return response.status_code == 200\n except (requests.RequestException, ConnectionError, OSError):\n return False\n \n def _find_running_mlflow(self):\n \"\"\"Find an existing MLflow UI in our port range. Returns port or None.\"\"\"\n import requests\n for port in range(self.port, self.port + 10):\n if not self.is_port_available(port):\n try:\n response = requests.get(f'http://localhost:{port}', timeout=2)\n if response.status_code == 200 and 'mlflow' in response.text.lower():\n return port\n except (requests.RequestException, ConnectionError, OSError):\n continue\n return None\n \n def find_available_port(self, start_port=5001):\n \"\"\"Find an available port starting from start_port.\"\"\"\n for port in range(start_port, start_port + 10):\n if self.is_port_available(port):\n return port\n return None\n \n def _log_tail(self, max_chars=2000):\n \"\"\"End of the server's stderr log, or '' when there is none.\"\"\"\n try:\n return Path(self._log_path).read_text(errors='replace')[-max_chars:].strip()\n except (OSError, TypeError):\n return ''\n\n def _remove_log(self):\n \"\"\"Delete the server's stderr log file, if any.\"\"\"\n if self._log_path is not None:\n try:\n os.remove(self._log_path)\n except OSError:\n pass\n self._log_path = None\n\n def start_ui(self, quiet=False):\n \"\"\"Start the MLflow UI, or show its link if it is already running.\n\n The server runs with this interpreter (``python -m mlflow ui``), so it uses the\n same MLflow install as the notebook. Its stderr goes to a temporary log file,\n and the end of that log is shown if the server fails to start.\n \"\"\"\n from IPython.display import display, HTML\n # Refresh in case the tracking URI was configured after __init__\n self.backend_store_uri = _resolve_backend_store_uri()\n\n if self._owns_process and self.process is not None and self.process.poll() is None:\n display(HTML(f'''\n \n MLflow UI (Port {self.port})\n \n '''))\n return True\n\n existing_port = self._find_running_mlflow()\n if existing_port is not None:\n self.port = existing_port\n self._owns_process = False\n self._reusing_external = True\n if not quiet:\n display(HTML(f'''\n
\n
\n Reusing existing MLflow UI on port {self.port}\n
\n \n Open MLflow UI\n \n
\n
URL: http://localhost:{self.port}
\n
\n
\n '''))\n else:\n display(HTML(f'''\n \n MLflow UI (Port {self.port}, reused)\n \n '''))\n return True\n \n available_port = self.find_available_port(self.port)\n if available_port is None:\n if not quiet:\n display(HTML('
No available ports found (5001-5010)
'))\n return False\n \n self.port = available_port\n\n # stderr goes to a file rather than a pipe: nobody drains a pipe while the\n # server runs, so a full pipe buffer would block it.\n self._remove_log()\n log_file = tempfile.NamedTemporaryFile(\n mode='wb', prefix='mlflow_ui_', suffix='.log', delete=False)\n self._log_path = log_file.name\n try:\n with log_file:\n self.process = subprocess.Popen([\n sys.executable, '-m', 'mlflow', 'ui',\n '--host', self.host,\n '--port', str(self.port),\n '--backend-store-uri', self.backend_store_uri\n ], stdout=subprocess.DEVNULL, stderr=log_file,\n preexec_fn=_die_with_parent if _libc is not None else None)\n except Exception as e:\n self._remove_log()\n if not quiet:\n display(HTML(f'
Error: {html.escape(str(e))}
'))\n return False\n self._owns_process = True\n self._reusing_external = False\n\n max_wait = 10\n for i in range(max_wait):\n time.sleep(1)\n if self.is_mlflow_running():\n if quiet:\n display(HTML(f'''\n \n MLflow UI (Port {self.port})\n \n '''))\n else:\n display(HTML(f'''\n
\n
\n MLflow UI is running successfully!\n
\n \n Open MLflow UI\n \n
\n
URL: http://localhost:{self.port}
\n
\n
\n '''))\n return True\n if self.process.poll() is not None:\n break\n\n exited = self.process.poll() is not None\n log_tail = self._log_tail()\n if exited:\n self.process, self._owns_process = None, False\n self._remove_log()\n if not quiet:\n reason = 'the server exited' if exited else f'no response after {max_wait} s'\n details = (f'
{html.escape(log_tail)}
'\n if log_tail else '')\n display(HTML(f'
Failed to start MLflow UI ({reason})
{details}'))\n return False\n\n def _kill_process(self):\n \"\"\"Terminate the mlflow ui subprocess: terminate -> wait(timeout) -> kill.\n\n Quiet (no IPython display) so it is safe to run during interpreter shutdown,\n and idempotent. Used by both ``stop()`` and the atexit handler. Only acts on a\n process this manager started (``_owns_process``); never an externally-started UI.\n \"\"\"\n proc = self.process\n if proc is None or not self._owns_process:\n return\n self.process, self._owns_process = None, False\n try:\n proc.terminate()\n try:\n proc.wait(timeout=5)\n except subprocess.TimeoutExpired:\n proc.kill()\n proc.wait(timeout=5)\n except Exception:\n pass\n self._remove_log()\n\n def stop(self):\n \"\"\"Stop the MLflow UI server (terminate -> wait -> kill).\"\"\"\n from IPython.display import display, HTML\n if self.process and self._owns_process:\n self._kill_process()\n display(HTML('''\n
\n MLflow UI stopped\n
\n '''))\n elif self._reusing_external:\n self._reusing_external = False\n display(HTML('''\n
\n MLflow UI was started externally, not stopping it\n
\n '''))\n else:\n display(HTML('''\n
\n MLflow UI is not currently running\n
\n '''))\n \n def status(self):\n \"\"\"Check MLflow UI status.\"\"\"\n from IPython.display import display, HTML\n if self.is_mlflow_running():\n display(HTML(f'''\n
\n
MLflow UI is running
\n \n http://localhost:{self.port}\n \n
\n '''))\n else:\n display(HTML('''\n
\n
MLflow UI is not running
\n
\n Run mlflow_ui.start_ui() to start it.\n
\n
\n '''))" }, { "cell_type": "code", @@ -2020,6 +2336,55 @@ "metadata": {}, "outputs": [], "source": "#| hide\n# MLflowUIManager._kill_process must terminate -> wait -> kill (escalating on\n# timeout) and be idempotent -- verified without launching a real `mlflow ui` server.\nfrom unittest.mock import MagicMock\nfrom fastcore.test import test_eq\n\n# Case 1: process exits on terminate() within the timeout -> no kill(); idempotent.\n_m = MLflowUIManager()\n_proc = MagicMock()\n_proc.wait.return_value = 0\n_m.process, _m._owns_process = _proc, True\n_m._kill_process()\n_proc.terminate.assert_called_once()\n_proc.kill.assert_not_called()\ntest_eq(_m.process, None)\ntest_eq(_m._owns_process, False)\n_m._kill_process() # idempotent: no-op, no extra terminate/kill\n_proc.terminate.assert_called_once()\n\n# Case 2: process hangs past the timeout -> escalate to kill().\n_m2 = MLflowUIManager()\n_proc2 = MagicMock()\n_proc2.wait.side_effect = [subprocess.TimeoutExpired(cmd='mlflow', timeout=5), 0]\n_m2.process, _m2._owns_process = _proc2, True\n_m2._kill_process()\n_proc2.terminate.assert_called_once()\n_proc2.kill.assert_called_once()\n\n# An externally-started UI (not owned) is never killed.\n_m3 = MLflowUIManager()\n_proc3 = MagicMock()\n_m3.process, _m3._owns_process, _m3._reusing_external = _proc3, False, True\n_m3._kill_process()\n_proc3.terminate.assert_not_called()\n\nprint(\"MLflowUIManager._kill_process terminate/wait/kill tests passed!\")" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "mlflowuistarttest", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# start_ui keeps ownership of a server it started, launches mlflow with this\n", + "# interpreter, and shows the end of the server's stderr when it fails to start.\n", + "from unittest.mock import MagicMock, patch as _patch\n", + "from fastcore.test import test_eq\n", + "\n", + "# A second start_ui() while its own server runs only shows the link.\n", + "_m4 = MLflowUIManager()\n", + "_proc4 = MagicMock()\n", + "_proc4.poll.return_value = None\n", + "_m4.process, _m4._owns_process = _proc4, True\n", + "with _patch.object(_m4, '_find_running_mlflow') as _find, \\\n", + " _patch.object(subprocess, 'Popen') as _popen:\n", + " test_eq(_m4.start_ui(quiet=True), True)\n", + "_find.assert_not_called()\n", + "_popen.assert_not_called()\n", + "test_eq((_m4.process, _m4._owns_process, _m4._reusing_external), (_proc4, True, False))\n", + "\n", + "\n", + "def _fake_popen(cmd, stdout=None, stderr=None, preexec_fn=None):\n", + " stderr.write(b'Error: backend store \\n')\n", + " stderr.flush()\n", + " _proc = MagicMock()\n", + " _proc.poll.return_value = 1\n", + " return _proc\n", + "\n", + "\n", + "# A server that exits at startup is reported with its stderr, and is no longer owned.\n", + "_m5 = MLflowUIManager()\n", + "with _patch.object(_m5, '_find_running_mlflow', return_value=None), \\\n", + " _patch.object(_m5, 'find_available_port', return_value=5099), \\\n", + " _patch.object(_m5, 'is_mlflow_running', return_value=False), \\\n", + " _patch.object(subprocess, 'Popen', side_effect=_fake_popen) as _popen, \\\n", + " _patch.object(time, 'sleep'), \\\n", + " _patch('IPython.display.display') as _display:\n", + " test_eq(_m5.start_ui(), False)\n", + "test_eq(_popen.call_args.args[0][:4], [sys.executable, '-m', 'mlflow', 'ui'])\n", + "_shown = _display.call_args.args[0].data\n", + "assert 'the server exited' in _shown and 'backend store <unreachable>' in _shown, _shown\n", + "test_eq((_m5.process, _m5._owns_process, _m5._log_path), (None, False, None))\n" + ] } ], "metadata": { From 45a10b40c45361bca69dbd5a4ec666514ba67368 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:25:47 +0100 Subject: [PATCH 05/11] refactor(vision_patch)!: remove the unimplemented weighted sampler sampler_type='weighted' was accepted and documented but always raised NotImplementedError. PatchConfig now rejects it up front. MedPatchDataLoaders built directly from two loaders report apply_reorder, target_spacing and patch_config from the train loader's config, so MLflow logs the right values. Also removed dead code (an unreachable affine guard, an unused attribute and loop variable), moved imports to the top and updated docstrings and hidden test cells. Prediction file names, return types and the private methods used by the VS PACS engine are unchanged. BREAKING CHANGE: sampler_type='weighted' now raises ValueError. --- fastMONAI/_modidx.py | 21 +++-- fastMONAI/vision_patch.py | 97 +++++++++++---------- nbs/10_vision_patch.ipynb | 173 +++++++++++++++++++++----------------- 3 files changed, 155 insertions(+), 136 deletions(-) diff --git a/fastMONAI/_modidx.py b/fastMONAI/_modidx.py index 9184852..f5010a6 100644 --- a/fastMONAI/_modidx.py +++ b/fastMONAI/_modidx.py @@ -22,8 +22,6 @@ 'fastMONAI/dataset_info.py'), 'fastMONAI.dataset_info.MedDataset._visualize_single_case': ( 'dataset_info.html#meddataset._visualize_single_case', 'fastMONAI/dataset_info.py'), - 'fastMONAI.dataset_info.MedDataset.calculate_target_size': ( 'dataset_info.html#meddataset.calculate_target_size', - 'fastMONAI/dataset_info.py'), 'fastMONAI.dataset_info.MedDataset.fingerprint': ( 'dataset_info.html#meddataset.fingerprint', 'fastMONAI/dataset_info.py'), 'fastMONAI.dataset_info.MedDataset.get_size_statistics': ( 'dataset_info.html#meddataset.get_size_statistics', @@ -69,8 +67,6 @@ 'fastMONAI/external_data.py'), 'fastMONAI.external_data._dataset_available_locally': ( 'external_data.html#_dataset_available_locally', 'fastMONAI/external_data.py'), - 'fastMONAI.external_data._df_sort_and_add_columns': ( 'external_data.html#_df_sort_and_add_columns', - 'fastMONAI/external_data.py'), 'fastMONAI.external_data._process_ixi_xls': ( 'external_data.html#_process_ixi_xls', 'fastMONAI/external_data.py'), 'fastMONAI.external_data._process_medmnist_img': ( 'external_data.html#_process_medmnist_img', @@ -93,8 +89,10 @@ 'fastMONAI/utils.py'), 'fastMONAI.utils.MLflowUIManager._kill_process': ( 'utils.html#mlflowuimanager._kill_process', 'fastMONAI/utils.py'), - 'fastMONAI.utils.MLflowUIManager.check_mlflow_installed': ( 'utils.html#mlflowuimanager.check_mlflow_installed', - 'fastMONAI/utils.py'), + 'fastMONAI.utils.MLflowUIManager._log_tail': ( 'utils.html#mlflowuimanager._log_tail', + 'fastMONAI/utils.py'), + 'fastMONAI.utils.MLflowUIManager._remove_log': ( 'utils.html#mlflowuimanager._remove_log', + 'fastMONAI/utils.py'), 'fastMONAI.utils.MLflowUIManager.find_available_port': ( 'utils.html#mlflowuimanager.find_available_port', 'fastMONAI/utils.py'), 'fastMONAI.utils.MLflowUIManager.is_mlflow_running': ( 'utils.html#mlflowuimanager.is_mlflow_running', @@ -113,6 +111,8 @@ 'fastMONAI/utils.py'), 'fastMONAI.utils.ModelTrackingCallback._extract_training_params': ( 'utils.html#modeltrackingcallback._extract_training_params', 'fastMONAI/utils.py'), + 'fastMONAI.utils.ModelTrackingCallback._log_run_setup': ( 'utils.html#modeltrackingcallback._log_run_setup', + 'fastMONAI/utils.py'), 'fastMONAI.utils.ModelTrackingCallback._log_split_df': ( 'utils.html#modeltrackingcallback._log_split_df', 'fastMONAI/utils.py'), 'fastMONAI.utils.ModelTrackingCallback._save_model_artifacts': ( 'utils.html#modeltrackingcallback._save_model_artifacts', @@ -135,6 +135,7 @@ 'fastMONAI/utils.py'), 'fastMONAI.utils._actual_split_df': ('utils.html#_actual_split_df', 'fastMONAI/utils.py'), 'fastMONAI.utils._canonical_tracking_json': ('utils.html#_canonical_tracking_json', 'fastMONAI/utils.py'), + 'fastMONAI.utils._check_patch_threshold': ('utils.html#_check_patch_threshold', 'fastMONAI/utils.py'), 'fastMONAI.utils._detect_patch_workflow': ('utils.html#_detect_patch_workflow', 'fastMONAI/utils.py'), 'fastMONAI.utils._die_with_parent': ('utils.html#_die_with_parent', 'fastMONAI/utils.py'), 'fastMONAI.utils._download_model_artifact': ('utils.html#_download_model_artifact', 'fastMONAI/utils.py'), @@ -147,9 +148,11 @@ '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._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'), 'fastMONAI.utils._normalize_monai_kwargs': ('utils.html#_normalize_monai_kwargs', 'fastMONAI/utils.py'), + 'fastMONAI.utils._prefer_raw_column': ('utils.html#_prefer_raw_column', 'fastMONAI/utils.py'), 'fastMONAI.utils._read_preprocessing_manifest': ( 'utils.html#_read_preprocessing_manifest', 'fastMONAI/utils.py'), 'fastMONAI.utils._resolve_backend_store_uri': ( 'utils.html#_resolve_backend_store_uri', @@ -457,8 +460,12 @@ '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.py'), 'fastMONAI.vision_inference._resize_labels': ( 'vision_inference.html#_resize_labels', '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', @@ -551,8 +558,6 @@ 'fastMONAI/vision_metrics.py'), 'fastMONAI.vision_metrics.calculate_dsc': ( 'vision_metrics.html#calculate_dsc', 'fastMONAI/vision_metrics.py'), - 'fastMONAI.vision_metrics.calculate_haus': ( 'vision_metrics.html#calculate_haus', - 'fastMONAI/vision_metrics.py'), 'fastMONAI.vision_metrics.calculate_hd95': ( 'vision_metrics.html#calculate_hd95', 'fastMONAI/vision_metrics.py'), 'fastMONAI.vision_metrics.calculate_lesion_detection_rate': ( 'vision_metrics.html#calculate_lesion_detection_rate', diff --git a/fastMONAI/vision_patch.py b/fastMONAI/vision_patch.py index 5b530ae..2505e01 100644 --- a/fastMONAI/vision_patch.py +++ b/fastMONAI/vision_patch.py @@ -13,6 +13,7 @@ import pandas as pd import numpy as np import warnings +import numbers import copy as _copy import matplotlib.pyplot as plt from pathlib import Path @@ -21,10 +22,12 @@ from torch.utils.data import DataLoader from tqdm.auto import tqdm from fastai.data.all import * +# keep below the fastai star import, which exports fastcore's ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor 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 +from .vision_inference import _do_resize, _resize_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 @@ -240,7 +243,7 @@ class PatchConfig: - Int >= 1: pixel overlap (e.g., 48 = 48 pixel overlap) - List: per-dimension overlap in pixels samples_per_volume: Number of patches to extract per volume during training. - sampler_type: Type of sampler ('uniform', 'label', 'weighted'). + sampler_type: Type of sampler ('uniform' or 'label'). label_probabilities: For LabelSampler, dict mapping label values to probabilities. queue_length: Maximum number of patches to store in queue. queue_num_workers: Number of workers for parallel patch extraction. @@ -305,7 +308,7 @@ class PatchConfig: def __post_init__(self): """Validate configuration.""" - valid_samplers = ['uniform', 'label', 'weighted'] + valid_samplers = ['uniform', 'label'] if self.sampler_type not in valid_samplers: raise ValueError(f"sampler_type must be one of {valid_samplers}") @@ -511,13 +514,6 @@ def create_patch_sampler(config: PatchConfig) -> tio.data.PatchSampler: label_probabilities=config.label_probabilities ) - elif config.sampler_type == 'weighted': - raise NotImplementedError( - "WeightedSampler requires a pre-computed probability map which is not currently supported. " - "Use 'label' sampler with label_probabilities for weighted sampling based on segmentation labels, " - "or 'uniform' for random patch extraction." - ) - raise ValueError(f"Unknown sampler type: {config.sampler_type}") # %% ../nbs/10_vision_patch.ipynb #cell-14 @@ -536,7 +532,7 @@ class MedPatchDataLoader: raw TorchIO transforms. fastMONAI wrappers are automatically normalized to raw TorchIO for internal use. Mutually exclusive with gpu_augmentation. gpu_augmentation: GpuPatchAugmentation instance for GPU-batched augmentation. - Operates on [B,C,D,H,W] tensors already on GPU, avoiding per-sample CPU + Operates on [B, C, X, Y, Z] tensors already on GPU, avoiding per-sample CPU overhead. Mutually exclusive with patch_tfms. Training only. shuffle: Whether to shuffle subjects and patches. drop_last: Whether to drop last incomplete batch. @@ -594,7 +590,7 @@ def __init__( def __iter__(self): """Iterate over batches, yielding (image, mask) tuples.""" for batch in self._dl: - img = batch['image'][tio.DATA] # [B, C, H, W, D] + img = batch['image'][tio.DATA] # [B, C, X, Y, Z] has_mask = 'mask' in batch # Apply CPU patch transforms if provided (per-sample TorchIO loop) @@ -726,19 +722,17 @@ def _build_pre_patch_tfms(patch_config, apply_reorder, target_spacing, pre_patch def _stash_from_df_metadata(instance, img_col, mask_col, pre_patch_tfms, apply_reorder, - target_spacing, ensure_affine_consistency, patch_config, - train_df, valid_df): - """Record the de-facto metadata contract attributes consumed by research/vs_seg scripts. + target_spacing, patch_config, train_df, valid_df): + """Record the from_df inputs on the instance for MLflow logging and split tracking. - Stores the RAW pre_patch_tfms (not the normalized list), matching how the train_final*.py - scripts set dls._pre_patch_tfms manually. + Stores pre_patch_tfms as given (fastMONAI wrappers are not unwrapped to TorchIO). + Code that builds MedPatchDataLoaders directly can set the same attributes. """ instance._img_col = img_col instance._mask_col = mask_col instance._pre_patch_tfms = pre_patch_tfms instance._apply_reorder = apply_reorder instance._target_spacing = target_spacing - instance._ensure_affine_consistency = ensure_affine_consistency instance._patch_config = patch_config instance._train_source_df = train_df instance._valid_source_df = valid_df @@ -759,10 +753,10 @@ class MedPatchDataLoaders: for full-volume sliding window inference. Example: - >>> config = PatchConfig(patch_size=[96, 96, 96], apply_reorder=True, target_spacing=[0.5, 0.5, 0.5]) + >>> config = PatchConfig(patch_size=[96, 96, 96], apply_reorder=True, target_spacing=[0.5, 0.5, 0.5], + ... normalization=[ZNormalization(masking_method='foreground')]) >>> dls = MedPatchDataLoaders.from_df( - ... df, img_col='image', mask_col='label', valid_pct=0.2, - ... patch_config=config, pre_patch_tfms=[tio.ZNormalization()], bs=4 + ... df, img_col='image', mask_col='label', valid_pct=0.2, patch_config=config, bs=4 ... ) >>> learn = Learner(dls, model, loss_func=DiceLoss()) """ @@ -919,7 +913,7 @@ def from_df( instance = cls(train_dl, valid_dl, device) return _stash_from_df_metadata( instance, img_col, mask_col, _norm_tfms, _apply_reorder, _target_spacing, - ensure_affine_consistency, patch_config, train_df, valid_df + patch_config, train_df, valid_df ) @property @@ -954,18 +948,27 @@ def bs(self): @property def apply_reorder(self): - """Whether reordering to RAS+ is enabled.""" - return getattr(self, '_apply_reorder', False) + """Whether reordering to RAS+ is enabled (from from_df, else patch_config).""" + if hasattr(self, '_apply_reorder'): + return self._apply_reorder + config = self.patch_config + return config.apply_reorder if config is not None else False @property def target_spacing(self): - """Target voxel spacing for resampling.""" - return getattr(self, '_target_spacing', None) + """Target voxel spacing for resampling (from from_df, else patch_config).""" + if hasattr(self, '_target_spacing'): + return self._target_spacing + config = self.patch_config + return config.target_spacing if config is not None else None @property def patch_config(self): - """The PatchConfig used for this DataLoaders.""" - return getattr(self, '_patch_config', None) + """The PatchConfig used for this DataLoaders (from from_df, else the train loader's config).""" + config = getattr(self, '_patch_config', None) + if config is None: + config = getattr(getattr(self, '_train_dl', None), 'config', None) + return config @property def split_df(self): @@ -1024,7 +1027,10 @@ def cpu(self): def show_batch(self, dl_idx=0, max_n=6, figsize=None, channel=0, slice_index=None, anatomical_plane=0, overlay=False, voxel_size=None, **kwargs): - """Show a batch of patch samples for visualization.""" + """Show a batch of patch samples for visualization. + + Extra keyword arguments are accepted for compatibility with fastai's show_batch and ignored. + """ dl = self[dl_idx] x, y = dl.one_batch() @@ -1063,9 +1069,9 @@ def show_batch(self, dl_idx=0, max_n=6, figsize=None, channel=0, slice_idxs.extend([idx] * len(im_channels)) _voxel_size = voxel_size if voxel_size is not None else self.target_spacing - ctxs = [im.show(ax=ax, slice_index=idx, anatomical_plane=anatomical_plane, - voxel_size=_voxel_size) - for im, ax, idx in zip(imgs, flat_axs, slice_idxs)] + for im, ax, idx in zip(imgs, flat_axs, slice_idxs): + im.show(ax=ax, slice_index=idx, anatomical_plane=anatomical_plane, + voxel_size=_voxel_size) if overlay and has_mask: for mask, ax, idx in zip(masks_for_overlay, flat_axs, slice_idxs): @@ -1126,15 +1132,13 @@ def __del__(self): pass # %% ../nbs/10_vision_patch.ipynb #cell-17 -import numbers - def _normalize_patch_overlap(patch_overlap: 'int | float | list', patch_size: list) -> tuple: """Convert patch_overlap (fraction 0-1, pixel int, numpy scalar, or sequence) to a tuple of even pixel ints for TorchIO's GridSampler. Format conversion only; value validation (negatives, overlap >= patch_size) lives in PatchConfig.__post_init__(). """ # Scalar fraction (0 < x < 1); excludes 1.0 since 100% overlap creates step_size=0 (infinite patches) - if isinstance(patch_overlap, (int, float, numbers.Number)) and 0 < float(patch_overlap) < 1: + if isinstance(patch_overlap, numbers.Number) and 0 < float(patch_overlap) < 1: result = [] for ps in patch_size: pixels = int(int(ps) * float(patch_overlap)) @@ -1144,7 +1148,7 @@ def _normalize_patch_overlap(patch_overlap: 'int | float | list', patch_size: li return tuple(result) # Scalar int (incl. numpy scalars): pixel count - if isinstance(patch_overlap, (int, float, numbers.Number)): + if isinstance(patch_overlap, numbers.Number): val = int(patch_overlap) if val % 2 != 0: val = val - 1 if val > 0 else 0 @@ -1191,8 +1195,8 @@ def _resolve_tta(tta) -> tuple | None: def _logits_to_probs(logits: torch.Tensor) -> torch.Tensor: """Logits -> probabilities: sigmoid for 1 channel (binary), else softmax over the channel dim. - Upcasts to float32 (logits may be bfloat16 under AMP). Does NOT move to CPU; callers decide - placement (TTA keeps tensors on-device to flip back; the non-TTA path moves to CPU immediately). + Upcasts to float32 (logits may be bfloat16 under AMP). Does not move to CPU: callers sum the + passes (flips and models) on the device and move the average to CPU once. """ return torch.sigmoid(logits.float()) if logits.shape[1] == 1 else torch.softmax(logits.float(), dim=1) @@ -1201,8 +1205,9 @@ def _predict_patch_tta(models, patch_input, amp_context=None, flip_axes=_TTA_FLI """Mirror TTA (with optional model ensemble): average probabilities over `flip_axes` and all models. models is a single model or a list of models; a list soft-votes (averages probabilities). - Uses a running sum (2x memory, not 9x per model). Each pass: flip -> forward -> activate -> - flip back -> accumulate. Returns the averaged probability tensor [B, C, D, H, W] on CPU. + Keeps a running sum, so memory holds two probability tensors instead of one per pass + (8 per model with all flips). Each pass: flip -> forward -> activate -> flip back -> + accumulate. Returns the averaged probability tensor [B, C, X, Y, Z] on CPU. amp_context, if given, wraps each forward pass. """ if amp_context is None: amp_context = nullcontext() @@ -1466,7 +1471,6 @@ def _postprocess( f"models, or post-process each class separately." ) else: - from fastMONAI.vision_inference import keep_largest result = keep_largest(result.squeeze(0)).unsqueeze(0) if return_probabilities: @@ -1489,10 +1493,6 @@ def _postprocess( result = result.long() # Use original affine matrix for correct spatial alignment - if not (hasattr(prepared.org_img, 'affine') and prepared.org_img.affine is not None): - raise RuntimeError( - "org_img.affine not available. This should never happen - please report this bug." - ) affine = prepared.org_img.affine.copy() return result, affine @@ -1544,8 +1544,8 @@ def predict_mask_and_probabilities( Returns: ``(mask, probabilities)`` or ``(mask, probabilities, affine)`` when - ``return_affine=True``. The mask has shape ``[1, D, H, W]`` and dtype - ``torch.long``; probabilities have shape ``[C, D, H, W]`` and floating dtype. + ``return_affine=True``. The mask has shape ``[1, X, Y, Z]`` and dtype + ``torch.long``; probabilities have shape ``[C, X, Y, Z]`` and floating dtype. """ prepared = self._prepare_subject(img_path) output = self._run_inference(prepared, tta=tta) @@ -1563,9 +1563,6 @@ def to(self, device): return self # %% ../nbs/10_vision_patch.ipynb #cell-18 -from concurrent.futures import ThreadPoolExecutor - - def prediction_filename(input_path: str | Path) -> str: """Return the filename used by patch inference for a saved prediction.""" input_path = Path(input_path) diff --git a/nbs/10_vision_patch.ipynb b/nbs/10_vision_patch.ipynb index 55a79e5..cf1bd11 100644 --- a/nbs/10_vision_patch.ipynb +++ b/nbs/10_vision_patch.ipynb @@ -34,6 +34,7 @@ "import pandas as pd\n", "import numpy as np\n", "import warnings\n", + "import numbers\n", "import copy as _copy\n", "import matplotlib.pyplot as plt\n", "from pathlib import Path\n", @@ -42,10 +43,12 @@ "from torch.utils.data import DataLoader\n", "from tqdm.auto import tqdm\n", "from fastai.data.all import *\n", + "# keep below the fastai star import, which exports fastcore's ThreadPoolExecutor\n", + "from concurrent.futures import ThreadPoolExecutor\n", "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\n", + "from fastMONAI.vision_inference import _do_resize, _resize_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" @@ -262,6 +265,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Test _extract_tio_transform and normalize_patch_transforms\n", "from fastMONAI.vision_augmentation import RandomAffine, RandomGamma, RandomFlip, RandomNoise\n", "\n", @@ -338,7 +342,7 @@ " - Int >= 1: pixel overlap (e.g., 48 = 48 pixel overlap)\n", " - List: per-dimension overlap in pixels\n", " samples_per_volume: Number of patches to extract per volume during training.\n", - " sampler_type: Type of sampler ('uniform', 'label', 'weighted').\n", + " sampler_type: Type of sampler ('uniform' or 'label').\n", " label_probabilities: For LabelSampler, dict mapping label values to probabilities.\n", " queue_length: Maximum number of patches to store in queue.\n", " queue_num_workers: Number of workers for parallel patch extraction.\n", @@ -403,7 +407,7 @@ " \n", " def __post_init__(self):\n", " \"\"\"Validate configuration.\"\"\"\n", - " valid_samplers = ['uniform', 'label', 'weighted']\n", + " valid_samplers = ['uniform', 'label']\n", " if self.sampler_type not in valid_samplers:\n", " raise ValueError(f\"sampler_type must be one of {valid_samplers}\")\n", " \n", @@ -510,6 +514,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "config = PatchConfig(patch_size=[96, 96, 96], samples_per_volume=16)\n", "test_eq(config.patch_size, [96, 96, 96])\n", "test_eq(config.samples_per_volume, 16)\n", @@ -528,7 +533,7 @@ "test_eq(config2.apply_reorder, True)\n", "test_eq(config2.target_spacing, [0.5, 0.5, 0.5])\n", "\n", - "# preprocessed=True with preprocessing params: no warning\n", + "# preprocessed=True keeps the preprocessing params (inference still uses them)\n", "config3 = PatchConfig(\n", " patch_size=[96, 96, 96],\n", " apply_reorder=True,\n", @@ -537,20 +542,7 @@ ")\n", "test_eq(config3.preprocessed, True)\n", "test_eq(config3.apply_reorder, True)\n", - "test_eq(config3.target_spacing, [0.5, 0.5, 0.5])\n", - "\n", - "# preprocessed=True without preprocessing params does NOT warn\n", - "# (preprocessed=True still has effect: skips pre_patch_tfms during training)\n", - "with warnings.catch_warnings(record=True) as w:\n", - " warnings.simplefilter(\"always\")\n", - " config4 = PatchConfig(\n", - " patch_size=[96, 96, 96],\n", - " apply_reorder=False,\n", - " target_spacing=None,\n", - " preprocessed=True\n", - " )\n", - " preprocessed_warns = [x for x in w if 'preprocessed' in str(x.message).lower()]\n", - " test_eq(len(preprocessed_warns), 0)" + "test_eq(config3.target_spacing, [0.5, 0.5, 0.5])" ] }, { @@ -699,13 +691,6 @@ " label_probabilities=config.label_probabilities\n", " )\n", " \n", - " elif config.sampler_type == 'weighted':\n", - " raise NotImplementedError(\n", - " \"WeightedSampler requires a pre-computed probability map which is not currently supported. \"\n", - " \"Use 'label' sampler with label_probabilities for weighted sampling based on segmentation labels, \"\n", - " \"or 'uniform' for random patch extraction.\"\n", - " )\n", - " \n", " raise ValueError(f\"Unknown sampler type: {config.sampler_type}\")" ] }, @@ -716,14 +701,18 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "config = PatchConfig(patch_size=[64, 64, 64], sampler_type='uniform')\n", "sampler = create_patch_sampler(config)\n", "test_eq(type(sampler), tio.UniformSampler)\n", "\n", - "# WeightedSampler raises NotImplementedError\n", + "label_config = PatchConfig(patch_size=[64, 64, 64], sampler_type='label')\n", + "test_eq(type(create_patch_sampler(label_config)), _PatchLabelSampler)\n", + "\n", + "# Only the 'uniform' and 'label' samplers are supported\n", "from fastcore.test import test_fail\n", - "config_weighted = PatchConfig(patch_size=[64, 64, 64], sampler_type='weighted')\n", - "test_fail(lambda: create_patch_sampler(config_weighted), contains='WeightedSampler')" + "test_fail(lambda: PatchConfig(patch_size=[64, 64, 64], sampler_type='weighted'),\n", + " contains='sampler_type must be one of')" ] }, { @@ -733,6 +722,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "device = _get_default_device()\n", "test_eq(type(device), torch.device)\n", "\n", @@ -784,7 +774,7 @@ " raw TorchIO transforms. fastMONAI wrappers are automatically normalized\n", " to raw TorchIO for internal use. Mutually exclusive with gpu_augmentation.\n", " gpu_augmentation: GpuPatchAugmentation instance for GPU-batched augmentation.\n", - " Operates on [B,C,D,H,W] tensors already on GPU, avoiding per-sample CPU\n", + " Operates on [B, C, X, Y, Z] tensors already on GPU, avoiding per-sample CPU\n", " overhead. Mutually exclusive with patch_tfms. Training only.\n", " shuffle: Whether to shuffle subjects and patches.\n", " drop_last: Whether to drop last incomplete batch.\n", @@ -842,7 +832,7 @@ " def __iter__(self):\n", " \"\"\"Iterate over batches, yielding (image, mask) tuples.\"\"\"\n", " for batch in self._dl:\n", - " img = batch['image'][tio.DATA] # [B, C, H, W, D]\n", + " img = batch['image'][tio.DATA] # [B, C, X, Y, Z]\n", " has_mask = 'mask' in batch\n", "\n", " # Apply CPU patch transforms if provided (per-sample TorchIO loop)\n", @@ -982,19 +972,17 @@ "\n", "\n", "def _stash_from_df_metadata(instance, img_col, mask_col, pre_patch_tfms, apply_reorder,\n", - " target_spacing, ensure_affine_consistency, patch_config,\n", - " train_df, valid_df):\n", - " \"\"\"Record the de-facto metadata contract attributes consumed by research/vs_seg scripts.\n", + " target_spacing, patch_config, train_df, valid_df):\n", + " \"\"\"Record the from_df inputs on the instance for MLflow logging and split tracking.\n", "\n", - " Stores the RAW pre_patch_tfms (not the normalized list), matching how the train_final*.py\n", - " scripts set dls._pre_patch_tfms manually.\n", + " Stores pre_patch_tfms as given (fastMONAI wrappers are not unwrapped to TorchIO).\n", + " Code that builds MedPatchDataLoaders directly can set the same attributes.\n", " \"\"\"\n", " instance._img_col = img_col\n", " instance._mask_col = mask_col\n", " instance._pre_patch_tfms = pre_patch_tfms\n", " instance._apply_reorder = apply_reorder\n", " instance._target_spacing = target_spacing\n", - " instance._ensure_affine_consistency = ensure_affine_consistency\n", " instance._patch_config = patch_config\n", " instance._train_source_df = train_df\n", " instance._valid_source_df = valid_df\n", @@ -1015,10 +1003,10 @@ " for full-volume sliding window inference.\n", "\n", " Example:\n", - " >>> config = PatchConfig(patch_size=[96, 96, 96], apply_reorder=True, target_spacing=[0.5, 0.5, 0.5])\n", + " >>> config = PatchConfig(patch_size=[96, 96, 96], apply_reorder=True, target_spacing=[0.5, 0.5, 0.5],\n", + " ... normalization=[ZNormalization(masking_method='foreground')])\n", " >>> dls = MedPatchDataLoaders.from_df(\n", - " ... df, img_col='image', mask_col='label', valid_pct=0.2,\n", - " ... patch_config=config, pre_patch_tfms=[tio.ZNormalization()], bs=4\n", + " ... df, img_col='image', mask_col='label', valid_pct=0.2, patch_config=config, bs=4\n", " ... )\n", " >>> learn = Learner(dls, model, loss_func=DiceLoss())\n", " \"\"\"\n", @@ -1175,7 +1163,7 @@ " instance = cls(train_dl, valid_dl, device)\n", " return _stash_from_df_metadata(\n", " instance, img_col, mask_col, _norm_tfms, _apply_reorder, _target_spacing,\n", - " ensure_affine_consistency, patch_config, train_df, valid_df\n", + " patch_config, train_df, valid_df\n", " )\n", "\n", " @property\n", @@ -1210,18 +1198,27 @@ "\n", " @property\n", " def apply_reorder(self):\n", - " \"\"\"Whether reordering to RAS+ is enabled.\"\"\"\n", - " return getattr(self, '_apply_reorder', False)\n", + " \"\"\"Whether reordering to RAS+ is enabled (from from_df, else patch_config).\"\"\"\n", + " if hasattr(self, '_apply_reorder'):\n", + " return self._apply_reorder\n", + " config = self.patch_config\n", + " return config.apply_reorder if config is not None else False\n", "\n", " @property\n", " def target_spacing(self):\n", - " \"\"\"Target voxel spacing for resampling.\"\"\"\n", - " return getattr(self, '_target_spacing', None)\n", + " \"\"\"Target voxel spacing for resampling (from from_df, else patch_config).\"\"\"\n", + " if hasattr(self, '_target_spacing'):\n", + " return self._target_spacing\n", + " config = self.patch_config\n", + " return config.target_spacing if config is not None else None\n", "\n", " @property\n", " def patch_config(self):\n", - " \"\"\"The PatchConfig used for this DataLoaders.\"\"\"\n", - " return getattr(self, '_patch_config', None)\n", + " \"\"\"The PatchConfig used for this DataLoaders (from from_df, else the train loader's config).\"\"\"\n", + " config = getattr(self, '_patch_config', None)\n", + " if config is None:\n", + " config = getattr(getattr(self, '_train_dl', None), 'config', None)\n", + " return config\n", "\n", " @property\n", " def split_df(self):\n", @@ -1280,7 +1277,10 @@ " def show_batch(self, dl_idx=0, max_n=6, figsize=None, channel=0,\n", " slice_index=None, anatomical_plane=0, overlay=False,\n", " voxel_size=None, **kwargs):\n", - " \"\"\"Show a batch of patch samples for visualization.\"\"\"\n", + " \"\"\"Show a batch of patch samples for visualization.\n", + "\n", + " Extra keyword arguments are accepted for compatibility with fastai's show_batch and ignored.\n", + " \"\"\"\n", "\n", " dl = self[dl_idx]\n", " x, y = dl.one_batch()\n", @@ -1319,9 +1319,9 @@ " slice_idxs.extend([idx] * len(im_channels))\n", "\n", " _voxel_size = voxel_size if voxel_size is not None else self.target_spacing\n", - " ctxs = [im.show(ax=ax, slice_index=idx, anatomical_plane=anatomical_plane,\n", - " voxel_size=_voxel_size)\n", - " for im, ax, idx in zip(imgs, flat_axs, slice_idxs)]\n", + " for im, ax, idx in zip(imgs, flat_axs, slice_idxs):\n", + " im.show(ax=ax, slice_index=idx, anatomical_plane=anatomical_plane,\n", + " voxel_size=_voxel_size)\n", "\n", " if overlay and has_mask:\n", " for mask, ax, idx in zip(masks_for_overlay, flat_axs, slice_idxs):\n", @@ -1400,6 +1400,34 @@ "test_eq(_a, _b)" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "direct-dls-metadata-test", + "metadata": {}, + "outputs": [], + "source": [ + "#| hide\n", + "# MedPatchDataLoaders built directly reports its train loader's config\n", + "_direct_subject = tio.Subject(\n", + " image=tio.ScalarImage(tensor=torch.zeros(1, 16, 16, 16), affine=np.eye(4)),\n", + " mask=tio.LabelMap(tensor=torch.zeros(1, 16, 16, 16, dtype=torch.uint8), affine=np.eye(4)),\n", + ")\n", + "_direct_cfg = PatchConfig(patch_size=[16, 16, 16], target_spacing=[1.0, 1.0, 2.0], queue_num_workers=0)\n", + "_direct_dls = MedPatchDataLoaders(\n", + " MedPatchDataLoader(tio.SubjectsDataset([_direct_subject]), _direct_cfg, batch_size=1),\n", + " MedPatchDataLoader(tio.SubjectsDataset([_direct_subject]), _direct_cfg, batch_size=1, shuffle=False),\n", + " device=torch.device('cpu'),\n", + ")\n", + "assert _direct_dls.patch_config is _direct_cfg\n", + "test_eq(_direct_dls.apply_reorder, True)\n", + "test_eq(_direct_dls.target_spacing, [1.0, 1.0, 2.0])\n", + "# Values recorded by from_df (or set by hand) take precedence\n", + "_direct_dls._apply_reorder = False\n", + "test_eq(_direct_dls.apply_reorder, False)\n", + "_direct_dls.close()" + ] + }, { "cell_type": "code", "execution_count": null, @@ -1407,6 +1435,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "from fastMONAI.vision_augmentation import GpuPatchAugmentation\n", "\n", "# Both gpu_augmentation and patch_tfms -> ValueError\n", @@ -1443,15 +1472,13 @@ "outputs": [], "source": [ "#| export\n", - "import numbers\n", - "\n", "def _normalize_patch_overlap(patch_overlap: 'int | float | list', patch_size: list) -> tuple:\n", " \"\"\"Convert patch_overlap (fraction 0-1, pixel int, numpy scalar, or sequence) to a tuple of even pixel ints for TorchIO's GridSampler.\n", "\n", " Format conversion only; value validation (negatives, overlap >= patch_size) lives in PatchConfig.__post_init__().\n", " \"\"\"\n", " # Scalar fraction (0 < x < 1); excludes 1.0 since 100% overlap creates step_size=0 (infinite patches)\n", - " if isinstance(patch_overlap, (int, float, numbers.Number)) and 0 < float(patch_overlap) < 1:\n", + " if isinstance(patch_overlap, numbers.Number) and 0 < float(patch_overlap) < 1:\n", " result = []\n", " for ps in patch_size:\n", " pixels = int(int(ps) * float(patch_overlap))\n", @@ -1461,7 +1488,7 @@ " return tuple(result)\n", "\n", " # Scalar int (incl. numpy scalars): pixel count\n", - " if isinstance(patch_overlap, (int, float, numbers.Number)):\n", + " if isinstance(patch_overlap, numbers.Number):\n", " val = int(patch_overlap)\n", " if val % 2 != 0:\n", " val = val - 1 if val > 0 else 0\n", @@ -1508,8 +1535,8 @@ "def _logits_to_probs(logits: torch.Tensor) -> torch.Tensor:\n", " \"\"\"Logits -> probabilities: sigmoid for 1 channel (binary), else softmax over the channel dim.\n", "\n", - " Upcasts to float32 (logits may be bfloat16 under AMP). Does NOT move to CPU; callers decide\n", - " placement (TTA keeps tensors on-device to flip back; the non-TTA path moves to CPU immediately).\n", + " Upcasts to float32 (logits may be bfloat16 under AMP). Does not move to CPU: callers sum the\n", + " passes (flips and models) on the device and move the average to CPU once.\n", " \"\"\"\n", " return torch.sigmoid(logits.float()) if logits.shape[1] == 1 else torch.softmax(logits.float(), dim=1)\n", "\n", @@ -1518,8 +1545,9 @@ " \"\"\"Mirror TTA (with optional model ensemble): average probabilities over `flip_axes` and all models.\n", "\n", " models is a single model or a list of models; a list soft-votes (averages probabilities).\n", - " Uses a running sum (2x memory, not 9x per model). Each pass: flip -> forward -> activate ->\n", - " flip back -> accumulate. Returns the averaged probability tensor [B, C, D, H, W] on CPU.\n", + " Keeps a running sum, so memory holds two probability tensors instead of one per pass\n", + " (8 per model with all flips). Each pass: flip -> forward -> activate -> flip back ->\n", + " accumulate. Returns the averaged probability tensor [B, C, X, Y, Z] on CPU.\n", " amp_context, if given, wraps each forward pass.\n", " \"\"\"\n", " if amp_context is None: amp_context = nullcontext()\n", @@ -1783,7 +1811,6 @@ " f\"models, or post-process each class separately.\"\n", " )\n", " else:\n", - " from fastMONAI.vision_inference import keep_largest\n", " result = keep_largest(result.squeeze(0)).unsqueeze(0)\n", "\n", " if return_probabilities:\n", @@ -1806,10 +1833,6 @@ " result = result.long()\n", "\n", " # Use original affine matrix for correct spatial alignment\n", - " if not (hasattr(prepared.org_img, 'affine') and prepared.org_img.affine is not None):\n", - " raise RuntimeError(\n", - " \"org_img.affine not available. This should never happen - please report this bug.\"\n", - " )\n", " affine = prepared.org_img.affine.copy()\n", "\n", " return result, affine\n", @@ -1861,8 +1884,8 @@ "\n", " Returns:\n", " ``(mask, probabilities)`` or ``(mask, probabilities, affine)`` when\n", - " ``return_affine=True``. The mask has shape ``[1, D, H, W]`` and dtype\n", - " ``torch.long``; probabilities have shape ``[C, D, H, W]`` and floating dtype.\n", + " ``return_affine=True``. The mask has shape ``[1, X, Y, Z]`` and dtype\n", + " ``torch.long``; probabilities have shape ``[C, X, Y, Z]`` and floating dtype.\n", " \"\"\"\n", " prepared = self._prepare_subject(img_path)\n", " output = self._run_inference(prepared, tta=tta)\n", @@ -1887,6 +1910,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# PatchInferenceEngine must store a raw model as-is, not its internal .model\n", "from monai.networks.nets import UNet\n", "from monai.networks.layers import Norm\n", @@ -1917,9 +1941,6 @@ "outputs": [], "source": [ "#| export\n", - "from concurrent.futures import ThreadPoolExecutor\n", - "\n", - "\n", "def prediction_filename(input_path: str | Path) -> str:\n", " \"\"\"Return the filename used by patch inference for a saved prediction.\"\"\"\n", " input_path = Path(input_path)\n", @@ -2094,6 +2115,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Test _TTA_FLIP_AXES and _predict_patch_tta\n", "from itertools import combinations\n", "\n", @@ -2126,7 +2148,7 @@ "\n", "# Binary case (1 output channel -> sigmoid)\n", "model_bin = _SimpleConv(1).eval()\n", - "dummy_input = torch.randn(2, 1, 8, 8, 8) # [B=2, C=1, D, H, W]\n", + "dummy_input = torch.randn(2, 1, 8, 8, 8) # [B=2, C=1, X, Y, Z]\n", "with torch.no_grad():\n", " tta_out = _predict_patch_tta(model_bin, dummy_input)\n", "test_eq(tta_out.shape, torch.Size([2, 1, 8, 8, 8]))\n", @@ -2167,6 +2189,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Test _PreparedSubject and decomposed predict path\n", "import tempfile, os, nibabel as nib\n", "from monai.networks.nets import UNet\n", @@ -2259,6 +2282,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Safety-net tests (added before refactor to characterize current behavior)\n", "import tempfile as _tempfile, os as _os, nibabel as _nib\n", "from monai.networks.nets import UNet as _UNet\n", @@ -2327,7 +2351,7 @@ " test_eq(len(_dls._valid_source_df), 2)\n", " assert isinstance(_dls.train, MedPatchDataLoader) and isinstance(_dls.valid, MedPatchDataLoader)\n", " for _a in ['_img_col', '_mask_col', '_pre_patch_tfms', '_apply_reorder', '_target_spacing',\n", - " '_ensure_affine_consistency', '_patch_config', '_train_source_df', '_valid_source_df']:\n", + " '_patch_config', '_train_source_df', '_valid_source_df']:\n", " assert hasattr(_dls, _a), f'from_df must set {_a}'\n", " test_eq(_dls.patch_config is _cfg, True)\n", " test_eq(_dls.apply_reorder, True)\n", @@ -2733,6 +2757,7 @@ "metadata": {}, "outputs": [], "source": [ + "#| hide\n", "# Ensemble (soft-voting): PatchInferenceEngine and patch_inference accept a list of models.\n", "# Per-patch probabilities are averaged, then a single argmax/threshold is applied.\n", "import tempfile, os, nibabel as nib\n", @@ -2853,14 +2878,6 @@ "print(\"torch.compile unwrap tests passed!\")" ] }, - { - "cell_type": "markdown", - "id": "cell-19", - "metadata": {}, - "source": [ - "## Export" - ] - }, { "cell_type": "code", "execution_count": null, From c1a685ab53eb06af9a6ae4451eac4f6abff60ef3 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:25:47 +0100 Subject: [PATCH 06/11] docs: fix tutorial text, front matter and the readme logo - the README/index logo pointed at a URL that returns 404 - nbdev turns the whole title cell into front matter, so intro text in 12a, 12b, 12c, 11f and 11d never rendered; it now has its own cell - 11b passed splitter= to from_df, which ignored it and split randomly; it now uses valid_col - stale prose in 11b, 11c, 11d and 12c, a dead link, typos, references to the removed research/vs_seg folder, and the GPU augmentation note in 12a - CITATION.cff author name fields, CONTRIBUTING install command and commit convention --- CITATION.cff | 6 +- CONTRIBUTING.md | 3 +- README.md | 2 +- nbs/11b_tutorial_classification.ipynb | 10 ++-- nbs/11c_tutorial_regression.ipynb | 2 +- nbs/11d_tutorial_binary_segmentation.ipynb | 11 ++-- ...11e_tutorial_multiclass_segmentation.ipynb | 2 +- nbs/11f_tutorial_inference.ipynb | 23 ++++---- nbs/12a_tutorial_patch_training.ipynb | 34 ++++++++---- nbs/12b_tutorial_patch_cross_validation.ipynb | 55 ++++++++++--------- nbs/12c_tutorial_patch_inference.ipynb | 35 +++++++----- nbs/index.ipynb | 4 +- nbs/nbdev.yml | 2 +- 13 files changed, 107 insertions(+), 82 deletions(-) diff --git a/CITATION.cff b/CITATION.cff index a0d00de..e1c80aa 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -13,8 +13,8 @@ authors: and Mathematical Sciences, Western Norway University of Applied Sciences, Bergen, Norway orcid: 'https://orcid.org/0000-0002-0038-5540' - - given-names: Lundervold - family-names: Alexander Selvikvåg + - given-names: Alexander Selvikvåg + family-names: Lundervold email: allu@hvl.no affiliation: >- Department of Computer Science, Electrical Engineering @@ -38,7 +38,7 @@ license: Apache-2.0 preferred-citation: type: article authors: - - family-names: " Kaliyugarasan" + - family-names: "Kaliyugarasan" given-names: "Satheshkumar" orcid: "https://orcid.org/0000-0002-0038-5540" - family-names: "Lundervold" diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f0ca6b5..e0cd474 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,10 +2,11 @@ Contributions to the source code are greatly appreciated. If you find any issues or have suggestions for improvements, please open an issue in the corresponding GitHub repository. To contribute code changes, we encourage you to follow these [steps](https://docs.github.com/en/get-started/quickstart/contributing-to-projects). ## Getting Started -1. Install an editable version of the forked project: `pip install -e 'fastMONAI[dev]'` +1. Install an editable version of the forked project from inside the project folder: `pip install -e '.[dev]'` 2. Git hooks are used to run automatic scripts during each commit and merge to strip the notebooks of superfluous metadata and avoid merge conflicts. To set up Git hooks, run the following command inside the project folder: `nbdev_install_hooks` 3. Once you have installed the project and set up the Git hooks, you can start making changes to the notebooks. 4. Run: `nbdev_prepare` to build .py modules from notebooks and test them (if unit tests are written). +5. Write commit messages and pull request titles as [Conventional Commits](https://www.conventionalcommits.org/), for example `fix(vision_patch): handle empty masks`. See `AGENTS.md` for the allowed types. ## Contact If you have any questions or need further assistance, please contact us at [skaliyugarasan@hotmail.com]. We appreciate your interest and look forward to your contributions. \ No newline at end of file diff --git a/README.md b/README.md index bc7678f..edc6ffc 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ -![](https://raw.githubusercontent.com/skaliy/skaliy.github.io/main/assets/fastmonai_v1.png) +![](https://raw.githubusercontent.com/MMIV-ML/fastMONAI/main/nbs/images/fastmonai_v1.png) ![CI](https://github.com/MMIV-ML/fastMONAI/workflows/CI/badge.svg) [![Docs](https://github.com/MMIV-ML/fastMONAI/actions/workflows/deploy.yaml/badge.svg)](https://fastmonai.no) diff --git a/nbs/11b_tutorial_classification.ipynb b/nbs/11b_tutorial_classification.ipynb index f8fa469..caa0c24 100644 --- a/nbs/11b_tutorial_classification.ipynb +++ b/nbs/11b_tutorial_classification.ipynb @@ -221,9 +221,9 @@ "metadata": {}, "source": [ "Before feeding the data into a model, we must create a `DataLoaders` object for our dataset. There are several ways to get the data in `DataLoaders`. \n", - "In the following line, we call the ` ImageDataLoaders.from_df` factory method, which is the most basic way of building a `DataLoaders`. \n", + "In the following line, we call the `MedImageDataLoaders.from_df` factory method, which is the most basic way of building a `DataLoaders`. \n", "\n", - "Here, we pass the processed DataFrame, define the columns for the images `fn_col` and the labels `label_col`, some transforms `item_tfms`, voxel spacing `resample`, and the batch size `bs`. " + "Here, we pass the processed DataFrame, define the columns for the images `fn_col` and the labels `label_col`, the column that marks validation rows `valid_col`, some transforms `item_tfms`, the target voxel spacing `target_spacing`, and the batch size `bs`. " ] }, { @@ -234,7 +234,7 @@ "outputs": [], "source": [ "dls = MedImageDataLoaders.from_df(df, fn_col='img_path', label_col='labels',\n", - " splitter=ColSplitter('is_val'), item_tfms=item_tfms,\n", + " valid_col='is_val', item_tfms=item_tfms,\n", " target_spacing=1, bs=64)" ] }, @@ -566,7 +566,7 @@ "id": "b6cb52ae", "metadata": {}, "source": [ - "Class imbalance is a common challenge in medical datasets, and it is something we're facing in our example dataset. When dealing with a classification task on such imbalanced datasets, specific sampling techniques may be necessary. See the [Advanced](#section_advanced)) section for further details." + "Class imbalance is a common challenge in medical datasets, and it is something we're facing in our example dataset. When dealing with a classification task on such imbalanced datasets, specific sampling techniques may be necessary. See the [Advanced](#section_advanced) section for further details." ] }, { @@ -892,7 +892,7 @@ "id": "bcf87048", "metadata": {}, "source": [ - "Now we pass our processed DataFrame and the bath size to create a `DataLoaders` object. " + "Now we pass our processed DataFrame and the batch size to create a `DataLoaders` object. " ] }, { diff --git a/nbs/11c_tutorial_regression.ipynb b/nbs/11c_tutorial_regression.ipynb index f4eca9a..7043d68 100644 --- a/nbs/11c_tutorial_regression.ipynb +++ b/nbs/11c_tutorial_regression.ipynb @@ -6,7 +6,7 @@ "metadata": {}, "source": [ "# Regression\n", - "> We will use the same data used in the classification tutorial for this task (the IXI Dataset). The approach for regression tasks is nearly identical. Therefore, take a look at the classification tutorial for explanations of the various cells." + "> We use the IXI dataset to predict age from T1-weighted brain MR images. The approach for regression tasks is nearly identical to classification, so take a look at the classification tutorial for explanations of the various cells." ] }, { diff --git a/nbs/11d_tutorial_binary_segmentation.ipynb b/nbs/11d_tutorial_binary_segmentation.ipynb index 05530aa..7821739 100644 --- a/nbs/11d_tutorial_binary_segmentation.ipynb +++ b/nbs/11d_tutorial_binary_segmentation.ipynb @@ -7,8 +7,7 @@ "source": [ "# Binary semantic segmentation\n", "\n", - "> Image segmentation, or semantic segmentation, is the task of classifying each pixel within an image with a corresponding class. \n", - "For the binary semantic segmentation task, we will use the Heart MRI dataset from the Medical Segmentation Decathlon challenge (http://medicaldecathlon.com/) provided by King’s College London. The target is to segment the left atrium. The challenge with this dataset is that it is small and has considerable variability. " + "> Image segmentation, or semantic segmentation, is the task of classifying each pixel within an image with a corresponding class. For the binary semantic segmentation task, we will use the Heart MRI dataset from the Medical Segmentation Decathlon challenge (http://medicaldecathlon.com/) provided by King’s College London. The target is to segment the left atrium. The challenge with this dataset is that it is small and has considerable variability. " ] }, { @@ -143,7 +142,7 @@ "id": "5ac3d920-b0b8-4b96-974c-fa5aeb51c7b7", "metadata": {}, "source": [ - "Split the labled data into training and test" + "Split the labeled data into training and test" ] }, { @@ -587,7 +586,7 @@ "In the following line, we import an enhanced version of 3D UNet from MONAI. In addition, we import the Dice coefficient (Dice) loss, a commonly used loss function in medical image segmentation tasks: \n", "\\begin{equation*}\n", "\\begin{aligned}\n", - "loss = 1- \\frac{y \\cap \\hat{y}}{\\left | y \\right |+\\left | \\hat{y} \\right |}\n", + "loss = 1- \\frac{2\\left | y \\cap \\hat{y} \\right |}{\\left | y \\right |+\\left | \\hat{y} \\right |}\n", "\\end{aligned}\n", "\\end{equation*}" ] @@ -740,7 +739,7 @@ "id": "1b5c5998-4aee-4ce8-9d8f-17af176ec5f3", "metadata": {}, "source": [ - "As shown in https://walkwithfastai.com/Segmentation, with the ranger optimizer, we want to use a different fit function, called `fit_flat_cos`" + "With the ranger optimizer, we use a different fit function, [`fit_flat_cos`](https://docs.fast.ai/callback.schedule.html#learner.fit_flat_cos), which keeps the learning rate flat before annealing it with a cosine schedule." ] }, { @@ -919,7 +918,7 @@ "id": "88abda21-4d86-457e-ad4d-ed5e800ce303", "metadata": {}, "source": [ - "### Test data (with labels to evaulate the model)\n", + "### Test data (with labels to evaluate the model)\n", "Evaluate the performance of the selected model on unseen data. \n", "It’s important to not touch this data until you have fine tuned your model to get an unbiased evaluation!" ] diff --git a/nbs/11e_tutorial_multiclass_segmentation.ipynb b/nbs/11e_tutorial_multiclass_segmentation.ipynb index e935df3..ec08ebe 100644 --- a/nbs/11e_tutorial_multiclass_segmentation.ipynb +++ b/nbs/11e_tutorial_multiclass_segmentation.ipynb @@ -142,7 +142,7 @@ "id": "1cee363e-6852-4e1e-90af-79ee1c49bcd8", "metadata": {}, "source": [ - "Split the labled data into training and test" + "Split the labeled data into training and test" ] }, { diff --git a/nbs/11f_tutorial_inference.ipynb b/nbs/11f_tutorial_inference.ipynb index 6097bef..30fb3ad 100644 --- a/nbs/11f_tutorial_inference.ipynb +++ b/nbs/11f_tutorial_inference.ipynb @@ -5,12 +5,7 @@ "id": "04e4d97b-d133-4c8d-8c56-8d1742cc3d5e", "metadata": {}, "source": [ - "# Inference with a safe model artifact\n", - "\n", - "This tutorial loads the registered architecture, Safetensors weights, preprocessing,\n", - "and output-decode contract. It does not require a serialized Learner. For Task02_Heart,\n", - "the best_model.safetensors artifact is produced by the binary-segmentation tutorial with\n", - "the monai.unet_inner model spec and binary segmentation output contract." + "# Inference with a safe model artifact" ] }, { @@ -24,6 +19,17 @@ "---" ] }, + { + "cell_type": "markdown", + "id": "intro", + "metadata": {}, + "source": [ + "This tutorial loads the registered architecture, Safetensors weights, preprocessing,\n", + "and output-decode contract. It does not require a serialized Learner. For Task02_Heart,\n", + "the best_model.safetensors artifact is produced by the binary-segmentation tutorial with\n", + "the monai.unet_inner model spec and binary segmentation output contract." + ] + }, { "cell_type": "code", "execution_count": null, @@ -31,7 +37,6 @@ "metadata": {}, "outputs": [], "source": [ - "#all_skip\n", "from fastMONAI.vision_all import *\n", "from monai.apps import DecathlonDataset" ] @@ -49,9 +54,7 @@ " \"spleen\": \"Task09_Spleen\"\n", "}\n", "\n", - "task = tasks[\"heart\"]\n", - "\n", - "model_artifact_path = Path(f\"model_artifacts/{task}\")" + "task = tasks[\"heart\"]" ] }, { diff --git a/nbs/12a_tutorial_patch_training.ipynb b/nbs/12a_tutorial_patch_training.ipynb index 6da55d3..ad5ea8d 100644 --- a/nbs/12a_tutorial_patch_training.ipynb +++ b/nbs/12a_tutorial_patch_training.ipynb @@ -7,15 +7,7 @@ "source": [ "# Patch-Based Binary Segmentation\n", "\n", - "> Patch-based training for 3D medical image segmentation using fastMONAI's `MedPatchDataLoaders` with lazy loading, where volumes are loaded on-demand, keeping memory usage constant regardless of dataset size.\n", - "\n", - "This tutorial demonstrates when and how to use patch-based training:\n", - "\n", - "- **Large volumes**: When full images don't fit in GPU memory\n", - "- **Memory constraints**: Constant ~150 MB memory usage regardless of dataset size\n", - "- **Class imbalance**: Foreground-weighted sampling for better learning on small structures\n", - "\n", - "We use the same Heart MRI dataset from the [binary segmentation tutorial](11d_tutorial_binary_segmentation.ipynb) for direct comparison." + "> Patch-based training for 3D medical image segmentation using fastMONAI's `MedPatchDataLoaders` with lazy loading, where volumes are loaded on-demand, keeping memory usage constant regardless of dataset size." ] }, { @@ -37,6 +29,20 @@ "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/nbs/12a_tutorial_patch_training.ipynb)" ] }, + { + "cell_type": "markdown", + "id": "cell-intro", + "metadata": {}, + "source": [ + "This tutorial demonstrates when and how to use patch-based training:\n", + "\n", + "- **Large volumes**: When full images don't fit in GPU memory\n", + "- **Memory constraints**: Constant ~150 MB memory usage regardless of dataset size\n", + "- **Class imbalance**: Foreground-weighted sampling for better learning on small structures\n", + "\n", + "We use the same Heart MRI dataset from the [binary segmentation tutorial](11d_tutorial_binary_segmentation.ipynb) for direct comparison." + ] + }, { "cell_type": "code", "execution_count": null, @@ -306,9 +312,13 @@ "~77 ms for the equivalent MONAI per-sample-on-GPU pipeline (~25x) and ~287 ms for the TorchIO\n", "CPU path (~100x).\n", "\n", - "It is an *approximate* reimplementation of the `patch_tfms` transforms (voxel-space affine,\n", - "per-axis flip, image-only anisotropy, sign-preserving gamma), It approximates the CPU path rather than matching it exactly. Use `gpu_augmentation` instead of\n", - "`patch_tfms` (they are mutually exclusive)." + "It takes the same parameters as `patch_tfms`, and each value means the same thing (per-axis\n", + "values follow array axis order, lengths are in voxels), but the output is not bit-identical.\n", + "Rotations act in normalised patch coordinates, so on a non-cubic patch a rotation between two\n", + "axes of different length is stretched by their length ratio. Blur uses one isotropic sigma per\n", + "sample, anisotropy resamples with `F.interpolate`, and gamma draws one exponent per sample\n", + "instead of one per channel. Use `gpu_augmentation` instead of `patch_tfms` (they are mutually\n", + "exclusive)." ] }, { diff --git a/nbs/12b_tutorial_patch_cross_validation.ipynb b/nbs/12b_tutorial_patch_cross_validation.ipynb index 9d31c63..77d5239 100644 --- a/nbs/12b_tutorial_patch_cross_validation.ipynb +++ b/nbs/12b_tutorial_patch_cross_validation.ipynb @@ -7,26 +7,7 @@ "source": [ "# Patch-Based Cross-Validation and Final Model Training\n", "\n", - "> Cross-validation gives a more reliable estimate of how well a model generalizes than a single\n", - "> train/validation split, which can vary a lot on the small datasets common in medical imaging.\n", - "> K-fold CV trains *k* models, each holding out a different fold. Every case is scored once, by\n", - "> the model that did not train on it, and we report the mean and standard deviation across all\n", - "> held-out cases. The spread shows how much performance varies from case to case.\n", - "\n", - "This tutorial has two parts:\n", - "\n", - "- Part 1, cross-validation (estimate): run 5-fold CV to estimate segmentation performance as a\n", - " mean and std across all held-out cases. None of these fold models is shipped; they exist to\n", - " measure the method.\n", - "- Part 2, final model (deploy): once we trust the estimate, we train a single model on all the\n", - " data (no fold held out) and export it for deployment.\n", - "\n", - "So you cross-validate to measure, then retrain on everything to deploy.\n", - "\n", - "We use the public Decathlon Heart dataset (`Task02_Heart`, left-atrium\n", - "segmentation) and the patch-based workflow. If you are new to patch-based training,\n", - "start with [Patch-Based Binary Segmentation](12a_tutorial_patch_training.ipynb), whose\n", - "setup (download, `MedDataset` analysis, `PatchConfig`) we reuse here." + "> Cross-validation gives a more reliable estimate of how well a model generalizes than a single train/validation split, which can vary a lot on the small datasets common in medical imaging." ] }, { @@ -48,6 +29,31 @@ "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/nbs/12b_tutorial_patch_cross_validation.ipynb)" ] }, + { + "cell_type": "markdown", + "id": "c00-intro", + "metadata": {}, + "source": [ + "K-fold CV trains *k* models, each holding out a different fold. Every case is scored once, by\n", + "the model that did not train on it, and we report the mean and standard deviation across all\n", + "held-out cases. The spread shows how much performance varies from case to case.\n", + "\n", + "This tutorial has two parts:\n", + "\n", + "- Part 1, cross-validation (estimate): run 5-fold CV to estimate segmentation performance as a\n", + " mean and std across all held-out cases. None of these fold models is shipped; they exist to\n", + " measure the method.\n", + "- Part 2, final model (deploy): once we trust the estimate, we train a single model on all the\n", + " data (no fold held out) and export it for deployment.\n", + "\n", + "So you cross-validate to measure, then retrain on everything to deploy.\n", + "\n", + "We use the public Decathlon Heart dataset (`Task02_Heart`, left-atrium\n", + "segmentation) and the patch-based workflow. If you are new to patch-based training,\n", + "start with [Patch-Based Binary Segmentation](12a_tutorial_patch_training.ipynb), whose\n", + "setup (download, `MedDataset` analysis, `PatchConfig`) we reuse here." + ] + }, { "cell_type": "code", "execution_count": null, @@ -92,8 +98,7 @@ "### Configuration\n", "\n", "We use reduced settings so the notebook runs quickly as a demo. Production runs use many\n", - "more epochs. The reference script `research/vs_seg/patch_based_dev/train_cv.py` trains each fold\n", - "for 500 epochs." + "more epochs, for example 500 per fold." ] }, { @@ -320,7 +325,7 @@ "> roughly homogeneous. Real datasets often need:\n", ">\n", "> - **`StratifiedKFold`** to balance a property across folds, for example binning tumor volume\n", - "> into quartiles so each fold has a similar tumor-size distribution (as the VS dataset does).\n", + "> into quartiles so each fold has a similar tumor-size distribution.\n", "> - **`GroupKFold`** to keep all scans from the same patient in the same fold, preventing\n", "> leakage when a subject contributes multiple volumes." ] @@ -394,8 +399,8 @@ "The cross-validation estimate is the **mean +/- std across all held-out cases** (pooled over\n", "the five folds) for each metric; the per-fold DSC block below is where you see fold-to-fold\n", "variation. We exclude non-finite `hd95_mm` / `assd_mm` values (which occur when exactly one of\n", - "the prediction or ground truth is empty for a case) from those means, mirroring\n", - "`aggregate_results` in `train_cv.py`." + "the prediction or ground truth is empty for a case) from those means and report how\n", + "many cases were excluded." ] }, { diff --git a/nbs/12c_tutorial_patch_inference.ipynb b/nbs/12c_tutorial_patch_inference.ipynb index 20b0163..e46b76c 100644 --- a/nbs/12c_tutorial_patch_inference.ipynb +++ b/nbs/12c_tutorial_patch_inference.ipynb @@ -7,17 +7,7 @@ "source": [ "# Patch-Based Inference: Deploying on New Scans\n", "\n", - "> Sliding-window inference for 3D medical image segmentation using fastMONAI's `PatchInferenceEngine`.\n", - "\n", - "This tutorial deploys the model trained on all labeled data and runs it on **new, unlabeled scans** (the Decathlon Heart test set). Because the model was trained on every labeled case, there is no held-out ground truth to score against here, so we predict and visualize rather than compute metrics. The performance estimate comes from the cross-validation in the prerequisite notebook.\n", - "\n", - "The patch-based inference pipeline:\n", - "\n", - "1. **GridSampler**: Extracts overlapping patches from each input volume\n", - "2. **Model**: Predicts on each patch\n", - "3. **GridAggregator**: Reconstructs the full volume with Hann windowing for smooth boundaries\n", - "\n", - "**Prerequisites**: Run [12b_tutorial_patch_cross_validation.ipynb](12b_tutorial_patch_cross_validation.ipynb) first to cross-validate, train the final model on all data, and save `models/final_model.safetensors`. Its inference configuration is embedded in the artifact." + "> Sliding-window inference for 3D medical image segmentation using fastMONAI's `PatchInferenceEngine`." ] }, { @@ -39,6 +29,22 @@ "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/nbs/12c_tutorial_patch_inference.ipynb)" ] }, + { + "cell_type": "markdown", + "id": "cell-intro", + "metadata": {}, + "source": [ + "This tutorial deploys the model trained on all labeled data and runs it on **new, unlabeled scans** (the Decathlon Heart test set). Because the model was trained on every labeled case, there is no held-out ground truth to score against here, so we predict and visualize rather than compute metrics. The performance estimate comes from the cross-validation in the prerequisite notebook.\n", + "\n", + "The patch-based inference pipeline:\n", + "\n", + "1. **GridSampler**: Extracts overlapping patches from each input volume\n", + "2. **Model**: Predicts on each patch\n", + "3. **GridAggregator**: Reconstructs the full volume with Hann windowing for smooth boundaries\n", + "\n", + "**Prerequisites**: Run [12b_tutorial_patch_cross_validation.ipynb](12b_tutorial_patch_cross_validation.ipynb) first to cross-validate, train the final model on all data, and save `models/final_model.safetensors`. Its inference configuration is embedded in the artifact." + ] + }, { "cell_type": "code", "execution_count": null, @@ -387,7 +393,7 @@ ], "source": [ "predictions = patch_inference(\n", - " learner=model, # the engine also accepts a raw model (here learn.model)\n", + " learner=model, # a plain PyTorch model; a fastai Learner also works\n", " config=patch_config,\n", " file_paths=test_imgs,\n", " save_dir='predictions/patch_heart',\n", @@ -420,7 +426,7 @@ "source": [ "### Load the explicitly selected fold models\n", "\n", - "Each fold run from 12b logged `model/best_model.safetensors` to MLflow. Declare the exact run ID for every model you intentionally want in the ensemble. `find_model_artifacts` downloads those artifacts without guessing membership from an experiment or filename pattern. The embedded inference contracts are checked against the all-data model before any model is loaded." + "Each fold run from 12b logged `model/best_model.safetensors` to MLflow. Declare the exact run ID for every model you intentionally want in the ensemble. `find_model_artifacts` downloads those artifacts without guessing membership from an experiment or filename pattern. Each model's embedded inference contract is compared with the first selected model's before any model is loaded." ] }, { @@ -497,7 +503,8 @@ "source": [ "if fold_models is not None:\n", " # Pass the list of models: the engine averages their softmax probabilities per patch\n", - " # (soft voting), then applies a single argmax + keep_largest_component, as in training.\n", + " # (soft voting), then applies a single argmax and, if enabled in the config,\n", + " # keep_largest_component post-processing.\n", " ensemble_predictions = patch_inference(\n", " fold_models,\n", " config=ensemble_patch_config,\n", diff --git a/nbs/index.ipynb b/nbs/index.ipynb index 110f1c2..8108752 100644 --- a/nbs/index.ipynb +++ b/nbs/index.ipynb @@ -6,7 +6,7 @@ "metadata": {}, "source": [ "# Overview\n", - "> " + "> A low-code deep learning library for 3D medical image analysis." ] }, { @@ -14,7 +14,7 @@ "id": "578d9fc6", "metadata": {}, "source": [ - "![](https://raw.githubusercontent.com/skaliy/skaliy.github.io/main/assets/fastmonai_v1.png)" + "![](https://raw.githubusercontent.com/MMIV-ML/fastMONAI/main/nbs/images/fastmonai_v1.png)" ] }, { diff --git a/nbs/nbdev.yml b/nbs/nbdev.yml index 3c73201..c759e28 100644 --- a/nbs/nbdev.yml +++ b/nbs/nbdev.yml @@ -4,6 +4,6 @@ project: website: title: "fastMONAI" site-url: "https://fastmonai.no/" - description: "fastMONAI library" + description: "A low-code deep learning library for 3D medical image analysis, built on fastai, MONAI and TorchIO" repo-branch: main repo-url: "https://github.com/MMIV-ML/fastMONAI" From f0cd00120af02952c802e12e9592af62032e42ed Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:26:03 +0100 Subject: [PATCH 07/11] build: raise the xlrd floor and drop unused dev requirements pandas needs xlrd>=2.0.1 for read_excel, which the IXI helpers use. tabulate is not imported anywhere, and the PyPI quarto package is a thin wrapper whose script shadows the real quarto CLI (it broke nbdev_readme locally); nbdev installs the CLI itself. The allowed metadata settings were parsed as a key literally named "True", so leaving them empty keeps nbdev_clean behaviour the same. Also a real package description. --- settings.ini | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/settings.ini b/settings.ini index 6516943..e4928aa 100644 --- a/settings.ini +++ b/settings.ini @@ -8,8 +8,8 @@ min_python = 3.10 version = 0.12.0 ### OPTIONAL ### -requirements = fastai>=2.8.6 monai>=1.5.2 torchio>=1.2.1 xlrd>=1.2.0 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 -dev_requirements = ipywidgets nbdev<3 execnb<0.2 fastcore<2 tabulate quarto pytest +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 +dev_requirements = ipywidgets nbdev<3 execnb<0.2 fastcore<2 pytest ### nbdev ### nbs_path = nbs @@ -26,15 +26,15 @@ jupyter_hooks = True clean_ids = False black_formatting = False readme_nb = index.ipynb -allowed_metadata_keys = True -allowed_cell_metadata_keys = True +allowed_metadata_keys = +allowed_cell_metadata_keys = ### PyPI ### audience = Developers author = Satheshkumar Kaliyugarasan author_email = skaliyugarasan@hotmail.com copyright = 2022 onwards, Mohn Medical Imaging and Visualization Centre -description = fastMONAI library +description = A low-code deep learning library for 3D medical image analysis, built on fastai, MONAI and TorchIO keywords = deep learning, medical imaging language = English license = apache2 From c9805ba58b9fc89c2a4b9a225910452916f6fb3e Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:26:03 +0100 Subject: [PATCH 08/11] ci: run tests once per pull request commit push and pull_request both fired for every commit on a PR branch. Pushes now trigger CI only on main. pytest comes from the dev extras the nbdev-ci step installs, so the separate install is gone. --- .github/workflows/test.yaml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 91506e1..f09c657 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -1,5 +1,9 @@ name: CI -on: [workflow_dispatch, pull_request, push] +on: + workflow_dispatch: + pull_request: + push: + branches: [main] jobs: test: @@ -10,5 +14,4 @@ jobs: env: PYTHONPATH: research:research/vestibular_schwannoma run: | - python -m pip install pytest python -m pytest -q research/vestibular_schwannoma/tests From e77019ab219528280a8dd0674a2a6a1a6a9f6447 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:26:03 +0100 Subject: [PATCH 09/11] chore: fix gitignore exceptions and ignore local files The IXI test image exception was root-anchored and never matched nbs/images/. MLflow database backups and the local todo.html were not ignored, so `git add -A` could have committed them. Removed duplicate rules, and tracked research/.gitignore, whose rule keeps scripts in presentation build folders out of git. --- .gitignore | 8 +++----- research/.gitignore | 2 ++ 2 files changed, 5 insertions(+), 5 deletions(-) create mode 100644 research/.gitignore diff --git a/.gitignore b/.gitignore index ea17901..1d00941 100644 --- a/.gitignore +++ b/.gitignore @@ -3,8 +3,6 @@ _docs/ *.bak .gitattributes .last_checked -.gitconfig -*.bak *.log *~ ~* @@ -66,7 +64,6 @@ coverage.xml *.pot # Django stuff: -*.log local_settings.py # Flask stuff: @@ -178,7 +175,7 @@ _proc/ # MLflow mlruns/ -mlruns.db +mlruns.db* mlflow.db # Reports @@ -189,7 +186,7 @@ reports/ *.nii *.nii.gz # Exception: keep test image used in vision_core -!images/IXI002-Guys-0828-T1.nii.gz +!nbs/images/IXI002-Guys-0828-T1.nii.gz # research/ holds patient data, so it denies by default and allows code back by # type. PACS exports name DICOM files without an extension, so the *.dcm rule @@ -228,3 +225,4 @@ research/vs_seg/ # Local working notes TODO.md +todo.html diff --git a/research/.gitignore b/research/.gitignore new file mode 100644 index 0000000..e8d07de --- /dev/null +++ b/research/.gitignore @@ -0,0 +1,2 @@ +# Temporary presentation build workspaces +/.pptx-build-*/ From f5dbdece36d5c2e0e86e46aff801581d1a47d901 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:26:03 +0100 Subject: [PATCH 10/11] fix(research): repair the spine inference script and refresh research docs - spine/inference_script.py used an undefined name for the save path - spine notebook 00 used an inplace fillna that does nothing in pandas 3 - the VS deployment README test command skipped 24 of 75 tests - the PACS guide named the old vs-seg image and listed outdated files - broken links and Colab badges, a dead Hugging Face Space claim, stale version notes, a hard-coded home path, and the misspelled contributor surname in the VS label docs (the released v1 zip keeps the old spelling) - removed unused aliases and re-exports in the VS workflow package Nothing that ships in the VS PACS container changed. --- research/README.md | 3 ++- research/RESEARCH_PACS_DEPLOYMENT.md | 16 +++++++++++++--- research/endometrial_cancer/README.md | 8 ++++---- .../endometrial_cancer/nbs/02_ec_inference.ipynb | 2 +- .../nbs/04_ec_inference_multi.ipynb | 2 +- ...-spine-segmentation-training-evaluation.ipynb | 2 +- research/spine/README.md | 4 ++-- research/spine/inference_script.py | 2 +- research/vestibular_schwannoma/README.md | 7 ++++--- research/vestibular_schwannoma/data/README.md | 6 +++--- .../data/corrected_labels/LICENSE.txt | 2 +- .../data/corrected_labels/README.md | 2 +- .../deployment/pacs/README.md | 4 ++-- .../vestibular_schwannoma/workflow/config.py | 6 +----- .../workflow/data_preparation.py | 2 +- .../workflow/data_sources.py | 8 -------- 16 files changed, 38 insertions(+), 38 deletions(-) diff --git a/research/README.md b/research/README.md index 84cf5b2..5293459 100644 --- a/research/README.md +++ b/research/README.md @@ -11,4 +11,5 @@ Projects targeting the Research PACS should follow |:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | [00-spine-segmentation-training-evaluation](https://nbviewer.org/github/MMIV-ML/fastMONAI/blob/main/research/spine/00-spine-segmentation-training-evaluation.ipynb)
Notebook to reproduce our test results and facilitate training on your own data. | [![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/research/spine/00-spine-segmentation-training-evaluation.ipynb) | | [01-spine-segmentation-inference](https://nbviewer.org/github/MMIV-ML/fastMONAI/blob/main/research/spine/01-spine-segmentation-inference.ipynb)
Notebook that shows how to run inference on your own data with our trained models. | [![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/research/spine/01-spine-segmentation-inference.ipynb) | -| [02-ec-inference.ipynb](https://nbviewer.org/github/MMIV-ML/fastMONAI/blob/main/research/endometrial_cancer/02-ec-inference.ipynb)
Notebook that shows how to run inference on your own data with our trained model. | [![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/research/endometrial_cancer/02-ec-inference.ipynb) | \ No newline at end of file +| [02_ec_inference](https://nbviewer.org/github/MMIV-ML/fastMONAI/blob/main/research/endometrial_cancer/nbs/02_ec_inference.ipynb)
Notebook that shows how to run inference on your own data with our trained model. | [![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/research/endometrial_cancer/nbs/02_ec_inference.ipynb) | +| [vestibular_schwannoma](vestibular_schwannoma/README.md)
Five-fold cross-validation, inference and Research PACS deployment for vestibular schwannoma segmentation. | Local setup, see the project README | diff --git a/research/RESEARCH_PACS_DEPLOYMENT.md b/research/RESEARCH_PACS_DEPLOYMENT.md index 09d9bba..74cb5c2 100644 --- a/research/RESEARCH_PACS_DEPLOYMENT.md +++ b/research/RESEARCH_PACS_DEPLOYMENT.md @@ -13,7 +13,7 @@ project's values. |-------------|---------|--------------------| | `` | Path to your local fastMONAI checkout | `~/ml_projects/fastMONAI` | | `` | Research project folder under `research/` | `vestibular_schwannoma` | -| `` | Docker image name you build | `vs-seg` | +| `` | Docker image name you build | `fastmonai-vsseg` | | `` | Registered model key | `dynunet` | | `` | UTC build identifier | `20260812T140500Z` | @@ -113,7 +113,7 @@ docker build --pull \ . ``` -For vestibular schwannoma, substitute `vs-seg` and use the project's +For vestibular schwannoma, substitute `fastmonai-vsseg` and use the project's `deployment/pacs/` directory. The date must be a Docker tag, not only part of the exported archive filename. An archive @@ -256,11 +256,16 @@ identity, while `SoftwareVersions` records the producing fastMONAI version. ## Project Structure +The vestibular schwannoma project is laid out as follows. The `redcap_*` and `pr2mask_*` files +are specific to its REDCap and pr2mask outputs. + ``` research// ├── README.md +├── data/ ├── deployment/ │ └── pacs/ +│ ├── .dockerignore │ ├── .ror/ │ │ └── virt/ │ │ └── Dockerfile @@ -271,10 +276,15 @@ research// │ ├── deployment_bundle.py │ ├── deployment_models.py │ ├── dicom_output.py -│ ├── prepare_model_bundle.py │ ├── pacs_inference.py +│ ├── parallel_inference.py +│ ├── prepare_model_bundle.py +│ ├── redcap_output.py +│ ├── redcap_model_instances.json +│ ├── pr2mask_dictionary.csv │ └── model_bundles// # generated locally and ignored by Git ├── notebooks/ +├── scripts/ ├── workflow/ └── tests/ ``` diff --git a/research/endometrial_cancer/README.md b/research/endometrial_cancer/README.md index 2d48b65..bcea226 100644 --- a/research/endometrial_cancer/README.md +++ b/research/endometrial_cancer/README.md @@ -2,8 +2,8 @@ The primary objective of this repository is to reproduce the results reported in [Automated segmentation of endometrial cancer on MR images using deep learning](https://link.springer.com/content/pdf/10.1038/s41598-020-80068-9.pdf) and to integrate this model into research PACS (see [Results for VIBE](#results-for-vibe) section). In addition, we have looked at improving the segmentation performance using multi-sequence MR images (T2w, VIBE, and ADC) (see [Results for multi-sequence (T2, VIBE, and ADC)](#results-for-multi-sequence-t2-vibe-and-adc) section). -The trained weights and exported learner are available on [Huggin Face](https://huggingface.co/skaliy/endometrial_cancer_segmentation). -Test our model live with the Gradio app for VIBE on [Hugging Face Spaces](https://skaliy-endometrial-cancer-segmentation-app.hf.space). To run the application on your local system, follow the instructions in the [Getting started](#getting-started) section. +The trained weights and exported learner are available on [Hugging Face](https://huggingface.co/skaliy/endometrial_cancer_segmentation). +A Gradio app for VIBE is included in `src/app.py`. To run it on your local system, follow the instructions in the [Getting started](#getting-started) section. ## Structure overview ``` @@ -44,11 +44,11 @@ chmod +x inference_script.py python inference_script.py IMG_PATH ``` -If you choose to use `app.py`, please execute the following command in the terminal: +If you choose to use `app.py`, install `gradio` and `opencv-python` (they are not fastMONAI dependencies) and execute the following command in the terminal: ```bash gradio app.py ``` -The predicted maske will be stored in folder `ec_pred` in the `src` folder. +The predicted mask will be stored in folder `ec_pred` in the `src` folder. ## Results for VIBE Note that our results are not directly comparable with the results reported in [study](https://link.springer.com/content/pdf/10.1038/s41598-020-80068-9.pdf), as we opted to use the test set for validation to allocate more data to training. Unlike the approach detailed in the study, we refrained from post-processing steps, such as retaining only the largest object. Predictions from new test cases indicate that this method could occasionally eliminate the tumor. Below is the box plot showcasing predictions on the validation set: diff --git a/research/endometrial_cancer/nbs/02_ec_inference.ipynb b/research/endometrial_cancer/nbs/02_ec_inference.ipynb index 5da5c75..54dec5e 100644 --- a/research/endometrial_cancer/nbs/02_ec_inference.ipynb +++ b/research/endometrial_cancer/nbs/02_ec_inference.ipynb @@ -25,7 +25,7 @@ "id": "628fdd07-948b-419f-b496-5fdcaeb13400", "metadata": {}, "source": [ - "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/master/research/endometrial_cancer/02-ec-inference.ipynb)" + "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/research/endometrial_cancer/nbs/02_ec_inference.ipynb)" ] }, { diff --git a/research/endometrial_cancer/nbs/04_ec_inference_multi.ipynb b/research/endometrial_cancer/nbs/04_ec_inference_multi.ipynb index b4d5e25..9924712 100644 --- a/research/endometrial_cancer/nbs/04_ec_inference_multi.ipynb +++ b/research/endometrial_cancer/nbs/04_ec_inference_multi.ipynb @@ -25,7 +25,7 @@ "id": "6c97c13e-6701-4df7-93c9-0885df93560e", "metadata": {}, "source": [ - "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/master/research/endometrial_cancer/04-ec-inference-multi.ipynb)" + "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MMIV-ML/fastMONAI/blob/main/research/endometrial_cancer/nbs/04_ec_inference_multi.ipynb)" ] }, { diff --git a/research/spine/00-spine-segmentation-training-evaluation.ipynb b/research/spine/00-spine-segmentation-training-evaluation.ipynb index fd56949..f166237 100644 --- a/research/spine/00-spine-segmentation-training-evaluation.ipynb +++ b/research/spine/00-spine-segmentation-training-evaluation.ipynb @@ -647,7 +647,7 @@ "\n", "# Concatenate the original DataFrame and the test DataFrame\n", "df = pd.concat([df, df_test], ignore_index=True)\n", - "df['is_test'].fillna(False, inplace=True)" + "df['is_test'] = df['is_test'].fillna(False)" ] }, { diff --git a/research/spine/README.md b/research/spine/README.md index bff2d37..0de2185 100644 --- a/research/spine/README.md +++ b/research/spine/README.md @@ -1,10 +1,10 @@ # Multi-Center CNN-based spine segmentation from T2w MRI using small amounts of data The source code, designed to reproduce our test results and facilitate training and running inference on your own data. -The trained weights and exported learners are available on [Huggin Face](https://huggingface.co/skaliy/spine-segmentation). +The trained weights and exported learners are available on [Hugging Face](https://huggingface.co/skaliy/spine-segmentation). ## Getting started ### Quick Start -Yyou can run the notebooks directly in Google Colab. +You can run the notebooks directly in Google Colab. ### Local Setup (Optional) 1. Clone this repository: diff --git a/research/spine/inference_script.py b/research/spine/inference_script.py index a277121..ac41476 100755 --- a/research/spine/inference_script.py +++ b/research/spine/inference_script.py @@ -23,7 +23,7 @@ # Set file name from command line argument img_fn = args.fn -save_fn = fn.split('.nii')[0] + '_pred.nii.gz' +save_fn = img_fn.split('.nii')[0] + '_pred.nii.gz' #pred_items org_img, input_img, org_size = med_img_reader(img_fn, apply_reorder, target_spacing, only_tensor=False) diff --git a/research/vestibular_schwannoma/README.md b/research/vestibular_schwannoma/README.md index b0ea454..2da925b 100644 --- a/research/vestibular_schwannoma/README.md +++ b/research/vestibular_schwannoma/README.md @@ -16,7 +16,7 @@ with five-fold cross-validation, optional all-data training, inference, and PACS ## Setup and data -Use fastMONAI 0.10.1 or a matching development checkout. From the fastMONAI repository root: +Use fastMONAI 0.12.0 or a matching development checkout. From the fastMONAI repository root: ```bash pip install -e '.[dev]' @@ -25,8 +25,9 @@ pip install -e '.[dev]' Download and prepare the 344-case dataset with `scripts/prepare_data.py`; see [data/README.md](data/README.md) for scripted acquisition, corrected labels, and reference verification. The default prepared layout is `../nii_data/`, matching -the shared CSV. UNet and DynUNet use MONAI. The training notebook provides -setup instructions for optional SegMamba support. +the shared CSV. UNet and DynUNet use MONAI. SegMamba is optional; the install +commands are in `SEGMAMBA_INSTALL_HINT` in `workflow/models.py`, and training prints +them when SegMamba is requested but not installed. ## Training diff --git a/research/vestibular_schwannoma/data/README.md b/research/vestibular_schwannoma/data/README.md index 0f2f9e5..96b11b2 100644 --- a/research/vestibular_schwannoma/data/README.md +++ b/research/vestibular_schwannoma/data/README.md @@ -140,7 +140,7 @@ The bundled masks are used automatically if they have not already been extracted ```bash python scripts/prepare_data.py prepare \ - --from-existing /home/sathiesh/ml_projects/VS_Seg \ + --from-existing /path/to/VS_Seg \ --output data/prepared ``` @@ -186,14 +186,14 @@ Correction cases: **4, 5, 29, 35, 43, 44, 45, 47, 55, 56, 61, 63, 71, 74, 76, 96, 119**. These are new NL masks drawn for scans whose original tumour annotations contained errors. -Contributors: **Njål Lura** and **Satheshkumar Kaliyugarassan**. +Contributors: **Njål Lura** and **Satheshkumar Kaliyugarasan**. The release uses **CC BY 4.0** and asks readers to cite the accompanying study paper when available. The [`corrected_labels/`](corrected_labels/) folder contains one ZIP, a [short README](corrected_labels/README.md), and a [license notice](corrected_labels/LICENSE.txt). These small masks are versioned with the preparation code on GitHub. Model -weights are distributed separately on Hugging Face. Use a tagged repository +weights are not part of this release. Use a tagged repository release to identify the code, case index, and masks used for a study. Maintainer packaging command: diff --git a/research/vestibular_schwannoma/data/corrected_labels/LICENSE.txt b/research/vestibular_schwannoma/data/corrected_labels/LICENSE.txt index 56c91f2..db1c9e0 100644 --- a/research/vestibular_schwannoma/data/corrected_labels/LICENSE.txt +++ b/research/vestibular_schwannoma/data/corrected_labels/LICENSE.txt @@ -5,7 +5,7 @@ This license permits redistribution and adaptation, including commercial use, with appropriate credit, a license link, and identification of changes. Consult the linked license for its full terms. -Contributors: Njål Lura, Satheshkumar Kaliyugarassan. +Contributors: Njål Lura, Satheshkumar Kaliyugarasan. Image source: Shapey et al. (2021), Vestibular-Schwannoma-SEG version 2, The Cancer Imaging Archive, DOI: 10.7937/TCIA.9YTJ-5Q73 (CC BY 4.0). These masks are new annotations for scans with errors in the original annotations. diff --git a/research/vestibular_schwannoma/data/corrected_labels/README.md b/research/vestibular_schwannoma/data/corrected_labels/README.md index fc859c3..f515097 100644 --- a/research/vestibular_schwannoma/data/corrected_labels/README.md +++ b/research/vestibular_schwannoma/data/corrected_labels/README.md @@ -7,7 +7,7 @@ license: cc-by-4.0 17 newly drawn tumour masks for Queen Square scans whose original annotations contained errors. These are the NL masks used in our study. -**Contributors:** Njål Lura and Satheshkumar Kaliyugarassan. +**Contributors:** Njål Lura and Satheshkumar Kaliyugarasan. Download [vs_corrected_labels_v1.zip](vs_corrected_labels_v1.zip). It contains binary NIfTI masks in native CE-T1 geometry (**0: background, 1: tumour**), a case diff --git a/research/vestibular_schwannoma/deployment/pacs/README.md b/research/vestibular_schwannoma/deployment/pacs/README.md index f574396..c03d777 100644 --- a/research/vestibular_schwannoma/deployment/pacs/README.md +++ b/research/vestibular_schwannoma/deployment/pacs/README.md @@ -40,11 +40,11 @@ Published series keep pr2mask's UIDs. ## Test and build -Run the deployment tests first: +Run the deployment tests first (the test discovery runs from `research/`): ```bash conda activate fastmonai -python -m unittest discover -s ../../tests/deployment -p 'test_*.py' +(cd ../../.. && python -m unittest discover -s vestibular_schwannoma/tests/deployment -t .) bash -n entrypoint.sh ``` diff --git a/research/vestibular_schwannoma/workflow/config.py b/research/vestibular_schwannoma/workflow/config.py index a7286ef..6330409 100644 --- a/research/vestibular_schwannoma/workflow/config.py +++ b/research/vestibular_schwannoma/workflow/config.py @@ -11,11 +11,7 @@ make_output_spec, ) from .models import TRAINING_MODEL_CONFIGS -from .run_selection import ( - CROSS_VALIDATION_FOLDS, - INFERENCE_RUN_IDS_FILENAME as INFERENCE_RUN_IDS_FILENAME, - INFERENCE_RUN_IDS_SCHEMA as INFERENCE_RUN_IDS_SCHEMA, -) +from .run_selection import CROSS_VALIDATION_FOLDS VS_OUTPUT_SPEC = make_output_spec("multiclass_segmentation", classes=2) diff --git a/research/vestibular_schwannoma/workflow/data_preparation.py b/research/vestibular_schwannoma/workflow/data_preparation.py index c9f0441..efe08ae 100644 --- a/research/vestibular_schwannoma/workflow/data_preparation.py +++ b/research/vestibular_schwannoma/workflow/data_preparation.py @@ -29,7 +29,7 @@ f"vs_gk_{i}" for i in [4, 5, 29, 35, 43, 44, 45, 47, 55, 56, 61, 63, 71, 74, 76, 96, 119] ) -LABEL_CONTRIBUTORS = ("Njål Lura", "Satheshkumar Kaliyugarassan") +LABEL_CONTRIBUTORS = ("Njål Lura", "Satheshkumar Kaliyugarasan") LABEL_RELEASE_LICENSE = "CC BY 4.0" LABEL_RELEASE_LICENSE_URL = "https://creativecommons.org/licenses/by/4.0/" COMPONENT_SIGNATURES = { diff --git a/research/vestibular_schwannoma/workflow/data_sources.py b/research/vestibular_schwannoma/workflow/data_sources.py index beae40c..ddca302 100644 --- a/research/vestibular_schwannoma/workflow/data_sources.py +++ b/research/vestibular_schwannoma/workflow/data_sources.py @@ -111,11 +111,6 @@ def asset(self, name: str) -> SourceAsset: ), ) -# Short aliases for callers that use cohort names rather than repository names. -TCIA_VS_SEG = TCIA_VESTIBULAR_SCHWANNOMA_SEG -QUEEN_SQUARE = TCIA_VESTIBULAR_SCHWANNOMA_SEG -TILBURG = CROSSMODA_2022 - DATASET_SOURCES = { TCIA_VESTIBULAR_SCHWANNOMA_SEG.key: TCIA_VESTIBULAR_SCHWANNOMA_SEG, CROSSMODA_2022.key: CROSSMODA_2022, @@ -126,9 +121,6 @@ def asset(self, name: str) -> SourceAsset: "CROSSMODA_2022", "DATASET_SOURCES", "DatasetSource", - "QUEEN_SQUARE", "SourceAsset", "TCIA_VESTIBULAR_SCHWANNOMA_SEG", - "TCIA_VS_SEG", - "TILBURG", ] From ce6a60846da433d3d5def86c46c5010a9df54549 Mon Sep 17 00:00:00 2001 From: Sathiesh Date: Tue, 27 Oct 2026 09:40:27 +0100 Subject: [PATCH 11/11] build: bump version to 0.13.0 Minor bump because this release removes public API and changes what GPU patch augmentation does with per-axis degrees, translation and scales. --- 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 ea370a8..f23a6b3 100644 --- a/fastMONAI/__init__.py +++ b/fastMONAI/__init__.py @@ -1 +1 @@ -__version__ = "0.12.0" +__version__ = "0.13.0" diff --git a/settings.ini b/settings.ini index e4928aa..ecf38f5 100644 --- a/settings.ini +++ b/settings.ini @@ -5,7 +5,7 @@ ### Python Library ### lib_name = fastMONAI min_python = 3.10 -version = 0.12.0 +version = 0.13.0 ### 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