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 f890c15..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/ @@ -158,4 +140,5 @@ test.txt derivatives_seg robert_test.py poetry.lock -*.DS_Store \ No newline at end of file +*.DS_Store +*DS_Store* diff --git a/MIGRATION.md b/MIGRATION.md new file mode 100644 index 0000000..0b07326 --- /dev/null +++ b/MIGRATION.md @@ -0,0 +1,97 @@ +# 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) | +| *(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: + +```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. +- **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/README.md b/README.md index 02d092f..db87385 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: - -1. Activate your venv -2. Run `spineps -h` to see the arguments +After installation (`pip install spineps`, or `pip install -e .` from a local clone), the `spineps` command is +available in your venv: -### 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 @@ -188,26 +194,48 @@ 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) | -| -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) | +| --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. #### 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 @@ -242,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 | | :--- | --------- | -| -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) | +| --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) | + +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`. -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 @@ -260,9 +297,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) @@ -287,11 +323,22 @@ In the subregion segmentation: | 100 | Vertebra_Disc | | 26 | Sacrum | +CT only +| 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 | + 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: -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 +356,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: + +```python +import spineps + +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. -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()`. +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`. -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! +> **Upgrading from 1.x?** See [MIGRATION.md](MIGRATION.md) for the renamed CLI flags, functions and classes. ## Authorship 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/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/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 c65c133..54cef55 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -77,18 +77,20 @@ 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-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, not accepted by `sample`.) + 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 +99,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..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/) @@ -28,6 +29,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 @@ -37,7 +39,15 @@ 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 +``` + +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 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/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/docs/modules/pipeline.md b/docs/modules/pipeline.md index 96477f2..396d5b9 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 @@ -43,25 +43,47 @@ 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 +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. 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/pyproject.toml b/pyproject.toml index 112465f..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' @@ -25,16 +26,26 @@ pytorch-lightning = "^2.0.8" torchmetrics = "^1.1.2" tqdm = "^4.66.1" einops= "^0.6.1" -nnunetv2 = "^2.4.2" -TPTBox = "*" -antspyx = "0.4.2" +TPTBox = "^0.7.5" +antspyx = "0.6.3" rich = "^13.6.0" monai="^1.3.0" -acvl-utils = "0.2" 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" }, ] +blosc2 = "*" +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" }, +] + [tool.poetry.group.dev.dependencies] @@ -42,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] @@ -87,14 +92,6 @@ exclude = [ "dist", "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/.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/__init__.py b/spineps/__init__.py index cb36017..2c50a54 100755 --- a/spineps/__init__.py +++ b/spineps/__init__.py @@ -1,10 +1,33 @@ -"""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 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 + +__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..47e9d18 --- /dev/null +++ b/spineps/api.py @@ -0,0 +1,234 @@ +"""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. +# 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 +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: 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: + 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: 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: + 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/architectures/pl_unet.py b/spineps/architectures/pl_unet.py index 7868c0e..85ff6da 100755 --- a/spineps/architectures/pl_unet.py +++ b/spineps/architectures/pl_unet.py @@ -8,7 +8,6 @@ import torch import torchmetrics.functional as mF from torch import nn -from torch.optim import lr_scheduler from spineps.architectures.unet3D import Unet3D @@ -37,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/config.py b/spineps/config.py new file mode 100644 index 0000000..6347863 --- /dev/null +++ b/spineps/config.py @@ -0,0 +1,90 @@ +"""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 + step_size: float | None = None + + 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, + "proc_sem_step_size": self.step_size, + } + + +@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 + amp: bool = False + 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, + "proc_inst_amp": self.amp, + "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/entrypoint.py b/spineps/entrypoint.py index 0349962..cc1b30a 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") @@ -39,58 +39,90 @@ 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).", + ) + 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).", ) - # proc_lab_force_no_tl_anomaly parser.add_argument( - "-non4", + "--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. 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", action="store_true", - help="Does not apply n4 bias field correction", + 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("-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( + "--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") return parser @@ -109,43 +141,33 @@ 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 ) ########################### ########################### - 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 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", + "--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 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", + "--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 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) @@ -153,54 +175,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. 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", + "--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 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", + "--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 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", + "--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", @@ -209,9 +222,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": @@ -225,7 +237,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, @@ -235,17 +247,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() @@ -264,6 +279,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 = { @@ -271,9 +289,7 @@ def run_sample(opt: Namespace): "model_semantic": model_semantic, "model_instance": model_instance, "model_labeling": model_labeling, - "derivative_name": opt.der_name, - # - # "save_uncertainty_image": opt.save_unc_img, + "derivative_name": opt.derivative_name, "save_softmax_logits": opt.save_softmax_logits, "save_debug_data": opt.save_debug, "save_modelres_mask": opt.save_modelres_mask, @@ -281,9 +297,12 @@ 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, + "proc_inst_amp": opt.amp, + "proc_sem_step_size": opt.step_size, "ignore_compatibility_issues": opt.ignore_inference_compatibility, "verbose": opt.verbose, } @@ -295,16 +314,16 @@ 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}, ) 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 @@ -325,11 +344,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 +363,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,17 +376,16 @@ 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, "model_semantic": model_semantic, "model_instance": model_instance, "model_labeling": model_labeling, - "rawdata_name": opt.raw_name, - "derivative_name": opt.der_name, - # - # "save_uncertainty_image": opt.save_unc_img, + "rawdata_name": opt.rawdata_name, + "derivative_name": opt.derivative_name, "save_modelres_mask": opt.save_modelres_mask, "save_softmax_logits": opt.save_softmax_logits, "save_debug_data": opt.save_debug, @@ -377,10 +397,14 @@ 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, + "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/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/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 8940982..cbb11bd 100755 --- a/spineps/get_models.py +++ b/spineps/get_models.py @@ -4,14 +4,13 @@ import os from pathlib import Path -from typing import Optional, Union from TPTBox import Log_Type, No_Logger from tqdm import tqdm 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 +32,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. @@ -42,18 +41,21 @@ 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) -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: @@ -61,15 +63,16 @@ 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 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) -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: @@ -77,10 +80,11 @@ 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 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,14 +100,15 @@ 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) -_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]: @@ -149,7 +154,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. @@ -167,12 +172,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 @@ -207,12 +213,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: @@ -223,7 +229,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 @@ -235,16 +241,12 @@ 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. 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 @@ -260,10 +262,7 @@ 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[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..925f5c5 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: @@ -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 9d3b4c9..722df9a 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 ( @@ -21,7 +20,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 +48,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, @@ -57,6 +56,8 @@ 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, + proc_inst_amp: bool = False, verbose: bool = False, ) -> tuple[NII | None, ErrCode]: """Build a per-vertebra instance mask from a subregion semantic segmentation. @@ -67,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. @@ -75,6 +76,10 @@ 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. + 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: @@ -100,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) @@ -119,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() # @@ -146,6 +147,8 @@ 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, + amp=proc_inst_amp, verbose=verbose, ) if vert_predictions is None: @@ -196,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 @@ -580,13 +581,15 @@ 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, 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, + 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. @@ -598,13 +601,17 @@ 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. 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. + 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: @@ -637,22 +644,22 @@ 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) - 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} 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) # 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 +675,22 @@ 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, + amp=amp, + 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,29 +703,19 @@ 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])] + 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) - # 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 - 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 @@ -937,7 +940,6 @@ def find_prediction_couple( agreement = 0 if len(couple) > 0: - # print(couple) agreement = 0 for c in couple: agreement += dices[c] @@ -978,6 +980,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(): @@ -985,10 +988,9 @@ 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) m = 1 if take_no_overlap else 2 # m = min(max(1, np.max(combine)), 2) # type:ignore combine[combine < m] = 0 @@ -998,9 +1000,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) @@ -1024,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 e1a0b2c..f253521 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -15,7 +15,6 @@ 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, @@ -300,6 +299,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, @@ -447,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, @@ -542,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, ( @@ -567,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( @@ -593,6 +669,16 @@ def find_vert_path_from_predictions( vert_w, ) + 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, + 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( @@ -616,6 +702,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 @@ -644,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 71f48af..bd6143f 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 @@ -226,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: @@ -366,7 +363,7 @@ 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) -> tuple[np.ndarray, np.ndarray]: +def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=False, verbose=True) -> tuple[np.ndarray, np.ndarray]: """Attach intervertebral-disc and endplate instance labels and split endplates into superior/inferior. Reorients both masks to PIR, computes each vertebra corpus center of mass along the inferior-superior axis, then assigns @@ -388,7 +385,9 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, verbose=True) -> tu 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 + 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() @@ -448,9 +447,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 +466,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 +481,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,7 +564,23 @@ def find_nearest_lower(seq, x) -> float: return max(values_lower) -def label_instance_top_to_bottom(vert_nii: NII, labeling_offset: int = 0) -> tuple[NII, np.ndarray]: +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, 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 @@ -548,8 +591,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() @@ -599,7 +643,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) @@ -618,14 +662,12 @@ 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) 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}", ) vert_nii.set_array_(vert_arr) diff --git a/spineps/phase_pre.py b/spineps/phase_pre.py index 641ac8a..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,27 +37,27 @@ 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. """ - 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/phase_semantic.py b/spineps/phase_semantic.py index 1e3605b..39d1638 100755 --- a/spineps/phase_semantic.py +++ b/spineps/phase_semantic.py @@ -2,12 +2,11 @@ from __future__ import annotations -# from utils.predictor import nnUNetPredictor import numpy as np 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,12 +25,13 @@ 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, 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. @@ -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 @@ -50,6 +50,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,12 +63,12 @@ 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, ) # 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...") @@ -217,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 8c9cc1c..38bb090 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -1,9 +1,10 @@ -"""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 import os from abc import ABC, abstractmethod +from contextlib import nullcontext from pathlib import Path import numpy as np @@ -29,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 @@ -66,10 +67,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 @@ -171,7 +173,10 @@ def segment_scan( logits). """ 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 @@ -244,6 +249,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. @@ -260,6 +308,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. @@ -317,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: @@ -339,8 +396,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, @@ -420,8 +477,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. @@ -461,12 +518,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: @@ -482,8 +543,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. @@ -496,36 +557,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] + + 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. - arr = input_nii_.get_seg_array().astype(np.int16) + 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 = from_numpy(arr) - n_classes = self.predictor.network.channels + 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. + 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 - - # channel-wise 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_pipeline.py b/spineps/seg_pipeline.py index 4a1f3b4..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 @@ -11,7 +10,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 +50,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 +62,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 01597f9..3954fec 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -19,24 +19,30 @@ 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 +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( +def process_dataset( # noqa: C901 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", 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, @@ -53,11 +59,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 @@ -69,18 +78,19 @@ 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, ): """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". @@ -109,12 +119,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. @@ -125,6 +141,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. """ @@ -133,7 +152,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") @@ -149,6 +169,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): @@ -198,15 +227,13 @@ 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, 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, @@ -221,11 +248,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, @@ -266,14 +296,12 @@ 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", - # - # save_uncertainty_image: bool = False, save_modelres_mask: bool = False, save_softmax_logits: bool = False, save_debug_data: bool = False, @@ -296,11 +324,14 @@ def process_img_nii( # 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, @@ -319,6 +350,8 @@ def process_img_nii( # noqa: C901 return_output_instead_of_save: bool = False, timing=False, verbose: bool = False, + _nii=None, + _dataset_id_ct_crop=100, ) -> tuple[dict[str, Path], ErrCode]: """Runs the SPINEPS framework over one nifty. @@ -327,8 +360,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. @@ -359,12 +392,18 @@ def process_img_nii( # 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. 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. 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. @@ -393,7 +432,12 @@ def process_img_nii( # 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"] @@ -423,12 +467,12 @@ def process_img_nii( # 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 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 @@ -447,12 +491,20 @@ 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 timing: logger.print(f"Loading files took: {perf_counter() - start_time2:.2f} seconds", Log_Type.OK, verbose=log_inference_time) start_time2 = perf_counter() + 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 + logger.print("Input image", input_nii.zoom, input_nii.orientation, input_nii.shape) + # First stage if not out_spine_raw.exists() or override_semantic: resolution_range = model_semantic.inference_config.resolution_range @@ -472,7 +524,13 @@ def process_img_nii( # 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 @@ -514,6 +572,7 @@ def process_img_nii( # 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 @@ -553,6 +612,8 @@ 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, + proc_inst_amp=proc_inst_amp, ) if errcode != ErrCode.OK: logger.print(f"Vert Mask creation failed with errcode {errcode}", Log_Type.FAIL) @@ -659,18 +720,14 @@ def process_img_nii( # 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) @@ -690,6 +747,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. @@ -775,7 +833,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, ) diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index 122e313..6b60769 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -2,23 +2,23 @@ from __future__ import annotations -# 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 +from spineps.seg_model import SegmentationModel from spineps.seg_pipeline import logger +# 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( 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,31 +28,16 @@ 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. """ 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( - model: Segmentation_Model, + model: SegmentationModel, mod_pair: Modality_Pair, verbose: bool = True, ) -> bool: @@ -62,7 +47,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 +101,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 +115,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. @@ -190,7 +175,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: 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/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/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/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/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/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 3d9758a..101b867 100755 --- a/spineps/utils/proc_functions.py +++ b/spineps/utils/proc_functions.py @@ -46,9 +46,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() mask[mask < threshold] = 0 mask[mask != 0] = 1 @@ -56,7 +53,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.to_ants(), spline_param=spline_param) if norm != -1: n4 *= norm / n4.max() if dtype2nii: @@ -132,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 @@ -156,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): @@ -173,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/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_api.py b/unit_tests/test_api.py new file mode 100644 index 0000000..7765234 --- /dev/null +++ b/unit_tests/test_api.py @@ -0,0 +1,120 @@ +# Call 'python -m unittest' on this folder +"""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, 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, 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): + 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() diff --git a/unit_tests/test_architectures.py b/unit_tests/test_architectures.py index 4245c6d..583b062 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 @@ -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_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..c6371dd 100644 --- a/unit_tests/test_disc_labels.py +++ b/unit_tests/test_disc_labels.py @@ -1,11 +1,15 @@ -# 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 __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_entrypoint.py b/unit_tests/test_entrypoint.py index 4400b3c..d0a4746 100644 --- a/unit_tests/test_entrypoint.py +++ b/unit_tests/test_entrypoint.py @@ -1,15 +1,17 @@ -# 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 __future__ import annotations import argparse +import contextlib +import io +import sys import unittest -from pathlib import Path +from unittest import mock from TPTBox import No_Logger -from typing_extensions import Self from spineps import entrypoint @@ -19,3 +21,60 @@ 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_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"): + 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") + 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_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 b2bf433..715add0 100644 --- a/unit_tests/test_filepaths.py +++ b/unit_tests/test_filepaths.py @@ -1,15 +1,16 @@ -# 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 __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 SegmentationModel, check_available_models, get_actual_model from spineps.utils.filepaths import ( filepath_model, get_mri_segmentor_models_dir, @@ -25,18 +26,18 @@ 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): 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 @@ -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 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 78d2b4c..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 @@ -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, SegmentationModel, SegmentationModelUnet3D logger = No_Logger() @@ -56,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() @@ -146,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 @@ -165,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( @@ -200,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( @@ -223,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) @@ -318,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 @@ -348,5 +349,141 @@ 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)) -> SegmentationModelUnet3D: + """A real SegmentationModelUnet3D 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 = SegmentationModelUnet3D(__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()) + + 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) + + +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 = 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 + 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() 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 deleted file mode 100644 index 78c5a0e..0000000 --- a/unit_tests/test_path.py +++ /dev/null @@ -1,97 +0,0 @@ -# Call 'python -m unittest' on this folder # noqa: INP001 -# 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) 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 14f8cb5..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 @@ -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_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 diff --git a/unit_tests/test_semantic.py b/unit_tests/test_semantic.py index e882749..9a012df 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(