From 8b017cc266b9cb3ab05dd8d65951cf4ba7b10b25 Mon Sep 17 00:00:00 2001 From: Robert Graf Date: Mon, 10 Nov 2025 17:09:28 +0000 Subject: [PATCH 01/29] bug fixes, add more paramter options - robert --- .gitignore | 4 +++- spineps/architectures/pl_unet.py | 1 - spineps/phase_post.py | 25 +++++++++++++++---------- spineps/seg_model.py | 6 +++++- spineps/seg_run.py | 8 ++++++-- spineps/seg_utils.py | 5 ++--- 6 files changed, 31 insertions(+), 18 deletions(-) diff --git a/.gitignore b/.gitignore index 8281c7d..107518b 100755 --- a/.gitignore +++ b/.gitignore @@ -157,4 +157,6 @@ lightning_logs/ test.txt derivatives_seg robert_test.py -poetry.lock \ No newline at end of file +poetry.lock + +*DS_Store* \ No newline at end of file diff --git a/spineps/architectures/pl_unet.py b/spineps/architectures/pl_unet.py index ca4415f..1faad73 100755 --- a/spineps/architectures/pl_unet.py +++ b/spineps/architectures/pl_unet.py @@ -6,7 +6,6 @@ import torch import torchmetrics.functional as mF # noqa: N812 from torch import nn -from torch.optim import lr_scheduler from spineps.architectures.unet3D import Unet3D diff --git a/spineps/phase_post.py b/spineps/phase_post.py index ce0ed98..149c9ff 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -236,12 +236,15 @@ def assign_missing_cc( return target_arr, reference_arr, deletion_map -def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, verbose=True): +def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=False, verbose=True): # PIR orientation = whole_vert_nii.orientation vert_t = whole_vert_nii.reorient() seg_t = seg_nii.reorient() - vert_labels = [t for t in vert_t.unique() if t <= 26 or t == 28] # without zero + if include_sacrum: + vert_labels = [t for t in vert_t.unique() if t < 40] # without zero + else: + vert_labels = [t for t in vert_t.unique() if t <= 26 or t == 28] # without zero vert_arr = vert_t.get_seg_array() subreg_arr = seg_t.get_seg_array() @@ -392,18 +395,20 @@ def label_instance_top_to_bottom(vert_nii: NII, labeling_offset: int = 0): return vert_nii, vert_nii.unique() -def assign_vertebra_inconsistency(seg_nii: NII, vert_nii: NII): +def assign_vertebra_inconsistency(seg_nii: NII, vert_nii: NII, verbose=False, locs=None): + if locs is None: + locs = [ + Location.Superior_Articular_Left, + Location.Superior_Articular_Right, + Location.Inferior_Articular_Left, + Location.Inferior_Articular_Right, + ] seg_nii.assert_affine(shape=vert_nii.shape) seg_arr = seg_nii.get_seg_array() vert_arr = vert_nii.get_seg_array() # assign inconsistent substructures - for loc in [ - Location.Superior_Articular_Left, - Location.Superior_Articular_Right, - Location.Inferior_Articular_Left, - Location.Inferior_Articular_Right, - ]: + for loc in locs: value = loc.value subreg_l = np_extract_label(seg_arr, value, inplace=False) # type:ignore @@ -436,7 +441,7 @@ def assign_vertebra_inconsistency(seg_nii: NII, vert_nii: NII): vert_arr[cc_map == 1] = to_label logger.print( - f"set cc to {to_label}, with volume decision {gt_volume}, based on {biggest_volume}, {second_volume}", verbose=False + f"set cc to {to_label}, with volume decision {gt_volume}, based on {biggest_volume}, {second_volume}", verbose=verbose ) vert_nii.set_array_(vert_arr) diff --git a/spineps/seg_model.py b/spineps/seg_model.py index 3a8c86f..a24cc45 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -123,7 +123,10 @@ def segment_scan( dict[OutputType, NII]: _description_ """ if self.predictor is None: - self.load() + try: + self.load() + except FileNotFoundError: + self.load(folds=["all"]) assert self.predictor is not None, "self.predictor == None after load(). Error!" # Check if input matches expectation @@ -333,6 +336,7 @@ def run( target = from_numpy(arr) target[target == 26] = 0 + target[target >= 10] = 0 do_backup = False # channel-wise diff --git a/spineps/seg_run.py b/spineps/seg_run.py index 6be60bd..58192b8 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -290,6 +290,8 @@ def process_img_nii( # noqa: C901 return_output_instead_of_save: bool = False, crop=None, verbose: bool = False, + _nii=None, + _end_after_subreg=False, ) -> tuple[dict[str, Path], ErrCode]: """Runs the SPINEPS framework over one nifty @@ -386,11 +388,12 @@ def process_img_nii( # noqa: C901 with logger: if verbose: model_semantic.logger.default_verbose = True - input_nii = img_ref.open_nii() + input_nii = _nii if _nii is not None else img_ref.open_nii() input_nii.seg = False input_nii_ = input_nii.copy() if crop is not None: try: + logger.print("Input image manuel crop", crop, "from", input_nii.shape) input_nii = input_nii.apply_crop(crop) except Exception: pass @@ -445,7 +448,8 @@ def process_img_nii( # noqa: C901 logger.print("Subreg Mask already exists. Set -override_subreg to create it anew") seg_nii_modelres = NII.load(out_spine_raw, seg=True) print("seg_nii", seg_nii_modelres.zoom, seg_nii_modelres.orientation, seg_nii_modelres.shape) - + if _end_after_subreg: + return output_paths, ErrCode.OK # Second stage if not out_vert_raw.exists() or override_instance: whole_vert_nii, errcode = predict_instance_mask( diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index 6dda7de..ab66a88 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -3,8 +3,7 @@ # from utils.predictor import nnUNetPredictor from typing import Union -import nibabel as nib -from TPTBox import BIDS_FILE, NII, ZOOMS, Log_Type +from TPTBox import BIDS_FILE, ZOOMS, Log_Type from spineps.seg_enums import Acquisition, Modality from spineps.seg_model import Segmentation_Model @@ -157,7 +156,7 @@ def check_input_model_compatibility( img_nii = img_ref.open_nii() if img_nii.get_plane() not in ["iso", *allowed_acq]: - logger_texts.append(f"input get_plane() is not 'iso' or one of the expected {allowed_acq}.") + logger_texts.append(f"input {img_nii.get_plane()=} is not 'iso' or one of the expected {allowed_acq}.") compatible = False if len(logger_texts) > 1: From 79d67771d701cd76eb4cda8b664afb82a04183d0 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 15:39:31 +0000 Subject: [PATCH 02/29] added prepare vertexact --- spineps/phase_labeling.py | 46 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/spineps/phase_labeling.py b/spineps/phase_labeling.py index 0223ff6..5ad239f 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -271,6 +271,41 @@ def prepare_vert( return softmax_values +def prepare_vertexact( + vert_softmax_values: np.ndarray, + gaussian_sigma: float = 0.85, + gaussian_radius: int = 2, + gaussian_regionwise: bool = True, +) -> np.ndarray: + """Smooth and normalize a per-vertebra-class softmax vector. + + Optionally applies a 1-D Gaussian filter (either per spinal region or across all classes) and then normalizes to sum to 1. + + Args: + vert_softmax_values (np.ndarray): Length-``VERT_CLASSES`` per-class softmax values. + gaussian_sigma (float): Gaussian smoothing sigma; 0 disables smoothing. + gaussian_radius (int): Half-width of the Gaussian kernel. + gaussian_regionwise (bool): If True, smooth each spinal region independently instead of across the whole vector. + + Returns: + np.ndarray: The smoothed, sum-normalized per-class vector. + """ + # gaussian region-wise + softmax_values = vert_softmax_values.copy() + softmax_values[18] += softmax_values[19] # add t13 to t12 + softmax_values[24] += softmax_values[25] # add l6 to l5 + # remove 19 and 25 from the entire array, because they are not real classes and would mess up the smoothing + softmax_values = np.delete(softmax_values, [19, 25], axis=0) + if gaussian_sigma > 0.0: + if gaussian_regionwise: + for s in [CERV, THOR, LUMB]: + softmax_values[s] = gaussian_filter1d(softmax_values[s], sigma=gaussian_sigma, mode="nearest", radius=gaussian_radius) + else: + softmax_values = gaussian_filter1d(softmax_values, sigma=gaussian_sigma, mode="nearest", radius=gaussian_radius) + softmax_values /= np.sum(softmax_values) + DIVIDE_BY_ZERO_OFFSET + return softmax_values + + def prepare_vertgrp( vertgrp_softmax_values: np.ndarray, gaussian_sigma: float = 0.85, @@ -562,6 +597,16 @@ def find_vert_path_from_predictions( vert_w, ) + vertex_softmax_output = k["soft"]["VERTEX"] if "VERTEX" in predict_keys else np.zeros(len(VertExact)) + vertex_values = np.multiply( + prepare_vertexact( + vertex_softmax_output, + gaussian_sigma=vert_gaussian_sigma, + gaussian_regionwise=vert_gaussian_regionwise, + ), + vert_w, + ) + vertgrp_softmax_output = k["soft"]["VERTGRP"] if "VERTGRP" in predict_keys else np.zeros(len(VertGroup)) vertgrp_values = np.multiply( prepare_vertgrp( @@ -585,6 +630,7 @@ def find_vert_path_from_predictions( # # add region and vert final_vert_pred = np.add(region_values, vert_values) + final_vert_pred = np.add(final_vert_pred, vertex_values) final_vert_pred = np.add(final_vert_pred, vertgrp_values) # normalize final_vert_pred /= np.sum(final_vert_pred) + DIVIDE_BY_ZERO_OFFSET From 54d1113e3cc72ad3d7833a906692aa356af56b76 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 19:43:37 +0000 Subject: [PATCH 03/29] fix: prepare_vertexact VERTEX-head fallback used a wrong-sized vector prepare_vertexact() expects a 26-element VertExactClass vector (T13 at index 19, L6 at index 25, merged then deleted), but the VERTEX-head fallback in find_vert_path_from_predictions allocated np.zeros(len(VertExact))=24, crashing the whole labeling path with 'IndexError: index 24 is out of bounds for axis 0 with size 24'. Use len(VertExactClass) for the fallback. Fixes the two failing Test_Labeling_Inference_Mocked tests. Co-Authored-By: Claude Opus 4.8 --- spineps/phase_labeling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spineps/phase_labeling.py b/spineps/phase_labeling.py index 5ad239f..65fc413 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -597,7 +597,7 @@ def find_vert_path_from_predictions( vert_w, ) - vertex_softmax_output = k["soft"]["VERTEX"] if "VERTEX" in predict_keys else np.zeros(len(VertExact)) + vertex_softmax_output = k["soft"]["VERTEX"] if "VERTEX" in predict_keys else np.zeros(len(VertExactClass)) vertex_values = np.multiply( prepare_vertexact( vertex_softmax_output, From ba0cc3cb04f55c6b43d1230503c1ef4f355d726d Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 19:43:48 +0000 Subject: [PATCH 04/29] refactor: replace user-input asserts with descriptive exceptions + cleanups Phase 1 usability/error-proofing pass: - Raise FileNotFoundError/NotADirectoryError/ValueError/RuntimeError at user-input boundaries instead of bare asserts (which vanish under python -O): entrypoint, filepaths, seg_model, get_models, inference_api, seg_run. Docstring Raises sections updated to match. - Clearer model errors: download failures -> RuntimeError naming model id + url; no models installed -> FileNotFoundError; unknown id -> KeyError listing available options. - Remove leftover print(opt) debug dump; fix stale '-model_vert' message; fix typos; document the '-ms auto' option; drop dead model_instance=='auto' branch. - Route stray print() through the shared logger. - Add 6 tests covering the new exception behavior. Co-Authored-By: Claude Opus 4.8 --- spineps/entrypoint.py | 49 +++++++++++++++++++--------------- spineps/get_models.py | 25 ++++++++++------- spineps/phase_post.py | 2 +- spineps/seg_model.py | 15 +++++++---- spineps/seg_run.py | 3 ++- spineps/utils/filepaths.py | 21 +++++++++++---- spineps/utils/inference_api.py | 3 ++- unit_tests/test_entrypoint.py | 19 +++++++++++-- unit_tests/test_filepaths.py | 34 +++++++++++++++++++++-- 9 files changed, 123 insertions(+), 48 deletions(-) diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index 0349962..7268fca 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -109,7 +109,7 @@ def entry_point(): main_parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) cmdparsers = main_parser.add_subparsers(title="cmd", help="Possible subcommands", dest="cmd", required=True) parser_sample = cmdparsers.add_parser( - "sample", help="Process a single image nifty", formatter_class=argparse.ArgumentDefaultsHelpFormatter + "sample", help="Process a single NIfTI image", formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser_dataset = cmdparsers.add_parser( "dataset", help="Process a whole dataset", formatter_class=argparse.ArgumentDefaultsHelpFormatter @@ -125,7 +125,7 @@ def entry_point(): required=True, # choices=modelids_semantic, metavar="", - help="The model used for the semantic segmentation. You can also pass an absolute path the model folder", + help="The model used for the semantic segmentation. You can also pass an absolute path to the model folder", ) parser_sample.add_argument( "-model_instance", @@ -135,7 +135,7 @@ def entry_point(): # required=True, # choices=modelids_instance, metavar="", - help="The model used for the vertebra instance segmentation. You can also pass an absolute path the model folder", + help="The model used for the vertebra instance segmentation. You can also pass an absolute path to the model folder", ) parser_sample.add_argument( "-model_labeling", @@ -145,7 +145,7 @@ def entry_point(): # required=True, # choices=modelids_instance, metavar="", - help="The model used for the vertebra labeling classification. You can also pass an absolute path the model folder", + help="The model used for the vertebra labeling classification. You can also pass an absolute path to the model folder", ) parser_sample = parser_arguments(parser_sample) @@ -163,7 +163,7 @@ def entry_point(): default="t2w", # choices=model_subreg_choices, metavar="", - help="The model used for the subregion segmentation. You can also pass an absolute path the model folder", + help="The model used for the subregion segmentation. Pass 'auto' to auto-select a model by modality, or an absolute path to the model folder", ) parser_dataset.add_argument( "-model_instance", @@ -172,7 +172,7 @@ def entry_point(): default="instance", # choices=model_vert_choices, metavar="", - help="The model used for the vertebra segmentation. You can also pass an absolute path the model folder", + help="The model used for the vertebra segmentation. You can also pass an absolute path to the model folder", ) parser_dataset.add_argument( "-model_labeling", @@ -182,7 +182,7 @@ def entry_point(): # required=True, # choices=modelids_instance, metavar="", - help="The model used for the vertebra labeling classification. You can also pass an absolute path the model folder", + help="The model used for the vertebra labeling classification. You can also pass an absolute path to the model folder", ) parser_dataset.add_argument( "-ignore_bids_filter", @@ -209,9 +209,8 @@ def entry_point(): # ########################### opt = main_parser.parse_args() - print(opt) - print() - # print(opt) + if opt.verbose: + logger.print("Parsed arguments:", opt) if opt.cmd == "sample": run_sample(opt) elif opt.cmd == "dataset": @@ -235,17 +234,20 @@ def run_sample(opt: Namespace): int: ``1`` on completion. Raises: - AssertionError: If the input path's parent directory is missing, only a filename was given, or the - input file does not exist. + ValueError: If only a filename was given instead of a path to the file. + FileNotFoundError: If the input path's parent directory is missing, or the input file does not exist. """ input_path = Path(opt.input).absolute() dataset = str(input_path.parent) - assert os.path.exists(dataset), f"-input parent does not exist, got {dataset}" # noqa: PTH110 - assert dataset != "", f"-input you only gave a filename, not a direction to the file, got {input_path}" + if dataset == "": + raise ValueError(f"-input you only gave a filename, not a path to the file, got {input_path}") + if not os.path.exists(dataset): # noqa: PTH110 + raise FileNotFoundError(f"-input parent directory does not exist, got {dataset}") input_path = str(input_path) if not input_path.endswith(".nii.gz"): input_path += ".nii.gz" - assert os.path.isfile(input_path), f"-input does not exist or is not a file, got {input_path}" # noqa: PTH113 + if not os.path.isfile(input_path): # noqa: PTH113 + raise FileNotFoundError(f"-input does not exist or is not a file, got {input_path}") # model semantic if "/" in str(opt.model_semantic): model_semantic = get_actual_model(opt.model_semantic, use_cpu=opt.cpu).load() @@ -325,11 +327,15 @@ def run_dataset(opt: Namespace): int: ``1`` on completion. Raises: - AssertionError: If the directory does not exist, is not a directory, or no instance model is resolved. + FileNotFoundError: If the directory does not exist. + NotADirectoryError: If the given path is not a directory. + ValueError: If no instance model could be resolved. """ input_dir = Path(opt.directory) - assert input_dir.exists(), f"-input does not exist, {input_dir}" - assert input_dir.is_dir(), f"-input is not a directory, got {input_dir}" + if not input_dir.exists(): + raise FileNotFoundError(f"-directory does not exist, got {input_dir}") + if not input_dir.is_dir(): + raise NotADirectoryError(f"-directory is not a directory, got {input_dir}") # Model semantic if opt.model_semantic == "auto": @@ -340,9 +346,7 @@ def run_dataset(opt: Namespace): model_semantic = get_semantic_model(opt.model_semantic, use_cpu=opt.cpu).load() # Model Instance - if opt.model_instance == "auto": - model_instance = None - elif "/" in str(opt.model_instance): + if "/" in str(opt.model_instance): model_instance = get_actual_model(opt.model_instance, use_cpu=opt.cpu).load() else: model_instance = get_instance_model(opt.model_instance, use_cpu=opt.cpu).load() @@ -355,7 +359,8 @@ def run_dataset(opt: Namespace): else: model_labeling = get_labeling_model(opt.model_labeling, use_cpu=opt.cpu).load() - assert model_instance is not None, "-model_vert was None" + if model_instance is None: + raise ValueError("-model_instance/-mv resolved to None; pass a valid instance model id or path") kwargs = { "dataset_path": input_dir, diff --git a/spineps/get_models.py b/spineps/get_models.py index 8940982..47ac279 100755 --- a/spineps/get_models.py +++ b/spineps/get_models.py @@ -42,14 +42,17 @@ def _get_model_by_name( possible_keys = list(modelid2folder.keys()) if len(possible_keys) == 0: logger.print(_NO_MODELS_AVAILABLE_MSG.format(kind=kind), Log_Type.FAIL) - raise KeyError(model_name) + raise FileNotFoundError(_NO_MODELS_AVAILABLE_MSG.format(kind=kind)) if model_name not in possible_keys: logger.print(f"Model with name {model_name} does not exist, options are {possible_keys}", Log_Type.FAIL) - raise KeyError(model_name) + raise KeyError(f"Model '{model_name}' does not exist. Available {kind} models: {possible_keys}") config_path = modelid2folder[model_name] if str(config_path).startswith("http"): - # Resolve HTTP - config_path = download_if_missing(model_name, config_path, phase=phase) + # Resolve HTTP (download the model weights on first use) + try: + config_path = download_if_missing(model_name, config_path, phase=phase) + except Exception as e: + raise RuntimeError(f"Failed to download model '{model_name}' from {config_path}: {e}") from e return get_actual_model(config_path, **kwargs) @@ -64,7 +67,8 @@ def get_semantic_model(model_name: str, **kwargs) -> Segmentation_Model: Segmentation_Model: The instantiated semantic model. Raises: - KeyError: If no model with the given name is available. + KeyError: If the given model name is not among the available models. + FileNotFoundError: If no models of this kind are installed at all. """ return _get_model_by_name(model_name, modelid2folder_semantic(), SpinepsPhase.SEMANTIC, "semantic", **kwargs) @@ -80,7 +84,8 @@ def get_instance_model(model_name: str, **kwargs) -> Segmentation_Model: Segmentation_Model: The instantiated instance model. Raises: - KeyError: If no model with the given name is available. + KeyError: If the given model name is not among the available models. + FileNotFoundError: If no models of this kind are installed at all. """ return _get_model_by_name(model_name, modelid2folder_instance(), SpinepsPhase.INSTANCE, "instance", **kwargs) @@ -96,7 +101,8 @@ def get_labeling_model(model_name: str, **kwargs) -> VertLabelingClassifier: VertLabelingClassifier: The instantiated labeling classifier. Raises: - KeyError: If no model with the given name is available. + KeyError: If the given model name is not among the available models. + FileNotFoundError: If no models of this kind are installed at all. """ return _get_model_by_name(model_name, modelid2folder_labeling(), SpinepsPhase.LABELING, "labeling", **kwargs) @@ -167,12 +173,13 @@ def check_available_models( id-to-folder maps. Raises: - AssertionError: If models_folder does not exist. + FileNotFoundError: If models_folder does not exist. """ logger.print("Check available models...") if isinstance(models_folder, str): models_folder = Path(models_folder) - assert models_folder.exists(), f"models_folder {models_folder} does not exist" + if not models_folder.exists(): + raise FileNotFoundError(f"models_folder {models_folder} does not exist") config_paths = search_path(models_folder, query="**/inference_config.json", suppress=True) global _modelid2folder_semantic, _modelid2folder_instance, _modelid2folder_labeling # noqa: PLW0603 diff --git a/spineps/phase_post.py b/spineps/phase_post.py index 71f48af..2e37c62 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -599,7 +599,7 @@ def assign_vertebra_inconsistency( try: subreg_cc, _ = np_connected_components(subreg_l, label_ref=1) except AssertionError as e: - print(f"Got error {e}, skip") + logger.print(f"Connected-components failed, skipping remaining inconsistency assignment: {e}", Log_Type.WARNING) break cc_labels = np_unique(subreg_cc) diff --git a/spineps/seg_model.py b/spineps/seg_model.py index 6159bdc..04d986f 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -66,10 +66,11 @@ def __init__( default_allow_tqdm (bool, optional): If true, shows a progress bar while segmenting. Defaults to True. Raises: - AssertionError: If model_folder does not exist. + FileNotFoundError: If model_folder does not exist. """ self.name: str = "" - assert Path(model_folder).exists(), f"model_folder does not exist, got {model_folder}" + if not Path(model_folder).exists(): + raise FileNotFoundError(f"model_folder does not exist, got {model_folder}") self.logger = No_Logger() self.use_cpu = use_cpu @@ -461,12 +462,16 @@ def load(self, folds: tuple[str, ...] | None = None) -> Self: # noqa: ARG002 Self: This model with its 3D U-Net predictor loaded and moved to the selected device. Raises: - AssertionError: If exactly one checkpoint file is not found in the model folder. + FileNotFoundError: If the model folder is missing or does not contain exactly one checkpoint file. """ - assert os.path.exists(self.model_folder) # noqa: PTH110 + if not os.path.exists(self.model_folder): # noqa: PTH110 + raise FileNotFoundError(f"model_folder does not exist, got {self.model_folder}") chktpath = search_path(self.model_folder, "**/*weights*.ckpt") - assert len(chktpath) == 1, chktpath + if len(chktpath) != 1: + raise FileNotFoundError( + f"expected exactly one '*weights*.ckpt' checkpoint in {self.model_folder}, found {len(chktpath)}: {chktpath}" + ) try: model = PLNet.load_from_checkpoint(checkpoint_path=chktpath[0], weights_only=False) except RuntimeError: diff --git a/spineps/seg_run.py b/spineps/seg_run.py index 01597f9..6bd6310 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -133,7 +133,8 @@ def process_dataset( # INITIALIZATION if not isinstance(modalities, list): modalities = [modalities] - assert len(modalities) > 0, "you must specifiy the modalities to be segmented!" + if len(modalities) == 0: + raise ValueError("you must specify the modalities to be segmented!") if snapshot_copy_folder is True: snapshot_copy_folder = dataset_path.joinpath("snaps_seg") diff --git a/spineps/utils/filepaths.py b/spineps/utils/filepaths.py index 89380f7..9b582d3 100755 --- a/spineps/utils/filepaths.py +++ b/spineps/utils/filepaths.py @@ -7,6 +7,10 @@ from itertools import chain from pathlib import Path +from TPTBox import No_Logger + +logger = No_Logger(prefix="filepaths") + spineps_environment_path_override = None # Path( # "/DATA/NAS/ongoing_projects/hendrik/mri_usage/models/" # ) # None # You can put an absolute path to the model weights here instead of using environment variable @@ -19,6 +23,10 @@ def get_mri_segmentor_models_dir() -> Path: Returns: Path: Path to the overall models folder + + Raises: + RuntimeError: If no models directory could be determined from the environment variable, override or backup. + FileNotFoundError: If the resolved models directory does not exist. """ folder_path = ( os.environ.get("SPINEPS_SEGMENTOR_MODELS") @@ -28,11 +36,14 @@ def get_mri_segmentor_models_dir() -> Path: if folder_path is None and spineps_environment_path_backup is not None: folder_path = spineps_environment_path_backup - assert folder_path is not None, ( - "Environment variable 'SPINEPS_SEGMENTOR_MODELS' is not defined. Setup the environment variable as stated in the readme or set the override in utils.filepaths.py" - ) + if folder_path is None: + raise RuntimeError( + "Environment variable 'SPINEPS_SEGMENTOR_MODELS' is not defined. Setup the environment variable as stated " + "in the readme or set the override in utils.filepaths.py" + ) folder_path = Path(folder_path) - assert folder_path.exists(), f"Environment variable 'SPINEPS_SEGMENTOR_MODELS' = {folder_path} does not exist" + if not folder_path.exists(): + raise FileNotFoundError(f"Environment variable 'SPINEPS_SEGMENTOR_MODELS' = {folder_path} does not exist") return folder_path @@ -75,7 +86,7 @@ def search_path(basepath: str | Path, query: str, verbose: bool = False, suppres basepath = str(basepath) if not basepath.endswith("/"): basepath += "/" - print(f"search_path: in {basepath}{query}") if verbose else None + logger.print(f"search_path: in {basepath}{query}", verbose=verbose) paths = sorted(chain(list(Path(f"{basepath}").glob(f"{query}")))) if len(paths) == 0 and not suppress: warnings.warn(f"did not find any paths in {basepath}{query}", UserWarning, stacklevel=1) diff --git a/spineps/utils/inference_api.py b/spineps/utils/inference_api.py index 1aaab86..7d61415 100755 --- a/spineps/utils/inference_api.py +++ b/spineps/utils/inference_api.py @@ -58,7 +58,8 @@ def load_inf_model( else: device = torch.device("mps") - assert model_folder.exists(), f"model-folder not found: got path {model_folder}" + if not model_folder.exists(): + raise FileNotFoundError(f"model-folder not found: got path {model_folder}") predictor = nnUNetPredictor( tile_step_size=step_size, diff --git a/unit_tests/test_entrypoint.py b/unit_tests/test_entrypoint.py index 4400b3c..c0ce729 100644 --- a/unit_tests/test_entrypoint.py +++ b/unit_tests/test_entrypoint.py @@ -6,10 +6,8 @@ import argparse import unittest -from pathlib import Path from TPTBox import No_Logger -from typing_extensions import Self from spineps import entrypoint @@ -19,3 +17,20 @@ class Test_EntryPoint(unittest.TestCase): def test_normal(self): entrypoint.parser_arguments(argparse.ArgumentParser()) + + def test_run_sample_missing_parent_raises(self): + # parent directory of the input does not exist -> FileNotFoundError (not a bare AssertionError) + opt = argparse.Namespace(input="/this/path/does/not/exist/scan.nii.gz") + with self.assertRaises(FileNotFoundError): + entrypoint.run_sample(opt) + + def test_run_dataset_missing_directory_raises(self): + opt = argparse.Namespace(directory="/this/path/does/not/exist") + with self.assertRaises(FileNotFoundError): + entrypoint.run_dataset(opt) + + def test_run_dataset_file_not_directory_raises(self): + # an existing path that is a file, not a directory -> NotADirectoryError + opt = argparse.Namespace(directory=__file__) + with self.assertRaises(NotADirectoryError): + entrypoint.run_dataset(opt) diff --git a/unit_tests/test_filepaths.py b/unit_tests/test_filepaths.py index b2bf433..0d8b974 100644 --- a/unit_tests/test_filepaths.py +++ b/unit_tests/test_filepaths.py @@ -5,11 +5,12 @@ from __future__ import annotations import os +import tempfile import unittest from pathlib import Path import spineps -from spineps.get_models import Segmentation_Model, get_actual_model +from spineps.get_models import Segmentation_Model, check_available_models, get_actual_model from spineps.utils.filepaths import ( filepath_model, get_mri_segmentor_models_dir, @@ -54,5 +55,34 @@ def test_env_path(self): self.assertEqual(p, spineps_environment_path_override) else: self.assertEqual(str(p) + "/", os.environ.get("SPINEPS_SEGMENTOR_MODELS")) - except AssertionError as e: + except (AssertionError, FileNotFoundError, RuntimeError) as e: print(e) + + +class Test_filepaths_errors(unittest.TestCase): + """User-input boundaries should raise descriptive exceptions, not bare AssertionError.""" + + def test_check_available_models_missing_dir(self): + with self.assertRaises(FileNotFoundError): + check_available_models("/this/path/does/not/exist/models") + + def test_get_actual_model_without_config(self): + with tempfile.TemporaryDirectory() as d, self.assertRaises(FileNotFoundError): + get_actual_model(d) + + def test_models_dir_nonexistent_env_path(self): + import spineps.utils.filepaths as fp + + old_env = os.environ.get("SPINEPS_SEGMENTOR_MODELS") + old_override = fp.spineps_environment_path_override + try: + fp.spineps_environment_path_override = None + os.environ["SPINEPS_SEGMENTOR_MODELS"] = "/this/path/does/not/exist/models" + with self.assertRaises(FileNotFoundError): + fp.get_mri_segmentor_models_dir() + finally: + fp.spineps_environment_path_override = old_override + if old_env is None: + os.environ.pop("SPINEPS_SEGMENTOR_MODELS", None) + else: + os.environ["SPINEPS_SEGMENTOR_MODELS"] = old_env From a6e6bbf73b4fad9aa44627e4c9c6c78ba0d9fe9d Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 20:03:55 +0000 Subject: [PATCH 05/29] perf: batched instance inference (one forward per chunk of cutouts) The instance phase ran the 3D U-Net once per vertebra cutout (7-33 sequential GPU forward passes per case). Every cutout has the same fixed cutout_size, so they now stack into batched forward passes. - Segmentation_Model_Unet3D.run_batch stacks up to batch_size equally shaped cutouts into one forward under torch.inference_mode() (also a free memory/speed win); run() delegates to run_batch so single- and batched-input paths share one implementation and are identical by construction. Falls back to one-by-one on CUDA OOM. - segment_scan_batch: base class loops segment_scan (unchanged for non-batched models); Unet3D overrides with the batched path. - collect_vertebra_predictions gathers all cutouts then runs one batched call. New proc_inst_batch_size knob (default 4), threaded through predict_instance_mask and process_img_nii. - Tests prove run_batch == looping run, segment_scan_batch == segment_scan, batch-size invariance, and that batching really reduces forward calls (5 cutouts @ bs=2 -> 3 forwards). Numerically identical in fp32. Co-Authored-By: Claude Opus 4.8 --- spineps/phase_instance.py | 48 +++---- spineps/seg_model.py | 188 +++++++++++++++++++++++----- spineps/seg_run.py | 4 + unit_tests/test_inference_mocked.py | 97 +++++++++++++- 4 files changed, 285 insertions(+), 52 deletions(-) diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index 9d3b4c9..b918343 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -57,6 +57,7 @@ def predict_instance_mask( proc_corpus_clean: bool = True, proc_inst_clean_small_cc_artifacts: bool = True, proc_inst_largest_k_cc: int = 0, + proc_inst_batch_size: int = 4, verbose: bool = False, ) -> tuple[NII | None, ErrCode]: """Build a per-vertebra instance mask from a subregion semantic segmentation. @@ -75,6 +76,8 @@ def predict_instance_mask( proc_corpus_clean (bool, optional): Whether to clean small corpus connected-component artifacts. Defaults to True. proc_inst_clean_small_cc_artifacts (bool, optional): Whether to delete small instance artifacts. Defaults to True. proc_inst_largest_k_cc (int, optional): Keep only the largest k connected components per cutout label; 0 disables. Defaults to 0. + proc_inst_batch_size (int, optional): Number of cutouts run through the instance model per batched forward pass. + Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: @@ -146,6 +149,7 @@ def predict_instance_mask( process_detect_and_solve_merged_corpi=proc_detect_and_solve_merged_corpi, proc_inst_largest_k_cc=proc_inst_largest_k_cc, proc_inst_fill_holes=False, + instance_batch_size=proc_inst_batch_size, verbose=verbose, ) if vert_predictions is None: @@ -587,6 +591,7 @@ def collect_vertebra_predictions( proc_inst_largest_k_cc: int = 0, process_detect_and_solve_merged_corpi: bool = True, proc_inst_fill_holes: bool = False, + instance_batch_size: int = 4, verbose: bool = False, ) -> tuple[np.ndarray | None, list[str], int]: """Run the instance model on a cutout around each corpus center of mass and collect per-label predictions. @@ -605,6 +610,8 @@ def collect_vertebra_predictions( proc_inst_largest_k_cc (int, optional): Keep only the largest k connected components per cutout label; 0 disables. Defaults to 0. process_detect_and_solve_merged_corpi (bool, optional): Whether to detect and split merged corpora. Defaults to True. proc_inst_fill_holes (bool, optional): Whether to fill holes in each cutout prediction. Defaults to False. + instance_batch_size (int, optional): Number of cutouts run through the instance model per batched forward + pass. Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: @@ -651,8 +658,12 @@ def collect_vertebra_predictions( # seg_nii_for_cut is constant across the loop; read its array once instead of copying it per centroid. seg_arr_c = seg_nii_for_cut.get_seg_array() - # iterate over sorted coms and segment vertebra from subreg - for com_idx, com in enumerate(tqdm(corpus_coms, desc=logger._get_logger_prefix() + " Vertebra Body predictions")): + # First gather every cutout, then run them through the model in batched forward passes instead of one GPU call + # per centroid. Every cutout has the same fixed cutout_size, so they stack cleanly into one batch and the batched + # result is identical (in fp32) to predicting each cutout on its own. + cut_niis: list[NII] = [] + cut_meta: list[tuple[int, tuple, tuple, tuple]] = [] # (com_idx, com, cutout_coords, paddings) + for com_idx, com in enumerate(tqdm(corpus_coms, desc=logger._get_logger_prefix() + " Vertebra Body cutouts")): # Shift the com until there is a segmentation there (to account for mishaps in the com calculation) seg_at_com = seg_arr_c[int(com[0])][int(com[1])][int(com[2])] != 0 orig_com = (com[0], com[1], com[2]) @@ -668,16 +679,21 @@ def collect_vertebra_predictions( arr_cut, cutout_coords, paddings = np_calc_crop_around_centerpoint(com, seg_arr_c, cutout_size) cut_nii = seg_nii_for_cut.set_array(arr_cut, verbose=False).reorient_() debug_data[f"inst_cutout_vert_nii_{com_idx}_cut"] = cut_nii - results = model.segment_scan( - cut_nii, - resample_to_recommended=False, - pad_size=0, - resample_output_to_input_space=False, - verbose=False, - ) + cut_niis.append(cut_nii) + cut_meta.append((com_idx, com, cutout_coords, paddings)) + + # Batched inference: one forward per chunk of `instance_batch_size` cutouts instead of one per cutout. + batched_results = model.segment_scan_batch( + cut_niis, + resample_to_recommended=False, + pad_size=0, + resample_output_to_input_space=False, + batch_size=instance_batch_size, + verbose=False, + ) + + for (com_idx, com, cutout_coords, paddings), results in zip(cut_meta, batched_results): vert_cut_nii = results[OutputType.seg].reorient_() - # print("vert_cut_nii", vert_cut_nii.shape) - # logger.print(f"Done {com_idx}") debug_data[f"inst_cutout_vert_nii_{com_idx}_pred"] = vert_cut_nii.copy() vert_cut_nii = post_process_single_3vert_prediction( vert_cut_nii, @@ -690,22 +706,12 @@ def collect_vertebra_predictions( cutout_sizes = tuple(cutout_coords[i].stop - cutout_coords[i].start for i in range(len(cutout_coords))) pad_cutout = tuple(slice(paddings[i][0], paddings[i][0] + cutout_sizes[i]) for i in range(len(paddings))) - # print("cutout_sizes", cutout_sizes) - # print("pad_cutout", pad_cutout) arr = vert_cut_nii.get_seg_array() vert_predict_map = vert_predict_template.copy() vert_predict_map[cutout_coords] = arr[pad_cutout] - # vert_predict_map[com_idx][cutout_coords][vert_predict_map[com_idx][cutout_coords] == 0] = arr[pad_cutout][ - # vert_predict_map[com_idx][cutout_coords] == 0 - # ] seg_at_com = vert_predict_map[int(com[0])][int(com[1])][int(com[2])] if seg_at_com == 0: logger.print("Zero at cutout center, mistake", Log_Type.WARNING) - # if seg_at_com != 0: - # # before (1) is above - # shift = 2 - seg_at_com - # vert_predictions[com_idx][vert_predictions[com_idx] != 0] = vert_predictions[com_idx][vert_predictions[com_idx] != 0] + shift - # debug_data[f"vert_nii_{com_idx}_proc2"] = seg_nii_for_cut.set_array(vert_predictions[com_idx][cutout_coords]) for l in vert_labels: vert_l_map = vert_predict_map.copy() vert_l_map[vert_l_map != l] = 0 diff --git a/spineps/seg_model.py b/spineps/seg_model.py index 04d986f..0dcf6cd 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -4,6 +4,7 @@ import os from abc import ABC, abstractmethod +from contextlib import nullcontext from pathlib import Path import numpy as np @@ -245,6 +246,49 @@ def segment_scan( self.print("Segmenting done!") return result + def segment_scan_batch( + self, + input_images: list[Image_Reference | dict[InputType, Image_Reference]], + pad_size: int = 0, + step_size: float | None = None, + resample_to_recommended: bool = True, + resample_output_to_input_space: bool = True, + batch_size: int = 1, # noqa: ARG002 - only honored by batched overrides + amp: bool = False, # noqa: ARG002 - only honored by batched overrides + verbose: bool = False, + ) -> list[dict[OutputType, NII | None]]: + """Segments a list of inputs and returns one result per input, in order. + + The default implementation simply calls :meth:`segment_scan` on each input. Subclasses that can run a single + batched forward pass over equally shaped inputs (e.g. the 3D U-Net instance model) override this for speed. + ``batch_size`` and ``amp`` are only honored by such batched overrides. + + Args: + input_images: The inputs to segment, each in the form accepted by :meth:`segment_scan`. + pad_size (int, optional): Padding added in each dimension, removed again from the output. Defaults to 0. + step_size (float | None, optional): Sliding-window tile step size; if None, uses the config default. Defaults to None. + resample_to_recommended (bool, optional): If true, rescales each input to the model's recommended zoom. Defaults to True. + resample_output_to_input_space (bool, optional): If true, resamples and pads the outputs back to the input space. + Defaults to True. + batch_size (int, optional): Maximum number of inputs per forward pass (batched overrides only). Defaults to 1. + amp (bool, optional): Run the forward pass under autocast (batched overrides only). Defaults to False. + verbose (bool, optional): If true, prints verbose information. Defaults to False. + + Returns: + list[dict[OutputType, NII | None]]: One result mapping per input, in the same order as ``input_images``. + """ + return [ + self.segment_scan( + img, + pad_size=pad_size, + step_size=step_size, + resample_to_recommended=resample_to_recommended, + resample_output_to_input_space=resample_output_to_input_space, + verbose=verbose, + ) + for img in input_images + ] + def modalities(self) -> list[Modality]: """Returns the modalities this model supports. @@ -487,8 +531,8 @@ def load(self, folds: tuple[str, ...] | None = None) -> Self: # noqa: ARG002 def run(self, input_nii: list[NII], verbose: bool = False) -> dict[OutputType, NII | None]: """Runs the 3D U-Net on a single input segmentation mask. - Converts the input mask to a network tensor (one-hot encoded for the multi-channel network, or intensity-normalized - for the legacy single-channel network), runs the forward pass and returns the per-voxel argmax class as a mask. + Thin wrapper around :meth:`run_batch` (batch of one) so the single- and batched-input paths share a single + implementation and always produce identical results. Args: input_nii (list[NII]): A single-element list containing the input segmentation mask. @@ -501,36 +545,120 @@ def run(self, input_nii: list[NII], verbose: bool = False) -> dict[OutputType, N AssertionError: If more than one input is provided. """ assert len(input_nii) == 1, "Unet3D does not support more than one input" - input_nii_ = input_nii[0] + return self.run_batch(input_nii, verbose=verbose)[0] - arr = input_nii_.get_seg_array().astype(np.int16) + def run_batch( + self, + input_nii: list[NII], + batch_size: int = 4, + amp: bool = False, + verbose: bool = False, + ) -> list[dict[OutputType, NII | None]]: + """Runs the 3D U-Net on a list of equally shaped cutout masks using batched forward passes. - target = from_numpy(arr) - n_classes = self.predictor.network.channels + Converts each cutout to a network tensor (one-hot for the multi-channel network, intensity-normalized for the + legacy single-channel network), stacks up to ``batch_size`` of them into a single forward pass and returns the + per-voxel argmax class for each. Each cutout is processed independently, so (in fp32) the result for every + cutout is identical to calling :meth:`run` on it on its own. If a batched forward runs out of GPU memory, it + transparently falls back to processing that chunk one cutout at a time. - target[target >= n_classes] = 0 + Args: + input_nii (list[NII]): Cutout masks to segment; all must share the same shape. + batch_size (int, optional): Maximum number of cutouts per forward pass. Defaults to 4. + amp (bool, optional): If true, runs the forward pass under CUDA autocast (faster, may slightly change + values). Defaults to False. + verbose (bool, optional): If true, prints verbose information. Defaults to False. - # channel-wise + Returns: + list[dict[OutputType, NII | None]]: One ``{OutputType.seg: mask}`` mapping per input, in order. + """ + if self.predictor is None: + self.load() + assert self.predictor is not None, "self.predictor == None after load(). Error!" + n_classes = self.predictor.network.channels + batch_size = max(1, batch_size) + results: list[dict[OutputType, NII | None]] = [None] * len(input_nii) # type: ignore[list-item] + with torch.inference_mode(): + for start in range(0, len(input_nii), batch_size): + chunk = input_nii[start : start + batch_size] + tensors = [self._network_input_from_nii(nii, n_classes) for nii in chunk] + try: + pred_cls = self._forward_argmax(torch.stack(tensors).to(self.device), amp=amp) + except RuntimeError as e: # pragma: no cover - depends on available GPU memory + if len(tensors) == 1 or "out of memory" not in str(e).lower(): + raise + self.print( + f"Out of memory on a batch of {len(tensors)} cutouts, falling back to one-by-one", + Log_Type.WARNING, + ) + if self.device.type == "cuda": + torch.cuda.empty_cache() + pred_cls = np.concatenate([self._forward_argmax(t.unsqueeze(0).to(self.device), amp=amp) for t in tensors], axis=0) + for offset, nii in enumerate(chunk): + results[start + offset] = {OutputType.seg: nii.set_array(pred_cls[offset])} + self.print("Batched segmentation done!", verbose=verbose) + return results + + def _network_input_from_nii(self, input_nii: NII, n_classes: int) -> torch.Tensor: + """Builds the (channels, *spatial) float input tensor for one cutout mask.""" + arr = input_nii.get_seg_array().astype(np.int16) + target = from_numpy(arr) + target[target >= n_classes] = 0 if n_classes != 1: - targetc = target.to(torch.int64) - targetc = F.one_hot(targetc, num_classes=n_classes) - targetc = targetc.permute(3, 0, 1, 2) - targetc = targetc.unsqueeze(0) - targetc = targetc.to(torch.float32) - logits = self.predictor.forward(targetc.to(self.device)) - else: - # legacy version - target = target.to(torch.float32) - target /= LEGACY_LABEL_NORMALIZATION - target = target.unsqueeze(0) - target = target.unsqueeze(0) - logits = self.predictor.forward(target.to(self.device)) - soft_max = torch.nn.Softmax(dim=1) - pred_x = soft_max(logits) - _, pred_cls = torch.max(pred_x, 1) - del logits - del pred_x - pred_cls = pred_cls.detach().cpu().numpy()[0] - seg_nii: NII = input_nii_.set_array(pred_cls) - self.print("out", seg_nii.zoom, seg_nii.orientation, seg_nii.shape) if verbose else None - return {OutputType.seg: seg_nii} + return F.one_hot(target.to(torch.int64), num_classes=n_classes).permute(3, 0, 1, 2).to(torch.float32) + # legacy single-channel network + return (target.to(torch.float32) / LEGACY_LABEL_NORMALIZATION).unsqueeze(0) + + def _forward_argmax(self, batch: torch.Tensor, amp: bool = False) -> np.ndarray: + """Runs the network on a (batch, channels, *spatial) tensor and returns the per-voxel argmax classes on CPU.""" + context = torch.autocast(self.device.type) if amp and self.device.type == "cuda" else nullcontext() + with context: + logits = self.predictor.forward(batch) + pred_cls = torch.argmax(torch.softmax(logits.float(), dim=1), dim=1) + return pred_cls.detach().cpu().numpy() + + def segment_scan_batch( + self, + input_images: list[Image_Reference | dict[InputType, Image_Reference]], + pad_size: int = 0, + step_size: float | None = None, # noqa: ARG002 - the 3D U-Net has no sliding window + resample_to_recommended: bool = True, + resample_output_to_input_space: bool = True, + batch_size: int = 4, + amp: bool = False, + verbose: bool = False, + ) -> list[dict[OutputType, NII | None]]: + """Batched :meth:`segment_scan` for the 3D U-Net. + + Preprocesses each input (optional padding, reorientation, rescaling) exactly as :meth:`segment_scan`, runs them + through batched forward passes and maps the outputs back into the input space. Equivalent to calling + :meth:`segment_scan` per input, but with a single forward per ``batch_size`` inputs instead of one per input. + """ + if self.predictor is None: + self.load() + input_type = self.inference_config.expected_inputs[0] + prepared: list[NII] = [] + metas: list[tuple] = [] + for img in input_images: + nii = to_nii(img, seg=input_type == InputType.seg) + if pad_size > 0: + nii.set_array_(np.pad(nii.get_array(), pad_size, mode="edge")) + orig_shape = nii.shape + nii.reorient_(self.inference_config.model_expected_orientation, verbose=self.logger) + if resample_to_recommended: + nii.rescale_(self.calc_recommended_resampling_zoom(nii.zoom), verbose=self.logger) + prepared.append(nii) + metas.append((orig_shape, img)) + results = self.run_batch(prepared, batch_size=batch_size, amp=amp, verbose=verbose) + for result, (orig_shape, img) in zip(results, metas): + for output_type, out_nii in result.items(): + if not isinstance(out_nii, NII): + continue + if resample_output_to_input_space: + out_nii.resample_from_to_(img) + out_nii.pad_to(orig_shape, inplace=True) + if output_type == OutputType.seg: + out_nii.map_labels_(self.inference_config.segmentation_labels, verbose=self.logger) + if pad_size > 0: + out_nii.set_array_(out_nii.get_array()[pad_size:-pad_size, pad_size:-pad_size, pad_size:-pad_size]) + return results diff --git a/spineps/seg_run.py b/spineps/seg_run.py index 6bd6310..1bd14de 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -302,6 +302,7 @@ def process_img_nii( # noqa: C901 proc_inst_clean_small_cc_artifacts: bool = True, proc_inst_largest_k_cc: int = 0, proc_inst_detect_and_solve_merged_corpi: bool = True, + proc_inst_batch_size: int = 4, vertebra_instance_labeling_offset=2, # Labeling proc_lab_force_no_tl_anomaly: bool = False, @@ -366,6 +367,8 @@ def process_img_nii( # noqa: C901 proc_inst_largest_k_cc (int, optional): If greater than 0, keeps only the largest k connected components of the instance mask. Defaults to 0. proc_inst_detect_and_solve_merged_corpi (bool, optional): If true, detects and splits merged vertebra corpi. Defaults to True. + proc_inst_batch_size (int, optional): Number of vertebra cutouts run through the instance model per batched forward + pass. Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. vertebra_instance_labeling_offset (int, optional): Offset applied when mapping instance ids to vertebra labels (set to 1 for CT models that include C1). Defaults to 2. proc_lab_force_no_tl_anomaly (bool, optional): If true, forces the labeling to assume no thoracolumbar transition anomaly. @@ -554,6 +557,7 @@ def process_img_nii( # noqa: C901 proc_corpus_clean=proc_inst_corpus_clean, proc_inst_clean_small_cc_artifacts=proc_inst_clean_small_cc_artifacts, proc_inst_largest_k_cc=proc_inst_largest_k_cc, + proc_inst_batch_size=proc_inst_batch_size, ) if errcode != ErrCode.OK: logger.print(f"Vert Mask creation failed with errcode {errcode}", Log_Type.FAIL) diff --git a/unit_tests/test_inference_mocked.py b/unit_tests/test_inference_mocked.py index 78d2b4c..533cf0f 100644 --- a/unit_tests/test_inference_mocked.py +++ b/unit_tests/test_inference_mocked.py @@ -16,6 +16,7 @@ from typing import ClassVar from unittest.mock import MagicMock +import nibabel as nib import numpy as np import torch from TPTBox import NII, No_Logger @@ -26,7 +27,7 @@ from spineps.phase_instance import predict_instance_mask from spineps.phase_labeling import VERT_CLASSES, perform_labeling_step, run_model_for_vert_labeling from spineps.seg_enums import ErrCode, OutputType -from spineps.seg_model import Segmentation_Inference_Config, Segmentation_Model +from spineps.seg_model import Segmentation_Inference_Config, Segmentation_Model, Segmentation_Model_Unet3D logger = No_Logger() @@ -348,5 +349,99 @@ def test_single_coarser_axis_does_not_match(self): self.assertFalse(a.same_modelzoom_as_model(b, (1.0, 1.0, 1.0))) +class _FakeUnetNetwork: + def __init__(self, channels: int) -> None: + self.channels = channels + + +class _FakeUnetPredictor: + """Identity 'network': returns its input as logits, so softmax+argmax recovers the one-hot input label.""" + + def __init__(self, channels: int) -> None: + self.network = _FakeUnetNetwork(channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +def _make_unet3d_test_model(n_classes: int = 4, cutout: tuple[int, int, int] = (6, 6, 6)) -> Segmentation_Model_Unet3D: + """A real Segmentation_Model_Unet3D wired to an identity predictor (no weights, runs on CPU).""" + config = Segmentation_Inference_Config( + logger=No_Logger(), + modality=["SEG"], + acquisition="sag", + log_name="TestUnet3D", + modeltype="unet", + model_expected_orientation=("P", "I", "R"), + available_folds=1, + inference_augmentation=False, + resolution_range=[1.5, 1.5, 1.5], + default_step_size=0.5, + labels={1: 1, 2: 2, 3: 3}, + expected_inputs=["seg"], + cutout_size=cutout, + ) + model = Segmentation_Model_Unet3D(__file__, config, default_verbose=False, default_allow_tqdm=False) + model.predictor = _FakeUnetPredictor(n_classes) + model.device = torch.device("cpu") + return model + + +def _seg_nii(arr: np.ndarray) -> NII: + return NII(nib.Nifti1Image(arr.astype(np.uint8), affine=np.eye(4)), seg=True) + + +class Test_Unet3D_Batching(unittest.TestCase): + """The batched instance path must produce exactly the same masks as predicting each cutout on its own.""" + + def _cutouts(self, n: int = 5) -> list[NII]: + rng = np.random.default_rng(0) + return [_seg_nii(rng.integers(0, 4, size=(6, 6, 6))) for _ in range(n)] + + def test_run_batch_matches_run(self): + # run_batch over a list must equal calling run() on each cutout individually (golden oracle). + model = _make_unet3d_test_model() + cutouts = self._cutouts() + batched = model.run_batch(cutouts, batch_size=2) + self.assertEqual(len(batched), len(cutouts)) + for cut, res in zip(cutouts, batched): + single = model.run([cut]) + np.testing.assert_array_equal(res[OutputType.seg].get_seg_array(), single[OutputType.seg].get_seg_array()) + # distinct inputs must yield distinct outputs (guards against the batch collapsing to one result) + outs = [r[OutputType.seg].get_seg_array() for r in batched] + self.assertTrue(any(not np.array_equal(outs[0], o) for o in outs[1:])) + + def test_run_batch_size_invariant(self): + # The result must not depend on how cutouts are chunked into forward passes. + model = _make_unet3d_test_model() + cutouts = self._cutouts() + for r1, r9 in zip(model.run_batch(cutouts, batch_size=1), model.run_batch(cutouts, batch_size=9)): + np.testing.assert_array_equal(r1[OutputType.seg].get_seg_array(), r9[OutputType.seg].get_seg_array()) + + def test_run_batch_actually_batches_forward_calls(self): + # 5 cutouts at batch_size 2 must take ceil(5/2)=3 forward passes, not 5 (the whole point of batching). + model = _make_unet3d_test_model() + real_forward = model.predictor.forward + calls = {"n": 0} + + def counting_forward(x: torch.Tensor) -> torch.Tensor: + calls["n"] += 1 + return real_forward(x) + + model.predictor.forward = counting_forward + model.run_batch(self._cutouts(5), batch_size=2) + self.assertEqual(calls["n"], 3) + + def test_segment_scan_batch_matches_segment_scan(self): + # The way the instance phase calls it: no resampling, batched == sequential segment_scan. + model = _make_unet3d_test_model() + cutouts = self._cutouts() + kwargs = {"resample_to_recommended": False, "pad_size": 0, "resample_output_to_input_space": False} + batched = model.segment_scan_batch(cutouts, batch_size=3, **kwargs) + for cut, res in zip(cutouts, batched): + single = model.segment_scan(cut, **kwargs) + np.testing.assert_array_equal(res[OutputType.seg].get_seg_array(), single[OutputType.seg].get_seg_array()) + + if __name__ == "__main__": unittest.main() From 0e526f78bb5756a6f5715623965cef660f62bf6d Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 20:18:59 +0000 Subject: [PATCH 06/29] refactor!: rename process_img_nii -> segment_image and Segmentation_Model* -> PascalCase Clean-break (2.0) naming: the snake_case class names and the implementation-detail 'img_nii' function name are replaced with PEP8 PascalCase classes and a clearer function name, consistent with VertLabelingClassifier. - process_img_nii -> segment_image - Segmentation_Model -> SegmentationModel; Segmentation_Model_NNunet -> SegmentationModelNNunet; Segmentation_Model_Unet3D -> SegmentationModelUnet3D All call sites, imports, type hints, __init__ exports and tests updated. No aliases (breaking). 178 tests pass. Co-Authored-By: Claude Opus 4.8 --- spineps/__init__.py | 4 ++-- spineps/entrypoint.py | 8 ++++---- spineps/example/helper_parallel.py | 4 ++-- spineps/example/template_roll_out.py | 4 ++-- spineps/get_models.py | 24 ++++++++++++------------ spineps/lab_model.py | 6 +++--- spineps/phase_instance.py | 10 +++++----- spineps/phase_semantic.py | 6 +++--- spineps/seg_model.py | 12 ++++++------ spineps/seg_pipeline.py | 6 +++--- spineps/seg_run.py | 24 ++++++++++++------------ spineps/seg_utils.py | 14 +++++++------- unit_tests/test_filepaths.py | 6 +++--- unit_tests/test_inference_mocked.py | 24 ++++++++++++------------ unit_tests/test_proc_functions.py | 2 +- unit_tests/test_semantic.py | 10 +++++----- 16 files changed, 82 insertions(+), 82 deletions(-) diff --git a/spineps/__init__.py b/spineps/__init__.py index cb36017..32504bb 100755 --- a/spineps/__init__.py +++ b/spineps/__init__.py @@ -6,5 +6,5 @@ from spineps.phase_labeling import perform_labeling_step from spineps.phase_post import phase_postprocess_combined from spineps.phase_semantic import predict_semantic_mask -from spineps.seg_model import Segmentation_Model -from spineps.seg_run import process_dataset, process_img_nii +from spineps.seg_model import SegmentationModel +from spineps.seg_run import process_dataset, segment_image diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index 7268fca..1a43d97 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -20,7 +20,7 @@ modelid2folder_labeling, modelid2folder_semantic, ) -from spineps.seg_run import process_dataset, process_img_nii +from spineps.seg_run import process_dataset, segment_image from spineps.utils.citation_reminder import citation_reminder logger = No_Logger(prefix="Init") @@ -224,7 +224,7 @@ def run_sample(opt: Namespace): """Run the full segmentation pipeline on a single input NIfTI file. Loads the requested semantic, instance and (optional) labeling models, wraps the input as a - ``BIDS_FILE`` and calls :func:`process_img_nii`, optionally under a cProfiler. + ``BIDS_FILE`` and calls :func:`segment_image`, optionally under a cProfiler. Args: opt (Namespace): Parsed CLI arguments from the ``sample`` subcommand (input path, model ids/paths, @@ -302,11 +302,11 @@ def run_sample(opt: Namespace): info={"desc": "cprofile", "mod": bids_sample.format, "ses": timestamp}, ) with cProfile.Profile() as pr: - process_img_nii(**kwargs) + segment_image(**kwargs) pr.dump_stats(cprofile_out) logger.print(f"Saved cprofile log into {cprofile_out}", Log_Type.SAVE) else: - process_img_nii(**kwargs) + segment_image(**kwargs) logger.print(f"Sample took: {perf_counter() - start_time} seconds") return 1 diff --git a/spineps/example/helper_parallel.py b/spineps/example/helper_parallel.py index 064df95..820fad1 100755 --- a/spineps/example/helper_parallel.py +++ b/spineps/example/helper_parallel.py @@ -14,7 +14,7 @@ from TPTBox import BIDS_FILE # noqa: E402 from spineps.get_models import get_instance_model, get_semantic_model # noqa: E402 -from spineps.seg_run import process_img_nii # noqa: E402 +from spineps.seg_run import segment_image # noqa: E402 # Example # python /spineps/example/helper_parallel.py -i PATH/TO/IMG.nii.gz -ds DATASET-PATH -der derivatives -ms [t1w,t2w,vibe] -mv instance @@ -36,7 +36,7 @@ mv = get_instance_model(opt.mv) if opt.snap is not None: Path(opt.snap).mkdir(exist_ok=True, parents=True) - process_img_nii( + segment_image( img_ref=input_bids_file, derivative_name=opt.der, model_semantic=ms, diff --git a/spineps/example/template_roll_out.py b/spineps/example/template_roll_out.py index 4d1be94..4e79069 100755 --- a/spineps/example/template_roll_out.py +++ b/spineps/example/template_roll_out.py @@ -15,7 +15,7 @@ from TPTBox import BIDS_FILE, NII, POI, BIDS_Global_info, No_Logger # noqa: E402 from spineps.get_models import get_instance_model, get_semantic_model # noqa: E402 -from spineps.seg_run import ErrCode, process_img_nii # noqa: E402 +from spineps.seg_run import ErrCode, segment_image # noqa: E402 # INPUT in_ds = Path("DATASET_PATH") # TODO change this to the path to your dataset folder @@ -80,7 +80,7 @@ def injection_function(seg_nii: NII): start_time = time.perf_counter() ref: BIDS_FILE = f["T2w_part-inphase"][0] if "T2w_part-inphase" in f else f["T2w"][0] # Call to the pipeline - output_paths, errcode = process_img_nii( + output_paths, errcode = segment_image( img_ref=ref, derivative_name=der, model_semantic=model_semantic, diff --git a/spineps/get_models.py b/spineps/get_models.py index 47ac279..31ed7c1 100755 --- a/spineps/get_models.py +++ b/spineps/get_models.py @@ -11,7 +11,7 @@ from spineps.lab_model import VertLabelingClassifier from spineps.seg_enums import Modality, ModelType, SpinepsPhase -from spineps.seg_model import Segmentation_Model, Segmentation_Model_NNunet, Segmentation_Model_Unet3D +from spineps.seg_model import SegmentationModel, SegmentationModelNNunet, SegmentationModelUnet3D from spineps.utils.auto_download import download_if_missing, instances, labeling, semantic from spineps.utils.filepaths import get_mri_segmentor_models_dir, search_path from spineps.utils.seg_modelconfig import load_inference_config @@ -33,7 +33,7 @@ def _get_model_by_name( phase: SpinepsPhase, kind: str, **kwargs, -) -> Segmentation_Model | VertLabelingClassifier: +) -> SegmentationModel | VertLabelingClassifier: """Looks up a model by name in a model-id-to-folder map and instantiates it. Shared implementation behind get_semantic_model / get_instance_model / get_labeling_model. @@ -56,7 +56,7 @@ def _get_model_by_name( return get_actual_model(config_path, **kwargs) -def get_semantic_model(model_name: str, **kwargs) -> Segmentation_Model: +def get_semantic_model(model_name: str, **kwargs) -> SegmentationModel: """Finds and returns a semantic (subregion) model by name. Args: @@ -64,7 +64,7 @@ def get_semantic_model(model_name: str, **kwargs) -> Segmentation_Model: **kwargs: Extra keyword arguments forwarded to the model constructor. Returns: - Segmentation_Model: The instantiated semantic model. + SegmentationModel: The instantiated semantic model. Raises: KeyError: If the given model name is not among the available models. @@ -73,7 +73,7 @@ def get_semantic_model(model_name: str, **kwargs) -> Segmentation_Model: return _get_model_by_name(model_name, modelid2folder_semantic(), SpinepsPhase.SEMANTIC, "semantic", **kwargs) -def get_instance_model(model_name: str, **kwargs) -> Segmentation_Model: +def get_instance_model(model_name: str, **kwargs) -> SegmentationModel: """Finds and returns an instance (vertebra) model by name. Args: @@ -81,7 +81,7 @@ def get_instance_model(model_name: str, **kwargs) -> Segmentation_Model: **kwargs: Extra keyword arguments forwarded to the model constructor. Returns: - Segmentation_Model: The instantiated instance model. + SegmentationModel: The instantiated instance model. Raises: KeyError: If the given model name is not among the available models. @@ -214,12 +214,12 @@ def modeltype2class(modeltype: ModelType) -> type: NotImplementedError: If the model type is not supported. Returns: - type: The class to instantiate (Segmentation_Model_NNunet, Segmentation_Model_Unet3D or VertLabelingClassifier). + type: The class to instantiate (SegmentationModelNNunet, SegmentationModelUnet3D or VertLabelingClassifier). """ if modeltype == ModelType.nnunet: - return Segmentation_Model_NNunet + return SegmentationModelNNunet elif modeltype == ModelType.unet: - return Segmentation_Model_Unet3D + return SegmentationModelUnet3D elif modeltype == ModelType.classifier: return VertLabelingClassifier else: @@ -230,7 +230,7 @@ def get_actual_model( in_config: str | Path, use_cpu: bool = False, **kwargs, -) -> Segmentation_Model | VertLabelingClassifier: +) -> SegmentationModel | VertLabelingClassifier: """Creates and returns the appropriate model from a given inference config path. Accepts either a path to an inference_config.json file or a folder containing exactly one such file (searched @@ -242,7 +242,7 @@ def get_actual_model( **kwargs: Extra keyword arguments forwarded to the model constructor. Returns: - Segmentation_Model | VertLabelingClassifier: The instantiated model. + SegmentationModel | VertLabelingClassifier: The instantiated model. Raises: FileNotFoundError: If no inference_config.json is found in the given folder. @@ -272,5 +272,5 @@ def get_actual_model( # in_dir = base inference_config = load_inference_config(str(in_dir)) - modeltype: type[Segmentation_Model] = modeltype2class(inference_config.modeltype) + modeltype: type[SegmentationModel] = modeltype2class(inference_config.modeltype) return modeltype(model_folder=in_config, inference_config=inference_config, use_cpu=use_cpu, **kwargs) diff --git a/spineps/lab_model.py b/spineps/lab_model.py index 8769b81..a14df6a 100755 --- a/spineps/lab_model.py +++ b/spineps/lab_model.py @@ -15,7 +15,7 @@ from spineps.architectures.pl_densenet import PLClassifier from spineps.seg_enums import OutputType -from spineps.seg_model import Segmentation_Inference_Config, Segmentation_Model +from spineps.seg_model import Segmentation_Inference_Config, SegmentationModel from spineps.utils.filepaths import search_path logger = No_Logger(prefix="VertLabelingClassifier") @@ -73,12 +73,12 @@ def rotate_patch_sagitally(patch: np.ndarray, angle: float, msk: bool = False, c return rotated_patch -class VertLabelingClassifier(Segmentation_Model): +class VertLabelingClassifier(SegmentationModel): """Classifier that assigns anatomical labels to individual vertebrae. For each vertebra a patch is cropped around its center of mass, optionally rotated to align with the spine axis, normalized and center-cropped to a fixed size, then passed through a DenseNet (PLClassifier) that outputs per-head - softmax predictions. Although it subclasses Segmentation_Model to reuse config loading, it does not perform voxel + softmax predictions. Although it subclasses SegmentationModel to reuse config loading, it does not perform voxel segmentation (run/segment_scan are not implemented). Attributes: diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index b918343..ba5fa71 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -21,7 +21,7 @@ from tqdm import tqdm from spineps.seg_enums import ErrCode, OutputType -from spineps.seg_model import Segmentation_Model +from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import IVD_LABEL_OFFSET, logger from spineps.utils.proc_functions import clean_cc_artifacts from spineps.utils.resolution import ( @@ -49,7 +49,7 @@ def predict_instance_mask( seg_nii: NII, - model: Segmentation_Model, + model: SegmentationModel, debug_data: dict, pad_size: int = 0, proc_inst_fill_3d_holes: bool = True, @@ -68,7 +68,7 @@ def predict_instance_mask( Args: seg_nii (NII): Subregion (semantic) segmentation mask used as input. - model (Segmentation_Model): Instance model producing the per-vertebra-body cutout predictions. + model (SegmentationModel): Instance model producing the per-vertebra-body cutout predictions. debug_data (dict): Dictionary for collecting intermediate results across the pipeline. pad_size (int, optional): Edge padding added before processing and removed afterwards. Defaults to 0. proc_inst_fill_3d_holes (bool, optional): Whether to fill 3D holes in the final vertebra mask. Defaults to True. @@ -584,7 +584,7 @@ def split_by_plane( def collect_vertebra_predictions( seg_nii: NII, - model: Segmentation_Model, + model: SegmentationModel, corpus_size_cleaning: int, cutout_size: tuple[int, int, int], debug_data: dict, @@ -603,7 +603,7 @@ def collect_vertebra_predictions( Args: seg_nii (NII): Subregion semantic mask used for the cutouts. - model (Segmentation_Model): Instance model producing the three-vertebra-body cutout predictions. + model (SegmentationModel): Instance model producing the three-vertebra-body cutout predictions. corpus_size_cleaning (int): Voxel threshold for cleaning small corpus artifacts when finding coms; 0 disables. cutout_size: Cutout window size (per axis) extracted around each center of mass. debug_data (dict): Dictionary for collecting per-cutout intermediate results. diff --git a/spineps/phase_semantic.py b/spineps/phase_semantic.py index 1e3605b..8768e74 100755 --- a/spineps/phase_semantic.py +++ b/spineps/phase_semantic.py @@ -7,7 +7,7 @@ from TPTBox import NII, Location, Log_Type from spineps.seg_enums import ErrCode, OutputType -from spineps.seg_model import Segmentation_Model +from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import fill_holes_labels, logger from spineps.utils.proc_functions import clean_cc_artifacts from spineps.utils.resolution import REFERENCE_VOXEL_VOLUME_MM3, REFERENCE_ZOOM, mm3_to_voxels, mm_to_voxels @@ -26,7 +26,7 @@ def predict_semantic_mask( mri_nii: NII, - model: Segmentation_Model, + model: SegmentationModel, debug_data: dict, proc_fill_3d_holes: bool = True, proc_clean_beyond_largest_bounding_box: bool = True, @@ -42,7 +42,7 @@ def predict_semantic_mask( Args: mri_nii (NII): Input grayscale MRI image (intensities must start at 0). - model (Segmentation_Model): Model used to produce the semantic segmentation. + model (SegmentationModel): Model used to produce the semantic segmentation. debug_data (dict): Dictionary for collecting intermediate results (e.g. the raw segmentation). proc_fill_3d_holes (bool, optional): Whether to fill 3D holes in the output mask. Defaults to True. proc_clean_beyond_largest_bounding_box (bool, optional): Whether to keep only connected components within diff --git a/spineps/seg_model.py b/spineps/seg_model.py index 0dcf6cd..e944c00 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -1,4 +1,4 @@ -"""Segmentation model abstractions: the abstract Segmentation_Model and its nnU-Net and Unet3D subclasses.""" +"""Segmentation model abstractions: the abstract SegmentationModel and its nnU-Net and Unet3D subclasses.""" from __future__ import annotations @@ -30,7 +30,7 @@ LEGACY_LABEL_NORMALIZATION = 9 -class Segmentation_Model(ABC): +class SegmentationModel(ABC): """Abstract base class wrapping a segmentation network together with its inference configuration. Subclasses implement load() and run() for a concrete backend (e.g. nnU-Net or Unet3D). The class handles input @@ -384,8 +384,8 @@ def __repr__(self) -> str: return str(self) -class Segmentation_Model_NNunet(Segmentation_Model): - """Segmentation_Model backed by an nnU-Net predictor.""" +class SegmentationModelNNunet(SegmentationModel): + """SegmentationModel backed by an nnU-Net predictor.""" def __init__( self, @@ -465,8 +465,8 @@ def run( return {OutputType.seg: seg_nii, OutputType.softmax_logits: softmax_logits} -class Segmentation_Model_Unet3D(Segmentation_Model): - """Segmentation_Model backed by a single-input 3D U-Net (PyTorch Lightning PLNet). +class SegmentationModelUnet3D(SegmentationModel): + """SegmentationModel backed by a single-input 3D U-Net (PyTorch Lightning PLNet). Used as the instance (vertebra) model: it takes a segmentation mask as input and refines it into the vertebra instance output. Supports both the current multi-channel network and a legacy single-channel network. diff --git a/spineps/seg_pipeline.py b/spineps/seg_pipeline.py index 4a1f3b4..7958de5 100755 --- a/spineps/seg_pipeline.py +++ b/spineps/seg_pipeline.py @@ -11,7 +11,7 @@ from TPTBox.core import poi from TPTBox.logger.log_file import format_time_short, get_time -from spineps.seg_model import Segmentation_Model +from spineps.seg_model import SegmentationModel logger = No_Logger(prefix="SPINEPS") @@ -51,7 +51,7 @@ def predict_centroids_from_both( vert_nii_cleaned: NII, seg_nii: NII, - models: list[Segmentation_Model | None], + models: list[SegmentationModel | None], parameter: dict[str, Any], ) -> poi.POI: """Calculate the centroids of each vertebra corpus using both the semantic and instance masks. @@ -63,7 +63,7 @@ def predict_centroids_from_both( Args: vert_nii_cleaned (NII): Cleaned vertebra instance segmentation mask. seg_nii (NII): Subregion semantic segmentation mask. - models (list[Segmentation_Model | None]): Models used in the pipeline, recorded in the centroid metadata. + models (list[SegmentationModel | None]): Models used in the pipeline, recorded in the centroid metadata. parameter (dict[str, Any]): Pipeline parameters to record on the centroid metadata. Returns: diff --git a/spineps/seg_run.py b/spineps/seg_run.py index 1bd14de..56ea705 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -19,7 +19,7 @@ from spineps.phase_pre import compute_crop, preprocess_input from spineps.phase_semantic import predict_semantic_mask from spineps.seg_enums import Acquisition, ErrCode, Modality -from spineps.seg_model import Segmentation_Model +from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import logger, predict_centroids_from_both from spineps.seg_utils import Modality_Pair, check_input_model_compatibility, check_model_modality_acquisition, find_best_matching_model from spineps.utils.citation_reminder import citation_reminder @@ -28,8 +28,8 @@ @citation_reminder def process_dataset( dataset_path: Path, - model_instance: Segmentation_Model, - model_semantic: list[Segmentation_Model] | Segmentation_Model | None = None, + model_instance: SegmentationModel, + model_semantic: list[SegmentationModel] | SegmentationModel | None = None, model_labeling: VertLabelingClassifier | None = None, # rawdata_name: str = "rawdata", @@ -75,12 +75,12 @@ def process_dataset( """Runs the SPINEPS framework over a whole BIDS-conform dataset. Iterates over every subject in the BIDS dataset, queries the matching scans for each requested modality pair and runs - process_img_nii on each, producing semantic (subregion), vertebra (instance) and centroid outputs plus a snapshot. + segment_image on each, producing semantic (subregion), vertebra (instance) and centroid outputs plus a snapshot. Args: dataset_path (Path): Path to the BIDS dataset. - model_instance (Segmentation_Model): Model for the vertebra (instance) segmentation. - model_semantic (list[Segmentation_Model] | Segmentation_Model | None, optional): Models for the subregion (semantic) + model_instance (SegmentationModel): Model for the vertebra (instance) segmentation. + model_semantic (list[SegmentationModel] | SegmentationModel | None, optional): Models for the subregion (semantic) segmentation, one per modality pair. If None, attempts to find a matching model for each modality. Defaults to None. model_labeling (VertLabelingClassifier | None, optional): Classifier used to label the vertebra instances. Defaults to None. rawdata_name (str, optional): Name of the rawdata folder. Defaults to "rawdata". @@ -199,7 +199,7 @@ def process_dataset( q.filter("acq", lambda x: x in allowed_acq, required=False) # noqa: B023 scans = q.loop_list(sort=True) # TODO make it family to allow for multi-inputs for s in scans: - output_paths, errcode = process_img_nii( + output_paths, errcode = segment_image( img_ref=s, model_semantic=model, model_instance=model_instance, @@ -267,10 +267,10 @@ def process_dataset( @citation_reminder -def process_img_nii( # noqa: C901 +def segment_image( # noqa: C901 img_ref: BIDS_FILE, - model_semantic: Segmentation_Model, - model_instance: Segmentation_Model, + model_semantic: SegmentationModel, + model_instance: SegmentationModel, model_labeling: VertLabelingClassifier | None = None, derivative_name: str = "derivatives_seg", # @@ -329,8 +329,8 @@ def process_img_nii( # noqa: C901 Args: img_ref (BIDS_FILE): Input BIDS_FILE referencing the image to segment. - model_semantic (Segmentation_Model): Model for the subregion (semantic) segmentation. - model_instance (Segmentation_Model): Model for the vertebra (instance) segmentation. + model_semantic (SegmentationModel): Model for the subregion (semantic) segmentation. + model_instance (SegmentationModel): Model for the vertebra (instance) segmentation. model_labeling (VertLabelingClassifier | None, optional): Classifier used to label the vertebra instances. Defaults to None. derivative_name (str, optional): Name of the derivatives output folder. Defaults to "derivatives_seg". save_modelres_mask (bool, optional): If true, additionally saves the semantic mask in the resolution of the model. diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index 122e313..4bfb358 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -9,7 +9,7 @@ from TPTBox import BIDS_FILE, NII, ZOOMS, Log_Type from spineps.seg_enums import Acquisition, Modality -from spineps.seg_model import Segmentation_Model +from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import logger Modality_Pair = tuple[Union[list[Modality], Modality], Acquisition] @@ -18,7 +18,7 @@ def find_best_matching_model( modality_pair: Modality_Pair, expected_resolution: ZOOMS | None, # actual resolution here? -) -> Segmentation_Model: +) -> SegmentationModel: """Select the segmentation model best matching a modality/acquisition pair and resolution. Not yet implemented: intended to iterate over model configs and pick the one best matching the requested resolution. @@ -28,7 +28,7 @@ def find_best_matching_model( expected_resolution (ZOOMS | None): The desired voxel resolution, or None. Returns: - Segmentation_Model: The best-matching model (once implemented). + SegmentationModel: The best-matching model (once implemented). Raises: NotImplementedError: Always, as this function is not yet implemented; also for an unmapped modality pair. @@ -52,7 +52,7 @@ def find_best_matching_model( def check_model_modality_acquisition( - model: Segmentation_Model, + model: SegmentationModel, mod_pair: Modality_Pair, verbose: bool = True, ) -> bool: @@ -62,7 +62,7 @@ def check_model_modality_acquisition( mismatch when ``verbose`` is True. Args: - model (Segmentation_Model): The model to check. + model (SegmentationModel): The model to check. mod_pair (Modality_Pair): The required ``(modality(ies), acquisition)`` pair. verbose (bool): If True, log a warning when incompatible. @@ -116,7 +116,7 @@ def add_ignore_text(logger_texts: list[str]) -> None: def check_input_model_compatibility( img_ref: BIDS_FILE, - model: Segmentation_Model, + model: SegmentationModel, ignore_modality: bool = False, ignore_acquisition: bool = False, ignore_labelkey: bool = False, @@ -130,7 +130,7 @@ def check_input_model_compatibility( Args: img_ref (BIDS_FILE): Reference to the input image file. - model (Segmentation_Model): The model to check against. + model (SegmentationModel): The model to check against. ignore_modality (bool): If True, tolerate a modality/format mismatch. ignore_acquisition (bool): If True, tolerate an acquisition mismatch. ignore_labelkey (bool): If True, tolerate an unexpected ``label`` key in the filename. diff --git a/unit_tests/test_filepaths.py b/unit_tests/test_filepaths.py index 0d8b974..1dc128b 100644 --- a/unit_tests/test_filepaths.py +++ b/unit_tests/test_filepaths.py @@ -10,7 +10,7 @@ from pathlib import Path import spineps -from spineps.get_models import Segmentation_Model, check_available_models, get_actual_model +from spineps.get_models import SegmentationModel, check_available_models, get_actual_model from spineps.utils.filepaths import ( filepath_model, get_mri_segmentor_models_dir, @@ -26,9 +26,9 @@ class Test_filepaths(unittest.TestCase): # mv_p = "/DATA/NAS/ongoing_projects/hendrik/nako-segmentation/nnUNet/unet3d_result/nakospider_highres_shiftposi" # model_dir = "/DATA/NAS/ongoing_projects/hendrik/nako-segmentation/nnUNet/" # ms = get_segmentation_model(in_config=filepath_model(ms_p, model_dir=model_dir)) - # self.assertTrue(isinstance(ms, Segmentation_Model)) + # self.assertTrue(isinstance(ms, SegmentationModel)) # mv = get_segmentation_model(in_config=filepath_model(mv_p, model_dir=model_dir)) - # self.assertTrue(isinstance(mv, Segmentation_Model)) + # self.assertTrue(isinstance(mv, SegmentationModel)) # self.assertTrue(True) def test_search_path_simple(self): diff --git a/unit_tests/test_inference_mocked.py b/unit_tests/test_inference_mocked.py index 533cf0f..5c3ae4a 100644 --- a/unit_tests/test_inference_mocked.py +++ b/unit_tests/test_inference_mocked.py @@ -27,7 +27,7 @@ from spineps.phase_instance import predict_instance_mask from spineps.phase_labeling import VERT_CLASSES, perform_labeling_step, run_model_for_vert_labeling from spineps.seg_enums import ErrCode, OutputType -from spineps.seg_model import Segmentation_Inference_Config, Segmentation_Model, Segmentation_Model_Unet3D +from spineps.seg_model import Segmentation_Inference_Config, SegmentationModel, SegmentationModelUnet3D logger = No_Logger() @@ -57,8 +57,8 @@ def _dummy_seg_config(cutout_size: tuple[int, int, int] = (48, 48, 32)) -> Segme ) -class Segmentation_Model_Dummy(Segmentation_Model): - """A Segmentation_Model whose load() installs a dummy predictor instead of real weights.""" +class SegmentationModelDummy(SegmentationModel): + """A SegmentationModel whose load() installs a dummy predictor instead of real weights.""" def __init__(self, cutout_size: tuple[int, int, int] = (48, 48, 32)) -> None: self.logger = No_Logger() @@ -147,7 +147,7 @@ def test_perform_labeling_step_relabels(self): class Test_Segment_Scan_Mocked(unittest.TestCase): def test_segment_scan_padding_round_trip(self): mri, _subreg, _vert, _label = get_test_mri() - model = Segmentation_Model_Dummy() + model = SegmentationModelDummy() # run() echoes its (padded, reoriented) input back as the segmentation. model.run = MagicMock(side_effect=lambda input_nii, verbose=False: {OutputType.seg: input_nii[0], OutputType.softmax_logits: None}) # noqa: ARG005 @@ -166,7 +166,7 @@ def test_segment_scan_padding_round_trip(self): def test_segment_scan_without_resample_back(self): mri, subreg, _vert, _label = get_test_mri() - model = Segmentation_Model_Dummy() + model = SegmentationModelDummy() model.run = MagicMock(return_value={OutputType.seg: subreg.copy(), OutputType.softmax_logits: None}) result = model.segment_scan( @@ -201,7 +201,7 @@ def _fake_segment_scan(cut_nii: NII, **kwargs): # noqa: ARG004 def test_predict_instance_mask_runs(self): _mri, subreg, _vert, _label = get_test_mri() - model = Segmentation_Model_Dummy(cutout_size=(48, 48, 32)) + model = SegmentationModelDummy(cutout_size=(48, 48, 32)) model.segment_scan = MagicMock(side_effect=self._fake_segment_scan) whole_vert_nii, errcode = predict_instance_mask( @@ -224,7 +224,7 @@ def test_predict_instance_mask_empty_without_corpus(self): # Remove the corpus-border label (49) so the instance phase has nothing to work with. no_corpus = subreg.copy() no_corpus[no_corpus == 49] = 0 - model = Segmentation_Model_Dummy() + model = SegmentationModelDummy() model.segment_scan = MagicMock(side_effect=self._fake_segment_scan) result, errcode = predict_instance_mask(no_corpus, model, debug_data={}, proc_corpus_clean=False) @@ -319,9 +319,9 @@ def test_run_all_seg_instances_full_path(self): class Test_Same_Modelzoom(unittest.TestCase): @staticmethod - def _model_with_zoom(zoom: tuple[float, float, float]) -> Segmentation_Model_Dummy: + def _model_with_zoom(zoom: tuple[float, float, float]) -> SegmentationModelDummy: # A fixed (length-3) resolution_range makes calc_recommended_resampling_zoom return it verbatim. - model = Segmentation_Model_Dummy() + model = SegmentationModelDummy() model.inference_config.resolution_range = tuple(zoom) return model @@ -364,8 +364,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -def _make_unet3d_test_model(n_classes: int = 4, cutout: tuple[int, int, int] = (6, 6, 6)) -> Segmentation_Model_Unet3D: - """A real Segmentation_Model_Unet3D wired to an identity predictor (no weights, runs on CPU).""" +def _make_unet3d_test_model(n_classes: int = 4, cutout: tuple[int, int, int] = (6, 6, 6)) -> SegmentationModelUnet3D: + """A real SegmentationModelUnet3D wired to an identity predictor (no weights, runs on CPU).""" config = Segmentation_Inference_Config( logger=No_Logger(), modality=["SEG"], @@ -381,7 +381,7 @@ def _make_unet3d_test_model(n_classes: int = 4, cutout: tuple[int, int, int] = ( expected_inputs=["seg"], cutout_size=cutout, ) - model = Segmentation_Model_Unet3D(__file__, config, default_verbose=False, default_allow_tqdm=False) + model = SegmentationModelUnet3D(__file__, config, default_verbose=False, default_allow_tqdm=False) model.predictor = _FakeUnetPredictor(n_classes) model.device = torch.device("cpu") return model diff --git a/unit_tests/test_proc_functions.py b/unit_tests/test_proc_functions.py index 14f8cb5..fac80dc 100644 --- a/unit_tests/test_proc_functions.py +++ b/unit_tests/test_proc_functions.py @@ -12,7 +12,7 @@ from TPTBox.tests.test_utils import get_test_mri import spineps -from spineps.get_models import Segmentation_Model, get_actual_model +from spineps.get_models import SegmentationModel, get_actual_model from spineps.utils.compat import zip_strict from spineps.utils.proc_functions import clean_cc_artifacts, connected_components_3d, n4_bias diff --git a/unit_tests/test_semantic.py b/unit_tests/test_semantic.py index 3e5dea4..95d04aa 100644 --- a/unit_tests/test_semantic.py +++ b/unit_tests/test_semantic.py @@ -16,7 +16,7 @@ from spineps.phase_pre import preprocess_input from spineps.phase_semantic import predict_semantic_mask from spineps.seg_enums import ErrCode, OutputType -from spineps.seg_model import Segmentation_Inference_Config, Segmentation_Model, run_inference +from spineps.seg_model import Segmentation_Inference_Config, SegmentationModel, run_inference from spineps.seg_utils import check_input_model_compatibility, check_model_modality_acquisition logger = No_Logger() @@ -30,7 +30,7 @@ def predict_single_npy_array(self, arr: np.ndarray): return arr -class Segmentation_Model_Dummy(Segmentation_Model): +class SegmentationModelDummy(SegmentationModel): def __init__( self, model_folder: str | Path = __file__, @@ -77,7 +77,7 @@ def test_compatibility(self): mri, _subreg, _vert, _label = get_test_mri() input_path = get_tests_dir().joinpath("sample_mri", "sub-mri_label-6_T2w.nii.gz") - model = Segmentation_Model_Dummy() + model = SegmentationModelDummy() print(input_path) bf = BIDS_FILE(input_path, dataset=input_path.parent) @@ -114,7 +114,7 @@ def test_phase_preprocess(self): def test_segment_scan(self): mri, subreg, _vert, _label = get_test_mri() - model = Segmentation_Model_Dummy() + model = SegmentationModelDummy() model.run = MagicMock(return_value={OutputType.seg: subreg, OutputType.softmax_logits: None}) debug_data = {} seg_nii, _softmax_logits, errcode = predict_semantic_mask( @@ -135,7 +135,7 @@ def test_segment_scan(self): def test_run_inference(self): mri, subreg, _vert, _label = get_test_mri() - model = Segmentation_Model_Dummy().load() + model = SegmentationModelDummy().load() s_arr = subreg.get_seg_array() model.predictor.predict_single_npy_array = MagicMock(return_value=(s_arr, s_arr[np.newaxis, :])) From fc0aa37c6780deeea7286bd56f7f878b384e8f07 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 20:25:54 +0000 Subject: [PATCH 07/29] feat: high-level spineps.segment() API + grouped config dataclasses Users no longer have to load three models and pass ~45 flat keyword arguments to segment_image. - spineps.segment(image, ...) and a reusable SpinepsPipeline (loads models once for batches), returning a small SpinepsResult (errcode + in-memory masks or written paths). Accepts a path, an in-memory NII, or a BIDS_FILE. - spineps/config.py: SemanticConfig / InstanceConfig / LabelingConfig / PostConfig group the proc_* flags; .to_kwargs() maps back to the flat segment_image args (pipeline unchanged). - citation_reminder now uses functools.wraps, so decorated functions keep their real signatures (better IDE/docs/introspection). - Export new names from __init__ with explicit __all__. 12 new tests, 190 pass. Co-Authored-By: Claude Opus 4.8 --- spineps/__init__.py | 25 +++- spineps/api.py | 226 +++++++++++++++++++++++++++++ spineps/config.py | 86 +++++++++++ spineps/utils/citation_reminder.py | 2 + unit_tests/test_api.py | 118 +++++++++++++++ 5 files changed, 456 insertions(+), 1 deletion(-) create mode 100644 spineps/api.py create mode 100644 spineps/config.py create mode 100644 unit_tests/test_api.py diff --git a/spineps/__init__.py b/spineps/__init__.py index 32504bb..2c50a54 100755 --- a/spineps/__init__.py +++ b/spineps/__init__.py @@ -1,5 +1,7 @@ -"""SPINEPS: spine MRI segmentation package exposing the model loaders, pipeline phases and run entry points.""" +"""SPINEPS: spine MRI segmentation package exposing the high-level API, model loaders and pipeline phases.""" +from spineps.api import SpinepsPipeline, SpinepsResult, segment +from spineps.config import InstanceConfig, LabelingConfig, PostConfig, SemanticConfig from spineps.entrypoint import entry_point from spineps.get_models import get_instance_model, get_labeling_model, get_semantic_model from spineps.phase_instance import predict_instance_mask @@ -8,3 +10,24 @@ from spineps.phase_semantic import predict_semantic_mask from spineps.seg_model import SegmentationModel from spineps.seg_run import process_dataset, segment_image + +__all__ = [ + "InstanceConfig", + "LabelingConfig", + "PostConfig", + "SegmentationModel", + "SemanticConfig", + "SpinepsPipeline", + "SpinepsResult", + "entry_point", + "get_instance_model", + "get_labeling_model", + "get_semantic_model", + "perform_labeling_step", + "phase_postprocess_combined", + "predict_instance_mask", + "predict_semantic_mask", + "process_dataset", + "segment", + "segment_image", +] diff --git a/spineps/api.py b/spineps/api.py new file mode 100644 index 0000000..92268dc --- /dev/null +++ b/spineps/api.py @@ -0,0 +1,226 @@ +"""High-level, one-call API for SPINEPS. + +The pipeline functions in :mod:`spineps.seg_run` are powerful but verbose: you must load three models yourself and pass +many flat keyword arguments. This module wraps that into a single :func:`segment` call (and a reusable +:class:`SpinepsPipeline` for batches), returning a small :class:`SpinepsResult`. + +Example:: + + import spineps + + result = spineps.segment("sub-01_T2w.nii.gz") # saves a BIDS derivatives folder next to the input + result = spineps.segment(nii, output_in_memory=True) # returns the masks in memory + if result.success: + seg, vert = result.semantic, result.vertebra +""" + +from __future__ import annotations + +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Union + +from TPTBox import BIDS_FILE, NII, POI + +from spineps.config import InstanceConfig, LabelingConfig, PostConfig, SemanticConfig +from spineps.get_models import get_actual_model, get_instance_model, get_labeling_model, get_semantic_model +from spineps.phase_labeling import VertLabelingClassifier +from spineps.seg_enums import ErrCode +from spineps.seg_model import SegmentationModel +from spineps.seg_run import segment_image + +# Accepted forms for an input image and a model argument. +ImageInput = Union[str, Path, NII, BIDS_FILE] +ModelInput = Union[str, Path, SegmentationModel, VertLabelingClassifier] + + +@dataclass +class SpinepsResult: + """Result of a single :func:`segment` call. + + Attributes: + errcode (ErrCode): Outcome of the run (``ErrCode.OK`` / ``ErrCode.ALL_DONE`` mean success). + semantic (NII | None): Subregion (semantic) mask, when run with ``output_in_memory=True``. + vertebra (NII | None): Vertebra (instance) mask, when run with ``output_in_memory=True``. + centroids (POI | None): Computed centroids, when run with ``output_in_memory=True``. + output_paths (dict[str, Path] | None): Written output file paths, when run with ``output_in_memory=False``. + """ + + errcode: ErrCode + semantic: NII | None = None + vertebra: NII | None = None + centroids: POI | None = None + output_paths: dict[str, Path] | None = None + + @property + def success(self) -> bool: + """True if the run completed (freshly processed or already done).""" + return self.errcode in (ErrCode.OK, ErrCode.ALL_DONE) + + +def _resolve_model(model: ModelInput | None, getter, use_cpu: bool): + """Turns a model id / path / already-loaded model into a loaded model object (or None).""" + if model is None or model == "none": + return None + if isinstance(model, (SegmentationModel, VertLabelingClassifier)): + return model + if "/" in str(model): # treat anything path-like as an explicit model folder + return get_actual_model(model, use_cpu=use_cpu).load() + return getter(str(model), use_cpu=use_cpu).load() + + +@contextmanager +def _as_bids_file(image: ImageInput): + """Yields a ``BIDS_FILE`` for the given image, writing an in-memory NII to a temporary file if needed.""" + if isinstance(image, BIDS_FILE): + yield image + elif isinstance(image, NII): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "sub-spineps_T2w.nii.gz" + image.save(path, verbose=False) + yield BIDS_FILE(str(path), dataset=tmp, verbose=False) + else: + path = Path(image).absolute() + if not str(path).endswith(".nii.gz"): + raise ValueError(f"image must be a .nii.gz file, got {path}") + if not path.is_file(): + raise FileNotFoundError(f"image does not exist or is not a file, got {path}") + yield BIDS_FILE(str(path), dataset=str(path.parent), verbose=False) + + +def _to_result(raw: tuple) -> SpinepsResult: + """Normalizes the two return shapes of ``segment_image`` into a SpinepsResult.""" + if len(raw) == 4: # in-memory mode, fully processed: (seg, vert, centroids, errcode) + seg_nii, vert_nii, centroids, errcode = raw + return SpinepsResult(errcode=errcode, semantic=seg_nii, vertebra=vert_nii, centroids=centroids) + output_paths, errcode = raw # save mode or any early exit: (output_paths, errcode) + return SpinepsResult(errcode=errcode, output_paths=output_paths) + + +class SpinepsPipeline: + """Holds loaded SPINEPS models so repeated :meth:`segment` calls don't reload them. + + Args: + model_semantic: Semantic model id, path, or an already-loaded model. Defaults to ``"t2w"``. + model_instance: Instance model id, path, or an already-loaded model. Defaults to ``"instance"``. + model_labeling: Labeling model id/path/model, or ``None``/``"none"`` to skip labeling. Defaults to ``"t2w_labeling"``. + use_cpu: If true, runs on CPU instead of GPU. Defaults to False. + """ + + def __init__( + self, + model_semantic: ModelInput = "t2w", + model_instance: ModelInput = "instance", + model_labeling: ModelInput | None = "t2w_labeling", + *, + use_cpu: bool = False, + ) -> None: + self.model_semantic = _resolve_model(model_semantic, get_semantic_model, use_cpu) + self.model_instance = _resolve_model(model_instance, get_instance_model, use_cpu) + self.model_labeling = _resolve_model(model_labeling, get_labeling_model, use_cpu) + if self.model_semantic is None: + raise ValueError("model_semantic could not be resolved") + if self.model_instance is None: + raise ValueError("model_instance could not be resolved") + + def segment( + self, + image: ImageInput, + *, + output_in_memory: bool = False, + derivative_name: str = "derivatives_seg", + override: bool = False, + semantic: SemanticConfig | None = None, + instance: InstanceConfig | None = None, + labeling: LabelingConfig | None = None, + post: PostConfig | None = None, + verbose: bool = False, + ) -> SpinepsResult: + """Segments a single image with the already-loaded models. + + Args: + image: A path to a ``.nii.gz`` file, an in-memory ``NII``, or a ``BIDS_FILE``. An in-memory ``NII`` always + returns its results in memory. + output_in_memory: If true, returns the masks in memory instead of writing a derivatives folder. Defaults to False. + derivative_name: Name of the derivatives output folder (save mode only). Defaults to ``"derivatives_seg"``. + override: If true, recomputes and overwrites existing outputs. Defaults to False. + semantic, instance, labeling, post: Optional grouped config objects; unset groups use the pipeline defaults. + verbose: If true, prints verbose information. Defaults to False. + + Returns: + SpinepsResult: The outcome, carrying either in-memory masks or the written output paths. + """ + in_memory = output_in_memory or isinstance(image, NII) + kwargs: dict = {} + for cfg in (semantic, instance, labeling, post): + if cfg is not None: + kwargs.update(cfg.to_kwargs()) + with _as_bids_file(image) as img_ref: + raw = segment_image( + img_ref=img_ref, + model_semantic=self.model_semantic, + model_instance=self.model_instance, + model_labeling=self.model_labeling, + derivative_name=derivative_name, + override_semantic=override, + override_instance=override, + override_postpair=override, + override_ctd=override, + return_output_instead_of_save=in_memory, + verbose=verbose, + **kwargs, + ) + return _to_result(raw) + + +def segment( + image: ImageInput, + *, + model_semantic: ModelInput = "t2w", + model_instance: ModelInput = "instance", + model_labeling: ModelInput | None = "t2w_labeling", + use_cpu: bool = False, + output_in_memory: bool = False, + derivative_name: str = "derivatives_seg", + override: bool = False, + semantic: SemanticConfig | None = None, + instance: InstanceConfig | None = None, + labeling: LabelingConfig | None = None, + post: PostConfig | None = None, + verbose: bool = False, +) -> SpinepsResult: + """Segments a single image end to end in one call (loads the models, runs the pipeline). + + This is the easiest entry point: it resolves and loads the three models, wraps the input as needed and runs the + full pipeline. To segment many images, build a :class:`SpinepsPipeline` once and call its :meth:`~SpinepsPipeline.segment` + repeatedly so the models are loaded only once. + + Args: + image: A path to a ``.nii.gz`` file, an in-memory ``NII``, or a ``BIDS_FILE``. + model_semantic: Semantic model id, path, or loaded model. Defaults to ``"t2w"``. + model_instance: Instance model id, path, or loaded model. Defaults to ``"instance"``. + model_labeling: Labeling model id/path/model, or ``None``/``"none"`` to skip labeling. Defaults to ``"t2w_labeling"``. + use_cpu: If true, runs on CPU instead of GPU. Defaults to False. + output_in_memory: If true, returns the masks in memory instead of writing a derivatives folder. Defaults to False. + derivative_name: Name of the derivatives output folder (save mode only). Defaults to ``"derivatives_seg"``. + override: If true, recomputes and overwrites existing outputs. Defaults to False. + semantic, instance, labeling, post: Optional grouped config objects (see :mod:`spineps.config`). + verbose: If true, prints verbose information. Defaults to False. + + Returns: + SpinepsResult: The outcome, carrying either in-memory masks or the written output paths. + """ + pipeline = SpinepsPipeline(model_semantic, model_instance, model_labeling, use_cpu=use_cpu) + return pipeline.segment( + image, + output_in_memory=output_in_memory, + derivative_name=derivative_name, + override=override, + semantic=semantic, + instance=instance, + labeling=labeling, + post=post, + verbose=verbose, + ) diff --git a/spineps/config.py b/spineps/config.py new file mode 100644 index 0000000..af66f93 --- /dev/null +++ b/spineps/config.py @@ -0,0 +1,86 @@ +"""Grouped configuration objects for the SPINEPS pipeline. + +These dataclasses bundle the many flat ``proc_*`` keyword arguments of :func:`spineps.seg_run.segment_image` into a few +themed groups so the high-level :func:`spineps.segment` API stays readable. Each config maps back to the flat keyword +arguments via :meth:`to_kwargs`, so nothing about the underlying pipeline changes. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class SemanticConfig: + """Options for the semantic (subregion) segmentation phase.""" + + crop_input: bool = True + n4_bias_correction: bool = True + remove_inferior_beyond_canal: bool = False + clean_beyond_largest_bounding_box: bool = True + clean_small_cc_artifacts: bool = True + + def to_kwargs(self) -> dict: + """Returns the corresponding flat ``segment_image`` keyword arguments.""" + return { + "proc_sem_crop_input": self.crop_input, + "proc_sem_n4_bias_correction": self.n4_bias_correction, + "proc_sem_remove_inferior_beyond_canal": self.remove_inferior_beyond_canal, + "proc_sem_clean_beyond_largest_bounding_box": self.clean_beyond_largest_bounding_box, + "proc_sem_clean_small_cc_artifacts": self.clean_small_cc_artifacts, + } + + +@dataclass +class InstanceConfig: + """Options for the vertebra (instance) segmentation phase.""" + + corpus_clean: bool = True + clean_small_cc_artifacts: bool = True + largest_k_cc: int = 0 + detect_and_solve_merged_corpi: bool = True + batch_size: int = 4 + labeling_offset: int = 2 + + def to_kwargs(self) -> dict: + """Returns the corresponding flat ``segment_image`` keyword arguments.""" + return { + "proc_inst_corpus_clean": self.corpus_clean, + "proc_inst_clean_small_cc_artifacts": self.clean_small_cc_artifacts, + "proc_inst_largest_k_cc": self.largest_k_cc, + "proc_inst_detect_and_solve_merged_corpi": self.detect_and_solve_merged_corpi, + "proc_inst_batch_size": self.batch_size, + "vertebra_instance_labeling_offset": self.labeling_offset, + } + + +@dataclass +class LabelingConfig: + """Options for the vertebra-labeling phase.""" + + force_no_tl_anomaly: bool = False + + def to_kwargs(self) -> dict: + """Returns the corresponding flat ``segment_image`` keyword arguments.""" + return {"proc_lab_force_no_tl_anomaly": self.force_no_tl_anomaly} + + +@dataclass +class PostConfig: + """Options for the combined post-processing phase (semantic + instance reconciliation).""" + + fill_3d_holes: bool = True + assign_missing_cc: bool = True + assign_missing_cc_fast: bool = False + clean_inst_by_sem: bool = True + vertebra_inconsistency: bool = True + + def to_kwargs(self) -> dict: + """Returns the corresponding flat ``segment_image`` keyword arguments.""" + return { + "proc_fill_3d_holes": self.fill_3d_holes, + "proc_assign_missing_cc": self.assign_missing_cc, + "proc_assign_missing_cc_fast": self.assign_missing_cc_fast, + "proc_clean_inst_by_sem": self.clean_inst_by_sem, + "proc_vertebra_inconsistency": self.vertebra_inconsistency, + } diff --git a/spineps/utils/citation_reminder.py b/spineps/utils/citation_reminder.py index 7160443..42cf08c 100644 --- a/spineps/utils/citation_reminder.py +++ b/spineps/utils/citation_reminder.py @@ -3,6 +3,7 @@ from __future__ import annotations import atexit +import functools import os from rich.console import Console @@ -17,6 +18,7 @@ def citation_reminder(func): """Decorator to remind users to cite SPINEPS.""" + @functools.wraps(func) def wrapper(*args, **kwargs): global has_reminded_citation # noqa: PLW0603 if not has_reminded_citation and os.environ.get("SPINEPS_TURN_OF_CITATION_REMINDER", "FALSE") != "TRUE": diff --git a/unit_tests/test_api.py b/unit_tests/test_api.py new file mode 100644 index 0000000..84e07d7 --- /dev/null +++ b/unit_tests/test_api.py @@ -0,0 +1,118 @@ +# Call 'python -m unittest' on this folder # noqa: INP001 +"""Tests for the high-level spineps.segment API, config dataclasses and result wrapping.""" + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from TPTBox import BIDS_FILE +from TPTBox.tests.test_utils import get_test_mri + +from spineps import InstanceConfig, LabelingConfig, PostConfig, SemanticConfig, segment +from spineps.api import SpinepsPipeline, _as_bids_file, _to_result +from spineps.seg_enums import ErrCode + + +class Test_Config_To_Kwargs(unittest.TestCase): + def test_semantic(self): + kw = SemanticConfig(crop_input=False, n4_bias_correction=False).to_kwargs() + self.assertFalse(kw["proc_sem_crop_input"]) + self.assertFalse(kw["proc_sem_n4_bias_correction"]) + + def test_instance(self): + kw = InstanceConfig(batch_size=8, largest_k_cc=3).to_kwargs() + self.assertEqual(kw["proc_inst_batch_size"], 8) + self.assertEqual(kw["proc_inst_largest_k_cc"], 3) + self.assertEqual(kw["vertebra_instance_labeling_offset"], 2) + + def test_labeling_and_post(self): + self.assertTrue(LabelingConfig(force_no_tl_anomaly=True).to_kwargs()["proc_lab_force_no_tl_anomaly"]) + self.assertFalse(PostConfig(fill_3d_holes=False).to_kwargs()["proc_fill_3d_holes"]) + + +class Test_To_Result(unittest.TestCase): + def test_save_mode_tuple(self): + r = _to_result(({"out_spine": Path("/tmp/x")}, ErrCode.OK)) + self.assertTrue(r.success) + self.assertIsNotNone(r.output_paths) + self.assertIsNone(r.semantic) + + def test_in_memory_tuple(self): + a, b, c = object(), object(), object() + r = _to_result((a, b, c, ErrCode.OK)) + self.assertIs(r.semantic, a) + self.assertIs(r.vertebra, b) + self.assertIs(r.centroids, c) + self.assertIsNone(r.output_paths) + + def test_failure_errcode_not_success(self): + self.assertFalse(_to_result(({}, ErrCode.COMPATIBILITY)).success) + + +class Test_As_Bids_File(unittest.TestCase): + def test_wraps_existing_path(self): + mri, *_ = get_test_mri() + with tempfile.TemporaryDirectory() as td: + path = Path(td) / "sub-x_T2w.nii.gz" + mri.save(path, verbose=False) + with _as_bids_file(str(path)) as bf: + self.assertIsInstance(bf, BIDS_FILE) + + def test_wraps_in_memory_nii(self): + mri, *_ = get_test_mri() + with _as_bids_file(mri) as bf: + self.assertIsInstance(bf, BIDS_FILE) + + def test_missing_path_raises(self): + with self.assertRaises(FileNotFoundError), _as_bids_file("/no/such/file.nii.gz"): + pass + + def test_non_niigz_raises(self): + with self.assertRaises(ValueError), _as_bids_file("/tmp/not_an_image.png"): + pass + + +class Test_Segment_Api(unittest.TestCase): + def test_segment_forwards_config_kwargs(self): + captured: dict = {} + + def fake_segment_image(**kwargs): + captured.update(kwargs) + return ({"out_spine": Path("/tmp/x")}, ErrCode.OK) + + mri, *_ = get_test_mri() + with ( + mock.patch("spineps.api._resolve_model", side_effect=lambda m, _getter, _cpu: m or "model"), + mock.patch("spineps.api.segment_image", side_effect=fake_segment_image), + ): + res = segment(mri, instance=InstanceConfig(batch_size=8), semantic=SemanticConfig(crop_input=False)) + self.assertEqual(res.errcode, ErrCode.OK) + self.assertEqual(captured["proc_inst_batch_size"], 8) + self.assertFalse(captured["proc_sem_crop_input"]) + # an in-memory NII input forces in-memory output + self.assertTrue(captured["return_output_instead_of_save"]) + + def test_pipeline_resolves_models_once(self): + calls = {"n": 0} + + def fake_resolve(model, _getter, _cpu): + calls["n"] += 1 + return model or "model" + + with ( + mock.patch("spineps.api._resolve_model", side_effect=fake_resolve), + mock.patch("spineps.api.segment_image", return_value=({"x": Path("/tmp/x")}, ErrCode.OK)), + ): + pipe = SpinepsPipeline("t2w", "instance", "t2w_labeling") + mri, *_ = get_test_mri() + pipe.segment(mri) + pipe.segment(mri) + # 3 models resolved once in __init__, never again per segment() call + self.assertEqual(calls["n"], 3) + + +if __name__ == "__main__": + unittest.main() From e2461ad8fd827d1d4f370d14353db29a0b75d87d Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 20:39:06 +0000 Subject: [PATCH 08/29] refactor!: kebab-case CLI flags, positive-polarity booleans, --batch-size; fix -h crash 2.0 CLI standardization (clean break): - All long flags are now --kebab-case (short aliases like -i/-ms/-mv kept); e.g. -model_semantic -> --model-semantic, -der_name -> --derivative-name, -raw_name -> --rawdata-name. - Negative-polarity flags replaced with positive BooleanOptionalAction pairs: -nocrop -> --crop/--no-crop, -non4 -> --n4/--no-n4, -no_tltv_labeling -> --enforce-12-thoracic/--no-enforce-12-thoracic. - New --batch-size/-bs knob exposing the instance batch size; added -mi alias for --model-instance. - Fix: empty metavar='' crashed 'spineps sample -h' / 'dataset -h' (argparse usage AssertionError); removed the empty metavars so help works again. - run_sample/run_dataset read the new dests. Tests for flag defaults/negation + a -h regression test. 192 pass. Co-Authored-By: Claude Opus 4.8 --- spineps/entrypoint.py | 132 ++++++++++++++++------------------ unit_tests/test_entrypoint.py | 30 ++++++++ 2 files changed, 93 insertions(+), 69 deletions(-) diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index 1a43d97..48799e7 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -39,58 +39,69 @@ def parser_arguments(parser: argparse.ArgumentParser): Returns: argparse.ArgumentParser: The same parser with the shared arguments registered. """ - parser.add_argument("-der_name", "-dn", type=str, default="derivatives_seg", metavar="", help="Name of the derivatives folder") - parser.add_argument("-save_debug", "-sd", action="store_true", help="Saves a lot of debug data and intermediate results") - # parser.add_argument("-save_unc_img", "-sui", action="store_true", help="Saves a uncertainty image from the subreg prediction") + parser.add_argument("--derivative-name", "-dn", type=str, default="derivatives_seg", help="Name of the derivatives folder") + parser.add_argument("--save-debug", "-sd", action="store_true", help="Saves a lot of debug data and intermediate results") parser.add_argument( - "-save_softmax_logits", "-ssl", action="store_true", help="Saves an .npz containing the softmax logit outputs of the semantic mask" + "--save-softmax-logits", "-ssl", action="store_true", help="Saves an .npz containing the softmax logit outputs of the semantic mask" ) parser.add_argument( - "-save_modelres_mask", + "--save-modelres-mask", "-smrm", action="store_true", help="If true, will additionally save the semantic mask in the resolution used by the model", ) - parser.add_argument("-override_semantic", "-os", action="store_true", help="Will override existing seg-spine files") + parser.add_argument("--override-semantic", "-os", action="store_true", help="Will override existing seg-spine files") parser.add_argument( - "-override_instance", "-oi", action="store_true", help="Will override existing seg-vert files (True if semantic mask changed)" + "--override-instance", "-oi", action="store_true", help="Will override existing seg-vert files (True if semantic mask changed)" ) parser.add_argument( - "-override_postpair", + "--override-postpair", "-opp", action="store_true", help="Will override existing cleaned files (True if either semantic or instance mask changed)", ) parser.add_argument( - "-override_ctd", "-oc", action="store_true", help="Will override existing centroid files (True if the instance mask changed)" + "--override-ctd", "-oc", action="store_true", help="Will override existing centroid files (True if the instance mask changed)" ) parser.add_argument( - "-ignore_inference_compatibility", + "--ignore-inference-compatibility", "-iic", action="store_true", help="Does not skip input masks that do not match the models modalities", ) parser.add_argument( - "-nocrop", - "-nc", - action="store_true", - help="Does not crop input before semantically segmenting. Can improve the segmentation a little but depending on size costs more computation time", + "--crop", + action=argparse.BooleanOptionalAction, + default=True, + dest="crop_input", + help="Crop the input to the spine before semantic segmentation. Use --no-crop to disable (can slightly improve the " + "segmentation but costs more computation time).", ) parser.add_argument( - "-no_tltv_labeling", - "-ntl", - action="store_true", - help="Enforces the labeling model to predict exactly 12 thoracic vertebrae", + "--n4", + action=argparse.BooleanOptionalAction, + default=True, + dest="n4", + help="Apply N4 bias field correction before semantic segmentation (MRI only). Use --no-n4 to disable (faster).", ) - # proc_lab_force_no_tl_anomaly parser.add_argument( - "-non4", - action="store_true", - help="Does not apply n4 bias field correction", + "--enforce-12-thoracic", + action=argparse.BooleanOptionalAction, + default=False, + dest="enforce_12_thoracic", + help="Force the labeling model to predict exactly 12 thoracic vertebrae (assume no thoracolumbar transition anomaly).", + ) + parser.add_argument( + "--batch-size", + "-bs", + type=int, + default=4, + help="Number of vertebra cutouts run through the instance model per batched forward pass. Higher is faster but uses " + "more GPU memory; falls back to one-by-one on out-of-memory.", ) - parser.add_argument("-cpu", action="store_true", help="Use CPU instead of GPU (will take way longer)") - parser.add_argument("-run_cprofiler", "-rcp", action="store_true", help="Runs a cprofiler over the entire action") - parser.add_argument("-verbose", "-v", action="store_true", help="Prints much more stuff, may fully clutter your terminal") + parser.add_argument("--cpu", action="store_true", help="Use CPU instead of GPU (will take way longer)") + parser.add_argument("--run-cprofiler", "-rcp", action="store_true", help="Runs a cprofiler over the entire action") + parser.add_argument("--verbose", "-v", action="store_true", help="Prints much more stuff, may fully clutter your terminal") return parser @@ -116,35 +127,25 @@ def entry_point(): ) ########################### ########################### - parser_sample.add_argument("-input", "-i", required=True, type=str, help="path to the input nifty file") + parser_sample.add_argument("--input", "-i", required=True, type=str, help="path to the input NIfTI file") parser_sample.add_argument( - "-model_semantic", + "--model-semantic", "-ms", - # type=str.lower, default=None, required=True, - # choices=modelids_semantic, - metavar="", help="The model used for the semantic segmentation. You can also pass an absolute path to the model folder", ) parser_sample.add_argument( - "-model_instance", + "--model-instance", "-mv", - # type=str.lower, + "-mi", default="instance", - # required=True, - # choices=modelids_instance, - metavar="", help="The model used for the vertebra instance segmentation. You can also pass an absolute path to the model folder", ) parser_sample.add_argument( - "-model_labeling", + "--model-labeling", "-ml", - # type=str.lower, default="t2w_labeling", - # required=True, - # choices=modelids_instance, - metavar="", help="The model used for the vertebra labeling classification. You can also pass an absolute path to the model folder", ) parser_sample = parser_arguments(parser_sample) @@ -153,54 +154,45 @@ def entry_point(): # # parser_dataset.add_argument( - "-directory", "-i", "-d", required=True, type=str, help="path to the input directory, preferably a BIDS dataset" + "--directory", "-i", "-d", required=True, type=str, help="path to the input directory, preferably a BIDS dataset" ) - parser_dataset.add_argument("-raw_name", "-rn", type=str, default="rawdata", metavar="", help="Name of the rawdata folder") + parser_dataset.add_argument("--rawdata-name", "-rn", type=str, default="rawdata", help="Name of the rawdata folder") parser_dataset.add_argument( - "-model_semantic", + "--model-semantic", "-ms", - # type=str.lower, default="t2w", - # choices=model_subreg_choices, - metavar="", help="The model used for the subregion segmentation. Pass 'auto' to auto-select a model by modality, or an absolute path to the model folder", ) parser_dataset.add_argument( - "-model_instance", + "--model-instance", "-mv", - # type=str.lower, + "-mi", default="instance", - # choices=model_vert_choices, - metavar="", help="The model used for the vertebra segmentation. You can also pass an absolute path to the model folder", ) parser_dataset.add_argument( - "-model_labeling", + "--model-labeling", "-ml", - # type=str.lower, default="t2w_labeling", - # required=True, - # choices=modelids_instance, - metavar="", help="The model used for the vertebra labeling classification. You can also pass an absolute path to the model folder", ) parser_dataset.add_argument( - "-ignore_bids_filter", + "--ignore-bids-filter", "-ibf", action="store_true", help="If true, will search the BIDS dataset without the strict filters. Use with care!", ) parser_dataset.add_argument( - "-ignore_model_compatibility", + "--ignore-model-compatibility", "-imc", action="store_true", help="If true, will not stop the pipeline to use the given models on unfitting input modalities", ) parser_dataset.add_argument( - "-save_log", "-sl", action="store_true", help="If true, saves the log into a separate folder in the dataset directory" + "--save-log", "-sl", action="store_true", help="If true, saves the log into a separate folder in the dataset directory" ) parser_dataset.add_argument( - "-save_snaps_folder", + "--save-snaps-folder", "-ssf", action="store_true", help="If true, saves the snapshots also in a separate folder in the dataset directory", @@ -273,7 +265,7 @@ def run_sample(opt: Namespace): "model_semantic": model_semantic, "model_instance": model_instance, "model_labeling": model_labeling, - "derivative_name": opt.der_name, + "derivative_name": opt.derivative_name, # # "save_uncertainty_image": opt.save_unc_img, "save_softmax_logits": opt.save_softmax_logits, @@ -283,9 +275,10 @@ def run_sample(opt: Namespace): "override_instance": opt.override_instance, "override_postpair": opt.override_postpair, "override_ctd": opt.override_ctd, - "proc_sem_crop_input": not opt.nocrop, - "proc_sem_n4_bias_correction": not opt.non4, - "proc_lab_force_no_tl_anomaly": opt.no_tltv_labeling, + "proc_sem_crop_input": opt.crop_input, + "proc_sem_n4_bias_correction": opt.n4, + "proc_lab_force_no_tl_anomaly": opt.enforce_12_thoracic, + "proc_inst_batch_size": opt.batch_size, "ignore_compatibility_issues": opt.ignore_inference_compatibility, "verbose": opt.verbose, } @@ -297,7 +290,7 @@ def run_sample(opt: Namespace): timestamp = format_time_short(get_time()) cprofile_out = bids_sample.get_changed_path( bids_format="log", - parent=opt.der_name, + parent=opt.derivative_name, file_type="log", info={"desc": "cprofile", "mod": bids_sample.format, "ses": timestamp}, ) @@ -367,8 +360,8 @@ def run_dataset(opt: Namespace): "model_semantic": model_semantic, "model_instance": model_instance, "model_labeling": model_labeling, - "rawdata_name": opt.raw_name, - "derivative_name": opt.der_name, + "rawdata_name": opt.rawdata_name, + "derivative_name": opt.derivative_name, # # "save_uncertainty_image": opt.save_unc_img, "save_modelres_mask": opt.save_modelres_mask, @@ -382,9 +375,10 @@ def run_dataset(opt: Namespace): "ignore_model_compatibility": opt.ignore_model_compatibility, "ignore_inference_compatibility": opt.ignore_inference_compatibility, "ignore_bids_filter": opt.ignore_bids_filter, - "proc_sem_crop_input": not opt.nocrop, - "proc_sem_n4_bias_correction": not opt.non4, - "proc_lab_force_no_tl_anomaly": opt.no_tltv_labeling, + "proc_sem_crop_input": opt.crop_input, + "proc_sem_n4_bias_correction": opt.n4, + "proc_lab_force_no_tl_anomaly": opt.enforce_12_thoracic, + "proc_inst_batch_size": opt.batch_size, "snapshot_copy_folder": opt.save_snaps_folder, "verbose": opt.verbose, } diff --git a/unit_tests/test_entrypoint.py b/unit_tests/test_entrypoint.py index c0ce729..47fcaf7 100644 --- a/unit_tests/test_entrypoint.py +++ b/unit_tests/test_entrypoint.py @@ -5,7 +5,11 @@ from __future__ import annotations import argparse +import contextlib +import io +import sys import unittest +from unittest import mock from TPTBox import No_Logger @@ -18,6 +22,32 @@ class Test_EntryPoint(unittest.TestCase): def test_normal(self): entrypoint.parser_arguments(argparse.ArgumentParser()) + def test_shared_flag_defaults_and_negation(self): + # positive-polarity boolean flags + the new batch-size knob + p = argparse.ArgumentParser() + entrypoint.parser_arguments(p) + defaults = p.parse_args([]) + self.assertTrue(defaults.crop_input) + self.assertTrue(defaults.n4) + self.assertFalse(defaults.enforce_12_thoracic) + self.assertEqual(defaults.batch_size, 4) + negated = p.parse_args(["--no-crop", "--no-n4", "--enforce-12-thoracic", "--batch-size", "8"]) + self.assertFalse(negated.crop_input) + self.assertFalse(negated.n4) + self.assertTrue(negated.enforce_12_thoracic) + self.assertEqual(negated.batch_size, 8) + + def test_subparser_help_builds(self): + # Regression: empty metavar="" used to crash `spineps -h` in argparse usage formatting. + for sub in ("sample", "dataset"): + with ( + self.assertRaises(SystemExit) as cm, + mock.patch.object(sys, "argv", ["spineps", sub, "-h"]), + contextlib.redirect_stdout(io.StringIO()), + ): + entrypoint.entry_point() + self.assertEqual(cm.exception.code, 0) + def test_run_sample_missing_parent_raises(self): # parent directory of the input does not exist -> FileNotFoundError (not a bare AssertionError) opt = argparse.Namespace(input="/this/path/does/not/exist/scan.nii.gz") From 4a945a8418bf7addb52fe922e23c1954d94c11ce Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 10 Jun 2026 20:42:49 +0000 Subject: [PATCH 09/29] docs: add MIGRATION.md and update README/docs for the 2.0 API and CLI - MIGRATION.md: complete old->new mapping for CLI flags, functions and classes, plus the new spineps.segment API. - README and docs: rename process_img_nii->segment_image and Segmentation_Model*->PascalCase, kebab-case all CLI flags/examples, and lead the Python usage with spineps.segment() / SpinepsPipeline / config objects. Co-Authored-By: Claude Opus 4.8 --- MIGRATION.md | 92 ++++++++++++++++++++++++++++++++++++++++ README.md | 43 ++++++++++++------- docs/getting-started.md | 44 +++++++++++++------ docs/index.md | 2 +- docs/modules/models.md | 8 ++-- docs/modules/pipeline.md | 30 +++++++++++-- 6 files changed, 181 insertions(+), 38 deletions(-) create mode 100644 MIGRATION.md diff --git a/MIGRATION.md b/MIGRATION.md new file mode 100644 index 0000000..7aa6060 --- /dev/null +++ b/MIGRATION.md @@ -0,0 +1,92 @@ +# Migrating to SPINEPS 2.0 + +SPINEPS 2.0 is a clean-break release: CLI flags, public functions and classes were renamed for clarity and +consistency, and a new one-call Python API was added. **No backward-compatible aliases are kept** — update your +scripts using the tables below. + +## TL;DR — the new one-call API + +The biggest change is that you no longer have to load three models and pass dozens of flags: + +```python +import spineps + +# Save a BIDS derivatives folder next to the input: +result = spineps.segment("sub-01_T2w.nii.gz") + +# Or get the masks in memory: +result = spineps.segment(nii, output_in_memory=True) +if result.success: + semantic, vertebra = result.semantic, result.vertebra + +# Segmenting many images? Load the models once: +from spineps import SpinepsPipeline +pipe = SpinepsPipeline(model_semantic="t2w", model_instance="instance") +for path in paths: + pipe.segment(path) +``` + +## CLI flags + +All long flags are now `--kebab-case` (short aliases are unchanged). Negative flags became positive on/off pairs. + +| Old | New | +| --- | --- | +| `-nocrop` / `-nc` | `--no-crop` (cropping is on by default; pass `--no-crop` to disable) | +| `-non4` | `--no-n4` (N4 is on by default; pass `--no-n4` to disable) | +| `-no_tltv_labeling` / `-ntl` | `--enforce-12-thoracic` | +| `-input` / `-i` | `--input` / `-i` | +| `-directory` / `-d` | `--directory` / `-d` | +| `-model_semantic` / `-ms` | `--model-semantic` / `-ms` | +| `-model_instance` / `-mv` | `--model-instance` / `-mv` / `-mi` | +| `-model_labeling` / `-ml` | `--model-labeling` / `-ml` | +| `-der_name` / `-dn` | `--derivative-name` / `-dn` | +| `-raw_name` / `-rn` | `--rawdata-name` / `-rn` | +| `-save_debug` / `-sd` | `--save-debug` / `-sd` | +| `-save_softmax_logits` / `-ssl` | `--save-softmax-logits` / `-ssl` | +| `-save_modelres_mask` / `-smrm` | `--save-modelres-mask` / `-smrm` | +| `-save_log` / `-sl` | `--save-log` / `-sl` | +| `-save_snaps_folder` / `-ssf` | `--save-snaps-folder` / `-ssf` | +| `-override_semantic` / `-os` | `--override-semantic` / `-os` | +| `-override_instance` / `-oi` | `--override-instance` / `-oi` | +| `-override_postpair` / `-opp` | `--override-postpair` / `-opp` | +| `-override_ctd` / `-oc` | `--override-ctd` / `-oc` | +| `-ignore_inference_compatibility` / `-iic` | `--ignore-inference-compatibility` / `-iic` | +| `-ignore_bids_filter` / `-ibf` | `--ignore-bids-filter` / `-ibf` | +| `-ignore_model_compatibility` / `-imc` | `--ignore-model-compatibility` / `-imc` | +| `-run_cprofiler` / `-rcp` | `--run-cprofiler` / `-rcp` | +| `-cpu` | `--cpu` | +| `-verbose` / `-v` | `--verbose` / `-v` | +| *(new)* | `--batch-size` / `-bs` — vertebra cutouts per batched forward pass (faster; default 4) | + +Example: + +```bash +# 1.x +spineps sample -i scan.nii.gz -model_semantic t2w -model_instance instance -nocrop -non4 + +# 2.0 +spineps sample -i scan.nii.gz --model-semantic t2w --model-instance instance --no-crop --no-n4 +``` + +## Python API + +| Old | New | +| --- | --- | +| `process_img_nii(...)` | `segment_image(...)` | +| `Segmentation_Model` | `SegmentationModel` | +| `Segmentation_Model_NNunet` | `SegmentationModelNNunet` | +| `Segmentation_Model_Unet3D` | `SegmentationModelUnet3D` | + +`process_dataset`, `get_semantic_model`, `get_instance_model`, `get_labeling_model`, `predict_semantic_mask`, +`predict_instance_mask` and the other phase functions keep their names. + +## New in 2.0 + +- **`spineps.segment(...)`**, **`SpinepsPipeline`**, **`SpinepsResult`** — the high-level API shown above. +- **Config objects** `SemanticConfig`, `InstanceConfig`, `LabelingConfig`, `PostConfig` group the many `proc_*` + flags. Pass them to `segment(...)`, e.g. `spineps.segment(path, instance=InstanceConfig(batch_size=8))`. +- **`--batch-size`** / `InstanceConfig.batch_size` — the instance model now runs cutouts in batched forward passes + (much faster on GPU); falls back to one-by-one on out-of-memory. +- Clearer errors: invalid paths / missing models now raise `FileNotFoundError` / `ValueError` instead of bare + `AssertionError`, and `spineps sample -h` / `dataset -h` no longer crash. diff --git a/README.md b/README.md index 02d092f..46fd168 100644 --- a/README.md +++ b/README.md @@ -165,7 +165,7 @@ If you **don't** set the environment variable, the pipeline will look into `spin 1. Activate your venv 2. Run `python entrypoint.py -h` to see the arguments. -3. For example, for a sample, run `python entrypoint.py sample -i -model_semantic -model_instance ` +3. For example, for a sample, run `python entrypoint.py sample -i --model-semantic --model-instance ` (replacing with the name of the model you want to use) ### Issues @@ -189,10 +189,10 @@ Processes a single nifty file, will create a derivatves folder next to the nifty | argument | explanation | | :--- | --------- | | -i | Absolute path to the single nifty file (.nii.gz) to be processed | -| -model_semantic , -ms | The model used for the semantic segmentation | -| -model_instance , -mv | The model used for the vertebra instance segmentation | -| -der_name , -dn | Name of the derivatives folder (default: derivatives_seg) | -| -save_debug, -sd | Saves a lot of debug data and intermediate results in a separate debug-labeled folder (default: False) | +| --model-semantic , -ms | The model used for the semantic segmentation | +| --model-instance , -mv | The model used for the vertebra instance segmentation | +| --derivative-name , -dn | Name of the derivatives folder (default: derivatives_seg) | +| --save-debug, -sd | Saves a lot of debug data and intermediate results in a separate debug-labeled folder (default: False) | | -save_unc_img, -sui | Saves a uncertainty image from the subreg prediction (default: False) | | -override_semantic, -os | Will override existing seg-spine files (default: False) | | -override_instance, -ov | Will override existing seg-vert files (default: False) | @@ -204,9 +204,9 @@ There are a lot more arguments, run `spineps sample -h` to see them. #### Example ```bash #T2w sagittal -spineps sample -ignore_bids_filter -ignore_inference_compatibility -i /path/sub-testsample_T2w.nii.gz -model_semantic t2w -model_instance instance +spineps sample --ignore-bids-filter --ignore-inference-compatibility -i /path/sub-testsample_T2w.nii.gz --model-semantic t2w --model-instance instance #T1w sagittal -spineps sample -ignore_bids_filter -ignore_inference_compatibility -i ~/path/sub-testsample_T1w.nii.gz -model_semantic t1w -model_instance instance +spineps sample --ignore-bids-filter --ignore-inference-compatibility -i ~/path/sub-testsample_T1w.nii.gz --model-semantic t1w --model-instance instance ``` @@ -245,11 +245,11 @@ To that end, we are using TPTBox (see https://github.com/Hendrik-code/TPTBox) It supports the same arguments as in sample mode (see table above), and additionally: | argument | explanation | | :--- | --------- | -| -raw_name, -rn | Sets the name of the rawdata folder of the dataset (default: "rawdata") -| -ignore_bids_filter, -ibf | If true, will search the BIDS dataset without the strict filters. Use with care! (default: False) | -| -ignore_model_compatibility, -imc | If true, will not stop the pipeline to use the given models on unfitting input modalities (default: False) | -| -save_log, -sl | If true, saves the log into a separate folder in the dataset directory (default: False) | -| -save_snaps_folder, -ssf | If true, additionally saves the snapshots in a separate folder in the dataset directory (default: False) | +| --rawdata-name, -rn | Sets the name of the rawdata folder of the dataset (default: "rawdata") +| --ignore-bids-filter, -ibf | If true, will search the BIDS dataset without the strict filters. Use with care! (default: False) | +| --ignore-model-compatibility, -imc | If true, will not stop the pipeline to use the given models on unfitting input modalities (default: False) | +| --save-log, -sl | If true, saves the log into a separate folder in the dataset directory (default: False) | +| --save-snaps-folder, -ssf | If true, additionally saves the snapshots in a separate folder in the dataset directory (default: False) | For a full list of arguments, call `spineps dataset -h` @@ -291,7 +291,7 @@ In the vertebra instance segmentation mask, each label X in [1, 25] are the uniq ## VERIDAH: -To run the vertebra labeling after segmentation, specify a -model_labeling model (similar to -model_semantic and -model_instance). +To run the vertebra labeling after segmentation, specify a --model-labeling model (similar to --model-semantic and --model-instance). If you use VERIDAH (labeling model) in addition to the segmentation models from SPINEPS, then a labeling model will run and give each vertebrae detected by SPINEPS a vertebra label. These are @@ -309,11 +309,22 @@ The labels 100+X still correspond to the vertebra's IVD and 200+X the respective ## Using the Code -If you want to call the code snippets yourself, start by initializing your models using `seg_model.get_segmentation_model()` giving it the absolute path to your model folder. +The easiest way to run SPINEPS from Python is the one-call `spineps.segment` API, which loads the models and runs the whole pipeline: -Depending on whether you want to process a single sample or a whole dataset, go into `seg_run.py` and run either `process_img_nii()` or `process_dataset()`. +```python +import spineps -If you want to perform even more detailed changes or code injections, see `process_img_nii()` as inspiration on how the underlaying functions work and behave. Treat with care! +result = spineps.segment("/path/to/sub-test_T2w.nii.gz") # saves a derivatives folder next to the input +result = spineps.segment(nii, output_in_memory=True) # or get the masks back in memory +``` + +To segment many images without reloading the models, use `SpinepsPipeline`; to group processing options, pass the +`SemanticConfig` / `InstanceConfig` / `LabelingConfig` / `PostConfig` objects. + +For full control, load the models yourself with `get_semantic_model()` / `get_instance_model()` and call +`segment_image()` (single image) or `process_dataset()` (whole dataset) from `spineps.seg_run`. + +> **Upgrading from 1.x?** See [MIGRATION.md](MIGRATION.md) for the renamed CLI flags, functions and classes. ## Authorship diff --git a/docs/getting-started.md b/docs/getting-started.md index c65c133..7065763 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -77,18 +77,18 @@ Segment a single scan: ```bash # T2w sagittal -spineps sample -ignore_bids_filter -ignore_inference_compatibility \ - -i /path/sub-testsample_T2w.nii.gz -model_semantic t2w -model_instance instance +spineps sample --ignore-bids-filter --ignore-inference-compatibility \ + -i /path/sub-testsample_T2w.nii.gz --model-semantic t2w --model-instance instance # T1w sagittal -spineps sample -ignore_bids_filter -ignore_inference_compatibility \ - -i /path/sub-testsample_T1w.nii.gz -model_semantic t1w -model_instance instance +spineps sample --ignore-bids-filter --ignore-inference-compatibility \ + -i /path/sub-testsample_T1w.nii.gz --model-semantic t1w --model-instance instance ``` Process a whole [BIDS](https://bids-specification.readthedocs.io/en/stable/) dataset: ```bash -spineps dataset -i /path/to/dataset -model_semantic t2w -model_instance instance +spineps dataset -i /path/to/dataset --model-semantic t2w --model-instance instance ``` ### Adding vertebra labels (VERIDAH) @@ -97,23 +97,41 @@ To assign anatomical vertebra labels after segmentation, additionally pass a lab ```bash spineps sample -i /path/sub-test_T2w.nii.gz \ - -model_semantic t2w -model_instance instance -model_labeling labeling + --model-semantic t2w --model-instance instance --model-labeling labeling ``` ### From Python +The easiest way is the one-call API, which loads the models and runs the whole pipeline: + +```python +import spineps + +result = spineps.segment("/path/to/sub-test_T2w.nii.gz") # saves a derivatives folder next to the input +if result.success: + print("done:", result.output_paths) +``` + +To keep the models loaded across many images, or to get the masks in memory: + +```python +from spineps import SpinepsPipeline, InstanceConfig + +pipe = SpinepsPipeline(model_semantic="t2w", model_instance="instance") +result = pipe.segment(nii, output_in_memory=True, instance=InstanceConfig(batch_size=8)) +semantic, vertebra = result.semantic, result.vertebra +``` + +For full control you can still call the underlying pipeline function directly: + ```python from TPTBox import BIDS_FILE -from spineps import get_semantic_model, get_instance_model, process_img_nii +from spineps import get_semantic_model, get_instance_model, segment_image semantic = get_semantic_model("t2w") instance = get_instance_model("instance") - -# img_ref is a TPTBox BIDS_FILE pointing at the input scan -img_ref = BIDS_FILE("sub-test_T2w.nii.gz", dataset="/path/to/dataset") - -process_img_nii( - img_ref, +segment_image( + BIDS_FILE("sub-test_T2w.nii.gz", dataset="/path/to/dataset"), model_semantic=semantic, model_instance=instance, derivative_name="derivatives_seg", diff --git a/docs/index.md b/docs/index.md index 858c541..22d74b2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -37,7 +37,7 @@ Given a sagittal MR scan, the pipeline: pip install spineps # Segment a single T2w sagittal scan -spineps sample -i /path/sub-test_T2w.nii.gz -model_semantic t2w -model_instance instance +spineps sample -i /path/sub-test_T2w.nii.gz --model-semantic t2w --model-instance instance ``` See [Getting Started](getting-started.md) for the full installation guide (including PyTorch setup and diff --git a/docs/modules/models.md b/docs/modules/models.md index 7916769..1513695 100644 --- a/docs/modules/models.md +++ b/docs/modules/models.md @@ -21,11 +21,11 @@ the weights if needed, reads the model's `inference_config.json` and instantiate The concrete model classes are defined in [`spineps.seg_model`](../api/models.md) and [`spineps.lab_model`](../api/models.md): -- **`Segmentation_Model`** — abstract base wrapping a network plus its inference configuration. It handles +- **`SegmentationModel`** — abstract base wrapping a network plus its inference configuration. It handles input preparation (reorientation, rescaling to the recommended zoom, padding), running the model and mapping outputs back into the input space. -- **`Segmentation_Model_NNunet`** — an nnU-Net backend. -- **`Segmentation_Model_Unet3D`** — a 3D U-Net backend. +- **`SegmentationModelNNunet`** — an nnU-Net backend. +- **`SegmentationModelUnet3D`** — a 3D U-Net backend. - **`VertLabelingClassifier`** — the VERIDAH vertebra-labeling classifier. The model type is selected from the `inference_config.json` via `ModelType` (see @@ -43,4 +43,4 @@ VERIDAH ("solving Enumeration Anomaly Aware Vertebra Labeling across Imaging Seq anatomical labels to the detected vertebra instances. It combines a per-vertebra classifier with a min-cost path solver that enforces a plausible ordering and is aware of enumeration anomalies such as a 13th thoracic vertebra (T13) or a sixth lumbar vertebra (L6). Enable it by passing a labeling model to -the pipeline (`-model_labeling` on the CLI, or `model_labeling=` in Python). +the pipeline (`--model-labeling` on the CLI, or `model_labeling=` in Python). diff --git a/docs/modules/pipeline.md b/docs/modules/pipeline.md index 96477f2..5dd8f3f 100644 --- a/docs/modules/pipeline.md +++ b/docs/modules/pipeline.md @@ -3,7 +3,7 @@ SPINEPS processes a scan in a sequence of phases orchestrated by the functions in [`spineps.seg_run`](../api/pipeline.md). There are two top-level entry points: -- **`process_img_nii`** — process a single image. +- **`segment_image`** — process a single image. - **`process_dataset`** — discover and process every suitable scan in a [BIDS](https://bids-specification.readthedocs.io/en/stable/) dataset. ## Two-phase approach @@ -47,21 +47,43 @@ For each processed scan SPINEPS writes a derivatives folder next to the input co ## Calling from Python +The high-level `spineps.segment` API loads the models and runs the pipeline in one call: + +```python +import spineps + +result = spineps.segment("sub-test_T2w.nii.gz") # saves a derivatives folder next to the input +result = spineps.segment(nii, output_in_memory=True) # or get the masks back in memory +``` + +Use `SpinepsPipeline` to load the models once and segment many images, and the `SemanticConfig` / `InstanceConfig` +/ `LabelingConfig` / `PostConfig` objects to group processing options: + +```python +from spineps import SpinepsPipeline, InstanceConfig + +pipe = SpinepsPipeline(model_semantic="t2w", model_instance="instance") +for path in paths: + pipe.segment(path, instance=InstanceConfig(batch_size=8)) +``` + +For full control, call the underlying `segment_image` directly with already-loaded models: + ```python from TPTBox import BIDS_FILE -from spineps import get_semantic_model, get_instance_model, process_img_nii +from spineps import get_semantic_model, get_instance_model, segment_image semantic = get_semantic_model("t2w") instance = get_instance_model("instance") -process_img_nii( +segment_image( BIDS_FILE("sub-test_T2w.nii.gz", dataset="/path/to/dataset"), model_semantic=semantic, model_instance=instance, ) ``` -`process_img_nii` exposes many `proc_*` flags to toggle individual processing steps (pre-processing, +`segment_image` exposes many `proc_*` flags to toggle individual processing steps (pre-processing, semantic/instance cleaning, hole filling, labeling, …). See the [Pipeline & Run API reference](../api/pipeline.md) for the full signature. From 542694499d18e93cd33c38fe2dff4eefa9eaa43f Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 11 Jun 2026 15:23:41 +0000 Subject: [PATCH 10/29] feat: wire --amp, --step-size, --tta/--no-tta speed knobs end to end Expose the remaining inference speed controls through the CLI, the config objects and the pipeline functions: - --amp / InstanceConfig.amp: run the instance forward under CUDA autocast. Threaded segment_image -> predict_instance_mask -> collect_vertebra_predictions -> segment_scan_batch. - --step-size / SemanticConfig.step_size: semantic nnU-Net sliding-window tile step. Threaded segment_image -> predict_semantic_mask -> segment_scan. - --tta / --no-tta (tri-state, default = model config): toggle test-time mirroring via new SegmentationModel.set_test_time_augmentation(), applied to the semantic model in run_sample/run_dataset. Also fixes a latent bug: process_dataset never forwarded proc_inst_batch_size to segment_image, so 'spineps dataset --batch-size' would have raised TypeError. process_dataset now accepts and forwards proc_inst_batch_size, proc_inst_amp and proc_sem_step_size. Tests for the new flags, config mappings and set_test_time_augmentation. 194 pass. Co-Authored-By: Claude Opus 4.8 --- MIGRATION.md | 5 +++++ spineps/config.py | 4 ++++ spineps/entrypoint.py | 29 +++++++++++++++++++++++++++++ spineps/phase_instance.py | 8 ++++++++ spineps/phase_semantic.py | 4 ++++ spineps/seg_model.py | 12 ++++++++++++ spineps/seg_run.py | 20 ++++++++++++++++++++ unit_tests/test_api.py | 6 ++++-- unit_tests/test_entrypoint.py | 14 ++++++++++++++ unit_tests/test_inference_mocked.py | 14 ++++++++++++++ 10 files changed, 114 insertions(+), 2 deletions(-) diff --git a/MIGRATION.md b/MIGRATION.md index 7aa6060..0b07326 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -58,6 +58,9 @@ All long flags are now `--kebab-case` (short aliases are unchanged). Negative fl | `-cpu` | `--cpu` | | `-verbose` / `-v` | `--verbose` / `-v` | | *(new)* | `--batch-size` / `-bs` — vertebra cutouts per batched forward pass (faster; default 4) | +| *(new)* | `--amp` — run the instance forward pass under CUDA autocast (faster, may slightly change the output) | +| *(new)* | `--step-size` — semantic model sliding-window tile step size (larger = faster, less accurate) | +| *(new)* | `--tta` / `--no-tta` — force test-time augmentation (mirroring) on/off for the semantic model | Example: @@ -88,5 +91,7 @@ spineps sample -i scan.nii.gz --model-semantic t2w --model-instance instance --n flags. Pass them to `segment(...)`, e.g. `spineps.segment(path, instance=InstanceConfig(batch_size=8))`. - **`--batch-size`** / `InstanceConfig.batch_size` — the instance model now runs cutouts in batched forward passes (much faster on GPU); falls back to one-by-one on out-of-memory. +- **More speed knobs**: `--amp` / `InstanceConfig.amp` (instance autocast), `--step-size` / `SemanticConfig.step_size` + (semantic sliding-window step), and `--tta` / `--no-tta` (toggle test-time mirroring; `SegmentationModel.set_test_time_augmentation(...)` in Python). - Clearer errors: invalid paths / missing models now raise `FileNotFoundError` / `ValueError` instead of bare `AssertionError`, and `spineps sample -h` / `dataset -h` no longer crash. diff --git a/spineps/config.py b/spineps/config.py index af66f93..6347863 100644 --- a/spineps/config.py +++ b/spineps/config.py @@ -19,6 +19,7 @@ class SemanticConfig: remove_inferior_beyond_canal: bool = False clean_beyond_largest_bounding_box: bool = True clean_small_cc_artifacts: bool = True + step_size: float | None = None def to_kwargs(self) -> dict: """Returns the corresponding flat ``segment_image`` keyword arguments.""" @@ -28,6 +29,7 @@ def to_kwargs(self) -> dict: "proc_sem_remove_inferior_beyond_canal": self.remove_inferior_beyond_canal, "proc_sem_clean_beyond_largest_bounding_box": self.clean_beyond_largest_bounding_box, "proc_sem_clean_small_cc_artifacts": self.clean_small_cc_artifacts, + "proc_sem_step_size": self.step_size, } @@ -40,6 +42,7 @@ class InstanceConfig: largest_k_cc: int = 0 detect_and_solve_merged_corpi: bool = True batch_size: int = 4 + amp: bool = False labeling_offset: int = 2 def to_kwargs(self) -> dict: @@ -50,6 +53,7 @@ def to_kwargs(self) -> dict: "proc_inst_largest_k_cc": self.largest_k_cc, "proc_inst_detect_and_solve_merged_corpi": self.detect_and_solve_merged_corpi, "proc_inst_batch_size": self.batch_size, + "proc_inst_amp": self.amp, "vertebra_instance_labeling_offset": self.labeling_offset, } diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index 48799e7..8db17ec 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -99,6 +99,25 @@ def parser_arguments(parser: argparse.ArgumentParser): help="Number of vertebra cutouts run through the instance model per batched forward pass. Higher is faster but uses " "more GPU memory; falls back to one-by-one on out-of-memory.", ) + parser.add_argument( + "--amp", + action="store_true", + help="Run the instance model's forward pass under CUDA autocast (faster, may slightly change the output).", + ) + parser.add_argument( + "--step-size", + type=float, + default=None, + help="Sliding-window tile step size for the semantic model (e.g. 0.7). Larger is faster but less accurate; " + "by default uses the model's configured value.", + ) + parser.add_argument( + "--tta", + action=argparse.BooleanOptionalAction, + default=None, + help="Force test-time augmentation (mirroring) on/off for the semantic model. Pass --no-tta to disable it for " + "a speed-up; by default uses the model's configured setting.", + ) parser.add_argument("--cpu", action="store_true", help="Use CPU instead of GPU (will take way longer)") parser.add_argument("--run-cprofiler", "-rcp", action="store_true", help="Runs a cprofiler over the entire action") parser.add_argument("--verbose", "-v", action="store_true", help="Prints much more stuff, may fully clutter your terminal") @@ -258,6 +277,9 @@ def run_sample(opt: Namespace): else: model_labeling = get_labeling_model(opt.model_labeling, use_cpu=opt.cpu).load() + if opt.tta is not None: + model_semantic.set_test_time_augmentation(opt.tta) + bids_sample = BIDS_FILE(input_path, dataset=dataset, verbose=True) kwargs = { @@ -279,6 +301,8 @@ def run_sample(opt: Namespace): "proc_sem_n4_bias_correction": opt.n4, "proc_lab_force_no_tl_anomaly": opt.enforce_12_thoracic, "proc_inst_batch_size": opt.batch_size, + "proc_inst_amp": opt.amp, + "proc_sem_step_size": opt.step_size, "ignore_compatibility_issues": opt.ignore_inference_compatibility, "verbose": opt.verbose, } @@ -355,6 +379,9 @@ def run_dataset(opt: Namespace): if model_instance is None: raise ValueError("-model_instance/-mv resolved to None; pass a valid instance model id or path") + if opt.tta is not None and model_semantic is not None: + model_semantic.set_test_time_augmentation(opt.tta) + kwargs = { "dataset_path": input_dir, "model_semantic": model_semantic, @@ -379,6 +406,8 @@ def run_dataset(opt: Namespace): "proc_sem_n4_bias_correction": opt.n4, "proc_lab_force_no_tl_anomaly": opt.enforce_12_thoracic, "proc_inst_batch_size": opt.batch_size, + "proc_inst_amp": opt.amp, + "proc_sem_step_size": opt.step_size, "snapshot_copy_folder": opt.save_snaps_folder, "verbose": opt.verbose, } diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index ba5fa71..f6db4f7 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -58,6 +58,7 @@ def predict_instance_mask( proc_inst_clean_small_cc_artifacts: bool = True, proc_inst_largest_k_cc: int = 0, proc_inst_batch_size: int = 4, + proc_inst_amp: bool = False, verbose: bool = False, ) -> tuple[NII | None, ErrCode]: """Build a per-vertebra instance mask from a subregion semantic segmentation. @@ -78,6 +79,8 @@ def predict_instance_mask( proc_inst_largest_k_cc (int, optional): Keep only the largest k connected components per cutout label; 0 disables. Defaults to 0. proc_inst_batch_size (int, optional): Number of cutouts run through the instance model per batched forward pass. Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. + proc_inst_amp (bool, optional): Run the instance forward pass under CUDA autocast (faster, may slightly change + values). Defaults to False. verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: @@ -150,6 +153,7 @@ def predict_instance_mask( proc_inst_largest_k_cc=proc_inst_largest_k_cc, proc_inst_fill_holes=False, instance_batch_size=proc_inst_batch_size, + amp=proc_inst_amp, verbose=verbose, ) if vert_predictions is None: @@ -592,6 +596,7 @@ def collect_vertebra_predictions( process_detect_and_solve_merged_corpi: bool = True, proc_inst_fill_holes: bool = False, instance_batch_size: int = 4, + amp: bool = False, verbose: bool = False, ) -> tuple[np.ndarray | None, list[str], int]: """Run the instance model on a cutout around each corpus center of mass and collect per-label predictions. @@ -612,6 +617,8 @@ def collect_vertebra_predictions( proc_inst_fill_holes (bool, optional): Whether to fill holes in each cutout prediction. Defaults to False. instance_batch_size (int, optional): Number of cutouts run through the instance model per batched forward pass. Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. + amp (bool, optional): Run the instance forward pass under CUDA autocast (faster, may slightly change values). + Defaults to False. verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: @@ -689,6 +696,7 @@ def collect_vertebra_predictions( pad_size=0, resample_output_to_input_space=False, batch_size=instance_batch_size, + amp=amp, verbose=False, ) diff --git a/spineps/phase_semantic.py b/spineps/phase_semantic.py index 8768e74..5aa350b 100755 --- a/spineps/phase_semantic.py +++ b/spineps/phase_semantic.py @@ -32,6 +32,7 @@ def predict_semantic_mask( proc_clean_beyond_largest_bounding_box: bool = True, proc_remove_inferior_beyond_canal: bool = False, proc_clean_small_cc_artifacts: bool = True, + step_size: float | None = None, verbose: bool = False, ) -> tuple[NII | None, NII | None, ErrCode]: """Predict the semantic (subregion) segmentation mask and run post-processing on it. @@ -50,6 +51,8 @@ def predict_semantic_mask( proc_remove_inferior_beyond_canal (bool, optional): Whether to remove non-sacrum structures below the spinal-canal height. Defaults to False. proc_clean_small_cc_artifacts (bool, optional): Whether to delete small connected-component artifacts. Defaults to True. + step_size (float | None, optional): Sliding-window tile step size for the model; larger is faster but less + accurate. If None, uses the model's configured default. Defaults to None. verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: @@ -61,6 +64,7 @@ def predict_semantic_mask( results = model.segment_scan( mri_nii, pad_size=0, + step_size=step_size, resample_to_recommended=True, resample_output_to_input_space=False, verbose=verbose, diff --git a/spineps/seg_model.py b/spineps/seg_model.py index e944c00..5533795 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -305,6 +305,18 @@ def acquisition(self) -> Acquisition: """ return self.inference_config.acquisition + def set_test_time_augmentation(self, enabled: bool) -> None: + """Enables or disables test-time augmentation (mirroring) for this model, if the backend supports it. + + Only the nnU-Net (semantic) backend uses mirroring; for other backends this is a no-op. Mirroring roughly + multiplies inference cost, so disabling it is a simple speed-up at a small accuracy cost. + + Args: + enabled (bool): Whether to use test-time mirroring augmentation. + """ + if self.predictor is not None and hasattr(self.predictor, "use_mirroring"): + self.predictor.use_mirroring = enabled + @abstractmethod def run(self, input_nii: list[NII], verbose: bool = False) -> dict[OutputType, NII | None]: """Runs the backend predictor on the prepared inputs. diff --git a/spineps/seg_run.py b/spineps/seg_run.py index 56ea705..bb64846 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -53,11 +53,14 @@ def process_dataset( proc_sem_remove_inferior_beyond_canal: bool = False, proc_sem_clean_beyond_largest_bounding_box: bool = True, proc_sem_clean_small_cc_artifacts: bool = True, + proc_sem_step_size: float | None = None, # Instance proc_inst_corpus_clean: bool = True, proc_inst_clean_small_cc_artifacts: bool = True, proc_inst_largest_k_cc: int = 0, proc_inst_detect_and_solve_merged_corpi: bool = True, + proc_inst_batch_size: int = 4, + proc_inst_amp: bool = False, # Labeling proc_lab_force_no_tl_anomaly: bool = False, # Both @@ -109,12 +112,18 @@ def process_dataset( bounding box. Defaults to True. proc_sem_clean_small_cc_artifacts (bool, optional): If true, removes small connected-component artifacts from the semantic mask. Defaults to True. + proc_sem_step_size (float | None, optional): Sliding-window tile step size for the semantic model; larger is + faster but less accurate. If None, uses the model's configured default. Defaults to None. proc_inst_corpus_clean (bool, optional): If true, cleans the vertebra corpus during instance processing. Defaults to True. proc_inst_clean_small_cc_artifacts (bool, optional): If true, removes small connected-component artifacts from the instance mask. Defaults to True. proc_inst_largest_k_cc (int, optional): If greater than 0, keeps only the largest k connected components of the instance mask. Defaults to 0. proc_inst_detect_and_solve_merged_corpi (bool, optional): If true, detects and splits merged vertebra corpi. Defaults to True. + proc_inst_batch_size (int, optional): Number of vertebra cutouts run through the instance model per batched forward + pass. Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. + proc_inst_amp (bool, optional): Run the instance forward pass under CUDA autocast (faster, may slightly change + values). Defaults to False. proc_lab_force_no_tl_anomaly (bool, optional): If true, forces the labeling to assume no thoracolumbar transition anomaly. Defaults to False. proc_fill_3d_holes (bool, optional): If true, fills 3D holes during post-processing. Defaults to True. @@ -222,11 +231,14 @@ def process_dataset( proc_sem_remove_inferior_beyond_canal=proc_sem_remove_inferior_beyond_canal, proc_sem_clean_beyond_largest_bounding_box=proc_sem_clean_beyond_largest_bounding_box, proc_sem_clean_small_cc_artifacts=proc_sem_clean_small_cc_artifacts, + proc_sem_step_size=proc_sem_step_size, proc_inst_detect_and_solve_merged_corpi=proc_inst_detect_and_solve_merged_corpi, proc_inst_corpus_clean=proc_inst_corpus_clean, proc_inst_clean_small_cc_artifacts=proc_inst_clean_small_cc_artifacts, proc_assign_missing_cc=proc_assign_missing_cc, proc_inst_largest_k_cc=proc_inst_largest_k_cc, + proc_inst_batch_size=proc_inst_batch_size, + proc_inst_amp=proc_inst_amp, proc_clean_inst_by_sem=proc_clean_inst_by_sem, proc_lab_force_no_tl_anomaly=proc_lab_force_no_tl_anomaly, proc_vertebra_inconsistency=proc_vertebra_inconsistency, @@ -297,12 +309,14 @@ def segment_image( # noqa: C901 proc_sem_remove_inferior_beyond_canal: bool = False, proc_sem_clean_beyond_largest_bounding_box: bool = True, proc_sem_clean_small_cc_artifacts: bool = True, + proc_sem_step_size: float | None = None, # Instance proc_inst_corpus_clean: bool = True, proc_inst_clean_small_cc_artifacts: bool = True, proc_inst_largest_k_cc: int = 0, proc_inst_detect_and_solve_merged_corpi: bool = True, proc_inst_batch_size: int = 4, + proc_inst_amp: bool = False, vertebra_instance_labeling_offset=2, # Labeling proc_lab_force_no_tl_anomaly: bool = False, @@ -361,6 +375,8 @@ def segment_image( # noqa: C901 bounding box. Defaults to True. proc_sem_clean_small_cc_artifacts (bool, optional): If true, removes small connected-component artifacts from the semantic mask. Defaults to True. + proc_sem_step_size (float | None, optional): Sliding-window tile step size for the semantic model; larger is + faster but less accurate. If None, uses the model's configured default. Defaults to None. proc_inst_corpus_clean (bool, optional): If true, cleans the vertebra corpus during instance processing. Defaults to True. proc_inst_clean_small_cc_artifacts (bool, optional): If true, removes small connected-component artifacts from the instance mask. Defaults to True. @@ -369,6 +385,8 @@ def segment_image( # noqa: C901 proc_inst_detect_and_solve_merged_corpi (bool, optional): If true, detects and splits merged vertebra corpi. Defaults to True. proc_inst_batch_size (int, optional): Number of vertebra cutouts run through the instance model per batched forward pass. Higher is faster but uses more GPU memory; falls back to one-by-one on out-of-memory. Defaults to 4. + proc_inst_amp (bool, optional): Run the instance forward pass under CUDA autocast (faster, may slightly change + values). Defaults to False. vertebra_instance_labeling_offset (int, optional): Offset applied when mapping instance ids to vertebra labels (set to 1 for CT models that include C1). Defaults to 2. proc_lab_force_no_tl_anomaly (bool, optional): If true, forces the labeling to assume no thoracolumbar transition anomaly. @@ -518,6 +536,7 @@ def segment_image( # noqa: C901 proc_clean_small_cc_artifacts=proc_sem_clean_small_cc_artifacts, proc_clean_beyond_largest_bounding_box=proc_sem_clean_beyond_largest_bounding_box, proc_remove_inferior_beyond_canal=proc_sem_remove_inferior_beyond_canal, + step_size=proc_sem_step_size, ) if errcode != ErrCode.OK: return output_paths, errcode @@ -558,6 +577,7 @@ def segment_image( # noqa: C901 proc_inst_clean_small_cc_artifacts=proc_inst_clean_small_cc_artifacts, proc_inst_largest_k_cc=proc_inst_largest_k_cc, proc_inst_batch_size=proc_inst_batch_size, + proc_inst_amp=proc_inst_amp, ) if errcode != ErrCode.OK: logger.print(f"Vert Mask creation failed with errcode {errcode}", Log_Type.FAIL) diff --git a/unit_tests/test_api.py b/unit_tests/test_api.py index 84e07d7..8a00e46 100644 --- a/unit_tests/test_api.py +++ b/unit_tests/test_api.py @@ -18,14 +18,16 @@ class Test_Config_To_Kwargs(unittest.TestCase): def test_semantic(self): - kw = SemanticConfig(crop_input=False, n4_bias_correction=False).to_kwargs() + kw = SemanticConfig(crop_input=False, n4_bias_correction=False, step_size=0.7).to_kwargs() self.assertFalse(kw["proc_sem_crop_input"]) self.assertFalse(kw["proc_sem_n4_bias_correction"]) + self.assertEqual(kw["proc_sem_step_size"], 0.7) def test_instance(self): - kw = InstanceConfig(batch_size=8, largest_k_cc=3).to_kwargs() + kw = InstanceConfig(batch_size=8, largest_k_cc=3, amp=True).to_kwargs() self.assertEqual(kw["proc_inst_batch_size"], 8) self.assertEqual(kw["proc_inst_largest_k_cc"], 3) + self.assertTrue(kw["proc_inst_amp"]) self.assertEqual(kw["vertebra_instance_labeling_offset"], 2) def test_labeling_and_post(self): diff --git a/unit_tests/test_entrypoint.py b/unit_tests/test_entrypoint.py index 47fcaf7..a866ab1 100644 --- a/unit_tests/test_entrypoint.py +++ b/unit_tests/test_entrypoint.py @@ -37,6 +37,20 @@ def test_shared_flag_defaults_and_negation(self): self.assertTrue(negated.enforce_12_thoracic) self.assertEqual(negated.batch_size, 8) + def test_speed_knob_flags(self): + # --amp / --step-size / --tta tri-state + p = argparse.ArgumentParser() + entrypoint.parser_arguments(p) + defaults = p.parse_args([]) + self.assertFalse(defaults.amp) + self.assertIsNone(defaults.step_size) + self.assertIsNone(defaults.tta) # None = use the model's configured setting + opts = p.parse_args(["--amp", "--step-size", "0.7", "--no-tta"]) + self.assertTrue(opts.amp) + self.assertEqual(opts.step_size, 0.7) + self.assertFalse(opts.tta) + self.assertTrue(p.parse_args(["--tta"]).tta) + def test_subparser_help_builds(self): # Regression: empty metavar="" used to crash `spineps -h` in argparse usage formatting. for sub in ("sample", "dataset"): diff --git a/unit_tests/test_inference_mocked.py b/unit_tests/test_inference_mocked.py index 5c3ae4a..f306566 100644 --- a/unit_tests/test_inference_mocked.py +++ b/unit_tests/test_inference_mocked.py @@ -442,6 +442,20 @@ def test_segment_scan_batch_matches_segment_scan(self): single = model.segment_scan(cut, **kwargs) np.testing.assert_array_equal(res[OutputType.seg].get_seg_array(), single[OutputType.seg].get_seg_array()) + def test_set_test_time_augmentation(self): + model = _make_unet3d_test_model() + # the instance predictor has no use_mirroring attribute -> no-op, must not raise + model.set_test_time_augmentation(False) + + class _MirroringPredictor: + use_mirroring = True + + model.predictor = _MirroringPredictor() + model.set_test_time_augmentation(False) + self.assertFalse(model.predictor.use_mirroring) + model.set_test_time_augmentation(True) + self.assertTrue(model.predictor.use_mirroring) + if __name__ == "__main__": unittest.main() From 330be2f1261b46cfef23f940af848c1707924845 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 11 Jun 2026 17:18:44 +0000 Subject: [PATCH 11/29] feat: apply --tta/--no-tta to auto-resolved dataset models Move the test-time-augmentation override into process_dataset (new 'tta' param) so it reaches both explicitly-passed and '--model-semantic auto' resolved semantic models, eager-loading them so the toggle lands on the predictor. run_dataset now passes tta through instead of only setting it on the pre-loaded model. Test: process_dataset(tta=False) applies set_test_time_augmentation on the semantic model. 195 pass. Co-Authored-By: Claude Opus 4.8 --- spineps/entrypoint.py | 4 +--- spineps/seg_run.py | 15 ++++++++++++++- unit_tests/test_inference_mocked.py | 28 ++++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index 8db17ec..da539ee 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -379,9 +379,6 @@ def run_dataset(opt: Namespace): if model_instance is None: raise ValueError("-model_instance/-mv resolved to None; pass a valid instance model id or path") - if opt.tta is not None and model_semantic is not None: - model_semantic.set_test_time_augmentation(opt.tta) - kwargs = { "dataset_path": input_dir, "model_semantic": model_semantic, @@ -409,6 +406,7 @@ def run_dataset(opt: Namespace): "proc_inst_amp": opt.amp, "proc_sem_step_size": opt.step_size, "snapshot_copy_folder": opt.save_snaps_folder, + "tta": opt.tta, "verbose": opt.verbose, } diff --git a/spineps/seg_run.py b/spineps/seg_run.py index bb64846..5515fe6 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -26,7 +26,7 @@ @citation_reminder -def process_dataset( +def process_dataset( # noqa: C901 dataset_path: Path, model_instance: SegmentationModel, model_semantic: list[SegmentationModel] | SegmentationModel | None = None, @@ -72,6 +72,7 @@ def process_dataset( ignore_model_compatibility: bool = False, ignore_inference_compatibility: bool = False, ignore_bids_filter: bool = False, + tta: bool | None = None, log_inference_time: bool = True, verbose: bool = False, ): @@ -134,6 +135,9 @@ def process_dataset( ignore_inference_compatibility (bool, optional): If true, ignores compatibility issues between models and individual inputs. Defaults to False. ignore_bids_filter (bool, optional): If true, disables the BIDS query filters and processes all niftys found. Defaults to False. + tta (bool | None, optional): If not None, forces test-time augmentation (mirroring) on/off for the semantic + model(s), covering both explicitly-passed and auto-resolved models. If None, uses each model's configured + setting. Defaults to None. log_inference_time (bool, optional): If true, logs the inference time of each step. Defaults to True. verbose (bool, optional): If true, prints verbose information. Defaults to False. """ @@ -159,6 +163,15 @@ def process_dataset( if not isinstance(model_semantic, list): model_semantic = [model_semantic] + # Optionally force test-time augmentation (mirroring) on/off for the semantic model(s); covers both + # explicitly-passed and auto-resolved models. Load eagerly so the toggle reaches the predictor. + if tta is not None: + for m in model_semantic: + if m is not None: + if m.predictor is None: + m.load() + m.set_test_time_augmentation(tta) + # check models and mod, acq tuples compatible = True for idx, mp in enumerate(modalities): diff --git a/unit_tests/test_inference_mocked.py b/unit_tests/test_inference_mocked.py index f306566..45b0b37 100644 --- a/unit_tests/test_inference_mocked.py +++ b/unit_tests/test_inference_mocked.py @@ -457,5 +457,33 @@ class _MirroringPredictor: self.assertTrue(model.predictor.use_mirroring) +class Test_Process_Dataset_TTA(unittest.TestCase): + """process_dataset must apply the tta override to its semantic models (explicit or auto-resolved).""" + + def test_tta_override_applied_to_semantic_model(self): + import tempfile + + from spineps.seg_enums import Acquisition, Modality + from spineps.seg_run import process_dataset + + model_sem = _make_unet3d_test_model() + model_inst = _make_unet3d_test_model() + calls: list[bool] = [] + model_sem.set_test_time_augmentation = lambda enabled: calls.append(enabled) # type: ignore[method-assign] + + with tempfile.TemporaryDirectory() as td: + # empty dataset -> the subject loop is a no-op, but the tta block runs before it + process_dataset( + dataset_path=Path(td), + model_instance=model_inst, + model_semantic=model_sem, + modalities=[(Modality.SEG, Acquisition.sag)], + tta=False, + save_log_data=False, + ignore_model_compatibility=True, + ) + self.assertEqual(calls, [False]) + + if __name__ == "__main__": unittest.main() From 210f0afcab39670967285bb74bf951b655e9cab6 Mon Sep 17 00:00:00 2001 From: ga84mun Date: Wed, 17 Jun 2026 12:04:48 +0000 Subject: [PATCH 12/29] add endplate support --- spineps/get_models.py | 2 +- spineps/phase_post.py | 144 +++++++++++++++++++++++++++--------------- 2 files changed, 95 insertions(+), 51 deletions(-) diff --git a/spineps/get_models.py b/spineps/get_models.py index 8940982..175792a 100755 --- a/spineps/get_models.py +++ b/spineps/get_models.py @@ -149,7 +149,7 @@ def modelid2folder_labeling() -> dict[str, Path | str]: def check_available_models( - models_folder: str | Path, verbose: bool = False + models_folder: str | Path, verbose: bool = True ) -> tuple[dict[str, Path | str], dict[str, Path | str], dict[str, Path | str]]: """Searches the given directory for models and sorts them into semantic, instance and labeling id-to-folder maps. diff --git a/spineps/phase_post.py b/spineps/phase_post.py index 71f48af..a634cba 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -448,9 +448,17 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, verbose=True) -> tu vert_arr[subreg_arr == Location.Vertebra_Disc.value] = subreg_ivd[subreg_arr == Location.Vertebra_Disc.value] n_eps = 0 n_eps_unique = 0 - if Location.Endplate.value in seg_t.unique(): + ep_labels = [ + Location.Endplate.value, + Location.Vertebral_Body_Endplate_Inferior.value, + Location.Vertebral_Body_Endplate_Superior.value, + ] + u = seg_t.unique() + has_split_endplates = Location.Vertebral_Body_Endplate_Inferior.value in u or Location.Vertebral_Body_Endplate_Superior.value in u + # FIXME Problem: For some reason Endplate are mapped to the IVD in MRI aka the superior endplate hat the IVD of vertebra above instead of below. + if Location.Endplate.value in u or has_split_endplates: # MAP Endplate - ep_cc = seg_t.get_connected_components(labels=Location.Endplate.value) + ep_cc = seg_t.get_connected_components(labels=ep_labels) ep_cc_n = len(ep_cc.unique()) ep_cc = ep_cc.get_seg_array() cc_ep_labelset = list(range(1, ep_cc_n + 1)) @@ -459,8 +467,14 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, verbose=True) -> tu for c in cc_ep_labelset: if c == 0: continue - com_y = np_center_of_mass(ep_cc == c)[1][1] # center_of_mass(c_l)[1] - nearest_lower = find_nearest_lower(coms_vert_y, com_y) + com_y = np_center_of_mass(ep_cc == c)[1][1] + nearest_lower = ( + find_nearest_lower(coms_vert_y, com_y) + if not has_split_endplates + or seg_t[ep_cc == c].max() + != Location.Vertebral_Body_Endplate_Superior.value # True for CT, so that is mapped to the vertebra instead of IVD + else find_nearest_higher(coms_vert_y, com_y) + ) label = next(i for i in coms_vert_dict if coms_vert_dict[i] == nearest_lower) mapping_ep_cc_to_vert_label[c] = label n_eps += 1 @@ -468,55 +482,69 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, verbose=True) -> tu subreg_ep = ep_cc.copy() n_eps_unique = len(np.unique(list(mapping_ep_cc_to_vert_label.values()))) subreg_ep = np_map_labels(subreg_ep, label_map=mapping_ep_cc_to_vert_label) - subreg_ep += ENDPLATE_LABEL_OFFSET + subreg_ep += ENDPLATE_LABEL_OFFSET if not has_split_endplates else 0 # True for CT, so that is mapped to the vertebra itself subreg_ep[subreg_ep == ENDPLATE_LABEL_OFFSET] = 0 - vert_arr[subreg_arr == Location.Endplate.value] = subreg_ep[subreg_arr == Location.Endplate.value] - vert_t.set_array_(vert_arr) - - # divide into upper and lower endplate - out = seg_t * 0 - pref = 1 - old_vol = -1 - # seg_t and vert_t are not modified in this loop, so compute these invariants once. - endplate_nii = seg_t.extract_label(Location.Endplate.value) - total = endplate_nii.sum() - vert_labels_to_split = vert_t.unique() - for dil in range(1, MAX_ENDPLATE_DILATION): - curr = out.extract_label([Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value]) - new_vol = curr.sum() - logger.print(rf"{new_vol / total * 100:.2f}% endplates detected", end="\r") if verbose else None - if old_vol == new_vol and old_vol != 0: - break - old_vol = new_vol - if total == new_vol: - logger.print("Found all Endplates ") - break - for i in vert_labels_to_split: - if i >= INSTANCE_LABEL_LIMIT: - break + if not has_split_endplates: + # This code sets the IDs to the respective IVD instead of vertebra disc! has_split_endplates is True for CT + vert_arr[subreg_arr == Location.Endplate.value] = subreg_ep[subreg_arr == Location.Endplate.value] + vert_t.set_array_(vert_arr) + # divide into upper and lower endplate + out = seg_t * 0 + pref = 1 + old_vol = -1 + # seg_t and vert_t are not modified in this loop, so compute these invariants once. + endplate_nii = seg_t.extract_label(ep_labels) + total = endplate_nii.sum() + vert_labels_to_split = vert_t.unique() + for dil in range(1, MAX_ENDPLATE_DILATION): curr = out.extract_label([Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value]) - v = vert_t.extract_label(i).dilate_msk(dil, verbose=False) - end = endplate_nii * v - end *= -curr + 1 # type: ignore - plates = vert_t * end - plates.map_labels_( - { - i + ENDPLATE_LABEL_OFFSET: Location.Vertebral_Body_Endplate_Inferior.value, - pref + ENDPLATE_LABEL_OFFSET: Location.Vertebral_Body_Endplate_Superior.value, - }, - verbose=False, - ) - out += plates - pref = i - curr = out.extract_label([Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value]) - end = seg_t.extract_label(Location.Endplate.value) - end *= -curr + 1 - # end += end.dilate_msk(3) - out += end * Location.Endplate.value - seg_t = out.extract_label( - [Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value, Location.Endplate.value] - ) + new_vol = curr.sum() + logger.print(rf"{new_vol / total * 100:.2f}% endplates detected", end="\r") if verbose else None + if old_vol == new_vol and old_vol != 0: + break + old_vol = new_vol + if total == new_vol: + logger.print("Found all Endplates ") + break + for i in vert_labels_to_split: + if i >= INSTANCE_LABEL_LIMIT: + break + curr = out.extract_label( + [Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value] + ) + v = vert_t.extract_label(i).dilate_msk(dil, verbose=False) + end = endplate_nii * v + end *= -curr + 1 # type: ignore + plates = vert_t * end + plates.map_labels_( + { + i + ENDPLATE_LABEL_OFFSET: Location.Vertebral_Body_Endplate_Inferior.value, + pref + ENDPLATE_LABEL_OFFSET: Location.Vertebral_Body_Endplate_Superior.value, + }, + verbose=False, + ) + out += plates + pref = i + curr = out.extract_label([Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value]) + + end = seg_t.extract_label(ep_labels) + end *= -curr + 1 + # end += end.dilate_msk(3) + out += end * Location.Endplate.value + seg_t = out.extract_label( + [Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value, Location.Endplate.value] + ) + else: + # Endplates are already split semantically. + # Assign endplate instance IDs while preserving the semantic labels. + ep_mask = np.isin( + subreg_arr, + [Location.Endplate.value, Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value], + ) + + vert_arr[ep_mask] = subreg_ep[ep_mask] + vert_t.set_array_(vert_arr) logger.print(f"Labeled {n_ivds} IVDs ({n_ivd_unique} unique), and {n_eps} Endplates ({n_eps_unique} unique)") return vert_t.set_array_(vert_arr).reorient_(orientation).get_seg_array(), seg_t.reorient_(orientation).get_seg_array() @@ -537,6 +565,22 @@ def find_nearest_lower(seq, x) -> float: return max(values_lower) +def find_nearest_higher(seq, x) -> float: + """Return the largest element of ``seq`` strictly smaller than ``x``, or the minimum if none exists. + + Args: + seq (Sequence[float]): Values to search. + x (float): Reference value. + + Returns: + float: The greatest element below ``x``, or ``min(seq)`` when no element is below ``x``. + """ + values_higher = [item for item in seq if item > x] + if len(values_higher) == 0: + return max(seq) + return min(values_higher) + + def label_instance_top_to_bottom(vert_nii: NII, labeling_offset: int = 0) -> tuple[NII, np.ndarray]: """Relabel vertebra instances consecutively from top to bottom by center-of-mass height. From c913a67e6912351e60ee9a2598ba63390b9b4d44 Mon Sep 17 00:00:00 2001 From: iback Date: Mon, 29 Jun 2026 13:17:10 +0000 Subject: [PATCH 13/29] build: require TPTBox ^0.7.3 (numpy 2); drop unused INP001 noqa; simplify test - pyproject: TPTBox "*" -> "^0.7.3" (the new TPTBox inference backend). 0.7.3 requires numpy>=2; the local env migration also needed connected-components-3d>=4 and statsmodels>=0.14.6 rebuilt for numpy 2 (cc3d>=4 technically exceeds TPTBox 0.7.3's own <4 pin but works). - Remove now-unused '# noqa: INP001' from unit_tests/*.py (unit_tests is a real package since __init__.py was added); kept in spineps/example/* which has no __init__.py. - test_inference_mocked: replace 'lambda enabled: calls.append(enabled)' with 'calls.append'. 195 tests pass under numpy 2.2.6 + TPTBox 0.7.3, ruff clean. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 2 +- unit_tests/test_api.py | 2 +- unit_tests/test_architectures.py | 2 +- unit_tests/test_compat.py | 2 +- unit_tests/test_disc_labels.py | 2 +- unit_tests/test_entrypoint.py | 2 +- unit_tests/test_enums.py | 2 +- unit_tests/test_filepaths.py | 2 +- unit_tests/test_find_min_cost_path.py | 2 +- unit_tests/test_generate_disc_labels_extra.py | 2 +- unit_tests/test_inference_mocked.py | 4 ++-- unit_tests/test_modelconfig.py | 2 +- unit_tests/test_models.py | 2 +- unit_tests/test_path.py | 2 +- unit_tests/test_phase_labeling.py | 2 +- unit_tests/test_postproc.py | 2 +- unit_tests/test_proc_functions.py | 2 +- unit_tests/test_pure_helpers.py | 2 +- unit_tests/test_resolution.py | 2 +- 19 files changed, 20 insertions(+), 20 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 112465f..6a278ae 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ torchmetrics = "^1.1.2" tqdm = "^4.66.1" einops= "^0.6.1" nnunetv2 = "^2.4.2" -TPTBox = "*" +TPTBox = "^0.7.3" antspyx = "0.4.2" rich = "^13.6.0" monai="^1.3.0" diff --git a/unit_tests/test_api.py b/unit_tests/test_api.py index 8a00e46..7765234 100644 --- a/unit_tests/test_api.py +++ b/unit_tests/test_api.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder """Tests for the high-level spineps.segment API, config dataclasses and result wrapping.""" from __future__ import annotations diff --git a/unit_tests/test_architectures.py b/unit_tests/test_architectures.py index 4245c6d..624e94d 100644 --- a/unit_tests/test_architectures.py +++ b/unit_tests/test_architectures.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_compat.py b/unit_tests/test_compat.py index 2223c21..53d3862 100644 --- a/unit_tests/test_compat.py +++ b/unit_tests/test_compat.py @@ -1,4 +1,4 @@ -# Call "python -m unittest" on this folder # noqa: INP001 +# Call "python -m unittest" on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_disc_labels.py b/unit_tests/test_disc_labels.py index 9a8ac9d..9616dda 100644 --- a/unit_tests/test_disc_labels.py +++ b/unit_tests/test_disc_labels.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_entrypoint.py b/unit_tests/test_entrypoint.py index a866ab1..d0a4746 100644 --- a/unit_tests/test_entrypoint.py +++ b/unit_tests/test_entrypoint.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_enums.py b/unit_tests/test_enums.py index 026913e..3f82233 100644 --- a/unit_tests/test_enums.py +++ b/unit_tests/test_enums.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_filepaths.py b/unit_tests/test_filepaths.py index 1dc128b..81583df 100644 --- a/unit_tests/test_filepaths.py +++ b/unit_tests/test_filepaths.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_find_min_cost_path.py b/unit_tests/test_find_min_cost_path.py index f05e4c6..ef95955 100644 --- a/unit_tests/test_find_min_cost_path.py +++ b/unit_tests/test_find_min_cost_path.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_generate_disc_labels_extra.py b/unit_tests/test_generate_disc_labels_extra.py index 8e12df1..acf066c 100644 --- a/unit_tests/test_generate_disc_labels_extra.py +++ b/unit_tests/test_generate_disc_labels_extra.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_inference_mocked.py b/unit_tests/test_inference_mocked.py index 45b0b37..aecab75 100644 --- a/unit_tests/test_inference_mocked.py +++ b/unit_tests/test_inference_mocked.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html @@ -469,7 +469,7 @@ def test_tta_override_applied_to_semantic_model(self): model_sem = _make_unet3d_test_model() model_inst = _make_unet3d_test_model() calls: list[bool] = [] - model_sem.set_test_time_augmentation = lambda enabled: calls.append(enabled) # type: ignore[method-assign] + model_sem.set_test_time_augmentation = calls.append # type: ignore[method-assign] with tempfile.TemporaryDirectory() as td: # empty dataset -> the subject loop is a no-op, but the tta block runs before it diff --git a/unit_tests/test_modelconfig.py b/unit_tests/test_modelconfig.py index 6665968..d915265 100644 --- a/unit_tests/test_modelconfig.py +++ b/unit_tests/test_modelconfig.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_models.py b/unit_tests/test_models.py index 3a6bf7d..3f8148f 100644 --- a/unit_tests/test_models.py +++ b/unit_tests/test_models.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_path.py b/unit_tests/test_path.py index 78c5a0e..e881824 100644 --- a/unit_tests/test_path.py +++ b/unit_tests/test_path.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_phase_labeling.py b/unit_tests/test_phase_labeling.py index c1fd0ee..5ad23c7 100644 --- a/unit_tests/test_phase_labeling.py +++ b/unit_tests/test_phase_labeling.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_postproc.py b/unit_tests/test_postproc.py index 3e8d5cf..4cca07a 100644 --- a/unit_tests/test_postproc.py +++ b/unit_tests/test_postproc.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_proc_functions.py b/unit_tests/test_proc_functions.py index fac80dc..e8741d4 100644 --- a/unit_tests/test_proc_functions.py +++ b/unit_tests/test_proc_functions.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_pure_helpers.py b/unit_tests/test_pure_helpers.py index c7edd61..4ad92e0 100644 --- a/unit_tests/test_pure_helpers.py +++ b/unit_tests/test_pure_helpers.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html diff --git a/unit_tests/test_resolution.py b/unit_tests/test_resolution.py index c7bf7ed..290df4a 100644 --- a/unit_tests/test_resolution.py +++ b/unit_tests/test_resolution.py @@ -1,4 +1,4 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 +# Call 'python -m unittest' on this folder # coverage run -m unittest # coverage report # coverage html From ae93caa252d8a75ef809a98cd0b4e1a530e21fe3 Mon Sep 17 00:00:00 2001 From: iback Date: Mon, 29 Jun 2026 13:26:38 +0000 Subject: [PATCH 14/29] docs: align docs/README with the 2.0 code; document the high-level API - Add an API reference page for the new high-level modules (spineps.api: segment/SpinepsPipeline/SpinepsResult, and spineps.config) and wire it into the nav; feature spineps.segment() in the index quick start. - README cleanup: single H1 title; drop the stale 'python entrypoint.py' usage (use 'spineps'); convert the argument table to the 2.0 kebab-case flags (remove the removed -save_unc_img, add --model-labeling/--batch-size); tidy the model-weights install steps. - Remove stale 'uncertainty image' output references (no longer exposed). - Fix two griffe docstring warnings in spineps.api (the grouped-config params). Docs build cleanly with mkdocs. Co-Authored-By: Claude Opus 4.8 --- README.md | 72 ++++++++++++++++++++++------------------ docs/api/high-level.md | 22 ++++++++++++ docs/index.md | 9 +++++ docs/modules/pipeline.md | 2 +- mkdocs.yml | 1 + spineps/api.py | 10 ++++-- 6 files changed, 80 insertions(+), 36 deletions(-) create mode 100644 docs/api/high-level.md diff --git a/README.md b/README.md index 46fd168..51fe947 100644 --- a/README.md +++ b/README.md @@ -7,11 +7,13 @@ [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Documentation Status](https://readthedocs.org/projects/spineps/badge/?version=latest)](https://spineps.readthedocs.io) -# SPINEPS – Automatic Whole Spine Segmentation of T2w MR images using a Two-Phase Approach to Multi-class Semantic and Instance Segmentation. -# and -# VERIDAH: Solving Enumeration Anomaly Aware Vertebra Labeling across Imaging Sequences +# SPINEPS -This is a segmentation pipeline to automatically, and robustly, segment the whole spine in T2w sagittal images. +**Automatic whole-spine segmentation of MR (and CT) images** — a two-phase approach to multi-class semantic and +instance segmentation, with **VERIDAH** ("Solving Enumeration Anomaly Aware Vertebra Labeling across Imaging +Sequences") for anatomical vertebra labeling. + +SPINEPS automatically and robustly segments the whole spine in sagittal images. ## NOW SUPPORTS BOTH CT AND T2W! There is a new release that finally supports both CT and T2W with completely independent, modality specific models. We are already working on completely modality/sequence robust version that works on everything. Stay tuned for that. @@ -121,12 +123,16 @@ This should throw no errors and return True ### Setup this package -You have to install the package to use it, even if you just want to locally use the code. -1. `cd` into the `spineps` folder and install it by running `pip install -e .` or using the `pyproject.toml` inside of the project folder. -2. If you want to use manual modelweights, download them from the corresponding release page. -3. Extract the downloaded modelweights folders into a folder of your choice (the "spineps/spineps/models" folders will be used as default), from now on referred to as your models folder. -This specified folder should have the following structure: -4. You don't need this, SPINEPS will automatically download the newest weights for you. +Install the package (required even for local use): + +```bash +cd spineps +pip install -e . +``` + +**Model weights download automatically on first use**, so you usually don't need to do anything else. To manage +weights manually instead, download them from the corresponding release page and extract each model folder into a +directory of your choice (default `spineps/spineps/models/`), structured like: ``` ├── @@ -138,7 +144,7 @@ This specified folder should have the following structure: ... ``` -3. You need to specify this models folder as argument when running. If you want to set it permanently, set the according environment variable in your `.bashrc` or `.zshrc` (whatever you are using). +Point SPINEPS at that directory via the `SPINEPS_SEGMENTOR_MODELS` environment variable (set it permanently in your `.bashrc`/`.zshrc`): ```bash export SPINEPS_SEGMENTOR_MODELS= ``` @@ -156,17 +162,17 @@ If you **don't** set the environment variable, the pipeline will look into `spin ## Usage -### Installed as package: +After installation (`pip install spineps`, or `pip install -e .` from a local clone), the `spineps` command is +available in your venv: -1. Activate your venv -2. Run `spineps -h` to see the arguments - -### Installed as local clone: - -1. Activate your venv -2. Run `python entrypoint.py -h` to see the arguments. -3. For example, for a sample, run `python entrypoint.py sample -i --model-semantic --model-instance ` -(replacing with the name of the model you want to use) +1. Activate your venv. +2. Run `spineps -h` for the subcommands, and `spineps sample -h` / `spineps dataset -h` for their arguments. +3. For example, to segment a single scan: +```bash +spineps sample -i --model-semantic --model-instance +``` +(replacing `` with the model you want to use). You can also call SPINEPS from Python — see +[Using the Code](#using-the-code). ### Issues @@ -189,15 +195,16 @@ Processes a single nifty file, will create a derivatves folder next to the nifty | argument | explanation | | :--- | --------- | | -i | Absolute path to the single nifty file (.nii.gz) to be processed | -| --model-semantic , -ms | The model used for the semantic segmentation | -| --model-instance , -mv | The model used for the vertebra instance segmentation | -| --derivative-name , -dn | Name of the derivatives folder (default: derivatives_seg) | -| --save-debug, -sd | Saves a lot of debug data and intermediate results in a separate debug-labeled folder (default: False) | -| -save_unc_img, -sui | Saves a uncertainty image from the subreg prediction (default: False) | -| -override_semantic, -os | Will override existing seg-spine files (default: False) | -| -override_instance, -ov | Will override existing seg-vert files (default: False) | -| -override_ctd, -oc | Will override existing centroid files (default: False) | -| -verbose, -v | Prints much more stuff, may fully clutter your terminal (default: False) | +| --model-semantic, -ms | The model used for the semantic segmentation | +| --model-instance, -mi | The model used for the vertebra instance segmentation | +| --model-labeling, -ml | The (optional) VERIDAH model used for vertebra labeling | +| --derivative-name, -dn | Name of the derivatives folder (default: derivatives_seg) | +| --save-debug, -sd | Saves debug data and intermediate results in a separate folder (default: False) | +| --override-semantic, -os | Override existing seg-spine files (default: False) | +| --override-instance, -oi | Override existing seg-vert files (default: False) | +| --override-ctd, -oc | Override existing centroid files (default: False) | +| --batch-size, -bs | Vertebra cutouts per batched forward pass; higher is faster but uses more GPU memory (default: 4) | +| --verbose, -v | Prints much more stuff, may fully clutter your terminal (default: False) | There are a lot more arguments, run `spineps sample -h` to see them. @@ -260,9 +267,8 @@ The pipeline segments in multiple steps: 1. Semantically segments 14 spinal structures (9 regions for vertebrae, Spinal Cord, Spinal Canal, Intervertebral Discs, Endplate, Sacrum) 2. From the vertebra regions, segment the different vertebrae as instance mask 3. Save the first as `seg-spine` mask, the second as `seg-vert` mask -4. It can save an uncertainty image for the semantic segmentation -5. From the two segmentations, calculates centroids for each vertebrae center point, endplate, and IVD and saves that into a .json -6. From the centroid and the segmentations, makes a snapshot showcasing the result as a .png +4. From the two segmentations, calculates centroids for each vertebrae center point, endplate, and IVD and saves that into a .json +5. From the centroid and the segmentations, makes a snapshot showcasing the result as a .png ![example_semantic](spineps/example/figures/example_semantic.png?raw=true) diff --git a/docs/api/high-level.md b/docs/api/high-level.md new file mode 100644 index 0000000..22a72e8 --- /dev/null +++ b/docs/api/high-level.md @@ -0,0 +1,22 @@ +# High-Level API + +The easiest way to use SPINEPS from Python: the one-call [`segment`][spineps.api.segment] function, a reusable +[`SpinepsPipeline`][spineps.api.SpinepsPipeline] (load the models once, segment many images), and grouped +configuration objects to toggle individual processing steps without long keyword-argument lists. + +```python +import spineps + +result = spineps.segment("sub-01_T2w.nii.gz") # saves a derivatives folder next to the input +result = spineps.segment(nii, output_in_memory=True) # or return the masks in memory +if result.success: + semantic, vertebra = result.semantic, result.vertebra +``` + +## spineps.api + +::: spineps.api + +## spineps.config + +::: spineps.config diff --git a/docs/index.md b/docs/index.md index 22d74b2..ffa3449 100644 --- a/docs/index.md +++ b/docs/index.md @@ -28,6 +28,7 @@ Given a sagittal MR scan, the pipeline: - [Pipeline](modules/pipeline.md) — how the two-phase pipeline is structured. - [Processing Phases](modules/phases.md) — pre-processing, semantic, instance, labeling and post-processing. - [Models & Labeling](modules/models.md) — model loading and the VERIDAH labeling model. +- [High-Level API](api/high-level.md) — the one-call `spineps.segment()` API and config objects. - [API Reference](api/pipeline.md) — full auto-generated API documentation. ## Quick start @@ -40,6 +41,14 @@ pip install spineps spineps sample -i /path/sub-test_T2w.nii.gz --model-semantic t2w --model-instance instance ``` +Or in one call from Python: + +```python +import spineps + +result = spineps.segment("/path/sub-test_T2w.nii.gz") # saves a derivatives folder next to the input +``` + See [Getting Started](getting-started.md) for the full installation guide (including PyTorch setup and model weights), and the [Pipeline](modules/pipeline.md) page for calling SPINEPS from Python. diff --git a/docs/modules/pipeline.md b/docs/modules/pipeline.md index 5dd8f3f..396d5b9 100644 --- a/docs/modules/pipeline.md +++ b/docs/modules/pipeline.md @@ -43,7 +43,7 @@ For each processed scan SPINEPS writes a derivatives folder next to the input co - a `seg-vert` mask (vertebra instance segmentation), - a centroid file (`.json`) with points of interest for each vertebra, endplate and disc, - a snapshot `.png` visualizing the result, -- optionally: an uncertainty image, the model-resolution masks, softmax logits and debug data. +- optionally: the model-resolution masks, softmax logits and debug data. ## Calling from Python diff --git a/mkdocs.yml b/mkdocs.yml index 624070c..f56bb59 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -69,6 +69,7 @@ nav: - Processing Phases: modules/phases.md - Models & Labeling: modules/models.md - API Reference: + - High-Level API: api/high-level.md - Pipeline & Run: api/pipeline.md - Models: api/models.md - Processing Phases: api/phases.md diff --git a/spineps/api.py b/spineps/api.py index 92268dc..bbdc003 100644 --- a/spineps/api.py +++ b/spineps/api.py @@ -146,7 +146,10 @@ def segment( output_in_memory: If true, returns the masks in memory instead of writing a derivatives folder. Defaults to False. derivative_name: Name of the derivatives output folder (save mode only). Defaults to ``"derivatives_seg"``. override: If true, recomputes and overwrites existing outputs. Defaults to False. - semantic, instance, labeling, post: Optional grouped config objects; unset groups use the pipeline defaults. + semantic: Optional SemanticConfig overriding the semantic-phase defaults. + instance: Optional InstanceConfig overriding the instance-phase defaults. + labeling: Optional LabelingConfig overriding the labeling-phase defaults. + post: Optional PostConfig overriding the post-processing defaults. verbose: If true, prints verbose information. Defaults to False. Returns: @@ -206,7 +209,10 @@ def segment( output_in_memory: If true, returns the masks in memory instead of writing a derivatives folder. Defaults to False. derivative_name: Name of the derivatives output folder (save mode only). Defaults to ``"derivatives_seg"``. override: If true, recomputes and overwrites existing outputs. Defaults to False. - semantic, instance, labeling, post: Optional grouped config objects (see :mod:`spineps.config`). + semantic: Optional SemanticConfig overriding the semantic-phase defaults (see :mod:`spineps.config`). + instance: Optional InstanceConfig overriding the instance-phase defaults. + labeling: Optional LabelingConfig overriding the labeling-phase defaults. + post: Optional PostConfig overriding the post-processing defaults. verbose: If true, prints verbose information. Defaults to False. Returns: From 743786e180d06a2b2c309c5deca9cf6f64e9294c Mon Sep 17 00:00:00 2001 From: ga84mun Date: Tue, 30 Jun 2026 12:37:27 +0000 Subject: [PATCH 15/29] push ants version --- README.md | 11 +++++++++++ pyproject.toml | 5 ++--- spineps/utils/proc_functions.py | 14 +++++++++++++- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 51fe947..d79aed2 100644 --- a/README.md +++ b/README.md @@ -293,6 +293,17 @@ In the subregion segmentation: | 100 | Vertebra_Disc | | 26 | Sacrum | +CT only +| Label | Structure | +| :---: | --------- | +| 51 | Dense | +| 70 | Sacrum_Sacral_Ala_Left | +| 71 | Sacrum_Sacral_Ala_Right | +| 72 | Sacrum_Posterior_Sacral_Elements | +| 73 | Sacrum_Body | +| 74 | Sacrum_Endplate | +| 80 | Metal | + In the vertebra instance segmentation mask, each label X in [1, 25] are the unique vertebrae, while 100+X are their corresponding IVD and 200+X their endplates. ## VERIDAH: diff --git a/pyproject.toml b/pyproject.toml index 6a278ae..d04693f 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,12 +25,11 @@ pytorch-lightning = "^2.0.8" torchmetrics = "^1.1.2" tqdm = "^4.66.1" einops= "^0.6.1" -nnunetv2 = "^2.4.2" +nnunetv2 = "^2.8.0" TPTBox = "^0.7.3" -antspyx = "0.4.2" +antspyx = "0.6.3" rich = "^13.6.0" monai="^1.3.0" -acvl-utils = "0.2" TypeSaveArgParse="^1.0.1" python-gdcm = [ { version = "==3.0.25", python = ">=3.9,<3.10" }, diff --git a/spineps/utils/proc_functions.py b/spineps/utils/proc_functions.py index 3d9758a..41ce2f9 100755 --- a/spineps/utils/proc_functions.py +++ b/spineps/utils/proc_functions.py @@ -21,6 +21,19 @@ MAX_VERTEBRA_INSTANCE_LABEL = 25 +def from_nibabel(nib): + try: + from ants.utils.convert_nibabel import from_nibabel + + return from_nibabel(nib) + except ModuleNotFoundError: + from ants.utils.nibabel_nifti_to_ants import ( + from_nibabel_nifti, + ) + + return from_nibabel_nifti(nib) + + def n4_bias( nii: NII, threshold: int = 60, @@ -46,7 +59,6 @@ def n4_bias( Returns: tuple[NII, NII]: The bias-corrected image and the binary foreground mask used for correction. """ - from ants.utils.convert_nibabel import from_nibabel # they keep renaming that thing. (version 0.4.2) # print("n4 bias", nii.dtype) mask = nii.get_array() From bf590ae13d73b58efc3d76173b7887e94aebafa4 Mon Sep 17 00:00:00 2001 From: ga84mun Date: Tue, 30 Jun 2026 12:39:02 +0000 Subject: [PATCH 16/29] remove unused imports --- spineps/phase_labeling.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/spineps/phase_labeling.py b/spineps/phase_labeling.py index e1a0b2c..3b94ea2 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -8,14 +8,12 @@ from spineps.architectures.read_labels import ( VertExact, - VertExactClass, VertGroup, VertRegion, VertRel, VertT13, vert_group_idx_to_exact_idx_dict, ) -from spineps.get_models import get_actual_model from spineps.lab_model import VertLabelingClassifier from spineps.utils.find_min_cost_path import ( DEFAULT_REGION_STARTS, From 9a7da9e5423cb0cfea2116ef9823e5cdb33563f0 Mon Sep 17 00:00:00 2001 From: ga84mun Date: Tue, 30 Jun 2026 14:14:34 +0000 Subject: [PATCH 17/29] update crop --- spineps/phase_pre.py | 26 +++++++++++++------------- spineps/seg_run.py | 20 ++++++++++++++++---- 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/spineps/phase_pre.py b/spineps/phase_pre.py index 641ac8a..a0a734e 100644 --- a/spineps/phase_pre.py +++ b/spineps/phase_pre.py @@ -57,22 +57,22 @@ def compute_crop( Returns: tuple[slice, slice, slice]: The crop slices around the segmented spine, with a ``VIBE_CROP_MARGIN_MM`` margin. """ - from TPTBox.core.vert_constants import Full_Body_Instance_Vibe + from TPTBox.core.vert_constants import Full_Body_Instance, Full_Body_Instance_Vibe from TPTBox.segmentation import run_vibeseg - if _has_logger_arg(run_vibeseg): - out = run_vibeseg(nii, out_file, dataset_id=dataset_id, ddevice=ddevice, gpu=gpu, max_folds=max_folds, logger=logger) - else: # backwards compatibility, can be removed if we force to a new version of TPTBox than 30.Apr.26 - out = run_vibeseg(nii, out_file, dataset_id=dataset_id, ddevice=ddevice, gpu=gpu, max_folds=max_folds) + out = run_vibeseg(nii, out_file, dataset_id=dataset_id, ddevice=ddevice, gpu=gpu, max_folds=max_folds, logger=logger) seg = to_nii(out, True) - seg.extract_label_( - [ - Full_Body_Instance_Vibe.IVD, - Full_Body_Instance_Vibe.vertebra_body, - Full_Body_Instance_Vibe.vertebra_posterior_elements, - Full_Body_Instance_Vibe.sacrum, - ] - ) + if dataset_id in range(30, 120): + seg.extract_label_( + [ + Full_Body_Instance_Vibe.IVD, + Full_Body_Instance_Vibe.vertebra_body, + Full_Body_Instance_Vibe.vertebra_posterior_elements, + Full_Body_Instance_Vibe.sacrum, + ] + ) + elif dataset_id in range(10, 20): + seg.extract_label_([Full_Body_Instance.ivd, Full_Body_Instance.vert_body, Full_Body_Instance.vert_post, Full_Body_Instance.sacrum]) return seg.compute_crop(0, dist=VIBE_CROP_MARGIN_MM / min(seg.zoom)) diff --git a/spineps/seg_run.py b/spineps/seg_run.py index ec0a3eb..aa646ee 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -349,7 +349,7 @@ def segment_image( # noqa: C901 timing=False, verbose: bool = False, _nii=None, - _end_after_subreg=False, + _dataset_id_ct_crop=100, ) -> tuple[dict[str, Path], ErrCode]: """Runs the SPINEPS framework over one nifty. @@ -430,7 +430,12 @@ def segment_image( # noqa: C901 input_format = img_ref.format output_paths = output_paths_from_input( - img_ref, derivative_name, snapshot_copy_folder, input_format=input_format, non_strict_mode=ignore_bids_filter + img_ref, + derivative_name, + snapshot_copy_folder, + input_format=input_format, + non_strict_mode=ignore_bids_filter, + _dataset_id_ct_crop=_dataset_id_ct_crop, ) out_spine = output_paths["out_spine"] out_spine_raw = output_paths["out_spine_raw"] @@ -517,7 +522,13 @@ def segment_image( # noqa: C901 "Compute spine crop with VIBESegmentator https://link.springer.com/article/10.1007/s00330-025-12035-9", Log_Type.OK ) out_vibeseg = output_paths["out_vibeseg"] - crop = compute_crop(input_nii, out_vibeseg, ddevice="cpu" if model_semantic.use_cpu else "cuda", logger=logger) + crop = compute_crop( + input_nii, + out_vibeseg, + dataset_id=_dataset_id_ct_crop, + ddevice="cpu" if model_semantic.use_cpu else "cuda", + logger=logger, + ) if timing: logger.print( f"Compute cropping took: {perf_counter() - start_time2:.2f} seconds", Log_Type.OK, verbose=log_inference_time @@ -738,6 +749,7 @@ def output_paths_from_input( snapshot_copy_folder: Path | str | None, input_format: str, non_strict_mode: bool = False, + _dataset_id_ct_crop=100, ) -> dict[str, Path]: """Derives all pipeline output file paths for a given input image. @@ -823,7 +835,7 @@ def output_paths_from_input( out_vibeseg = img_ref.get_changed_path( bids_format="msk", parent=derivative_name, - info={"seg": "VIBESeg-100", "mod": img_ref.format}, + info={"seg": f"VIBESeg-{_dataset_id_ct_crop}", "mod": img_ref.format}, non_strict_mode=non_strict_mode, make_parent=False, ) From d03eb1e6e24caea13ef0d3388188f5f4847c5885 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 06:54:58 +0000 Subject: [PATCH 18/29] adaptive nnunet version depending on python >=3.10 or not --- pyproject.toml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d04693f..2cb588f 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,6 @@ pytorch-lightning = "^2.0.8" torchmetrics = "^1.1.2" tqdm = "^4.66.1" einops= "^0.6.1" -nnunetv2 = "^2.8.0" TPTBox = "^0.7.3" antspyx = "0.6.3" rich = "^13.6.0" @@ -35,6 +34,11 @@ python-gdcm = [ { version = "==3.0.25", python = ">=3.9,<3.10" }, ] +nnunetv2 = [ + { version = "2.4.2", python = "<3.10" }, + { version = "^2.8.0", python = ">=3.10" } +] + [tool.poetry.group.dev.dependencies] pytest = "^7.4.4" From 3ae4ad276037df57007be28145439e420cd5ee7e Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 07:02:49 +0000 Subject: [PATCH 19/29] fix test --- spineps/phase_labeling.py | 1 + 1 file changed, 1 insertion(+) diff --git a/spineps/phase_labeling.py b/spineps/phase_labeling.py index f514f96..5c8f5b3 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -8,6 +8,7 @@ from spineps.architectures.read_labels import ( VertExact, + VertExactClass, VertGroup, VertRegion, VertRel, From 83301bba6e9de2f7f542fd1bd321b4d4fc4b5717 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 07:17:19 +0000 Subject: [PATCH 20/29] move to newest tptbox version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 2cb588f..b06c679 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ pytorch-lightning = "^2.0.8" torchmetrics = "^1.1.2" tqdm = "^4.66.1" einops= "^0.6.1" -TPTBox = "^0.7.3" +TPTBox = "^0.7.4" antspyx = "0.6.3" rich = "^13.6.0" monai="^1.3.0" From d91d47412bcf936296966e7bdd2febf66dd3bafd Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 07:44:20 +0000 Subject: [PATCH 21/29] move to newest tptbox version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b06c679..8dba25b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ pytorch-lightning = "^2.0.8" torchmetrics = "^1.1.2" tqdm = "^4.66.1" einops= "^0.6.1" -TPTBox = "^0.7.4" +TPTBox = "^0.7.5" antspyx = "0.6.3" rich = "^13.6.0" monai="^1.3.0" From 819baaef02f3561918817dda74119e0edd905194 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 08:01:11 +0000 Subject: [PATCH 22/29] fix versioning issue --- pyproject.toml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8dba25b..0836865 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,13 +33,16 @@ TypeSaveArgParse="^1.0.1" python-gdcm = [ { version = "==3.0.25", python = ">=3.9,<3.10" }, ] - +blosc2 = [ + { version = "^3.3.2", python = "<3.10" } +] nnunetv2 = [ { version = "2.4.2", python = "<3.10" }, { version = "^2.8.0", python = ">=3.10" } ] + [tool.poetry.group.dev.dependencies] pytest = "^7.4.4" pre-commit = "*" From c7097856d5b3c821d2cf7e3cb12bee097b9f95a5 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 08:05:39 +0000 Subject: [PATCH 23/29] changed to arbitrary blosc2 version --- pyproject.toml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0836865..fd8ce00 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,9 +33,7 @@ TypeSaveArgParse="^1.0.1" python-gdcm = [ { version = "==3.0.25", python = ">=3.9,<3.10" }, ] -blosc2 = [ - { version = "^3.3.2", python = "<3.10" } -] +blosc2 = "*" nnunetv2 = [ { version = "2.4.2", python = "<3.10" }, { version = "^2.8.0", python = ">=3.10" } From ecbed444457fab1d1395602959850daaf8026b54 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 08:21:11 +0000 Subject: [PATCH 24/29] ruff happy --- spineps/phase_post.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/spineps/phase_post.py b/spineps/phase_post.py index d8965ca..7b8955e 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -388,10 +388,9 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=Fals orientation = whole_vert_nii.orientation vert_t = whole_vert_nii.reorient() seg_t = seg_nii.reorient() - if include_sacrum: - vert_labels = [t for t in vert_t.unique() if t < 40] # without zero - else: - vert_labels = [t for t in vert_t.unique() if t <= 26 or t == 28] # without zero + vert_labels = ( + [t for t in vert_t.unique() if t < 40] if include_sacrum else [t for t in vert_t.unique() if t <= 26 or t == 28] + ) # without zero vert_arr = vert_t.get_seg_array() subreg_arr = seg_t.get_seg_array() @@ -584,7 +583,7 @@ def find_nearest_higher(seq, x) -> float: return min(values_higher) -def label_instance_top_to_bottom(vert_nii: NII, labeling_offset: int = 0) -> tuple[NII, np.ndarray]: +def label_instance_top_to_bottom(vert_nii: NII, labeling_offset: int = 0) -> tuple[NII, list[int]]: """Relabel vertebra instances consecutively from top to bottom by center-of-mass height. Reorients to PIR, sorts the instances by their center of mass along the inferior-superior axis, and assigns consecutive @@ -595,8 +594,9 @@ def label_instance_top_to_bottom(vert_nii: NII, labeling_offset: int = 0) -> tup labeling_offset (int): Offset added to the consecutive labels. Returns: - tuple[NII, np.ndarray]: The relabeled instance mask and its array of unique labels. + tuple[NII, list[int]]: The relabeled instance mask and its list of unique labels. """ + ori = vert_nii.orientation vert_nii.reorient_() vert_arr = vert_nii.get_seg_array() @@ -672,7 +672,7 @@ def assign_vertebra_inconsistency( vert_arr[cc_map == 1] = to_label logger.print( - f"set cc to {to_label}, with volume decision {gt_volume}, based on {biggest_volume}, {second_volume}", verbose=verbose + f"set cc to {to_label}, with volume decision {gt_volume}, based on {biggest_volume}, {second_volume}", ) vert_nii.set_array_(vert_arr) From a04bc9ed1a5102ca5831f61999142c0137c6b13a Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 1 Jul 2026 08:40:53 +0000 Subject: [PATCH 25/29] some cleanup to newest tptbox, also acvl-utils --- pyproject.toml | 4 ++++ spineps/utils/proc_functions.py | 15 +-------------- 2 files changed, 5 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fd8ce00..3d8ca84 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,10 @@ nnunetv2 = [ { version = "2.4.2", python = "<3.10" }, { version = "^2.8.0", python = ">=3.10" } ] +acvl-utils = [ + { version = "0.2", python = "<3.10" }, + { version = "*", python = ">=3.10" }, +] diff --git a/spineps/utils/proc_functions.py b/spineps/utils/proc_functions.py index 41ce2f9..3987897 100755 --- a/spineps/utils/proc_functions.py +++ b/spineps/utils/proc_functions.py @@ -21,19 +21,6 @@ MAX_VERTEBRA_INSTANCE_LABEL = 25 -def from_nibabel(nib): - try: - from ants.utils.convert_nibabel import from_nibabel - - return from_nibabel(nib) - except ModuleNotFoundError: - from ants.utils.nibabel_nifti_to_ants import ( - from_nibabel_nifti, - ) - - return from_nibabel_nifti(nib) - - def n4_bias( nii: NII, threshold: int = 60, @@ -68,7 +55,7 @@ def n4_bias( mask[slices] = 1 mask_nii = nii.set_array(mask) mask_nii.seg = True - n4: NII = nii.n4_bias_field_correction(threshold=0, mask=from_nibabel(mask_nii.nii), spline_param=spline_param) + n4: NII = nii.n4_bias_field_correction(threshold=0, mask=mask_nii, spline_param=spline_param) if norm != -1: n4 *= norm / n4.max() if dtype2nii: From 1b91ebb89cac2164d5a04dab0c7bf1e2f387e9fd Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 2 Jul 2026 08:09:40 +0000 Subject: [PATCH 26/29] attempt to reduce memory consumption --- README.md | 2 +- spineps/entrypoint.py | 4 +++- spineps/phase_instance.py | 25 +++++++++++-------------- spineps/seg_run.py | 10 +++++++++- 4 files changed, 24 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index d79aed2..0d89347 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ Processes a single nifty file, will create a derivatves folder next to the nifty | --override-semantic, -os | Override existing seg-spine files (default: False) | | --override-instance, -oi | Override existing seg-vert files (default: False) | | --override-ctd, -oc | Override existing centroid files (default: False) | -| --batch-size, -bs | Vertebra cutouts per batched forward pass; higher is faster but uses more GPU memory (default: 4) | +| --batch-size, -bs | Vertebra cutouts per batched forward pass; higher is faster but uses more GPU memory. Only affects GPU memory; host RAM usage in the instance phase scales with scan length/vertebra count instead (default: 4) | | --verbose, -v | Prints much more stuff, may fully clutter your terminal (default: False) | There are a lot more arguments, run `spineps sample -h` to see them. diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index da539ee..6799061 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -97,7 +97,9 @@ def parser_arguments(parser: argparse.ArgumentParser): type=int, default=4, help="Number of vertebra cutouts run through the instance model per batched forward pass. Higher is faster but uses " - "more GPU memory; falls back to one-by-one on out-of-memory.", + "more GPU memory; falls back to one-by-one on out-of-memory. Only affects GPU memory during the instance model's " + "forward pass -- host RAM usage in the instance phase scales with scan length and vertebra count instead, and is " + "unaffected by this setting.", ) parser.add_argument( "--amp", diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index f6db4f7..e36c365 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -652,8 +652,6 @@ def collect_vertebra_predictions( # which would needlessly inflate this n_coms x 3 x volume array and slow the Dice comparisons below). hierarchical_predictions = np.zeros((n_corpus_coms, 3, *shp), dtype=np.uint8) # print("hierarchical_predictions", hierarchical_predictions.shape) - vert_predict_template = np.zeros(shp, dtype=np.uint16) - # print("vert_predict_template", vert_predict_template.shape) # relabel to the labels expected by the model # {41: 1, 42: 2, 43: 3, 44: 4, 45: 5, 46: 6, 47: 7, 48: 8, 49: 9, 50: 9, Location.Dens_axis.value: 9} @@ -715,18 +713,18 @@ def collect_vertebra_predictions( cutout_sizes = tuple(cutout_coords[i].stop - cutout_coords[i].start for i in range(len(cutout_coords))) pad_cutout = tuple(slice(paddings[i][0], paddings[i][0] + cutout_sizes[i]) for i in range(len(paddings))) arr = vert_cut_nii.get_seg_array() - vert_predict_map = vert_predict_template.copy() - vert_predict_map[cutout_coords] = arr[pad_cutout] - seg_at_com = vert_predict_map[int(com[0])][int(com[1])][int(com[2])] + cutout_vals = arr[pad_cutout] + # Write straight into the (already fully-allocated) hierarchical_predictions slice instead of building + # full-volume-sized temporaries per vertebra/label: everything outside cutout_coords is 0 either way. + local_com = tuple(int(com[i]) - cutout_coords[i].start for i in range(3)) + seg_at_com = cutout_vals[local_com] if seg_at_com == 0: logger.print("Zero at cutout center, mistake", Log_Type.WARNING) for l in vert_labels: - vert_l_map = vert_predict_map.copy() - vert_l_map[vert_l_map != l] = 0 - vert_l_map[vert_l_map != 0] = 1 + mask = cutout_vals == l labelindex = l - 1 - if vert_l_map.max() > 0: - hierarchical_predictions[com_idx][labelindex] = vert_l_map + if mask.any(): + hierarchical_predictions[com_idx, labelindex][cutout_coords] = mask.astype(np.uint8) hierarchical_existing_predictions.append(str_id_com_label(com_idx, labelindex)) return hierarchical_predictions, hierarchical_existing_predictions, n_corpus_coms @@ -992,6 +990,7 @@ def merge_coupled_predictions( """ whole_vert_nii = seg_nii.copy() whole_vert_arr = np.zeros(whole_vert_nii.shape, dtype=np.uint16) # this is fixed segmentations from vert + combine = np.zeros(whole_vert_nii.shape, dtype=whole_vert_nii.dtype) # reused scratch buffer, reset per couple below idx = 1 for k, overall_agreement in coupled_predictions.items(): @@ -999,7 +998,7 @@ def merge_coupled_predictions( take_no_overlap = len(k) <= 2 if overall_agreement < 0.3 + 0.15 * (4 - len(k)): take_no_overlap = True - combine = np.zeros(whole_vert_nii.shape, dtype=whole_vert_nii.dtype) + combine.fill(0) for cid in k: combine += hierarchical_predictions[cid[0]][cid[1]] # print(combine.shape) @@ -1012,9 +1011,7 @@ def merge_coupled_predictions( if count_new == 0: logger.print("ZERO instance mask failure on vertebra instance creation", Log_Type.FAIL) return seg_nii, debug_data, ErrCode.EMPTY - fixed_n = combine.copy() - fixed_n[whole_vert_arr != 0] = 0 - count_cut = np_count_nonzero(fixed_n) + count_cut = np_count_nonzero((combine != 0) & (whole_vert_arr == 0)) relative_overlap = (count_new - count_cut) / count_new if relative_overlap > 0.6: logger.print(k, f" was skipped because it overlaps {round(relative_overlap, 4)} with established verts", verbose=verbose) diff --git a/spineps/seg_run.py b/spineps/seg_run.py index aa646ee..dd266ba 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -25,6 +25,13 @@ from spineps.utils.citation_reminder import citation_reminder +class _NoOpDebugDict(dict): + """Dict-like sink that discards writes; used to skip retaining debug data when it won't be saved.""" + + def __setitem__(self, key, value): + pass + + @citation_reminder def process_dataset( # noqa: C901 dataset_path: Path, @@ -465,7 +472,8 @@ def segment_image( # noqa: C901 return output_paths, ErrCode.ALL_DONE done_something = False - debug_data_run: dict[str, NII] = {} + # Avoid retaining full-volume debug copies for the whole run when they'll never be saved (see seg_run.py:699). + debug_data_run: dict[str, NII] = {} if save_debug_data else _NoOpDebugDict() if Modality.CT in model_semantic.modalities(): proc_normalize_input = False # Never normalize input if it is an CT From 993910e927e0b21fb49f9b00fb58520273858605 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 2 Jul 2026 11:50:36 +0000 Subject: [PATCH 27/29] update t newest PTBox, removed dead version-compat shims, to ants to tptbox, deleted vendored nnUNet predictor --- pyproject.toml | 7 - spineps/phase_instance.py | 1 - spineps/phase_post.py | 1 - spineps/phase_pre.py | 17 +- spineps/phase_semantic.py | 1 - spineps/seg_pipeline.py | 1 - spineps/seg_run.py | 20 +- spineps/seg_utils.py | 1 - spineps/utils/data_iterators.py | 112 ----- spineps/utils/default_preprocessor.py | 220 --------- spineps/utils/export_prediction.py | 129 ------ spineps/utils/get_network_from_plans.py | 98 ---- spineps/utils/plans_handler.py | 331 ------------- spineps/utils/predictor.py | 514 --------------------- spineps/utils/proc_functions.py | 2 +- spineps/utils/sliding_window_prediction.py | 75 --- unit_tests/test_filepaths.py | 4 +- 17 files changed, 12 insertions(+), 1522 deletions(-) delete mode 100755 spineps/utils/data_iterators.py delete mode 100755 spineps/utils/default_preprocessor.py delete mode 100755 spineps/utils/export_prediction.py delete mode 100755 spineps/utils/get_network_from_plans.py delete mode 100755 spineps/utils/plans_handler.py delete mode 100755 spineps/utils/predictor.py delete mode 100755 spineps/utils/sliding_window_prediction.py diff --git a/pyproject.toml b/pyproject.toml index 3d8ca84..602f156 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,13 +96,6 @@ exclude = [ "node_modules", "venv", "spineps/Unet3D", - "spineps/utils/data_iterators.py", - "spineps/utils/default_preprocessor.py", - "spineps/utils/export_prediction.py", - "spineps/utils/get_network_from_plans.py", - "spineps/utils/plans_handler.py", - "spineps/utils/predictor.py", - "spineps/utils/sliding_window_prediction.py", "spineps/utils/image.py", # vendored from spinalcordtoolbox ".toml", ] diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index e36c365..a9dcb99 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -2,7 +2,6 @@ from __future__ import annotations -# from utils.predictor import nnUNetPredictor import numpy as np from TPTBox import NII, Location, Log_Type from TPTBox.core.np_utils import ( diff --git a/spineps/phase_post.py b/spineps/phase_post.py index 7b8955e..e3485c8 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -2,7 +2,6 @@ from __future__ import annotations -# from utils.predictor import nnUNetPredictor import heapq import numpy as np diff --git a/spineps/phase_pre.py b/spineps/phase_pre.py index a0a734e..d72a5bc 100644 --- a/spineps/phase_pre.py +++ b/spineps/phase_pre.py @@ -2,9 +2,6 @@ from __future__ import annotations -import inspect - -# from utils.predictor import nnUNetPredictor from time import perf_counter from typing import TYPE_CHECKING, Literal @@ -25,18 +22,6 @@ VIBE_CROP_MARGIN_MM = 25 * min(REFERENCE_ZOOM) -def _has_logger_arg(func) -> bool: - """Check whether a callable accepts a ``logger`` keyword argument. - - Args: - func (Callable): The function whose signature is inspected. - - Returns: - bool: True if ``logger`` is among the function's parameters, else False. - """ - return "logger" in inspect.signature(func).parameters - - def compute_crop( nii: NII, out_file: str | Path, dataset_id=100, ddevice: Literal["cpu", "cuda", "mps"] = "cuda", gpu=0, max_folds=None, logger=None ) -> tuple[slice, slice, slice]: @@ -52,7 +37,7 @@ def compute_crop( ddevice (Literal["cpu", "cuda", "mps"], optional): Compute device for inference. Defaults to "cuda". gpu (int, optional): GPU index used when running on CUDA. Defaults to 0. max_folds (int | None, optional): Maximum number of model folds to ensemble. Defaults to None (all folds). - logger (optional): Logger forwarded to ``run_vibeseg`` when that version supports it. Defaults to None. + logger (optional): Logger forwarded to ``run_vibeseg``. Defaults to None. Returns: tuple[slice, slice, slice]: The crop slices around the segmented spine, with a ``VIBE_CROP_MARGIN_MM`` margin. diff --git a/spineps/phase_semantic.py b/spineps/phase_semantic.py index 5aa350b..c8f67c3 100755 --- a/spineps/phase_semantic.py +++ b/spineps/phase_semantic.py @@ -2,7 +2,6 @@ from __future__ import annotations -# from utils.predictor import nnUNetPredictor import numpy as np from TPTBox import NII, Location, Log_Type diff --git a/spineps/seg_pipeline.py b/spineps/seg_pipeline.py index 7958de5..df49115 100755 --- a/spineps/seg_pipeline.py +++ b/spineps/seg_pipeline.py @@ -2,7 +2,6 @@ from __future__ import annotations -# from utils.predictor import nnUNetPredictor import subprocess from typing import Any diff --git a/spineps/seg_run.py b/spineps/seg_run.py index dd266ba..b32826b 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -726,18 +726,14 @@ def segment_image( # noqa: C901 if snapshot_copy_folder is not None: out_snap = [out_snap, out_snap2] ctd = ctd.extract_subregion(Location.Vertebra_Corpus) - try: - mri_snapshot( # TODO update snapshot - img_ref, - vert_nii_clean, - ctd, - subreg_msk=seg_nii_clean, - out_path=out_snap, - mode="MRI" if img_ref.bids_format.lower() != "ct" else "CT", - ) - except Exception: - # Fall back for older TPTBox versions TODO remove later - mri_snapshot(img_ref, vert_nii_clean, ctd, subreg_msk=seg_nii_clean, out_path=out_snap) + mri_snapshot( + img_ref, + vert_nii_clean, + ctd, + subreg_msk=seg_nii_clean, + out_path=out_snap, + mode="MRI" if img_ref.bids_format.lower() != "ct" else "CT", + ) logger.print(f"Snapshot saved into {out_snap}", Log_Type.SAVE) if timing: logger.print(f"Snapshot took: {perf_counter() - start_time2:.2f} seconds", Log_Type.OK, verbose=log_inference_time) diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index 2adaca5..185ecca 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -2,7 +2,6 @@ from __future__ import annotations -# from utils.predictor import nnUNetPredictor from typing import Union from TPTBox import BIDS_FILE, ZOOMS, Log_Type diff --git a/spineps/utils/data_iterators.py b/spineps/utils/data_iterators.py deleted file mode 100755 index 1f9b0a2..0000000 --- a/spineps/utils/data_iterators.py +++ /dev/null @@ -1,112 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -from typing import Union, List - -import numpy as np -import torch -from batchgenerators.dataloading.data_loader import DataLoader - -from spineps.utils.default_preprocessor import DefaultPreprocessor -from spineps.utils.plans_handler import PlansManager, ConfigurationManager - - -class PreprocessAdapterFromNpy(DataLoader): - def __init__( - self, - list_of_images: List[np.ndarray], - list_of_segs_from_prev_stage: Union[List[np.ndarray], None], - list_of_image_properties: List[dict], - truncated_ofnames: Union[List[str], None], - plans_manager: PlansManager, - dataset_json: dict, - configuration_manager: ConfigurationManager, - num_threads_in_multithreaded: int = 1, - verbose: bool = False, - ): - preprocessor = DefaultPreprocessor(verbose=verbose) - self.preprocessor, self.plans_manager, self.configuration_manager, self.dataset_json, self.truncated_ofnames = ( - preprocessor, - plans_manager, - configuration_manager, - dataset_json, - truncated_ofnames, - ) - - self.label_manager = plans_manager.get_label_manager(dataset_json) - - if list_of_segs_from_prev_stage is None: - list_of_segs_from_prev_stage = [None] * len(list_of_images) - if truncated_ofnames is None: - truncated_ofnames = [None] * len(list_of_images) - - super().__init__( - list(zip(list_of_images, list_of_segs_from_prev_stage, list_of_image_properties, truncated_ofnames)), - 1, - num_threads_in_multithreaded, - seed_for_shuffle=1, - return_incomplete=True, - shuffle=False, - infinite=False, - sampling_probabilities=None, - ) - - self.indices = list(range(len(list_of_images))) - - def generate_train_batch(self): - idx = self.get_indices()[0] - image = self._data[idx][0] - seg_prev_stage = self._data[idx][1] - props = self._data[idx][2] - ofname = self._data[idx][3] - # if we have a segmentation from the previous stage we have to process it together with the images so that we - # can crop it appropriately (if needed). Otherwise it would just be resized to the shape of the data after - # preprocessing and then there might be misalignments - data, seg = self.preprocessor.run_case_npy( - image, seg_prev_stage, props, self.plans_manager, self.configuration_manager, self.dataset_json - ) - if seg_prev_stage is not None: - seg_onehot = convert_labelmap_to_one_hot(seg[0], self.label_manager.foreground_labels, data.dtype) - data = np.vstack((data, seg_onehot)) - - data = torch.from_numpy(data) - - return {"data": data, "data_properites": props, "ofile": ofname} - - -def convert_labelmap_to_one_hot( - segmentation: Union[np.ndarray, torch.Tensor], all_labels: Union[List, torch.Tensor, np.ndarray, tuple], output_dtype=None -) -> Union[np.ndarray, torch.Tensor]: - """ - if output_dtype is None then we use np.uint8/torch.uint8 - if input is torch.Tensor then output will be on the same device - - np.ndarray is faster than torch.Tensor - - if segmentation is torch.Tensor, this function will be faster if it is LongTensor. If it is somethine else we have - to cast which takes time. - - IMPORTANT: This function only works properly if your labels are consecutive integers, so something like 0, 1, 2, 3, ... - DO NOT use it with 0, 32, 123, 255, ... or whatever (fix your labels, yo) - """ - if isinstance(segmentation, torch.Tensor): - result = torch.zeros( - (len(all_labels), *segmentation.shape), - dtype=output_dtype if output_dtype is not None else torch.uint8, - device=segmentation.device, - ) - # variant 1, 2x faster than 2 - result.scatter_(0, segmentation[None].long(), 1) # why does this have to be long!? - # variant 2, slower than 1 - # for i, l in enumerate(all_labels): - # result[i] = segmentation == l - else: - result = np.zeros((len(all_labels), *segmentation.shape), dtype=output_dtype if output_dtype is not None else np.uint8) - # variant 1, fastest in my testing - for i, l in enumerate(all_labels): - result[i] = segmentation == l - # variant 2. Takes about twice as long so nah - # result = np.eye(len(all_labels))[segmentation].transpose((3, 0, 1, 2)) - return result diff --git a/spineps/utils/default_preprocessor.py b/spineps/utils/default_preprocessor.py deleted file mode 100755 index 0b867c2..0000000 --- a/spineps/utils/default_preprocessor.py +++ /dev/null @@ -1,220 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -from typing import Union, Tuple - -import numpy as np - -# from acvl_utils.miscellaneous.ptqdm import ptqdm -from batchgenerators.utilities.file_and_folder_operations import * -from acvl_utils.cropping_and_padding.bounding_boxes import get_bbox_from_mask, bounding_box_to_slice - -import nnunetv2 -from nnunetv2.utilities.find_class_by_name import recursive_find_python_class - -# from nnunetv2.preprocessing.cropping.cropping import crop_to_nonzero -from spineps.utils.plans_handler import PlansManager, ConfigurationManager - - -class DefaultPreprocessor(object): - def __init__(self, verbose: bool = True): - self.verbose = verbose - """ - Everything we need is in the plans. Those are given when run() is called - """ - - def run_case_npy( - self, - data: np.ndarray, - seg: Union[np.ndarray, None], - properties: dict, - plans_manager: PlansManager, - configuration_manager: ConfigurationManager, - dataset_json: Union[dict, str], - ): - # let's not mess up the inputs! - data = np.copy(data) - if seg is not None: - seg = np.copy(seg) - - has_seg = seg is not None - - # apply transpose_forward, this also needs to be applied to the spacing! - data = data.transpose([0, *[i + 1 for i in plans_manager.transpose_forward]]) - if seg is not None: - seg = seg.transpose([0, *[i + 1 for i in plans_manager.transpose_forward]]) - original_spacing = [properties["spacing"][i] for i in plans_manager.transpose_forward] - - # crop, remember to store size before cropping! - shape_before_cropping = data.shape[1:] - properties["shape_before_cropping"] = shape_before_cropping - # this command will generate a segmentation. This is important because of the nonzero mask which we may need - data, seg, bbox = crop_to_nonzero(data, seg) - properties["bbox_used_for_cropping"] = bbox - # print(data.shape, seg.shape) - properties["shape_after_cropping_and_before_resampling"] = data.shape[1:] - - # resample - target_spacing = configuration_manager.spacing # this should already be transposed - - if len(target_spacing) < len(data.shape[1:]): - # target spacing for 2d has 2 entries but the data and original_spacing have three because everything is 3d - # in 3d we do not change the spacing between slices - target_spacing = [original_spacing[0]] + target_spacing - new_shape = compute_new_shape(data.shape[1:], original_spacing, target_spacing) - - # normalize - # normalization MUST happen before resampling or we get huge problems with resampled nonzero masks no - # longer fitting the images perfectly! - data = self._normalize(data, seg, configuration_manager, plans_manager.foreground_intensity_properties_per_channel) - - # print('current shape', data.shape[1:], 'current_spacing', original_spacing, - # '\ntarget shape', new_shape, 'target_spacing', target_spacing) - old_shape = data.shape[1:] - data = configuration_manager.resampling_fn_data(data, new_shape, original_spacing, target_spacing) - seg = configuration_manager.resampling_fn_seg(seg, new_shape, original_spacing, target_spacing) - if self.verbose: - print( - f"old shape: {old_shape}, new_shape: {new_shape}, old_spacing: {original_spacing}, " - f"new_spacing: {target_spacing}, fn_data: {configuration_manager.resampling_fn_data}" - ) - - # if we have a segmentation, sample foreground locations for oversampling and add those to properties - if has_seg: - # reinstantiating LabelManager for each case is not ideal. We could replace the dataset_json argument - # with a LabelManager Instance in this function because that's all its used for. Dunno what's better. - # LabelManager is pretty light computation-wise. - label_manager = plans_manager.get_label_manager(dataset_json) - collect_for_this = label_manager.foreground_regions if label_manager.has_regions else label_manager.foreground_labels - - # when using the ignore label we want to sample only from annotated regions. Therefore we also need to - # collect samples uniformly from all classes (incl background) - if label_manager.has_ignore_label: - collect_for_this.append(label_manager.all_labels) - - # no need to filter background in regions because it is already filtered in handle_labels - # print(all_labels, regions) - properties["class_locations"] = self._sample_foreground_locations(seg, collect_for_this, verbose=self.verbose) - seg = self.modify_seg_fn(seg, plans_manager, dataset_json, configuration_manager) - if np.max(seg) > 127: - seg = seg.astype(np.int16) - else: - seg = seg.astype(np.int8) - return data, seg - - @staticmethod - def _sample_foreground_locations( - seg: np.ndarray, classes_or_regions: Union[List[int], List[Tuple[int, ...]]], seed: int = 1234, verbose: bool = False - ): - num_samples = 10000 - min_percent_coverage = 0.01 # at least 1% of the class voxels need to be selected, otherwise it may be too - # sparse - rndst = np.random.RandomState(seed) - class_locs = {} - for c in classes_or_regions: - k = c if not isinstance(c, list) else tuple(c) - if isinstance(c, (tuple, list)): - mask = seg == c[0] - for cc in c[1:]: - mask = mask | (seg == cc) - all_locs = np.argwhere(mask) - else: - all_locs = np.argwhere(seg == c) - if len(all_locs) == 0: - class_locs[k] = [] - continue - target_num_samples = min(num_samples, len(all_locs)) - target_num_samples = max(target_num_samples, int(np.ceil(len(all_locs) * min_percent_coverage))) - - selected = all_locs[rndst.choice(len(all_locs), target_num_samples, replace=False)] - class_locs[k] = selected - if verbose: - print(c, target_num_samples) - return class_locs - - def _normalize( - self, - data: np.ndarray, - seg: np.ndarray, - configuration_manager: ConfigurationManager, - foreground_intensity_properties_per_channel: dict, - ) -> np.ndarray: - for c in range(data.shape[0]): - scheme = configuration_manager.normalization_schemes[c] - normalizer_class = recursive_find_python_class( - join(nnunetv2.__path__[0], "preprocessing", "normalization"), scheme, "nnunetv2.preprocessing.normalization" - ) - if normalizer_class is None: - raise RuntimeError("Unable to locate class '%s' for normalization" % scheme) - normalizer = normalizer_class( - use_mask_for_norm=configuration_manager.use_mask_for_norm[c], - intensityproperties=foreground_intensity_properties_per_channel[str(c)], - ) - data[c] = normalizer.run(data[c], seg[0]) - return data - - def modify_seg_fn( - self, seg: np.ndarray, plans_manager: PlansManager, dataset_json: dict, configuration_manager: ConfigurationManager - ) -> np.ndarray: - # this function will be called at the end of self.run_case. Can be used to change the segmentation - # after resampling. Useful for experimenting with sparse annotations: I can introduce sparsity after resampling - # and don't have to create a new dataset each time I modify my experiments - return seg - - -def compute_new_shape( - old_shape: Union[Tuple[int, ...], List[int], np.ndarray], - old_spacing: Union[Tuple[float, ...], List[float], np.ndarray], - new_spacing: Union[Tuple[float, ...], List[float], np.ndarray], -) -> np.ndarray: - assert len(old_spacing) == len(old_shape) - assert len(old_shape) == len(new_spacing) - new_shape = np.array([int(round(i / j * k)) for i, j, k in zip(old_spacing, new_spacing, old_shape)]) - return new_shape - - -def crop_to_nonzero(data, seg=None, nonzero_label=-1): - """ - - :param data: - :param seg: - :param nonzero_label: this will be written into the segmentation map - :return: - """ - nonzero_mask = create_nonzero_mask(data) - bbox = get_bbox_from_mask(nonzero_mask) - - slicer = bounding_box_to_slice(bbox) - data = data[tuple([slice(None), *slicer])] - - if seg is not None: - seg = seg[tuple([slice(None), *slicer])] - - nonzero_mask = nonzero_mask[slicer][None] - if seg is not None: - seg[(seg == 0) & (~nonzero_mask)] = nonzero_label - else: - nonzero_mask = nonzero_mask.astype(np.int8) - nonzero_mask[nonzero_mask == 0] = nonzero_label - nonzero_mask[nonzero_mask > 0] = 0 - seg = nonzero_mask - return data, seg, bbox - - -def create_nonzero_mask(data): - """ - - :param data: - :return: the mask is True where the data is nonzero - """ - from scipy.ndimage import binary_fill_holes - - assert len(data.shape) == 4 or len(data.shape) == 3, "data must have shape (C, X, Y, Z) or shape (C, X, Y)" - nonzero_mask = np.zeros(data.shape[1:], dtype=bool) - for c in range(data.shape[0]): - this_mask = data[c] != 0 - nonzero_mask = nonzero_mask | this_mask - nonzero_mask = binary_fill_holes(nonzero_mask) - return nonzero_mask diff --git a/spineps/utils/export_prediction.py b/spineps/utils/export_prediction.py deleted file mode 100755 index 9d73e15..0000000 --- a/spineps/utils/export_prediction.py +++ /dev/null @@ -1,129 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -import os -from copy import deepcopy -from typing import Union, List - -import numpy as np -import torch -from acvl_utils.cropping_and_padding.bounding_boxes import bounding_box_to_slice -from batchgenerators.utilities.file_and_folder_operations import ( - load_json, - isfile, - save_pickle, -) - -from nnunetv2.utilities.label_handling.label_handling import LabelManager -from spineps.utils.plans_handler import PlansManager, ConfigurationManager - - -def convert_predicted_logits_to_segmentation_with_correct_shape( - predicted_logits: Union[torch.Tensor, np.ndarray], - plans_manager: PlansManager, - configuration_manager: ConfigurationManager, - label_manager: LabelManager, - properties_dict: dict, - return_probabilities: bool = False, - num_threads_torch: int = 8, -): - old_threads = torch.get_num_threads() - torch.set_num_threads(num_threads_torch) - - # resample to original shape - current_spacing = ( - configuration_manager.spacing - if len(configuration_manager.spacing) == len(properties_dict["shape_after_cropping_and_before_resampling"]) - else [properties_dict["spacing"][0], *configuration_manager.spacing] - ) - predicted_logits = configuration_manager.resampling_fn_probabilities( - predicted_logits, - properties_dict["shape_after_cropping_and_before_resampling"], - current_spacing, - properties_dict["spacing"], - ) - # return value of resampling_fn_probabilities can be ndarray or Tensor but that doesnt matter because - # apply_inference_nonlin will covnert to torch - predicted_probabilities = label_manager.apply_inference_nonlin(predicted_logits) - del predicted_logits - segmentation = label_manager.convert_probabilities_to_segmentation(predicted_probabilities) - - # segmentation may be torch.Tensor but we continue with numpy - if isinstance(segmentation, torch.Tensor): - segmentation = segmentation.cpu().numpy() - - # put segmentation in bbox (revert cropping) - segmentation_reverted_cropping = np.zeros( - properties_dict["shape_before_cropping"], - dtype=np.uint8 if len(label_manager.foreground_labels) < 255 else np.uint16, - ) - slicer = bounding_box_to_slice(properties_dict["bbox_used_for_cropping"]) - segmentation_reverted_cropping[slicer] = segmentation - del segmentation - - # revert transpose - segmentation_reverted_cropping = segmentation_reverted_cropping.transpose(plans_manager.transpose_backward) - if return_probabilities: - # revert cropping - predicted_probabilities = label_manager.revert_cropping_on_probabilities( - predicted_probabilities, - properties_dict["bbox_used_for_cropping"], - properties_dict["shape_before_cropping"], - ) - if isinstance(predicted_probabilities, torch.Tensor): - predicted_probabilities = predicted_probabilities.cpu().numpy() - # revert transpose - predicted_probabilities = predicted_probabilities.transpose([0] + [i + 1 for i in plans_manager.transpose_backward]) - torch.set_num_threads(old_threads) - return segmentation_reverted_cropping, predicted_probabilities - else: - torch.set_num_threads(old_threads) - return segmentation_reverted_cropping - - -def export_prediction_from_logits( - predicted_array_or_file: Union[np.ndarray, torch.Tensor], - properties_dict: dict, - configuration_manager: ConfigurationManager, - plans_manager: PlansManager, - dataset_json_dict_or_file: Union[dict, str], - output_file_truncated: str, - save_probabilities: bool = False, -): - # if isinstance(predicted_array_or_file, str): - # tmp = deepcopy(predicted_array_or_file) - # if predicted_array_or_file.endswith('.npy'): - # predicted_array_or_file = np.load(predicted_array_or_file) - # elif predicted_array_or_file.endswith('.npz'): - # predicted_array_or_file = np.load(predicted_array_or_file)['softmax'] - # os.remove(tmp) - if isinstance(dataset_json_dict_or_file, str): - dataset_json_dict_or_file = load_json(dataset_json_dict_or_file) - - label_manager = plans_manager.get_label_manager(dataset_json_dict_or_file) - ret = convert_predicted_logits_to_segmentation_with_correct_shape( - predicted_array_or_file, - plans_manager, - configuration_manager, - label_manager, - properties_dict, - return_probabilities=save_probabilities, - ) - del predicted_array_or_file - - # save - if save_probabilities: - segmentation_final, probabilities_final = ret - np.savez_compressed(output_file_truncated + ".npz", probabilities=probabilities_final) - save_pickle(properties_dict, output_file_truncated + ".pkl") - del probabilities_final, ret - else: - segmentation_final = ret - del ret - - rw = plans_manager.image_reader_writer_class() - out_fname = output_file_truncated + dataset_json_dict_or_file["file_ending"] - rw.write_seg(segmentation_final, out_fname, properties_dict) - print(f"Saved Segmentation into {out_fname}") diff --git a/spineps/utils/get_network_from_plans.py b/spineps/utils/get_network_from_plans.py deleted file mode 100755 index 4ad20f8..0000000 --- a/spineps/utils/get_network_from_plans.py +++ /dev/null @@ -1,98 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -from dynamic_network_architectures.architectures.unet import PlainConvUNet, ResidualEncoderUNet -from dynamic_network_architectures.building_blocks.helper import get_matching_instancenorm, convert_dim_to_conv_op -from dynamic_network_architectures.initialization.weight_init import init_last_bn_before_add_to_0 -from nnunetv2.utilities.network_initialization import InitWeights_He -from nnunetv2.utilities.plans_handling.plans_handler import ConfigurationManager, PlansManager -from torch import nn - - -def get_network_from_plans( - plans_manager: PlansManager, - dataset_json: dict, - configuration_manager: ConfigurationManager, - num_input_channels: int, - num_output_channels: int, - deep_supervision: bool = True, -): - """ - we may have to change this in the future to accommodate other plans -> network mappings - - num_input_channels can differ depending on whether we do cascade. Its best to make this info available in the - trainer rather than inferring it again from the plans here. - """ - if "architecture" in configuration_manager.configuration: - from nnunetv2.utilities.get_network_from_plans import get_network_from_plans as plans - - class_name = configuration_manager.configuration["architecture"]["network_class_name"] - kwargs = configuration_manager.configuration["architecture"]["arch_kwargs"] - _kw_requires_import = configuration_manager.configuration["architecture"]["_kw_requires_import"] - - return plans(class_name, kwargs, _kw_requires_import, num_input_channels, num_output_channels, deep_supervision=deep_supervision) - num_stages = len(configuration_manager.conv_kernel_sizes) - - dim = len(configuration_manager.conv_kernel_sizes[0]) - conv_op = convert_dim_to_conv_op(dim) - - label_manager = plans_manager.get_label_manager(dataset_json) - - segmentation_network_class_name = configuration_manager.UNet_class_name - mapping = {"PlainConvUNet": PlainConvUNet, "ResidualEncoderUNet": ResidualEncoderUNet} - kwargs = { - "PlainConvUNet": { - "conv_bias": True, - "norm_op": get_matching_instancenorm(conv_op), - "norm_op_kwargs": {"eps": 1e-5, "affine": True}, - "dropout_op": None, - "dropout_op_kwargs": None, - "nonlin": nn.LeakyReLU, - "nonlin_kwargs": {"inplace": True}, - }, - "ResidualEncoderUNet": { - "conv_bias": True, - "norm_op": get_matching_instancenorm(conv_op), - "norm_op_kwargs": {"eps": 1e-5, "affine": True}, - "dropout_op": None, - "dropout_op_kwargs": None, - "nonlin": nn.LeakyReLU, - "nonlin_kwargs": {"inplace": True}, - }, - } - assert segmentation_network_class_name in mapping.keys(), ( - "The network architecture specified by the plans file " - "is non-standard (maybe your own?). Yo'll have to dive " - "into either this " - "function (get_network_from_plans) or " - "the init of your nnUNetModule to accomodate that." - ) - network_class = mapping[segmentation_network_class_name] - - conv_or_blocks_per_stage = { - ( - "n_conv_per_stage" if network_class != ResidualEncoderUNet else "n_blocks_per_stage" - ): configuration_manager.n_conv_per_stage_encoder, - "n_conv_per_stage_decoder": configuration_manager.n_conv_per_stage_decoder, - } - # network class name!! - model = network_class( - input_channels=num_input_channels, - n_stages=num_stages, - features_per_stage=[ - min(configuration_manager.UNet_base_num_features * 2**i, configuration_manager.unet_max_num_features) for i in range(num_stages) - ], - conv_op=conv_op, - kernel_sizes=configuration_manager.conv_kernel_sizes, - strides=configuration_manager.pool_op_kernel_sizes, - num_classes=label_manager.num_segmentation_heads, - deep_supervision=deep_supervision, - **conv_or_blocks_per_stage, - **kwargs[segmentation_network_class_name] - ) - model.apply(InitWeights_He(1e-2)) - if network_class == ResidualEncoderUNet: - model.apply(init_last_bn_before_add_to_0) - return model diff --git a/spineps/utils/plans_handler.py b/spineps/utils/plans_handler.py deleted file mode 100755 index 2db0045..0000000 --- a/spineps/utils/plans_handler.py +++ /dev/null @@ -1,331 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -import dynamic_network_architectures -from copy import deepcopy -from functools import lru_cache, partial -from typing import Union, Tuple, List, Type, Callable - -import numpy as np -import torch - -from nnunetv2.preprocessing.resampling.utils import recursive_find_resampling_fn_by_name -from torch import nn - -import nnunetv2 -from batchgenerators.utilities.file_and_folder_operations import load_json, join - -from nnunetv2.imageio.reader_writer_registry import recursive_find_reader_writer_by_name -from nnunetv2.utilities.find_class_by_name import recursive_find_python_class -from nnunetv2.utilities.label_handling.label_handling import get_labelmanager_class_from_plans - - -# see https://adamj.eu/tech/2021/05/13/python-type-hints-how-to-fix-circular-imports/ -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from nnunetv2.utilities.label_handling.label_handling import LabelManager - from nnunetv2.imageio.base_reader_writer import BaseReaderWriter - from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor - from nnunetv2.experiment_planning.experiment_planners.default_experiment_planner import ExperimentPlanner - - -class ConfigurationManager(object): - def __init__(self, configuration_dict: dict): - self.configuration = configuration_dict - - def __repr__(self): - return self.configuration.__repr__() - - @property - def data_identifier(self) -> str: - return self.configuration["data_identifier"] - - @property - def preprocessor_name(self) -> str: - return self.configuration["preprocessor_name"] - - @property - @lru_cache(maxsize=1) - def preprocessor_class(self) -> Type[DefaultPreprocessor]: - preprocessor_class = recursive_find_python_class( - join(nnunetv2.__path__[0], "preprocessing"), self.preprocessor_name, current_module="nnunetv2.preprocessing" - ) - return preprocessor_class - - @property - def batch_size(self) -> int: - return self.configuration["batch_size"] - - @property - def patch_size(self) -> List[int]: - return self.configuration["patch_size"] - - @property - def median_image_size_in_voxels(self) -> List[int]: - return self.configuration["median_image_size_in_voxels"] - - @property - def spacing(self) -> List[float]: - return self.configuration["spacing"] - - @property - def normalization_schemes(self) -> List[str]: - return self.configuration["normalization_schemes"] - - @property - def use_mask_for_norm(self) -> List[bool]: - return self.configuration["use_mask_for_norm"] - - @property - def UNet_class_name(self) -> str: - return self.configuration["UNet_class_name"] - - @property - @lru_cache(maxsize=1) - def UNet_class(self) -> Type[nn.Module]: - unet_class = recursive_find_python_class( - join(dynamic_network_architectures.__path__[0], "architectures"), - self.UNet_class_name, - current_module="dynamic_network_architectures.architectures", - ) - if unet_class is None: - raise RuntimeError( - "The network architecture specified by the plans file " - "is non-standard (maybe your own?). Fix this by not using " - "ConfigurationManager.UNet_class to instantiate " - "it (probably just overwrite build_network_architecture of your trainer." - ) - return unet_class - - @property - def UNet_base_num_features(self) -> int: - return self.configuration["UNet_base_num_features"] - - @property - def n_conv_per_stage_encoder(self) -> List[int]: - return self.configuration["n_conv_per_stage_encoder"] - - @property - def n_conv_per_stage_decoder(self) -> List[int]: - return self.configuration["n_conv_per_stage_decoder"] - - @property - def num_pool_per_axis(self) -> List[int]: - return self.configuration["num_pool_per_axis"] - - @property - def pool_op_kernel_sizes(self) -> List[List[int]]: - return self.configuration["pool_op_kernel_sizes"] - - @property - def conv_kernel_sizes(self) -> List[List[int]]: - for k in ["conv_kernel_sizes", "kernel_sizes"]: - if k in self.configuration: - return self.configuration[k] - return self.configuration["conv_kernel_sizes"] - - @property - def unet_max_num_features(self) -> int: - return self.configuration["unet_max_num_features"] - - @property - @lru_cache(maxsize=1) - def resampling_fn_data( - self, - ) -> Callable[ - [ - Union[torch.Tensor, np.ndarray], - Union[Tuple[int, ...], List[int], np.ndarray], - Union[Tuple[float, ...], List[float], np.ndarray], - Union[Tuple[float, ...], List[float], np.ndarray], - ], - Union[torch.Tensor, np.ndarray], - ]: - fn = recursive_find_resampling_fn_by_name(self.configuration["resampling_fn_data"]) - fn = partial(fn, **self.configuration["resampling_fn_data_kwargs"]) - return fn - - @property - @lru_cache(maxsize=1) - def resampling_fn_probabilities( - self, - ) -> Callable[ - [ - Union[torch.Tensor, np.ndarray], - Union[Tuple[int, ...], List[int], np.ndarray], - Union[Tuple[float, ...], List[float], np.ndarray], - Union[Tuple[float, ...], List[float], np.ndarray], - ], - Union[torch.Tensor, np.ndarray], - ]: - fn = recursive_find_resampling_fn_by_name(self.configuration["resampling_fn_probabilities"]) - fn = partial(fn, **self.configuration["resampling_fn_probabilities_kwargs"]) - return fn - - @property - @lru_cache(maxsize=1) - def resampling_fn_seg( - self, - ) -> Callable[ - [ - Union[torch.Tensor, np.ndarray], - Union[Tuple[int, ...], List[int], np.ndarray], - Union[Tuple[float, ...], List[float], np.ndarray], - Union[Tuple[float, ...], List[float], np.ndarray], - ], - Union[torch.Tensor, np.ndarray], - ]: - fn = recursive_find_resampling_fn_by_name(self.configuration["resampling_fn_seg"]) - fn = partial(fn, **self.configuration["resampling_fn_seg_kwargs"]) - return fn - - @property - def batch_dice(self) -> bool: - return self.configuration["batch_dice"] - - @property - def next_stage_names(self) -> Union[List[str], None]: - ret = self.configuration.get("next_stage") - if ret is not None: - if isinstance(ret, str): - ret = [ret] - return ret - - @property - def previous_stage_name(self) -> Union[str, None]: - return self.configuration.get("previous_stage") - - -class PlansManager(object): - def __init__(self, plans_file_or_dict: Union[str, dict]): - """ - Why do we need this? - 1) resolve inheritance in configurations - 2) expose otherwise annoying stuff like getting the label manager or IO class from a string - 3) clearly expose the things that are in the plans instead of hiding them in a dict - 4) cache shit - - This class does not prevent you from going wild. You can still use the plans directly if you prefer - (PlansHandler.plans['key']) - """ - self.plans = plans_file_or_dict if isinstance(plans_file_or_dict, dict) else load_json(plans_file_or_dict) - - def __repr__(self): - return self.plans.__repr__() - - def _internal_resolve_configuration_inheritance(self, configuration_name: str, visited: Tuple[str, ...] = None) -> dict: - if configuration_name not in self.plans["configurations"].keys(): - raise ValueError( - f"The configuration {configuration_name} does not exist in the plans I have. Valid " - f'configuration names are {list(self.plans["configurations"].keys())}.' - ) - configuration = deepcopy(self.plans["configurations"][configuration_name]) - if "inherits_from" in configuration: - parent_config_name = configuration["inherits_from"] - - if visited is None: - visited = (configuration_name,) - else: - if parent_config_name in visited: - raise RuntimeError( - f"Circular dependency detected. The following configurations were visited " - f"while solving inheritance (in that order!): {visited}. " - f"Current configuration: {configuration_name}. Its parent configuration " - f"is {parent_config_name}." - ) - visited = (*visited, configuration_name) - - base_config = self._internal_resolve_configuration_inheritance(parent_config_name, visited) - base_config.update(configuration) - configuration = base_config - return configuration - - @lru_cache(maxsize=10) - def get_configuration(self, configuration_name: str): - if configuration_name not in self.plans["configurations"].keys(): - raise RuntimeError( - f"Requested configuration {configuration_name} not found in plans. " - f"Available configurations: {list(self.plans['configurations'].keys())}" - ) - - configuration_dict = self._internal_resolve_configuration_inheritance(configuration_name) - return ConfigurationManager(configuration_dict) - - @property - def dataset_name(self) -> str: - return self.plans["dataset_name"] - - @property - def plans_name(self) -> str: - return self.plans["plans_name"] - - @property - def original_median_spacing_after_transp(self) -> List[float]: - return self.plans["original_median_spacing_after_transp"] - - @property - def original_median_shape_after_transp(self) -> List[float]: - return self.plans["original_median_shape_after_transp"] - - @property - @lru_cache(maxsize=1) - def image_reader_writer_class(self) -> Type[BaseReaderWriter]: - return recursive_find_reader_writer_by_name(self.plans["image_reader_writer"]) - - @property - def transpose_forward(self) -> List[int]: - return self.plans["transpose_forward"] - - @property - def transpose_backward(self) -> List[int]: - return self.plans["transpose_backward"] - - @property - def available_configurations(self) -> List[str]: - return list(self.plans["configurations"].keys()) - - @property - @lru_cache(maxsize=1) - def experiment_planner_class(self) -> Type[ExperimentPlanner]: - planner_name = self.experiment_planner_name - experiment_planner = recursive_find_python_class( - join(nnunetv2.__path__[0], "experiment_planning"), planner_name, current_module="nnunetv2.experiment_planning" - ) - return experiment_planner - - @property - def experiment_planner_name(self) -> str: - return self.plans["experiment_planner_used"] - - @property - @lru_cache(maxsize=1) - def label_manager_class(self) -> Type[LabelManager]: - return get_labelmanager_class_from_plans(self.plans) - - def get_label_manager(self, dataset_json: dict, **kwargs) -> LabelManager: - return self.label_manager_class( - label_dict=dataset_json["labels"], regions_class_order=dataset_json.get("regions_class_order"), **kwargs - ) - - @property - def foreground_intensity_properties_per_channel(self) -> dict: - if "foreground_intensity_properties_per_channel" not in self.plans.keys(): - if "foreground_intensity_properties_by_modality" in self.plans.keys(): - return self.plans["foreground_intensity_properties_by_modality"] - return self.plans["foreground_intensity_properties_per_channel"] - - -if __name__ == "__main__": - from nnunetv2.paths import nnUNet_preprocessed - from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name - - plans = load_json(join(nnUNet_preprocessed, maybe_convert_to_dataset_name(3), "nnUNetPlans.json")) - # build new configuration that inherits from 3d_fullres - plans["configurations"]["3d_fullres_bs4"] = {"batch_size": 4, "inherits_from": "3d_fullres"} - # now get plans and configuration managers - plans_manager = PlansManager(plans) - configuration_manager = plans_manager.get_configuration("3d_fullres_bs4") - print(configuration_manager) # look for batch size 4 diff --git a/spineps/utils/predictor.py b/spineps/utils/predictor.py deleted file mode 100755 index 38094ca..0000000 --- a/spineps/utils/predictor.py +++ /dev/null @@ -1,514 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -import os -import traceback -from typing import Tuple, Union - -import numpy as np -import torch -from acvl_utils.cropping_and_padding.padding import pad_nd_image -from batchgenerators.utilities.file_and_folder_operations import join, load_json -from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels -from torch._dynamo import OptimizedModule -from tqdm import tqdm - -from spineps.utils.data_iterators import PreprocessAdapterFromNpy -from spineps.utils.export_prediction import convert_predicted_logits_to_segmentation_with_correct_shape, export_prediction_from_logits -from spineps.utils.get_network_from_plans import get_network_from_plans -from spineps.utils.plans_handler import PlansManager -from spineps.utils.sliding_window_prediction import compute_gaussian, compute_steps_for_sliding_window - - -class nnUNetPredictor(object): - def __init__( - self, - tile_step_size: float = 0.5, - use_gaussian: bool = True, - use_mirroring: bool = True, - perform_everything_on_gpu: bool = True, - device: torch.device = torch.device("cuda"), - verbose: bool = False, - verbose_preprocessing: bool = False, - allow_tqdm: bool = True, - ): - self.verbose = verbose - self.verbose_preprocessing = verbose_preprocessing - self.allow_tqdm = allow_tqdm - - ( - self.plans_manager, - self.configuration_manager, - self.list_of_parameters, - self.network, - self.dataset_json, - self.trainer_name, - self.allowed_mirroring_axes, - self.label_manager, - ) = (None, None, None, None, None, None, None, None) - - self.tile_step_size = tile_step_size - self.use_gaussian = use_gaussian - self.use_mirroring = use_mirroring - if device.type == "cuda": - torch.backends.cudnn.benchmark = True - device = torch.device(type="cuda", index=0) # set the desired GPU with CUDA_VISIBLE_DEVICES! - if device.type != "cuda" and perform_everything_on_gpu: - print("perform_everything_on_gpu=True is only supported for cuda devices! Setting this to False") - perform_everything_on_gpu = False - self.device = device - self.perform_everything_on_gpu = perform_everything_on_gpu - - def initialize_from_trained_model_folder( - self, - model_training_output_dir: str, - use_folds: Union[Tuple[Union[int, str], ...], None], - checkpoint_name: str = "checkpoint_final.pth", - cache_state_dicts: bool = True, - ): - """ - This is used when making predictions with a trained model - """ - if use_folds is None: - use_folds = ("0", "1", "2", "3", "4") - - dataset_json = load_json(join(model_training_output_dir, "dataset.json")) - plans = load_json(join(model_training_output_dir, "plans.json")) - plans_manager = PlansManager(plans) - - if isinstance(use_folds, str): - use_folds = [use_folds] - - parameters = [] - for i, f in enumerate(use_folds): - f = int(f) if f != "all" else f - checkpoint = torch.load( - join(model_training_output_dir, f"fold_{f}", checkpoint_name), - map_location=torch.device("cpu"), - weights_only=False, - ) - if i == 0: - trainer_name = checkpoint["trainer_name"] - configuration_name = checkpoint["init_args"]["configuration"] - inference_allowed_mirroring_axes = ( - checkpoint["inference_allowed_mirroring_axes"] if "inference_allowed_mirroring_axes" in checkpoint.keys() else None - ) - - parameters.append(checkpoint["network_weights"]) - - configuration_manager = plans_manager.get_configuration(configuration_name) - # restore network - num_input_channels = determine_num_input_channels(plans_manager, configuration_manager, dataset_json) - num_output_channels = len(dataset_json["labels"]) - if "ignore" in dataset_json["labels"]: - num_output_channels -= 1 - # num_input_channels = 1 - network = get_network_from_plans( - plans_manager, - dataset_json, - configuration_manager, - num_input_channels, - num_output_channels, - deep_supervision=False, - ) - self.plans_manager = plans_manager - self.configuration_manager = configuration_manager - self.list_of_parameters = parameters # Lists of model folds - self.network = network - self.dataset_json = dataset_json - self.trainer_name = trainer_name - self.allowed_mirroring_axes = inference_allowed_mirroring_axes - self.label_manager = plans_manager.get_label_manager(dataset_json) - if ( - ("nnUNet_compile" in os.environ.keys()) - and (os.environ["nnUNet_compile"].lower() in ("true", "1", "t")) - and not isinstance(self.network, OptimizedModule) - ): - print("compiling network") - self.network = torch.compile(self.network) - - self.loaded_networks = [] - if cache_state_dicts: - for params in self.list_of_parameters: - if not isinstance(self.network, OptimizedModule): - self.network.load_state_dict(params) - else: - self.network._orig_mod.load_state_dict(params) - if self.device.type == "cuda": - self.network.cuda() - self.network.eval() - self.loaded_networks.append(self.network) - # print(type(self.loaded_networks[0])) - - def predict_single_npy_array( - self, - input_image: np.ndarray, - image_properties: dict, - save_or_return_probabilities: bool = False, - ): - """ - image_properties must only have a 'spacing' key! - """ - segmentation_previous_stage: np.ndarray = None # Was previously a parameter - output_file_truncated: str = None # previously a parameter - - ppa = PreprocessAdapterFromNpy( - [input_image], - [segmentation_previous_stage], - [image_properties], - [output_file_truncated], - self.plans_manager, - self.dataset_json, - self.configuration_manager, - num_threads_in_multithreaded=1, - verbose=self.verbose, - ) - if self.verbose: - print("preprocessing") - dct = next(ppa) - - if self.verbose: - print("predicting") - - predicted_logits = self.predict_logits_from_preprocessed_data(dct["data"]) - predicted_logits.cpu() - prediction_stacked = None - - if self.verbose: - print("resampling to original shape") - if output_file_truncated is not None: - export_prediction_from_logits( - predicted_logits, - dct["data_properites"], - self.configuration_manager, - self.plans_manager, - self.dataset_json, - output_file_truncated, - save_or_return_probabilities, - ) - else: - ret = convert_predicted_logits_to_segmentation_with_correct_shape( - predicted_logits, - self.plans_manager, - self.configuration_manager, - self.label_manager, - dct["data_properites"], - return_probabilities=save_or_return_probabilities, - ) - if False and prediction_stacked is not None: - seg_stacked = np.stack( - list( - [ - convert_predicted_logits_to_segmentation_with_correct_shape( - r, - self.plans_manager, - self.configuration_manager, - self.label_manager, - dct["data_properites"], - return_probabilities=False, - ) - for r in prediction_stacked - ] - ) - ) - - if save_or_return_probabilities: - return ret[0], ret[1] - else: - return ret # , seg_stacked - - def predict_logits_from_preprocessed_data(self, data: torch.Tensor) -> torch.Tensor: - """ - IMPORTANT! IF YOU ARE RUNNING THE CASCADE, THE SEGMENTATION FROM THE PREVIOUS STAGE MUST ALREADY BE STACKED ON - TOP OF THE IMAGE AS ONE-HOT REPRESENTATION! SEE PreprocessAdapter ON HOW THIS SHOULD BE DONE! - - RETURNED LOGITS HAVE THE SHAPE OF THE INPUT. THEY MUST BE CONVERTED BACK TO THE ORIGINAL IMAGE SIZE. - SEE convert_predicted_logits_to_segmentation_with_correct_shape - """ - # USED - - # we have some code duplication here but this allows us to run with perform_everything_on_gpu=True as - # default and not have the entire program crash in case of GPU out of memory. Neat. That should make - # things a lot faster for some datasets. - original_perform_everything_on_gpu = self.perform_everything_on_gpu - with torch.no_grad(): - prediction = None - # prediction_stacked = None # TODO REMOVE (unecessary time!) - if self.perform_everything_on_gpu: - try: - for idx, params in enumerate(self.list_of_parameters): - network = None - if self.loaded_networks is not None: - network = self.loaded_networks[idx] - # messing with state dict names... - elif not isinstance(self.network, OptimizedModule): - self.network.load_state_dict(params) - else: - self.network._orig_mod.load_state_dict(params) - # print(type(self.network)) - - if prediction is None: - prediction = self.predict_sliding_window_return_logits(data, network=network) - # prediction_stacked = [prediction.to("cpu")] - else: - new_prediction = self.predict_sliding_window_return_logits(data, network=network) - prediction += new_prediction - # prediction_stacked.append(new_prediction.to("cpu")) - - if len(self.list_of_parameters) > 1: - prediction /= len(self.list_of_parameters) - # empty_cache(self.device) - - except RuntimeError: - print( - "Prediction with perform_everything_on_gpu=True failed due to insufficient GPU memory. " - "Falling back to perform_everything_on_gpu=False. Not a big deal, just slower..." - ) - print("Error:") - traceback.print_exc() - prediction = None - self.perform_everything_on_gpu = False - - # CPU version - if prediction is None: - for idx, params in enumerate(self.list_of_parameters): - network = None - if self.loaded_networks is not None: - network = self.loaded_networks[idx] - # messing with state dict names... - elif not isinstance(self.network, OptimizedModule): - self.network.load_state_dict(params) - else: - self.network._orig_mod.load_state_dict(params) - - if prediction is None: - prediction = self.predict_sliding_window_return_logits(data, network=network) - # prediction_stacked = [prediction.to("cpu")] - else: - new_prediction = self.predict_sliding_window_return_logits(data, network=network) - prediction += new_prediction - # prediction_stacked.append(new_prediction.to("cpu")) - - if len(self.list_of_parameters) > 1: - prediction /= len(self.list_of_parameters) - - print("Prediction done, transferring to CPU if needed") if self.verbose else None - prediction = prediction.to("cpu") - self.perform_everything_on_gpu = original_perform_everything_on_gpu - return prediction # , torch.stack(prediction_stacked) - - def _internal_get_sliding_window_slicers(self, image_size: Tuple[int, ...]): - # USED - slicers = [] - if len(self.configuration_manager.patch_size) < len(image_size): - assert len(self.configuration_manager.patch_size) == len(image_size) - 1, ( - "if tile_size has less entries than image_size, " - "len(tile_size) " - "must be one shorter than len(image_size) " - "(only dimension " - "discrepancy of 1 allowed)." - ) - steps = compute_steps_for_sliding_window( - image_size[1:], - self.configuration_manager.patch_size, - self.tile_step_size, - ) - if self.verbose: - print( - f"n_steps {image_size[0] * len(steps[0]) * len(steps[1])}, image size is" - f" {image_size}, tile_size {self.configuration_manager.patch_size}, " - f"tile_step_size {self.tile_step_size}\nsteps:\n{steps}" - ) - for d in range(image_size[0]): - for sx in steps[0]: - for sy in steps[1]: - slicers.append( - tuple( - [ - slice(None), - d, - *[ - slice(si, si + ti) - for si, ti in zip( - (sx, sy), - self.configuration_manager.patch_size, - ) - ], - ] - ) - ) - else: - steps = compute_steps_for_sliding_window(image_size, self.configuration_manager.patch_size, self.tile_step_size) - if self.verbose: - print( - f"n_steps {np.prod([len(i) for i in steps])}, image size is {image_size}, tile_size {self.configuration_manager.patch_size}, " - f"tile_step_size {self.tile_step_size}\nsteps:\n{steps}" - ) - for sx in steps[0]: - for sy in steps[1]: - for sz in steps[2]: - slicers.append( - tuple( - [ - slice(None), - *[ - slice(si, si + ti) - for si, ti in zip( - (sx, sy, sz), - self.configuration_manager.patch_size, - ) - ], - ] - ) - ) - return slicers - - def _internal_maybe_mirror_and_predict(self, x: torch.Tensor, network) -> torch.Tensor: - # USED - mirror_axes = self.allowed_mirroring_axes if self.use_mirroring else None - prediction = network(x) - - if mirror_axes is not None: - # check for invalid numbers in mirror_axes - # x should be 5d for 3d images and 4d for 2d. so the max value of mirror_axes cannot exceed len(x.shape) - 3 - assert max(mirror_axes) <= len(x.shape) - 3, "mirror_axes does not match the dimension of the input!" - - num_predictons = 2 ** len(mirror_axes) - if 0 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2,))), (2,)) - if 1 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (3,))), (3,)) - if 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (4,))), (4,)) - if 0 in mirror_axes and 1 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2, 3))), (2, 3)) - if 0 in mirror_axes and 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2, 4))), (2, 4)) - if 1 in mirror_axes and 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (3, 4))), (3, 4)) - if 0 in mirror_axes and 1 in mirror_axes and 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2, 3, 4))), (2, 3, 4)) - prediction /= num_predictons - return prediction - - def predict_sliding_window_return_logits(self, input_image: torch.Tensor, network=None) -> Union[np.ndarray, torch.Tensor]: - # USED - assert isinstance(input_image, torch.Tensor) - if network is None: - network = self.network - network.eval() - network = network.to(self.device) - - # self.network = self.network.to(self.device) - # self.network.eval() - - # empty_cache(self.device) - - # Autocast is a little bitch. - # If the device_type is 'cpu' then it's slow as heck on some CPUs (no auto bfloat16 support detection) - # and needs to be disabled. - # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False - # is set. Whyyyyyyy. (this is why we don't make use of enabled=False) - # So autocast will only be active if we have a cuda device. - with torch.no_grad(): - with torch.autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): - assert len(input_image.shape) == 4, "input_image must be a 4D np.ndarray or torch.Tensor (c, x, y, z)" - - if self.verbose: - print(f"Input shape: {input_image.shape}") - if self.verbose: - print("step_size:", self.tile_step_size) - if self.verbose: - print( - "mirror_axes:", - self.allowed_mirroring_axes if self.use_mirroring else None, - ) - - # if input_image is smaller than tile_size we need to pad it to tile_size. - data, slicer_revert_padding = pad_nd_image( - input_image, - self.configuration_manager.patch_size, - "constant", - {"value": 0}, - True, - None, - ) - - slicers = self._internal_get_sliding_window_slicers(data.shape[1:]) - - precision = torch.half if self.perform_everything_on_gpu else torch.float32 - - # preallocate results and num_predictions - results_device = self.device if self.perform_everything_on_gpu else torch.device("cpu") - if self.verbose: - print("preallocating arrays") - try: - data = data.to(self.device, dtype=precision) - predicted_logits = torch.zeros( - (self.label_manager.num_segmentation_heads, *data.shape[1:]), - dtype=precision, - device=results_device, - ) - n_predictions = torch.zeros(data.shape[1:], dtype=precision, device=results_device) - if self.use_gaussian: - gaussian = compute_gaussian( - tuple(self.configuration_manager.patch_size), - sigma_scale=1.0 / 8, - value_scaling_factor=1000, - device=results_device, - ) - except RuntimeError: - # sometimes the stuff is too large for GPUs. In that case fall back to CPU - results_device = torch.device("cpu") - data = data.to(results_device, dtype=precision) - predicted_logits = torch.zeros( - (self.label_manager.num_segmentation_heads, *data.shape[1:]), - dtype=precision, - device=results_device, - ) - n_predictions = torch.zeros(data.shape[1:], dtype=precision, device=results_device) - if self.use_gaussian: - gaussian = compute_gaussian( - tuple(self.configuration_manager.patch_size), - sigma_scale=1.0 / 8, - value_scaling_factor=1000, - device=results_device, - ) - finally: - empty_cache(self.device) - - if self.verbose: - print("running prediction") - for sl in tqdm(slicers, disable=not self.allow_tqdm): - workon = data[sl][None] - workon = workon.to(self.device, non_blocking=False) - - prediction = self._internal_maybe_mirror_and_predict(workon, network=network)[0].to(results_device) - - predicted_logits[sl] += prediction * gaussian if self.use_gaussian else prediction - n_predictions[sl[1:]] += gaussian if self.use_gaussian else 1 - - predicted_logits /= n_predictions - # empty_cache(self.device) - return predicted_logits[tuple([slice(None), *slicer_revert_padding[1:]])] - - -def empty_cache(device: torch.device): - if device.type == "cuda": - torch.cuda.empty_cache() - elif device.type == "mps": - from torch import mps - - mps.empty_cache() - else: - pass - - -class dummy_context(object): - def __enter__(self): - pass - - def __exit__(self, exc_type, exc_val, exc_tb): - pass diff --git a/spineps/utils/proc_functions.py b/spineps/utils/proc_functions.py index 3987897..052c6f0 100755 --- a/spineps/utils/proc_functions.py +++ b/spineps/utils/proc_functions.py @@ -55,7 +55,7 @@ def n4_bias( mask[slices] = 1 mask_nii = nii.set_array(mask) mask_nii.seg = True - n4: NII = nii.n4_bias_field_correction(threshold=0, mask=mask_nii, spline_param=spline_param) + n4: NII = nii.n4_bias_field_correction(threshold=0, mask=mask_nii.to_ants(), spline_param=spline_param) if norm != -1: n4 *= norm / n4.max() if dtype2nii: diff --git a/spineps/utils/sliding_window_prediction.py b/spineps/utils/sliding_window_prediction.py deleted file mode 100755 index 575a275..0000000 --- a/spineps/utils/sliding_window_prediction.py +++ /dev/null @@ -1,75 +0,0 @@ -# Adapted from https://github.com/MIC-DKFZ/nnUNet -# Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring -# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. -from __future__ import annotations - -from functools import lru_cache - -import numpy as np -import torch -from typing import Union, Tuple, List -from acvl_utils.cropping_and_padding.padding import pad_nd_image -from scipy.ndimage import gaussian_filter - - -@lru_cache(maxsize=2) -def compute_gaussian( - tile_size: Union[Tuple[int, ...], List[int]], - sigma_scale: float = 1.0 / 8, - value_scaling_factor: float = 1, - dtype=torch.float16, - device=torch.device("cuda", 0), -) -> torch.Tensor: - tmp = np.zeros(tile_size) - center_coords = [i // 2 for i in tile_size] - sigmas = [i * sigma_scale for i in tile_size] - tmp[tuple(center_coords)] = 1 - gaussian_importance_map = gaussian_filter(tmp, sigmas, 0, mode="constant", cval=0) - - gaussian_importance_map = torch.from_numpy(gaussian_importance_map).type(dtype).to(device) - - gaussian_importance_map = gaussian_importance_map / torch.max(gaussian_importance_map) * value_scaling_factor - gaussian_importance_map = gaussian_importance_map.type(dtype) - - # gaussian_importance_map cannot be 0, otherwise we may end up with nans! - gaussian_importance_map[gaussian_importance_map == 0] = torch.min(gaussian_importance_map[gaussian_importance_map != 0]) - - return gaussian_importance_map - - -def compute_steps_for_sliding_window(image_size: Tuple[int, ...], tile_size: Tuple[int, ...], tile_step_size: float) -> List[List[int]]: - assert [i >= j for i, j in zip(image_size, tile_size)], "image size must be as large or larger than patch_size" - assert 0 < tile_step_size <= 1, "step_size must be larger than 0 and smaller or equal to 1" - - # our step width is patch_size*step_size at most, but can be narrower. For example if we have image size of - # 110, patch size of 64 and step_size of 0.5, then we want to make 3 steps starting at coordinate 0, 23, 46 - target_step_sizes_in_voxels = [i * tile_step_size for i in tile_size] - - num_steps = [int(np.ceil((i - k) / j)) + 1 for i, j, k in zip(image_size, target_step_sizes_in_voxels, tile_size)] - - steps = [] - for dim in range(len(tile_size)): - # the highest step value for this dimension is - max_step_value = image_size[dim] - tile_size[dim] - if num_steps[dim] > 1: - actual_step_size = max_step_value / (num_steps[dim] - 1) - else: - actual_step_size = 99999999999 # does not matter because there is only one step at 0 - - steps_here = [int(np.round(actual_step_size * i)) for i in range(num_steps[dim])] - - steps.append(steps_here) - - return steps - - -if __name__ == "__main__": - a = torch.rand((4, 2, 32, 23)) - a_npy = a.numpy() - - a_padded = pad_nd_image(a, new_shape=(48, 27)) - a_npy_padded = pad_nd_image(a_npy, new_shape=(48, 27)) - assert all([i == j for i, j in zip(a_padded.shape, (4, 2, 48, 27))]) - assert all([i == j for i, j in zip(a_npy_padded.shape, (4, 2, 48, 27))]) - assert np.all(a_padded.numpy() == a_npy_padded) - print("Test Done") diff --git a/unit_tests/test_filepaths.py b/unit_tests/test_filepaths.py index 81583df..715add0 100644 --- a/unit_tests/test_filepaths.py +++ b/unit_tests/test_filepaths.py @@ -34,10 +34,10 @@ class Test_filepaths(unittest.TestCase): def test_search_path_simple(self): package_path = Path(spineps.__file__).parent print(package_path) - predictor_search = search_path(package_path, query="**/predictor.py") + predictor_search = search_path(package_path, query="**/seg_model.py") print(predictor_search) self.assertTrue(len(predictor_search) == 1) - self.assertEqual(predictor_search[0], package_path.joinpath("utils", "predictor.py")) + self.assertEqual(predictor_search[0], package_path.joinpath("seg_model.py")) def test_search_path_multi(self): package_path = Path(spineps.__file__).parent From 948ee11fdf12a3ad9f5095e0112f2bac437e9ffe Mon Sep 17 00:00:00 2001 From: iback Date: Fri, 3 Jul 2026 06:05:38 +0000 Subject: [PATCH 28/29] cleanup and sync --- .github/workflows/linter.yml | 4 +- .github/workflows/python-publish.yml | 10 +-- .github/workflows/tests.yml | 2 +- .gitignore | 20 +----- README.md | 48 +++++++++++--- docs/api/architectures.md | 18 ++++++ docs/api/utils.md | 7 ++ docs/getting-started.md | 6 +- docs/index.md | 9 +-- docs/modules/phases.md | 12 ++++ pyproject.toml | 14 ++-- spineps/.vscode/settings.json | 6 -- spineps/api.py | 5 +- spineps/architectures/pl_unet.py | 4 -- spineps/architectures/unet3D.py | 1 - spineps/architectures_new/unet2D.py | 1 - spineps/entrypoint.py | 4 -- spineps/example/get_gpu.py | 3 - spineps/get_models.py | 14 +--- spineps/lab_model.py | 10 --- spineps/phase_instance.py | 15 ----- spineps/phase_labeling.py | 75 +++++++++++++++------ spineps/phase_post.py | 4 -- spineps/phase_semantic.py | 2 - spineps/seg_enums.py | 4 +- spineps/seg_model.py | 3 - spineps/seg_run.py | 6 -- spineps/seg_utils.py | 19 +----- spineps/utils/auto_download.py | 11 ++-- spineps/utils/find_min_cost_path.py | 11 ---- spineps/utils/proc_functions.py | 5 -- unit_tests/test_architectures.py | 25 ++++++- unit_tests/test_disc_labels.py | 13 +++- unit_tests/test_path.py | 97 ---------------------------- 34 files changed, 204 insertions(+), 284 deletions(-) delete mode 100755 spineps/.vscode/settings.json delete mode 100644 unit_tests/test_path.py diff --git a/.github/workflows/linter.yml b/.github/workflows/linter.yml index 3e91f82..7133174 100644 --- a/.github/workflows/linter.yml +++ b/.github/workflows/linter.yml @@ -10,8 +10,8 @@ jobs: ruff: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 - - uses: actions/setup-python@v3 + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 - run: pip install ruff - run: ruff check . #- uses: chartboost/ruff-action@v1 diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 6eaa179..cb479c1 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -13,10 +13,10 @@ jobs: runs-on: ubuntu-latest steps: - - name: Checkout Repository - uses: actions/checkout@v3 + - name: Checkout Repository + uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.11' - name: Install dependencies @@ -30,5 +30,5 @@ jobs: env: TWINE_USERNAME: __token__ TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} - run: | - twine upload dist/*.whl + run: | + twine upload dist/* diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f0f179c..b34983c 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -30,7 +30,7 @@ jobs: run: | python -m pip install --upgrade pip python -m pip install poetry - python -m pip install flake8 pytest + python -m pip install pytest - name: Install dependancies run: | python -m poetry install diff --git a/.gitignore b/.gitignore index 980c0d0..546e626 100755 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,6 @@ nnUNet/* .vscode/ spineps/models/* -*.png # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] @@ -109,8 +108,6 @@ ENV/ .mypy_cache/ # project directories -code/bids -code/bids/* data data/* output/A/ @@ -121,22 +118,8 @@ unet/out*/ *.nii *.png *.csv -FashionMNIST/ -MNIST/ .idea/ -anomaly_logs/* -anomaly_logs -forward_weighting_tests -forward_weighting_tests/* evaluation_logs/ -corner_detection_test/ -corner_detection_test/* -paper/ -paper/* -code/Unet3D/lightning_logs/ -code/Unet3D/lightning_logs/* -code/Unet3D/L_logs_archive/ -code/Unet3D/L_logs_archive/* # model checkpoints *.pth @@ -146,10 +129,9 @@ code/Unet3D/L_logs_archive/* unet/*_seg/ *.jpg -*.txt !Notes/ +*.txt weights/ *.gif -*.png *.npy lightning_logs/ diff --git a/README.md b/README.md index 0d89347..db87385 100644 --- a/README.md +++ b/README.md @@ -194,16 +194,36 @@ Processes a single nifty file, will create a derivatves folder next to the nifty | argument | explanation | | :--- | --------- | -| -i | Absolute path to the single nifty file (.nii.gz) to be processed | -| --model-semantic, -ms | The model used for the semantic segmentation | -| --model-instance, -mi | The model used for the vertebra instance segmentation | -| --model-labeling, -ml | The (optional) VERIDAH model used for vertebra labeling | +| --input, -i | Absolute path to the single nifty file (.nii.gz) to be processed (required) | +| --model-semantic, -ms | The model used for the semantic segmentation (required) | +| --model-instance, -mv, -mi | The model used for the vertebra instance segmentation (default: instance) | +| --model-labeling, -ml | The (optional) VERIDAH model used for vertebra labeling (default: t2w_labeling) | + +Plus the common processing options below, shared with `dataset` mode. Run `spineps sample -h` for the full list +with defaults. + +#### Common processing options (both `sample` and `dataset`) + +| argument | explanation | +| :--- | --------- | | --derivative-name, -dn | Name of the derivatives folder (default: derivatives_seg) | | --save-debug, -sd | Saves debug data and intermediate results in a separate folder (default: False) | +| --save-softmax-logits, -ssl | Saves an .npz of the semantic model's raw softmax logits (default: False) | +| --save-modelres-mask, -smrm | Also saves the semantic mask at the model's native resolution (default: False) | | --override-semantic, -os | Override existing seg-spine files (default: False) | | --override-instance, -oi | Override existing seg-vert files (default: False) | +| --override-postpair, -opp | Override existing cleaned/paired files (default: False) | | --override-ctd, -oc | Override existing centroid files (default: False) | +| --ignore-inference-compatibility, -iic | Don't skip inputs whose modality doesn't match the models (default: False) | +| --crop / --no-crop | Crop the input to the spine before semantic segmentation (default: on) | +| --n4 / --no-n4 | N4 bias field correction before semantic segmentation, MRI only (default: on) | +| --enforce-12-thoracic | Force the labeling model to predict exactly 12 thoracic vertebrae (default: False) | | --batch-size, -bs | Vertebra cutouts per batched forward pass; higher is faster but uses more GPU memory. Only affects GPU memory; host RAM usage in the instance phase scales with scan length/vertebra count instead (default: 4) | +| --amp | Run the instance model's forward pass under CUDA autocast, faster but may slightly change output (default: False) | +| --step-size | Semantic model sliding-window tile step size; larger is faster but less accurate (default: model's own setting) | +| --tta / --no-tta | Force test-time mirroring augmentation on/off for the semantic model (default: model's own setting) | +| --cpu | Run on CPU instead of GPU, much slower (default: False) | +| --run-cprofiler, -rcp | Runs a cProfiler over the entire run (default: False) | | --verbose, -v | Prints much more stuff, may fully clutter your terminal (default: False) | There are a lot more arguments, run `spineps sample -h` to see them. @@ -211,10 +231,11 @@ There are a lot more arguments, run `spineps sample -h` to see them. #### Example ```bash #T2w sagittal -spineps sample --ignore-bids-filter --ignore-inference-compatibility -i /path/sub-testsample_T2w.nii.gz --model-semantic t2w --model-instance instance +spineps sample --ignore-inference-compatibility -i /path/sub-testsample_T2w.nii.gz --model-semantic t2w --model-instance instance #T1w sagittal -spineps sample --ignore-bids-filter --ignore-inference-compatibility -i ~/path/sub-testsample_T1w.nii.gz --model-semantic t1w --model-instance instance +spineps sample --ignore-inference-compatibility -i ~/path/sub-testsample_T1w.nii.gz --model-semantic t1w --model-instance instance ``` +(`--ignore-bids-filter` is a `dataset`-only option — see below — it isn't accepted by `sample`.) ### Dataset @@ -249,16 +270,25 @@ Meaning you can have some key-value pairs (like `sub-`) in the name. Those k To that end, we are using TPTBox (see https://github.com/Hendrik-code/TPTBox) -It supports the same arguments as in sample mode (see table above), and additionally: | argument | explanation | | :--- | --------- | +| --directory, -i, -d | Absolute path to the dataset directory, preferably a BIDS dataset (required) | +| --model-semantic, -ms | The model used for the semantic segmentation, or `auto` to select automatically by modality (default: t2w) | +| --model-instance, -mv, -mi | The model used for the vertebra instance segmentation (default: instance) | +| --model-labeling, -ml | The (optional) VERIDAH model used for vertebra labeling (default: t2w_labeling) | | --rawdata-name, -rn | Sets the name of the rawdata folder of the dataset (default: "rawdata") | --ignore-bids-filter, -ibf | If true, will search the BIDS dataset without the strict filters. Use with care! (default: False) | | --ignore-model-compatibility, -imc | If true, will not stop the pipeline to use the given models on unfitting input modalities (default: False) | | --save-log, -sl | If true, saves the log into a separate folder in the dataset directory (default: False) | | --save-snaps-folder, -ssf | If true, additionally saves the snapshots in a separate folder in the dataset directory (default: False) | -For a full list of arguments, call `spineps dataset -h` +It also accepts all of the [common processing options](#single-nifty) listed above (`--batch-size`, `--crop`, +`--n4`, `--amp`, etc.). For a full list of arguments, call `spineps dataset -h`. + +#### Example +```bash +spineps dataset --ignore-bids-filter -i /path/to/dataset-folder --model-semantic t2w --model-instance instance +``` ## Segmentation @@ -296,7 +326,7 @@ In the subregion segmentation: CT only | Label | Structure | | :---: | --------- | -| 51 | Dense | +| 51 | Dens_axis (odontoid process of C2) | | 70 | Sacrum_Sacral_Ala_Left | | 71 | Sacrum_Sacral_Ala_Right | | 72 | Sacrum_Posterior_Sacral_Elements | diff --git a/docs/api/architectures.md b/docs/api/architectures.md index f6cf9f4..6b9923d 100644 --- a/docs/api/architectures.md +++ b/docs/api/architectures.md @@ -17,3 +17,21 @@ Network architectures and the vertebra label definitions used by the models. ## spineps.architectures.unet3D ::: spineps.architectures.unet3D + +## spineps.architectures_new.pl_unet + +The newer PyTorch Lightning U-Net wrapper (2D or 3D), used by [`spineps.seg_model`](models.md). + +::: spineps.architectures_new.pl_unet + +## spineps.architectures_new.unet3D + +::: spineps.architectures_new.unet3D + +## spineps.architectures_new.unet2D + +::: spineps.architectures_new.unet2D + +## spineps.architectures_new.dice + +::: spineps.architectures_new.dice diff --git a/docs/api/utils.md b/docs/api/utils.md index 866b5de..6df1055 100644 --- a/docs/api/utils.md +++ b/docs/api/utils.md @@ -2,6 +2,13 @@ Image processing, the vertebra-labeling path solver, disc labeling and other helpers. +## spineps.utils.resolution + +Resolution-aware conversion of physical (mm) thresholds to voxels, used throughout the processing phases so +behaviour is consistent across MRI and CT resolutions. + +::: spineps.utils.resolution + ## spineps.utils.proc_functions ::: spineps.utils.proc_functions diff --git a/docs/getting-started.md b/docs/getting-started.md index 7065763..54cef55 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -77,14 +77,16 @@ Segment a single scan: ```bash # T2w sagittal -spineps sample --ignore-bids-filter --ignore-inference-compatibility \ +spineps sample --ignore-inference-compatibility \ -i /path/sub-testsample_T2w.nii.gz --model-semantic t2w --model-instance instance # T1w sagittal -spineps sample --ignore-bids-filter --ignore-inference-compatibility \ +spineps sample --ignore-inference-compatibility \ -i /path/sub-testsample_T1w.nii.gz --model-semantic t1w --model-instance instance ``` +(`--ignore-bids-filter` is a `dataset`-only option, not accepted by `sample`.) + Process a whole [BIDS](https://bids-specification.readthedocs.io/en/stable/) dataset: ```bash diff --git a/docs/index.md b/docs/index.md index ffa3449..90dce15 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,9 +1,10 @@ # SPINEPS -**SPINEPS** is a framework for out-of-the-box **whole-spine segmentation of MR images**. It segments the -spine in sagittal MR images (T2w, T1w and others) using a two-phase approach to multi-class **semantic** -and **instance** segmentation, and can additionally assign anatomical **vertebra labels** via the -**VERIDAH** labeling model. +**SPINEPS** is a framework for out-of-the-box **whole-spine segmentation of MR and CT images**. It segments +the spine in sagittal MR images (T2w, T1w and others) and CT scans — each with independent, +modality-specific models — using a two-phase approach to multi-class **semantic** and **instance** +segmentation, and can additionally assign anatomical **vertebra labels** via the **VERIDAH** labeling +model. [![Paper](https://img.shields.io/badge/Paper-10.1007-blue)](https://link.springer.com/article/10.1007/s00330-024-11155-y) [![PyPI version](https://badge.fury.io/py/spineps.svg)](https://pypi.python.org/pypi/spineps/) diff --git a/docs/modules/phases.md b/docs/modules/phases.md index 2bc5fa7..4cf828a 100644 --- a/docs/modules/phases.md +++ b/docs/modules/phases.md @@ -34,6 +34,18 @@ The semantic mask uses these labels: | 100 | Vertebra_Disc | | 26 | Sacrum | +CT models additionally use these labels: + +| Label | Structure | +| :---: | --------------------------------- | +| 51 | Dens_axis (odontoid process of C2) | +| 70 | Sacrum_Sacral_Ala_Left | +| 71 | Sacrum_Sacral_Ala_Right | +| 72 | Sacrum_Posterior_Sacral_Elements | +| 73 | Sacrum_Body | +| 74 | Sacrum_Endplate | +| 80 | Metal | + ## Instance phase — [`spineps.phase_instance`](../api/phases.md) Derives a per-vertebra instance mask from the vertebra subregions. For each corpus center of mass it diff --git a/pyproject.toml b/pyproject.toml index 602f156..ea1db17 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,8 @@ build-backend = "poetry_dynamic_versioning.backend" [tool.poetry] name = "SPINEPS" -version = "1.0.0" +# Fallback only: poetry-dynamic-versioning derives the real version from git tags at build time. +version = "2.0.0" description = "Framework for out-of-the box whole spine MRI segmentation." authors = ["Hendrik Möller "] repository = "https://github.com/Hendrik-code/spineps" @@ -12,7 +13,7 @@ homepage = "https://github.com/Hendrik-code/spineps" documentation = "https://spineps.readthedocs.io" license = "Apache License Version 2.0, January 2004" readme = "README.md" -exclude = ["models", "examples"] +exclude = ["spineps/models", "spineps/example"] [tool.poetry.scripts] spineps = 'spineps.entrypoint:entry_point' @@ -30,6 +31,8 @@ antspyx = "0.6.3" rich = "^13.6.0" monai="^1.3.0" TypeSaveArgParse="^1.0.1" +# nnunetv2/acvl-utils/blosc2/python-gdcm are transitive dependencies of TPTBox's nnU-Net inference path +# (TPTBox.segmentation.nnUnet_utils.inference_api); pinned directly here to control the resolved versions. python-gdcm = [ { version = "==3.0.25", python = ">=3.9,<3.10" }, ] @@ -50,12 +53,6 @@ pytest = "^7.4.4" pre-commit = "*" coverage = ">=7.0.1" pytest-mock = "^3.6.0" -pandas = "^2.1.0" -joblib = "^1.3.2" -future = "^0.18.3" -flake8 = ">=4.0.1" -#auxiliary = ">=0.1.0" -tqdm = ">=4.62.3" [tool.poetry.group.docs] @@ -95,7 +92,6 @@ exclude = [ "dist", "node_modules", "venv", - "spineps/Unet3D", "spineps/utils/image.py", # vendored from spinalcordtoolbox ".toml", ] diff --git a/spineps/.vscode/settings.json b/spineps/.vscode/settings.json deleted file mode 100755 index d99f2f3..0000000 --- a/spineps/.vscode/settings.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "[python]": { - "editor.defaultFormatter": "ms-python.black-formatter" - }, - "python.formatting.provider": "none" -} \ No newline at end of file diff --git a/spineps/api.py b/spineps/api.py index bbdc003..ab431ca 100644 --- a/spineps/api.py +++ b/spineps/api.py @@ -20,7 +20,6 @@ from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Union from TPTBox import BIDS_FILE, NII, POI @@ -32,8 +31,8 @@ from spineps.seg_run import segment_image # Accepted forms for an input image and a model argument. -ImageInput = Union[str, Path, NII, BIDS_FILE] -ModelInput = Union[str, Path, SegmentationModel, VertLabelingClassifier] +ImageInput = str | Path | NII | BIDS_FILE +ModelInput = str | Path | SegmentationModel | VertLabelingClassifier @dataclass diff --git a/spineps/architectures/pl_unet.py b/spineps/architectures/pl_unet.py index 141cc6d..85ff6da 100755 --- a/spineps/architectures/pl_unet.py +++ b/spineps/architectures/pl_unet.py @@ -36,10 +36,6 @@ def __init__(self, opt: Any = None, do2D: bool = False, *args: Any, **kwargs: An dim_mults = (1, 2, 4, 8) dim = 16 # 16 - # if opt.high_res: - # dim = 16 - # dim_mults = (2, 4, 8, 8) - self.network = nclass( dim=dim, dim_mults=dim_mults, diff --git a/spineps/architectures/unet3D.py b/spineps/architectures/unet3D.py index 203f90a..7a7b018 100755 --- a/spineps/architectures/unet3D.py +++ b/spineps/architectures/unet3D.py @@ -134,7 +134,6 @@ def forward( if self.first_forward: print("|", x.shape) - # time = None if time is None: time = torch.ones((1,), device=x.device) x = self.init_conv(x) diff --git a/spineps/architectures_new/unet2D.py b/spineps/architectures_new/unet2D.py index 61021d3..4d84da6 100644 --- a/spineps/architectures_new/unet2D.py +++ b/spineps/architectures_new/unet2D.py @@ -400,7 +400,6 @@ def __init__( self.channels = channels init_dim = default(init_dim, dim) - # print(init_dim, channels) self.init_conv = nn.Conv2d((channels + conditional_dimensions) * patch_size * patch_size, init_dim, 7, padding=3) dims = [init_dim, *(int(dim * m) for m in dim_mults)] diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index 6799061..cc1b30a 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -290,8 +290,6 @@ def run_sample(opt: Namespace): "model_instance": model_instance, "model_labeling": model_labeling, "derivative_name": opt.derivative_name, - # - # "save_uncertainty_image": opt.save_unc_img, "save_softmax_logits": opt.save_softmax_logits, "save_debug_data": opt.save_debug, "save_modelres_mask": opt.save_modelres_mask, @@ -388,8 +386,6 @@ def run_dataset(opt: Namespace): "model_labeling": model_labeling, "rawdata_name": opt.rawdata_name, "derivative_name": opt.derivative_name, - # - # "save_uncertainty_image": opt.save_unc_img, "save_modelres_mask": opt.save_modelres_mask, "save_softmax_logits": opt.save_softmax_logits, "save_debug_data": opt.save_debug, diff --git a/spineps/example/get_gpu.py b/spineps/example/get_gpu.py index cfe1a7a..5c18504 100644 --- a/spineps/example/get_gpu.py +++ b/spineps/example/get_gpu.py @@ -62,16 +62,13 @@ def get_free_gpus(blocked_gpus=None, max_load: float = 0.3, max_memory: float = Returns: list[int]: IDs of GPUs that are consistently available and not blocked. """ - # print("get_free_gpus") if blocked_gpus is None: blocked_gpus = {0: False, 1: False, 2: False, 3: False} cached_list = get_gpu(max_load=max_load, max_memory=max_memory) for _ in range(15): time.sleep(0.25) cached_list = intersection(cached_list, get_gpu()) - # print("result:", list(cached_list)) gpulist = [i for i in list(cached_list) if i not in blocked_gpus or blocked_gpus[i] is False] - # print("result:", gpulist) return gpulist diff --git a/spineps/get_models.py b/spineps/get_models.py index f127a92..cbb11bd 100755 --- a/spineps/get_models.py +++ b/spineps/get_models.py @@ -4,7 +4,6 @@ import os from pathlib import Path -from typing import Optional, Union from TPTBox import Log_Type, No_Logger from tqdm import tqdm @@ -107,9 +106,9 @@ def get_labeling_model(model_name: str, **kwargs) -> VertLabelingClassifier: return _get_model_by_name(model_name, modelid2folder_labeling(), SpinepsPhase.LABELING, "labeling", **kwargs) -_modelid2folder_semantic: Optional[dict[str, Union[Path, str]]] = None -_modelid2folder_instance: Optional[dict[str, Union[Path, str]]] = None -_modelid2folder_labeling: Optional[dict[str, Union[Path, str]]] = None +_modelid2folder_semantic: dict[str, Path | str] | None = None +_modelid2folder_instance: dict[str, Path | str] | None = None +_modelid2folder_labeling: dict[str, Path | str] | None = None def modelid2folder_semantic() -> dict[str, Path | str]: @@ -248,10 +247,6 @@ def get_actual_model( FileNotFoundError: If no inference_config.json is found in the given folder. AssertionError: If more than one inference_config.json is found in the given folder. """ - # if isinstance(in_config, MODELS): - # in_dir = filepath_model(in_config.value, model_dir=None) - # else: - in_dir = in_config if os.path.isdir(str(in_dir)): # noqa: PTH112 @@ -267,9 +262,6 @@ def get_actual_model( f"get_actual_model: found more than one inference_config.json in {in_dir}/**/*inference_config.json. Ambiguous behavior, please manually correct this by removing one of these.\nFound {path_search}" ) in_dir = path_search[0] - # else: - # base = filepath_model(in_config, model_dir=None) - # in_dir = base inference_config = load_inference_config(str(in_dir)) modeltype: type[SegmentationModel] = modeltype2class(inference_config.modeltype) diff --git a/spineps/lab_model.py b/spineps/lab_model.py index a14df6a..925f5c5 100755 --- a/spineps/lab_model.py +++ b/spineps/lab_model.py @@ -112,7 +112,6 @@ def __init__( """ super().__init__(model_folder, inference_config, use_cpu, default_verbose, default_allow_tqdm) assert len(self.inference_config.expected_inputs) == 1, "Unet3D cannot expect more than one input" - # self.model: PLClassifier = model self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.final_size: tuple[int, int, int] = DEFAULT_CLASSIFIER_INPUT_SIZE self.totensor = ToTensor() @@ -192,7 +191,6 @@ def from_modelfolder(cls, model_folder: str | Path): NotImplementedError: Always; use from_checkpoint_path instead. """ raise NotImplementedError() - # find checkpoint yourself, then load from checkpoitn path @classmethod def from_checkpoint_path(cls, checkpoint_path: str | Path) -> VertLabelingClassifier: @@ -212,9 +210,6 @@ def from_checkpoint_path(cls, checkpoint_path: str | Path) -> VertLabelingClassi if isinstance(checkpoint_path, str): checkpoint_path = Path(checkpoint_path) assert checkpoint_path.exists(), f"Checkpoint path does not exist: {checkpoint_path}" - # model = PLClassifier.load_from_checkpoint( - # str(checkpoint_path), - # ) d = cls(checkpoint_path.parent.parent) logger.print("Model loaded from", checkpoint_path, verbose=True) return d @@ -339,7 +334,6 @@ def run_given_center_pos( img_v = img.set_array(arr_cut).reorient_(("I", "P", "L")) seg_v = seg.set_array(sem_cut).reorient_(("I", "P", "L")) - # angle = 0 if angle is not None and angle != 0: arr_cut = rotate_patch_sagitally(img_v.get_array(), -angle, msk=False) sem_cut = rotate_patch_sagitally(seg_v.get_seg_array(), -angle, msk=True) @@ -358,8 +352,6 @@ def run_given_center_pos( img_v.set_array_(arr_cut).reorient_(ori) seg_v.set_array_(sem_cut).reorient_(ori) - # img_v.save("/DATA/NAS/ongoing_projects/hendrik/img_v.nii.gz") - # seg_v.save("/DATA/NAS/ongoing_projects/hendrik/seg_v.nii.gz") return self._run_array(img_v.get_array(), seg_v.get_seg_array()) # sem_cut def _run_nii(self, img_nii: NII): @@ -426,9 +418,7 @@ def _run_array(self, img_arr: np.ndarray, seg_arr: np.ndarray | None | torch.Ten # TODO seg channelwise and stuff model_input = d["img"] - # print(model_input.shape) model_input.unsqueeze_(0) - # print(model_input.shape) model_input = model_input.to(torch.float32) model_input = model_input.to(self.device) diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index a9dcb99..722df9a 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -105,11 +105,9 @@ def predict_instance_mask( # Padding? if pad_size > 0: - # logger.print(seg_nii_rdy.shape) arr = seg_nii_rdy.get_array() arr = np.pad(arr, pad_size, mode="edge") seg_nii_rdy.set_array_(arr) - # logger.print(seg_nii_rdy.shape) zms = seg_nii_rdy.zoom logger.print("zms", zms, verbose=verbose) @@ -124,10 +122,8 @@ def predict_instance_mask( uncropped_vert_mask = np.zeros(seg_nii_uncropped.shape, dtype=seg_nii_uncropped.dtype) logger.print("Vertebra uncropped_vert_mask empty", uncropped_vert_mask.shape, verbose=verbose) crop = seg_nii_rdy.compute_crop(dist=INSTANCE_CROP_MARGIN_MM / min(seg_nii_rdy.zoom)) - # logger.print("Crop", crop, verbose=verbose) seg_nii_rdy.apply_crop_(crop) logger.print(f"Crop down from {uncropped_vert_mask.shape} to {seg_nii_rdy.shape}", verbose=verbose) - # arr[crop] = X, then set nifty to arr logger.print("Vertebra seg_nii_rdy", seg_nii_rdy.zoom, seg_nii_rdy.orientation, seg_nii_rdy.shape, verbose=verbose) debug_data["inst_cropped_Subreg_nii_b"] = seg_nii_rdy.copy() # @@ -203,11 +199,9 @@ def predict_instance_mask( # Uncrop again if pad_size > 0: - # logger.print(whole_vert_nii_uncropped.shape) arr = whole_vert_nii_uncropped.get_array() arr = arr[pad_size:-pad_size, pad_size:-pad_size, pad_size:-pad_size] whole_vert_nii_uncropped.set_array_(arr) - # logger.print(whole_vert_nii_uncropped.shape) return whole_vert_nii_uncropped, ErrCode.OK @@ -650,13 +644,11 @@ def collect_vertebra_predictions( # Holds only binary {0, 1} per-label masks, so uint8 is sufficient (the source dtype can be wider, # which would needlessly inflate this n_coms x 3 x volume array and slow the Dice comparisons below). hierarchical_predictions = np.zeros((n_corpus_coms, 3, *shp), dtype=np.uint8) - # print("hierarchical_predictions", hierarchical_predictions.shape) # relabel to the labels expected by the model # {41: 1, 42: 2, 43: 3, 44: 4, 45: 5, 46: 6, 47: 7, 48: 8, 49: 9, 50: 9, Location.Dens_axis.value: 9} mapping = {int(a): int(b) for a, b in model.inference_config.mapping.items()} seg_nii_for_cut: NII = seg_nii.copy().extract_label(list(mapping.keys()), keep_label=True).map_labels_(mapping, verbose=False) - # print("seg_nii_for_cut", seg_nii_for_cut.shape) logger.print("Vertebra collect in", seg_nii.zoom, seg_nii.orientation, seg_nii.shape, verbose=verbose) @@ -948,7 +940,6 @@ def find_prediction_couple( agreement = 0 if len(couple) > 0: - # print(couple) agreement = 0 for c in couple: agreement += dices[c] @@ -1000,7 +991,6 @@ def merge_coupled_predictions( combine.fill(0) for cid in k: combine += hierarchical_predictions[cid[0]][cid[1]] - # print(combine.shape) m = 1 if take_no_overlap else 2 # m = min(max(1, np.max(combine)), 2) # type:ignore combine[combine < m] = 0 @@ -1034,10 +1024,5 @@ def merge_coupled_predictions( logger=logger, verbose=verbose, ) - # print("whole_vert_arr", whole_vert_arr.shape) - # print("seg_nii", seg_nii.shape) whole_vert_nii_proc = seg_nii.set_array(whole_vert_arr) - # print("whole_vert_arr_proc", whole_vert_arr_proc.shape) - # debug_data["whole_vert_arr_proc"] = seg_nii.set_array(whole_vert_arr) - # return seg_nii.set_array(whole_vert_arr, verbose=False).map_labels_(com_map, verbose=False), debug_data, ErrCode.OK return whole_vert_nii_proc, debug_data, ErrCode.OK diff --git a/spineps/phase_labeling.py b/spineps/phase_labeling.py index 5c8f5b3..f253521 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -481,6 +481,47 @@ def prepare_vertrel(vertrel_softmax_values: np.ndarray, gaussian_sigma: float = return softmax_values +def _print_labeling_weights( + predict_keys, + visible_w, + vert_w, + region_w, + vertrel_w, + vertgrp_w, + vertt13_w, + disable_c1, + boost_c2, + allow_cervical_skip, + region_gaussian_sigma, + vert_gaussian_sigma, + vert_gaussian_regionwise, + vertrel_gaussian_sigma, +) -> None: + """Prints the per-objective weights and gaussian-smoothing settings used by ``find_vert_path_from_predictions``.""" + if "FULLYVISIBLE" in predict_keys: + print("visible_w", visible_w) + if "VERT" in predict_keys: + print("vert_w", vert_w) + if "REGION" in predict_keys: + print("region_w", region_w) + if "VERTREL" in predict_keys: + print("vertrel_w", vertrel_w) + if "VERTGRP" in predict_keys: + print("vertgrp_w", vertgrp_w) + if "VERTT13" in predict_keys: + print("vertt13_w", vertt13_w) + print("disable_c1", disable_c1) + print("boost_c2", boost_c2) + print("allow_cervical_skip", allow_cervical_skip) + if "VERTREGION" in predict_keys: + print("region_gaussian_sigma", region_gaussian_sigma) + if "VERT" in predict_keys: + print("vert_gaussian_sigma", vert_gaussian_sigma) + print("vert_gaussian_regionwise", vert_gaussian_regionwise) + if "VERTREL" in predict_keys: + print("vertrel_gaussian_sigma", vertrel_gaussian_sigma) + + def find_vert_path_from_predictions( predictions, visible_w: float = 0.5, @@ -576,7 +617,6 @@ def find_vert_path_from_predictions( cost_matrix = np.zeros((n_vert, VERT_CLASSES)) relative_cost_matrix = np.zeros((n_vert, len(VertRel))) visible_chain = prepare_visible(predictions, visible_w) - # print(visible_chain) predict_keys = list(predictions[list(predictions.keys())[0]]["soft"].keys()) # noqa: RUF015 assert "VERT" in predict_keys or "VERTEXACT" in predict_keys or "VERTEX" in predict_keys or "VERTGRP" in predict_keys, ( @@ -601,21 +641,23 @@ def find_vert_path_from_predictions( ) if verbose: - print("visible_w", visible_w) if "FULLYVISIBLE" in predict_keys else None - print("vert_w", vert_w) if "VERT" in predict_keys else None - print("region_w", region_w) if "REGION" in predict_keys else None - print("vertrel_w", vertrel_w) if "VERTREL" in predict_keys else None - print("vertgrp_w", vertgrp_w) if "VERTGRP" in predict_keys else None - print("vertt13_w", vertt13_w) if "VERTT13" in predict_keys else None - print("disable_c1", disable_c1) - print("boost_c2", boost_c2) - print("allow_cervical_skip", allow_cervical_skip) - print("region_gaussian_sigma", region_gaussian_sigma) if "VERTREGION" in predict_keys else None - print("vert_gaussian_sigma", vert_gaussian_sigma) if "VERT" in predict_keys else None - print("vert_gaussian_regionwise", vert_gaussian_regionwise) if "VERT" in predict_keys else None - print("vertrel_gaussian_sigma", vertrel_gaussian_sigma) if "VERTREL" in predict_keys else None + _print_labeling_weights( + predict_keys, + visible_w=visible_w, + vert_w=vert_w, + region_w=region_w, + vertrel_w=vertrel_w, + vertgrp_w=vertgrp_w, + vertt13_w=vertt13_w, + disable_c1=disable_c1, + boost_c2=boost_c2, + allow_cervical_skip=allow_cervical_skip, + region_gaussian_sigma=region_gaussian_sigma, + vert_gaussian_sigma=vert_gaussian_sigma, + vert_gaussian_regionwise=vert_gaussian_regionwise, + vertrel_gaussian_sigma=vertrel_gaussian_sigma, + ) - # for idx, (_, k) in enumerate(predictions.items()): vert_softmax_output = k["soft"]["VERT"] if "VERT" in predict_keys else np.zeros(len(VertExact)) vert_values = np.multiply( @@ -689,9 +731,6 @@ def find_vert_path_from_predictions( cost_matrix = np.asarray(cost_matrix) # invert rel cost relative_cost_matrix = np.multiply(-relative_cost_matrix, vertrel_w) - # for i in range(len(relative_cost_matrix)): - # print(relative_cost_matrix[i]) - # if argmax_combined_cost_matrix_instead_of_path_algorithm: fcost = 0 min_costs_path = [[]] diff --git a/spineps/phase_post.py b/spineps/phase_post.py index e3485c8..bd6143f 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -225,8 +225,6 @@ def mask_cleaning_other( # make copy where both masks clean each other vert_arr_cleaned = whole_vert_nii.get_seg_array() subreg_vert_arr = subreg_vert_nii.get_seg_array() - # if dilation_fill: - # vert_arr_cleaned = np_dilate_msk(vert_arr_cleaned, label_ref=vert_labels, mm=5) # , mask=subreg_vert_arr subreg_arr = seg_nii.get_seg_array() if proc_assign_missing_cc: @@ -664,8 +662,6 @@ def assign_vertebra_inconsistency( biggest_volume = (k_keys_sorted[0], gt_volume[k_keys_sorted[0]]) second_volume = (k_keys_sorted[1], gt_volume[k_keys_sorted[1]]) - # print(biggest_volume, second_volume) - if biggest_volume[1] * ARTICULAR_DOMINANCE_RATIO > second_volume[1]: to_label = biggest_volume[0] - 1 # int(list(gt_volume.keys())[argmax] - 1) diff --git a/spineps/phase_semantic.py b/spineps/phase_semantic.py index c8f67c3..39d1638 100755 --- a/spineps/phase_semantic.py +++ b/spineps/phase_semantic.py @@ -69,7 +69,6 @@ def predict_semantic_mask( verbose=verbose, ) # type:ignore seg_nii = results[OutputType.seg] - # unc_nii = results.get(OutputType.unc, None) softmax_logits = results[OutputType.softmax_logits] logger.print("Post-process semantic mask...") @@ -220,7 +219,6 @@ def semantic_bounding_box_clean(seg_nii: NII) -> NII: seg_bin_clean_arr[largest_k_arr == 0] = 0 seg_arr = seg_nii.get_seg_array() - # logger.print(seg_nii.volumes()) seg_arr[seg_bin_clean_arr != 1] = 0 seg_nii.set_array_(seg_arr) seg_nii.reorient_(ori) diff --git a/spineps/seg_enums.py b/spineps/seg_enums.py index ad4c0da..d55931d 100755 --- a/spineps/seg_enums.py +++ b/spineps/seg_enums.py @@ -195,12 +195,10 @@ class InputType(Enum_Compare): class OutputType(Enum_Compare): - """Type of model output: a segmentation (``seg``), softmax logits or an uncertainty map (``unc``).""" + """Type of model output: a segmentation (``seg``) or softmax logits.""" seg = auto() - # seg_modelres = auto() softmax_logits = auto() - unc = auto() class ErrCode(Enum_Compare): diff --git a/spineps/seg_model.py b/spineps/seg_model.py index 0e44745..38bb090 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -377,9 +377,6 @@ def dict_representation(self) -> dict[str, str]: "aquisition": str(self.acquisition()), "resolution_range": str(self.inference_config.resolution_range), } - # if input_zms is not None: - # proc_zms = self.calc_recommended_resampling_zoom(input_zms) - # info["resolution_processed"] = str(proc_zms) return info def __str__(self) -> str: diff --git a/spineps/seg_run.py b/spineps/seg_run.py index b32826b..3954fec 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -43,7 +43,6 @@ def process_dataset( # noqa: C901 derivative_name: str = "derivatives_seg", modalities: list[Modality_Pair] | Modality_Pair = [(Modality.T2w, Acquisition.sag)], # noqa: B006 save_debug_data: bool = False, - # save_uncertainty_image: bool = False, save_modelres_mask: bool = False, save_softmax_logits: bool = False, save_log_data: bool = True, @@ -235,8 +234,6 @@ def process_dataset( # noqa: C901 model_labeling=model_labeling, # derivative_name=derivative_name, - # - # save_uncertainty_image=save_uncertainty_image, save_modelres_mask=save_modelres_mask, save_softmax_logits=save_softmax_logits, save_debug_data=save_debug_data, @@ -305,8 +302,6 @@ def segment_image( # noqa: C901 model_instance: SegmentationModel, model_labeling: VertLabelingClassifier | None = None, derivative_name: str = "derivatives_seg", - # - # save_uncertainty_image: bool = False, save_modelres_mask: bool = False, save_softmax_logits: bool = False, save_debug_data: bool = False, @@ -478,7 +473,6 @@ def segment_image( # noqa: C901 if Modality.CT in model_semantic.modalities(): proc_normalize_input = False # Never normalize input if it is an CT proc_sem_n4_bias_correction = False # n4_bias_correction is a MRI thing - # proc_assign_missing_cc_fast = True # TODO remove if model_semantic.inference_config.has_c1: vertebra_instance_labeling_offset = 1 diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index 185ecca..cb5bfcc 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -2,15 +2,13 @@ from __future__ import annotations -from typing import Union - from TPTBox import BIDS_FILE, ZOOMS, Log_Type from spineps.seg_enums import Acquisition, Modality from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import logger -Modality_Pair = tuple[Union[list[Modality], Modality], Acquisition] +Modality_Pair = tuple[list[Modality] | Modality, Acquisition] def find_best_matching_model( @@ -32,21 +30,6 @@ def find_best_matching_model( NotImplementedError: Always, as this function is not yet implemented; also for an unmapped modality pair. """ raise NotImplementedError("find_best_matching_model()") - logger.print(expected_resolution) - # TODO replace with automatic going through model configs to find best matching the resolution - mapping: dict = { - # (Modality.CT, Acquisition.sag): MODELS.CT_SEGMENTOR, - # (Modality.T2w, Acquisition.sag): MODELS.T2w_NAKOSPIDER_HIGHRES, - # (Modality.T1w, Acquisition.sag): MODELS.T1w_SEGMENTOR, - # (Modality.Vibe, Acquisition.ax): MODELS.VIBE_SEGMENTOR, - # (Modality.SEG, Acquisition.sag): MODELS.VERT_HIGHRES, - } - if isinstance(modality_pair[0], list) and len(modality_pair[0]) == 1: - modality_pair = (modality_pair[0][0], modality_pair[1]) - if modality_pair not in mapping: - raise NotImplementedError(str(modality_pair[0]), str(modality_pair[1])) - else: - return mapping[modality_pair] def check_model_modality_acquisition( diff --git a/spineps/utils/auto_download.py b/spineps/utils/auto_download.py index 835f1f1..e80c4f5 100644 --- a/spineps/utils/auto_download.py +++ b/spineps/utils/auto_download.py @@ -6,7 +6,6 @@ import urllib.request import zipfile from pathlib import Path -from typing import Union from TPTBox import Print_Logger from tqdm import tqdm @@ -28,17 +27,17 @@ SpinepsPhase.SEMANTIC.name + "_ct": current_highest_ct_version, } -instances: dict[str, Union[Path, str]] = { +instances: dict[str, Path | str] = { "instance": link + current_instance_highest_version + "/instance.zip", "ct_instance": link + current_highest_ct_version + "/CT_instance.zip", } -semantic: dict[str, Union[Path, str]] = { +semantic: dict[str, Path | str] = { "t2w": link + current_highest_version + "/t2w.zip", "t1w": link + current_highest_version + "/t1w.zip", "vibe": link + current_highest_version + "/vibe.zip", "ct": link + current_highest_ct_version + "/ct.zip", } -labeling: dict[str, Union[Path, str]] = { +labeling: dict[str, Path | str] = { "t2w_labeling": link + current_labeling_highest_version + "/labeling.zip", "ct_labeling": link + current_labeling_highest_version + "/ct_labeling.zip", } @@ -55,7 +54,7 @@ } -def download_if_missing(key: str, url: Union[Path, str], phase: SpinepsPhase) -> Path: +def download_if_missing(key: str, url: Path | str, phase: SpinepsPhase) -> Path: """Return the local model folder for a model, downloading and extracting its weights if absent. The target folder name combines the model's download name with the version resolved for its phase (and the @@ -77,7 +76,7 @@ def download_if_missing(key: str, url: Union[Path, str], phase: SpinepsPhase) -> return out_path -def download_weights(weights_url: Union[Path, str], out_path: Union[Path, str]) -> None: +def download_weights(weights_url: Path | str, out_path: Path | str) -> None: """Download a weights zip archive, extract it into ``out_path`` and remove the archive. Shows a progress bar during download. If the extracted archive nests its contents in an extra subfolder diff --git a/spineps/utils/find_min_cost_path.py b/spineps/utils/find_min_cost_path.py index 93c762f..25410d6 100644 --- a/spineps/utils/find_min_cost_path.py +++ b/spineps/utils/find_min_cost_path.py @@ -268,22 +268,11 @@ def minCostAlgo(r, c): logger.print(f"Setting vert {r}, label {c} to {cost_value}, {internal_to_real_path(p)}") return min_costs_path[r][c] - # def t13_cost(r, c, pnext, p, region_cur): - # cost_add = 0 - # if vertt13_cost is not None: - # vt13_cost = vertt13_cost[r][1] - # # print(r, c, p[-1][1], p[-2][1], internal_to_real_path(p)) - # if p[-1][1] == 18 and p[-2][1] == 18: - # print(f"Added F {vt13_cost} to {r}, {c}, {internal_to_real_path(p)}") - # cost_add += vt13_cost - # return cost_add - def t13_cost_single(r, c): cost_add = 0 if vertt13_cost is not None: vt13_cost = vertt13_cost[r][1] if c == 18: - # print(f"Added F {vt13_cost} to {r}, {c}") cost_add += vt13_cost return cost_add diff --git a/spineps/utils/proc_functions.py b/spineps/utils/proc_functions.py index 052c6f0..101b867 100755 --- a/spineps/utils/proc_functions.py +++ b/spineps/utils/proc_functions.py @@ -46,8 +46,6 @@ def n4_bias( Returns: tuple[NII, NII]: The bias-corrected image and the binary foreground mask used for correction. """ - - # print("n4 bias", nii.dtype) mask = nii.get_array() mask[mask < threshold] = 0 mask[mask != 0] = 1 @@ -131,7 +129,6 @@ def clean_cc_artifacts( cc_to_clean = {} for lidx, label in enumerate(tqdm(labels, desc=f"{logger._get_logger_prefix()} cleaning...", disable=not verbose)): - # print(l, subreg_cc_stats[l]["voxel_counts"]) idx = [i for i, v in enumerate(subreg_cc_stats[label]["voxel_counts"]) if v < cc_size_threshold[lidx] and v > 0] if len(idx) > 0: cc_to_clean[label] = idx @@ -155,7 +152,6 @@ def clean_cc_artifacts( dilated_m = np_dilate_msk(mask_cc_l, n_pixel=1) dilated_m[mask_cc_l != 0] = 0 neighbor_voxel_count = np_count_nonzero(dilated_m) - # print(subreg_cc_stats[label]) mult = mask_arr * dilated_m if np_count_nonzero(mult) <= int(neighbor_voxel_count * neighbor_factor_2_delete): @@ -172,7 +168,6 @@ def clean_cc_artifacts( newlabel = nlabels[np.argmax(volumes_values)] # type: ignore result_arr[mask_cc_l != 0] = newlabel logger.print(log_string + f"labeled as {newlabel}") if verbose else None - # print(labels, count) n_to_clean = {k: len(v) for k, v in cc_to_clean.items()} # By clearning: look at surrounding neighbor pixels. If too few, remove cc. otherwise, do majority voting if len(n_to_clean) != 0: diff --git a/unit_tests/test_architectures.py b/unit_tests/test_architectures.py index 624e94d..583b062 100644 --- a/unit_tests/test_architectures.py +++ b/unit_tests/test_architectures.py @@ -33,7 +33,7 @@ logger = No_Logger() -class Test_Labeling_Model_Dummy(VertLabelingClassifier): +class LabelingModelDummy(VertLabelingClassifier): def __init__( self, model_folder: str | Path = __file__, @@ -155,3 +155,26 @@ def test_simple_testing_case(self): _, entry = get_vert_entry(v, subject_info) label = objectives(entry) print(v, label) + + def test_objective(self): + objectives = Objectives( + [ + Target.FULLYVISIBLE, + Target.REGION, + Target.VERTREL, + Target.VERT, + ], + as_group=True, + ) + + entry_dict = { + "vert_exact": VertExact.L1, + "vert_region": VertRegion.LWS, + "vert_rel": VertRel.FIRST_LWK, + "vert_cut": True, + } + + label = objectives(entry_dict) + self.assertEqual(label["FULLYVISIBLE"], [1, 0]) + self.assertEqual(label["REGION"], [0, 0, 1]) + self.assertEqual(label["VERTREL"], [0, 0, 0, 0, 1, 0]) diff --git a/unit_tests/test_disc_labels.py b/unit_tests/test_disc_labels.py index 9616dda..c6371dd 100644 --- a/unit_tests/test_disc_labels.py +++ b/unit_tests/test_disc_labels.py @@ -4,8 +4,12 @@ # coverage html from __future__ import annotations +import contextlib +import io +import sys import unittest from pathlib import Path +from unittest import mock import numpy as np from TPTBox import No_Logger @@ -18,7 +22,14 @@ class Test_DiscLabels(unittest.TestCase): def test_main_without_args(self): - self.skipTest("Not implemented") + # --path-vert is required; omitting it must exit via argparse, not proceed. + with ( + self.assertRaises(SystemExit) as cm, + mock.patch.object(sys, "argv", ["generate_disc_labels"]), + contextlib.redirect_stderr(io.StringIO()), + ): + main() + self.assertEqual(cm.exception.code, 2) def test_image(self): img = Image(param=np.array([0, 0, 0, 0])) diff --git a/unit_tests/test_path.py b/unit_tests/test_path.py deleted file mode 100644 index e881824..0000000 --- a/unit_tests/test_path.py +++ /dev/null @@ -1,97 +0,0 @@ -# Call 'python -m unittest' on this folder -# coverage run -m unittest -# coverage report -# coverage html -from __future__ import annotations - -import unittest -from unittest.mock import MagicMock - -import numpy as np -from TPTBox.tests.test_utils import get_test_mri - -from spineps.architectures.read_labels import Objectives, Target, VertExact, VertRegion, VertRel -from spineps.phase_labeling import VertLabelingClassifier, perform_labeling_step -from spineps.utils.find_min_cost_path import find_most_probably_sequence - - -class VertLabelingClassifierDummy(VertLabelingClassifier): - def __init__(self): - pass - - -class Test_PathLabeling(unittest.TestCase): - def test_search_path_simple(self): - cost = np.array( - [ - # colum lables - # rows predictions - [0, 10, 0, 0, 0], - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - ], - dtype=int, - ) - rel_cost = np.array( - [ - # nothing, last0, first1, last2 - [0, 0, 0, 0], - [0, 10, 0, 0], - [0, 0, 0, 0], - [0, 0, 11, 0], - ], - dtype=int, - ) - rel_cost = -rel_cost - fcost, fpath, _min_costs_path = find_most_probably_sequence( - cost, - region_rel_cost=rel_cost, - regions=[0, 3], - softmax_temp=0, - softmax_cost=False, - ) - fcost_avg = fcost / len(fpath) - print() - print("Path cost", round(fcost, 3)) - print("Path", fpath) - self.assertEqual(round(fcost_avg, 3), -5.0) - self.assertEqual(fpath, [1, 2, 3, 4]) - - def test_search_path_relativeonly(self): - self.skipTest("Notimplemented") - - def test_search_path_complex(self): - self.skipTest("Notimplemented") - - def test_objective(self): - objectives = Objectives( - [ - Target.FULLYVISIBLE, - Target.REGION, - Target.VERTREL, - Target.VERT, - ], - as_group=True, - ) - - entry_dict = { - "vert_exact": VertExact.L1, - "vert_region": VertRegion.LWS, - "vert_rel": VertRel.FIRST_LWK, - "vert_cut": True, - } - - label = objectives(entry_dict) - print(label) - self.assertEqual(label["FULLYVISIBLE"], [1, 0]) - self.assertEqual(label["REGION"], [0, 0, 1]) - self.assertEqual(label["VERTREL"], [0, 0, 0, 0, 1, 0]) - - def test_labeling_easy(self): - self.skipTest("Notimplemented") - # mri, subreg, vert, label = get_test_mri() - # model = VertLabelingClassifierDummy() - # l = vert.unique() - # model.run_all_seg_instances = MagicMock(return_value=l) - # perform_labeling_step(model, mri, vert) From 557ebc3ad04c7231cbf2027461d4ac14607b25ab Mon Sep 17 00:00:00 2001 From: iback Date: Fri, 3 Jul 2026 07:55:21 +0000 Subject: [PATCH 29/29] fixed 3.9 generic alias issue --- spineps/api.py | 7 +++++-- spineps/seg_utils.py | 6 +++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/spineps/api.py b/spineps/api.py index ab431ca..47e9d18 100644 --- a/spineps/api.py +++ b/spineps/api.py @@ -20,6 +20,7 @@ from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path +from typing import Union from TPTBox import BIDS_FILE, NII, POI @@ -31,8 +32,10 @@ from spineps.seg_run import segment_image # Accepted forms for an input image and a model argument. -ImageInput = str | Path | NII | BIDS_FILE -ModelInput = str | Path | SegmentationModel | VertLabelingClassifier +# NOTE: these are plain runtime assignments, not annotations, so `from __future__ import annotations` does not +# defer them -- `X | Y` on bare types needs Python 3.10+, so this must stay `Union[...]` for the 3.9 floor. +ImageInput = Union[str, Path, NII, BIDS_FILE] +ModelInput = Union[str, Path, SegmentationModel, VertLabelingClassifier] @dataclass diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index cb5bfcc..6b60769 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -2,13 +2,17 @@ from __future__ import annotations +from typing import Union + from TPTBox import BIDS_FILE, ZOOMS, Log_Type from spineps.seg_enums import Acquisition, Modality from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import logger -Modality_Pair = tuple[list[Modality] | Modality, Acquisition] +# NOTE: plain runtime assignment, not an annotation, so `from __future__ import annotations` does not defer it -- +# `X | Y` on bare types needs Python 3.10+, so this must stay `Union[...]` for the 3.9 floor. +Modality_Pair = tuple[Union[list[Modality], Modality], Acquisition] def find_best_matching_model(