SPINEPS 2.0: usability, error-proofing, batched inference, segment() API + CLI cleanup#64
Open
Hendrik-code wants to merge 32 commits into
Open
SPINEPS 2.0: usability, error-proofing, batched inference, segment() API + CLI cleanup#64Hendrik-code wants to merge 32 commits into
Hendrik-code wants to merge 32 commits into
Conversation
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 <noreply@anthropic.com>
…eanups 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…odel* -> 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Hendrik-code
added
documentation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces a breaking SPINEPS 2.0 update focused on a higher-level Python API (spineps.segment + SpinepsPipeline), batched instance inference for performance, and CLI/API renames/cleanup with improved error handling and updated docs/migration guidance.
Changes:
- Added a new high-level Python API (
segment,SpinepsPipeline,SpinepsResult) plus grouped config dataclasses to simplify pipeline invocation. - Implemented batched 3D U-Net instance inference (
run_batch/segment_scan_batch) with CLI/API knobs (--batch-size,--amp,--step-size,--tta) and propagated these through pipeline functions. - Performed 2.0 breaking renames (e.g.,
process_img_nii→segment_image,Segmentation_Model*→SegmentationModel*), replaced assertion-based user input checks with explicit exceptions, and updated tests/docs accordingly.
Reviewed changes
Copilot reviewed 31 out of 31 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| unit_tests/test_semantic.py | Updates semantic tests to renamed model classes. |
| unit_tests/test_proc_functions.py | Adjusts imports to renamed model loader API. |
| unit_tests/test_inference_mocked.py | Adds extensive tests for batched Unet3D inference + updated model names. |
| unit_tests/test_filepaths.py | Adds/updates tests asserting exception-based error handling in filepaths/model resolution. |
| unit_tests/test_entrypoint.py | Adds CLI regression tests for boolean flag defaults/negation and help output. |
| unit_tests/test_api.py | New tests covering the new high-level segment API, configs, and result wrapping. |
| spineps/utils/inference_api.py | Converts missing-model-folder assertion into FileNotFoundError. |
| spineps/utils/filepaths.py | Adds logger and replaces prints/asserts with logged output and exceptions. |
| spineps/utils/citation_reminder.py | Preserves wrapped function metadata via functools.wraps. |
| spineps/seg_utils.py | Renames model types in function signatures/docs. |
| spineps/seg_run.py | Renames process_img_nii→segment_image, adds new knobs (step_size/batch/amp/tta) and forwards them. |
| spineps/seg_pipeline.py | Renames model types for centroid metadata plumbing. |
| spineps/seg_model.py | Renames model base classes and adds batched Unet3D inference and segment-scan batching. |
| spineps/phase_semantic.py | Adds semantic sliding-window step size plumbing into model inference. |
| spineps/phase_post.py | Replaces a stray print with structured logging. |
| spineps/phase_labeling.py | Adds prepare_vertexact and incorporates VERTEX head into labeling path scoring. |
| spineps/phase_instance.py | Switches instance cutout inference from per-cutout calls to batched model calls. |
| spineps/lab_model.py | Updates classifier base class rename and docs. |
| spineps/get_models.py | Improves missing-model/download error reporting and updates renamed model classes. |
| spineps/example/template_roll_out.py | Updates example to renamed segment_image API. |
| spineps/example/helper_parallel.py | Updates example to renamed segment_image API. |
| spineps/entrypoint.py | CLI cleanup: kebab-case long flags, boolean optional flags, new speed knobs, and exception-based validation. |
| spineps/config.py | New grouped config dataclasses mapping to flat proc_* kwargs. |
| spineps/api.py | New high-level one-call API + reusable pipeline wrapper and result normalization. |
| spineps/init.py | Exposes new public API objects and updates exports for renamed functions/classes. |
| README.md | Updates CLI examples and adds guidance for the new Python API and migration link. |
| MIGRATION.md | New migration guide documenting the 2.0 breaking renames and new knobs. |
| docs/modules/pipeline.md | Updates pipeline docs to reference segment_image and adds spineps.segment examples. |
| docs/modules/models.md | Updates model docs for renamed model classes and CLI flag naming. |
| docs/index.md | Updates quickstart CLI command to new kebab-case flags. |
| docs/getting-started.md | Updates installation/usage docs with new CLI flags and new Python API examples. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+293
to
+299
| # 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: |
Comment on lines
+654
to
+675
| 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]) |
Comment on lines
+624
to
+630
| 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() |
Comment on lines
72
to
93
| 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).", | ||
| ) | ||
| parser.add_argument( | ||
| "--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).", | ||
| ) |
Integrates main's faster TPTBox inference backend (local inference_api.py removed in favor of TPTBox.segmentation.nnUnet_utils.inference_api; run_inference now returns a 3-tuple) and main's new C1/C2 labeling-force logic into the 2.0 branch. Resolution: - spineps/utils/inference_api.py: accepted main's deletion (my assert->exception change there is moot; load_inf_model/run_inference now come from TPTBox). - seg_model.py: kept main's TPTBox import + NNunet 3-tuple run_inference unpacking alongside the batched-inference / segment_scan_batch / set_test_time_augmentation additions. - phase_labeling.py: kept both the prepare_vertexact/VERTEX fix and main's force_c1/force_c2 forcing. - test_semantic.py: kept main's 3-tuple run_inference test under the SegmentationModel rename. 195 tests pass, ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…lify 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 <noreply@anthropic.com>
- 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 <noreply@anthropic.com>
…tptbox, deleted vendored nnUNet predictor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
A clean-break 2.0 pass to make SPINEPS easier to use, harder to misuse, and faster. Full old→new mapping in MIGRATION.md.
Highlights
New high-level API —
spineps.segment(image, ...)one-call API returning a SpinepsResult;SpinepsPipeline(load once, segment many); config objectsSemanticConfig/InstanceConfig/LabelingConfig/PostConfig.
Performance — batched instance inference (one forward per --batch-size cutouts vs 7–33 sequential, under inference_mode), numerically identical to
sequential in fp32 with CUDA-OOM fallback; knobs --batch-size/--amp/--step-size/--tta.
Usability — user-input asserts → FileNotFoundError/ValueError/etc.; clearer model-download errors; typo/stray-print cleanup.
CLI (breaking) — all long flags --kebab-case (short aliases kept); --crop/--no-crop, --n4/--no-n4, --enforce-12-thoracic.
Renames (breaking) — process_img_nii→segment_image; Segmentation_Model→SegmentationModel.
Bug fixes found while testing (pre-existing)
spineps sample -h/dataset -h.dataset --batch-sizewould TypeError.Breaking release — bump major and tag v2.0.0 (dynamic versioning).