From 7774a121ef19bcdb446cb9c95dbcd54cf224ae91 Mon Sep 17 00:00:00 2001 From: Julien Jomier Date: Wed, 25 Feb 2026 13:18:32 -0500 Subject: [PATCH 1/6] Added pre-commit + CI --- .github/workflows/ci.yml | 26 ++++++++++++++++++++++++++ .gitignore | 2 +- .markdownlint.yaml | 7 +++++++ .pre-commit-config.yaml | 35 +++++++++++++++++++++++++++++++++++ pyproject.toml | 13 +++++++++++++ 5 files changed, 82 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/ci.yml create mode 100644 .markdownlint.yaml create mode 100644 .pre-commit-config.yaml create mode 100644 pyproject.toml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..3bbbd10 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,26 @@ +# CI for NV-Generate-CTMR: lint and format checks via pre-commit +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + pre-commit: + name: pre-commit + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install pre-commit + run: pip install pre-commit + + - name: Run pre-commit + run: pre-commit run --all-files diff --git a/.gitignore b/.gitignore index ab0f6c8..9ba23e0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,2 @@ *.pt -*.pyc \ No newline at end of file +*.pyc diff --git a/.markdownlint.yaml b/.markdownlint.yaml new file mode 100644 index 0000000..998dbbd --- /dev/null +++ b/.markdownlint.yaml @@ -0,0 +1,7 @@ +# Markdownlint config for NV-Generate-CTMR +# Relaxed for existing docs (READMEs with tables, HTML, long lines). +# Re-enable rules as you clean up docs or for new files. + +# Line length: allow long lines common in docs (tables, code, links) +MD013: + line_length: 700 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..dc42e8a --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,35 @@ +# Pre-commit hooks for NV-Generate-CTMR +# Install: pip install pre-commit && pre-commit install +# Run manually: pre-commit run --all-files + +repos: + # General file checks + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-json + - id: check-merge-conflict + - id: check-added-large-files + args: [--maxkb=1000] + - id: check-case-conflict + - id: debug-statements + + # Python linting and formatting (ruff) — fixes applied locally + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.8.4 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + # Markdown linting + - repo: https://github.com/igorshubovych/markdownlint-cli + rev: v0.38.0 + hooks: + - id: markdownlint + +ci: + autoupdate_commit_msg: "chore: pre-commit autoupdate" diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a0254ab --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,13 @@ +# Minimal config for tooling (pre-commit, ruff). NV-Generate-CTMR has no installable package. + +[tool.ruff] +target-version = "py311" +line-length = 150 +exclude = [".git", "__pycache__", "data", "figures", "assets", "*.ipynb"] + +[tool.ruff.lint] +select = ["E", "F", "I", "N", "W", "UP"] +ignore = ["E501"] # line length handled by formatter + +[tool.ruff.format] +quote-style = "double" From 3fb21f5cbf029fc5f84cdfd7db71484825a0fa23 Mon Sep 17 00:00:00 2001 From: Julien Jomier Date: Wed, 25 Feb 2026 13:19:12 -0500 Subject: [PATCH 2/6] Fixing README --- NV-Segment-CT/docs/data.md | 6 +- NV-Segment-CT/docs/finetune.md | 33 +++-- NV-Segment-CT/docs/inference.md | 47 ++++--- NV-Segment-CTMR/brain_t1_preprocess/README.md | 4 +- NV-Segment-CTMR/docs/README.md | 117 ++++++++++++------ NV-Segment-CTMR/docs/data.md | 6 +- NV-Segment-CTMR/docs/finetune.md | 33 +++-- NV-Segment-CTMR/docs/inference.md | 47 ++++--- README.md | 9 +- 9 files changed, 195 insertions(+), 107 deletions(-) diff --git a/NV-Segment-CT/docs/data.md b/NV-Segment-CT/docs/data.md index 45595a8..031ece4 100644 --- a/NV-Segment-CT/docs/data.md +++ b/NV-Segment-CT/docs/data.md @@ -1,5 +1,9 @@ -### Best practice to generate data list +# Data + +## Best practice to generate data list + User can use monai to generate the 5-fold data lists. Full exampls can be found in VISTA3D open source [codebase](https://github.com/Project-MONAI/VISTA/blob/main/vista3d/data/make_datalists.py) + ```python from monai.data.utils import partition_dataset from monai.bundle import ConfigParser diff --git a/NV-Segment-CT/docs/finetune.md b/NV-Segment-CT/docs/finetune.md index 855350a..7625d6e 100644 --- a/NV-Segment-CT/docs/finetune.md +++ b/NV-Segment-CT/docs/finetune.md @@ -1,38 +1,46 @@ -### Configurations +# Finetune configurations +## Configurations - -#### Best practice to set label_mapping +### Best practice to set label_mapping For a class that represent the same or similar class as the global index, directly map it to the global index. For example, "mouse left lung" (e.g. index 2 in the mouse dataset) can be mapped to the 28 "left lung upper lobe"(or 29 "left lung lower lobe") with [[2,28]]. After finetuning, 28 now represents "mouse left lung" and will be used for segmentation. If you want to segment 4 substructures of aorta, you can map one of the substructuress to 6 aorta and the rest to any value, [[1,6],[2,133],[3,134],[4,135]]. -``` + +```text NOTE: Do not map to global index value >= 255. `num_classes=255` in the config only represent the maximum mapping index, while the actual output class number only depends on your label_mapping definition. The 255 value in the inference output is also used to represent 'NaN' value. ``` -#### `val_at_start` +### `val_at_start` + Default `true`, VISTA3D will perform out-of-the-box segmentation before the training. Users can disable if the validation takes too long. +### `n_train_samples` and `n_val_samples` -#### `n_train_samples` and `n_val_samples` In `train_continual.json`, only `n_train_samples` and `n_val_samples` are used for training and validation. -#### `patch_size` +### `patch_size` + The patch size parameter is defined in `configs/train_continual.json`: `"patch_size": [128, 128, 128]`. For finetuning purposes, this value needs to be changed acccording to user's task and GPU memory. Usually a larger patch_size will give better final results. `[192,192,128]` is a good value for larger memory GPU. -#### `resample_to_spacing` +### `resample_to_spacing` + The resample_to_spacing parameter is defined in `configs/train_continual.json` and it represents the resolution the model will be trained on. The `1.5,1.5,1.5` mm default is suitable for large CT organs, but for other tasks, this value should be changed to achive the optimal performance. -#### Advanced user: `drop_label_prob` and `drop_point_prob` (in train.json) +### Advanced user: `drop_label_prob` and `drop_point_prob` (in train.json) + VISTA3D is trained to perform both automatic (class prompts) and interactive point segmentation. `drop_label_prob` and `drop_point_prob` means percentage to remove class prompts and point prompts during training respectively. If `drop_point_prob=1`, the model is only finetuning for automatic segmentation, while `drop_label_prob=1` means only finetuning for interactive segmentation. The VISTA3D foundation model is trained with interactive only (drop_label_prob=1) and then froze the point branch and trained with fully automatic segmentation (`drop_point_prob=1`). In this bundle, the training is simplified by jointly training with class prompts and point prompts and both of the drop ratio is set to 0.25. -``` + +```text NOTE: If user doesn't use interactive segmentation, set `drop_point_prob=1` and `drop_label_prob=0` in train.json might provide a faster and easier finetuning process. ``` -#### Other explanatory items + +### Other explanatory items + In `train.json`, `validate[evaluator][val_head]` can be `auto` and `point`. If `auto`, the validation results will be automatic segmentation. If `point`, the validation results will be sampling one positive point per object per patch. The validation scheme of combining auto and point is deprecated due to speed issue. @@ -44,7 +52,8 @@ to the VISTA3D global class index. `label_set` is used to identify the VISTA model classes for providing training prompts. `val_label_set` is used to identify the original training label classes for computing foreground/background mask during validation. The default configs for both variables are derived from the `label_mappings` config and include `[0]`: -``` + +```json "label_set": "$[0] + list(x[1] for x in @label_mappings#default)" "val_label_set": "$[0] + list(x[0] for x in @label_mappings#default)" ``` diff --git a/NV-Segment-CT/docs/inference.md b/NV-Segment-CT/docs/inference.md index 220a621..653f70a 100644 --- a/NV-Segment-CT/docs/inference.md +++ b/NV-Segment-CT/docs/inference.md @@ -1,10 +1,16 @@ +# Inference configurations + All the configurations for inference is stored in inference.json, change those parameters: -### `input_dict` + +## `input_dict` + `input_dict` defines the image to segment and the prompt for segmentation. -``` + +```json "input_dict": "$[{'image': '/data/Task09_Spleen/imagesTs/spleen_15.nii.gz', 'label_prompt':[1]}]", "input_dict": "$[{'image': '/data/Task09_Spleen/imagesTs/spleen_15.nii.gz', 'points':[[138,245,18], [271,343,27]], 'point_labels':[1,0]}]" ``` + - The input_dict must include the key `image` which contain the absolute path to the nii image file, and includes prompt keys of `label_prompt`, `points` and `point_labels`. - The `label_prompt` is a list of length `B`, which can perform `B` foreground objects segmentation, e.g. `[2,3,4,5]`. If `B>1`, Point prompts must NOT be provided. - The `points` is of shape `[N, 3]` like `[[x1,y1,z1],[x2,y2,z2],...[xN,yN,zN]]`, representing `N` point coordinates **IN THE ORIGINAL IMAGE SPACE** of a single foreground object. `point_labels` is a list of length [N] like [1,1,0,-1,...], which @@ -12,7 +18,7 @@ matches the `points`. 0 means background, 1 means foreground, -1 means ignoring - **B must be 1 if label_prompt and points are provided together**. The inferer only supports SINGLE OBJECT point click segmentatation. - If no prompt is provided, the model will use `everything_labels` to segment 117 classes: -```Python +```python list(set([i+1 for i in range(132)]) - set([2,16,18,20,21,23,24,25,26,27,128,129,130,131,132])) ``` @@ -20,10 +26,11 @@ list(set([i+1 for i in range(132)]) - set([2,16,18,20,21,23,24,25,26,27,128,129, - To specify a new class for zero-shot segmentation, set the `label_prompt` to a value between 133 and 254. Ensure that `points` and `point_labels` are also provided; otherwise, the inference result will be a tensor of zeros. ### `label_prompt` and `label_dict` + The `label_dict` defined in `configs/metadata.json` has in total 132 classes. However, there are 5 we do not support and we keep them due to legacy issue. So in total VISTA3D support 127 classes. -``` +```text "16, # prostate or uterus" since we already have "prostate" class, "18, # rectum", insufficient data or dataset excluded. "130, # liver tumor" already have hepatic tumor. @@ -34,34 +41,36 @@ VISTA3D support 127 classes. These 5 are excluded in the `everything_labels`. Another 7 tumor and vessel classes are also removed since they will overlap with other organs and make the output messy. To segment those 7 classes, we recommend users to directly set `label_prompt` to those indexes and avoid using them in `everything_labels`. For "Kidney", "Lung", "Bone" (class index [2, 20, 21]), VISTA3D did not directly use the class index for segmentation, but instead convert them to their subclass indexes as defined by `subclass` dict. For example, "2-Kidney" is converted to "14-Left Kidney" + "5-Right Kidney" since "2" is defined in `subclasss` dict. ### `resample_spacing` + The optimal inference resample spacing should be changed according to the task. For monkey data, a high resolution of [1,1,1] showed better automatic inference results. This spacing applies to both automatic and interactive segmentation. For zero-shot interactive segmentation for non-human CTs e.g. mouse CT or even rock/stone CT, using original resolution (set `resample_spacing` to [-1,-1,-1]) may give better interactive results. ### `use_point_window` + When user click a point, there is no need to perform whole image sliding window inference. Set "use_point_window" to true in the inference.json to enable this function. A window centered at the clicked points will be used for inference. All values outside of the window will set to be "NaN" unless "prev_mask" is passed to the inferer (255 is used to represent NaN). If no point click exists, this function will not be used. Notice if "use_point_window" is true and user provided point clicks, there will be obvious cut-off box artefacts. - ### Inference GPU benchmarks + Benchmarks on a 16GB V100 GPU with 400G system cpu memory. + | Volume size at 1.5x1.5x1.5 mm | 333x333x603 | 512x512x512 | 512x512x768 | 1024x1024x512 | 1024x1024x768 | | :---: | :---: | :---: | :---: | :---: | :---: | |RunTime| 1m07s | 2m09s | 3m25s| 9m20s| killed | +### Execute inference with the TensorRT model - -### Execute inference with the TensorRT model: - -``` +```bash python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']" ``` By default, the argument `head_trt_enabled` is set to `false` in `configs/inference_trt.json`. This means that the `class_head` module of the network will not be converted into a TensorRT model. Setting this to `true` may accelerate the process, but there are some limitations: -Since the `label_prompt` will be converted into a tensor and input into the `class_head` module, the batch size of this input tensor will equal the length of the original `label_prompt` list (if no prompt is provided, the length is 117). To make the TensorRT model work on the `class_head` module, you should set a suitable dynamic batch size range. The maximum dynamic batch size can be configured using the argument `max_prompt_size` in `configs/inference_trt.json`. If the length of the `label_prompt` list exceeds `max_prompt_size`, the engine will fall back to using the normal PyTorch model for inference. Setting a larger `max_prompt_size` can cover more input cases but may require more GPU memory (the default value is 4, which requires 16 GB of GPU memory). Therefore, please set it to a reasonable value according to your actual requirements. - +Since the `label_prompt` will be converted into a tensor and input into the `class_head` module, the batch size of this input tensor will equal the length of the original `label_prompt` list (if no prompt is provided, the length is 117). To make the TensorRT model work on the `class_head` module, you should set a suitable dynamic batch size range. The maximum dynamic batch size can be configured using the argument `max_prompt_size` in `configs/inference_trt.json`. If the length of the `label_prompt` list exceeds `max_prompt_size`, the engine will fall back to using the normal PyTorch model for inference. +Setting a larger `max_prompt_size` can cover more input cases but may require more GPU memory (the default value is 4, which requires 16 GB of GPU memory). Therefore, please set it to a reasonable value according to your actual requirements. ### TensorRT speedup + The `vista3d` bundle supports acceleration with TensorRT. The table below displays the speedup ratios observed on an A100 80G GPU. Please note for 32bit precision models, they are benchmarked with tf32 weight format. | method | torch_tf32(ms) | torch_amp(ms) | trt_tf32(ms) | trt_fp16(ms) | speedup amp | speedup tf32 | speedup fp16 | amp vs fp16| @@ -70,6 +79,7 @@ The `vista3d` bundle supports acceleration with TensorRT. The table below displa | end2end | 6740 | 5166 | 5242 | 3386 | 1.30 | 1.29 | 1.99 | 1.53 | Where: + - `model computation` means the speedup ratio of model's inference with a random input without preprocessing and postprocessing - `end2end` means run the bundle end-to-end with the TensorRT based model. - `torch_tf32` and `torch_amp` are for the PyTorch models with or without `amp` mode. @@ -78,10 +88,11 @@ Where: - `amp vs fp16` is the speedup ratio between the PyTorch amp model and the TensorRT float16 based model. This result is benchmarked under: - - TensorRT: 10.3.0+cuda12.6 - - Torch-TensorRT Version: 2.4.0 - - CPU Architecture: x86-64 - - OS: ubuntu 20.04 - - Python version:3.10.12 - - CUDA version: 12.6 - - GPU models and configuration: A100 80G + +- TensorRT: 10.3.0+cuda12.6 +- Torch-TensorRT Version: 2.4.0 +- CPU Architecture: x86-64 +- OS: ubuntu 20.04 +- Python version:3.10.12 +- CUDA version: 12.6 +- GPU models and configuration: A100 80G diff --git a/NV-Segment-CTMR/brain_t1_preprocess/README.md b/NV-Segment-CTMR/brain_t1_preprocess/README.md index 92dcc1e..766b81c 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/README.md +++ b/NV-Segment-CTMR/brain_t1_preprocess/README.md @@ -7,6 +7,7 @@ Visit the [official website](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/ `python3.8 -u run_synthstrip.py` ## STEP 2 MRI preprocessing + Affine align to the MNI template\ `python3.8 preprocess.py sub-01_T1w.nii.gz mni_icbm152_2009c_t1_1mm_masked_img.nii.gz output.nii.gz --save-preprocess output.preprocess.json` @@ -14,8 +15,9 @@ Affine alignment to the LUMIR template for consistency with the LUMIR dataset\ `python3.8 preprocess.py sub-01_T1w.nii.gz LUMIR_template.nii.gz output.nii.gz --save-preprocess output.preprocess.json` ## STEP 3 Revert to original space + Revert a processed image back to the original space (uses saved metadata)\ `python3.8 revert_preprocess.py output.nii.gz --out output_revert.nii.gz --meta output.preprocess.json` Revert a processed mask back to the original space\ -`python3.8 revert_preprocess.py output.nii.gz --out output_revert.nii.gz --mask output_lbl.nii.gz --mask-out output_lbl_revert.nii.gz --meta output.preprocess.json` \ No newline at end of file +`python3.8 revert_preprocess.py output.nii.gz --out output_revert.nii.gz --mask output_lbl.nii.gz --mask-out output_lbl_revert.nii.gz --meta output.preprocess.json` diff --git a/NV-Segment-CTMR/docs/README.md b/NV-Segment-CTMR/docs/README.md index e65679b..bc65f4b 100644 --- a/NV-Segment-CTMR/docs/README.md +++ b/NV-Segment-CTMR/docs/README.md @@ -1,12 +1,15 @@ # Model Overview -NV-Segment-CTMR is a unified CT and MRI segmentation foundation model. It is based on VISTA3D CT model and extended to both CT and MRI. Please refer to https://github.com/Project-MONAI/VISTA/tree/main/vista3d for more information. + +NV-Segment-CTMR is a unified CT and MRI segmentation foundation model. It is based on VISTA3D CT model and extended to both CT and MRI. Please refer to [VISTA3D repo](https://github.com/Project-MONAI/VISTA/tree/main/vista3d) for more information. ## Performance on held-out test set -
+![Benchmark CT](./benchmarkct.png) ![Benchmark MR](./benchmarkmr.png) + +## Quick Start + +### Installation -### Quick Start -#### Installation ```bash # Create and activate conda environment conda create -y -n vista3d-nv python=3.9 @@ -25,11 +28,13 @@ mkdir -p NV-Segment-CTMR/models wget -O NV-Segment-CTMR/models/model.pt https://huggingface.co/nvidia/NV-Segment-CTMR/resolve/main/vista3d_pretrained_model/model.pt ``` - ## Automatic Segmentation (support multi-gpu batch processing) -We defined 345 classes as in [label_dict.json](../configs/label_dict.json). It shows the label organ name, index, training dataset, modality and evaluation dice score. If a class only comes from CT training dataset, it may not perform well on MRI, but the actual performance will vary case by case. We support three types of segment everything: "CT_BODY", "MRI_BODY", and "MRI_BRAIN". "CT_BODY" is the previous VISTA3D bundle supported 132 CT classes. "MRI_BODY" shares the same 50 label classes as TotalsegmentatorMR. "MRI_BRAIN" is trained on skull stripped [LUMIR](https://github.com/JHU-MedImage-Reg/LUMIR_L2R) dataset and will segment brain MRI substructures. Preprocessing is needed. Follow [tutorials](https://github.com/junyuchen245/MIR/tree/main/tutorials/brain_MRI_preprocessing). The exact mapping for those three everything labels can be found in [metadata.json](../configs/metadata.json). + +We defined 345 classes as in [label_dict.json](../configs/label_dict.json). It shows the label organ name, index, training dataset, modality and evaluation dice score. If a class only comes from CT training dataset, it may not perform well on MRI, but the actual performance will vary case by case. We support three types of segment everything: "CT_BODY", "MRI_BODY", and "MRI_BRAIN". "CT_BODY" is the previous VISTA3D bundle supported 132 CT classes. "MRI_BODY" shares the same 50 label classes as TotalsegmentatorMR. "MRI_BRAIN" is trained on skull stripped [LUMIR](https://github.com/JHU-MedImage-Reg/LUMIR_L2R) dataset and will segment brain MRI substructures. +Preprocessing is needed. Follow [tutorials](https://github.com/junyuchen245/MIR/tree/main/tutorials/brain_MRI_preprocessing). The exact mapping for those three everything labels can be found in [metadata.json](../configs/metadata.json). ## Single image inference to segment everything (automatic) + The output will be saved to `output_dir/s0289/s0289_{output_postfix}{output_ext}`. By default the everything will be "CT_BODY". Add "MRI_BODY" to segment the MRI body classes. ```bash @@ -41,6 +46,7 @@ python -m monai.bundle run --config_file configs/inference.json --input_dict "{' ``` ## Single image inference to segment specific class (automatic) + The detailed automatic segmentation class index can be found [here](../configs/label_dict.json). ```bash @@ -51,6 +57,7 @@ python -m monai.bundle run --config_file configs/inference.json --input_dict "{' ## Batch inference with multiGPU support for segmenting everything (automatic) ### Single-GPU Batch Inference + ```bash # Make sure conda environment is activated conda activate vista3d-nv @@ -60,6 +67,7 @@ python -m monai.bundle run --config_file="['configs/inference.json', 'configs/ba ``` ### Multi-GPU Batch Inference + **Important**: Always activate your conda environment before running `torchrun`. If you don't, you may encounter `ModuleNotFoundError` because `torchrun` will use the system Python instead of your conda environment's Python. ```bash @@ -69,17 +77,17 @@ conda activate vista3d-nv # Automatic Batch segmentation for the whole folder with multi-gpu support # Change --nproc_per_node to match your number of GPUs torchrun --nproc_per_node=2 --nnodes=1 -m monai.bundle run --config_file="['configs/inference.json', 'configs/batch_inference.json', 'configs/mgpu_inference.json']" --input_dir="example/" --output_dir="example/" -``` - +``` `configs/batch_inference.json` by default runs the segment everything workflow (classes defined by `everything_labels`) on all (`*.nii.gz`) files in `input_dir`. This default is overridable by changing the input folder `input_dir`, or the input image name suffix `input_suffix`, or directly setting the list of filenames `input_list`. -``` +```text Note: if using the finetuned checkpoint and the finetuning label_mapping mapped to global index "2, 20, 21", remove the `subclass` dict from inference.json since those values defined in `subclass` will trigger the wrong subclass segmentation. ``` ## Brain MRI segmentation + For brain MRI segmentation, we only support T1 and require preprocessing. We provide a convenient bash script that handles all preprocessing steps automatically. ### Using the Brain Segmentation Script @@ -87,6 +95,7 @@ For brain MRI segmentation, we only support T1 and require preprocessing. We pro The script `brain_t1_preprocess/run_brain_segmentation.sh` automates the entire pipeline: skull stripping, preprocessing, segmentation, and reverting results back to original space. It also handles temporary file cleanup automatically. It is modified from [MIR tutorials](https://github.com/junyuchen245/MIR/tree/main/tutorials/brain_MRI_preprocessing). #### Single File Processing + ```bash # Process a single brain MRI file ./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz @@ -99,6 +108,7 @@ The script `brain_t1_preprocess/run_brain_segmentation.sh` automates the entire ``` #### Batch Processing + ```bash # Process all NIfTI files in a folder ./brain_t1_preprocess/run_brain_segmentation.sh --input_folder example/ --output_dir results/ @@ -108,6 +118,7 @@ The script `brain_t1_preprocess/run_brain_segmentation.sh` automates the entire ``` #### Script Options + - `--input FILE`: Single NIfTI file to segment - `--input_folder FOLDER`: Folder containing NIfTI files (batch mode) - `--output_dir DIR`: Output directory (default: `./eval`) @@ -119,6 +130,7 @@ The script `brain_t1_preprocess/run_brain_segmentation.sh` automates the entire **Note**: The script automatically activates the conda environment and handles all temporary file management. By default, temporary files are cleaned up after processing. Use `--keep-temp` if you need to inspect intermediate results. ### Manual Processing (Advanced) + If you need more control over individual steps, you can run them manually: ```bash @@ -146,7 +158,7 @@ python -m monai.bundle run --config_file configs/inference.json --input_dict "{' python brain_t1_preprocess/revert_preprocess.py $preprocess_tmp --out ${preprocess_tmp}.revert.nii.gz --mask eval/${file}_p/${file}_p_trans.nii.gz --mask-out eval/${file}_trans.nii.gz --meta $preprocess_meta ``` -## Execute inference with the TensorRT model: +## Execute inference with the TensorRT model ```bash # Make sure conda environment is activated @@ -157,11 +169,12 @@ python -m monai.bundle run --config_file "['configs/inference.json', 'configs/in For more details, please refer to [this](inference.md). +## Continual learning / Finetuning -# Continual learning / Finetuning +### Step1: Generate Data json file -## Step1: Generate Data json file Users need to provide a json data split for continuous learning (`configs/msd_task09_spleen_folds.json` from the [MSD](http://medicaldecathlon.com/) is provided as an example). The data split should meet the following format ('testing' labels are optional): + ```json { "training": [ @@ -176,14 +189,20 @@ Users need to provide a json data split for continuous learning (`configs/msd_ta ] } ``` + Example code for 5 fold cross-validation generation can be found [here](data.md) -``` + +```text Note the data is not the absolute path to the image and label file. The actual image file will be `os.path.join(dataset_dir, data["training"][item]["image"])`, where `dataset_dir` is defined in `configs/train_continual.json`. Also 5-fold cross-validation is not required! `fold=0` is defined in train.json, which means any data item with fold==0 will be used as validation and other fold will be used for training. So if you only have train/val split, you can manually set validation data with "fold": 0 in its datalist and the other to be training by setting "fold" to any number other than 0. ``` -## Step2: Changing hyperparameters + +### Step2: Changing hyperparameters + For continual learning, user can change `configs/train_continual.json`. More advanced users can change configurations in `configs/train.json`. Most hyperparameters are straighforward and user can tell based on their names. The users must manually change the following keys in `configs/train_continual.json`. + #### 1. `label_mappings` -``` + +```json "label_mappings": { "default": [ [ @@ -201,37 +220,46 @@ For continual learning, user can change `configs/train_continual.json`. More adv ] }, ``` -`index_1_in_user_data`,...,`index_N_in_user_data` is the class index value in the groundtruth that user tries to segment. `mapped_index_1`,...,`mapped_index_N` is the mapped index value that the bundle will output. You can make these two the same for finetuning, but we suggest finding the semantic relevant mappings from our unified [global label index](../configs/metadata.json). For example, "Spleen" in MSD09 groundtruth label is represented by 1, but "Spleen" is 3 in `docs/labels.json`. So by defining label mapping `[[1, 3]]`, VISTA3D can segment "Spleen" using its pretrained weights out-of-the-box, and can speed up the finetuning convergence speed. If you cannot find a relevant semantic label for your class, just use any value < `num_classes` defined in train_continue.json. For more details about this label_mapping, please read [this](finetune.md). + +`index_1_in_user_data`,...,`index_N_in_user_data` is the class index value in the groundtruth that user tries to segment. `mapped_index_1`,...,`mapped_index_N` is the mapped index value that the bundle will output. You can make these two the same for finetuning, but we suggest finding the semantic relevant mappings from our unified [global label index](../configs/metadata.json). For example, "Spleen" in MSD09 groundtruth label is represented by 1, but "Spleen" is 3 in `docs/labels.json`. So by defining label mapping `[[1, 3]]`, VISTA3D can segment "Spleen" using its pretrained weights out-of-the-box, +and can speed up the finetuning convergence speed. +If you cannot find a relevant semantic label for your class, just use any value < `num_classes` defined in train_continue.json. +For more details about this label_mapping, please read [this](finetune.md). #### 2. `data_list_file_path` and `dataset_dir` + Change `data_list_file_path` to the absolute path of your data json split. Change `dataset_dir` to the root folder that combines with the relative path in the data json split. -#### 3. Optional hyperparameters and details are [here](finetune.md). +#### 3. Optional hyperparameters and details are [here](finetune.md) + Hyperparameter finetuning is important and varies from task to task. ## Step3: Run finetuning + The hyperparameters in `configs/train_continual.json` will overwrite ones in `configs/train.json`. Configs in the back will overide the previous ones if they have the same key. Single-GPU: + ```bash # Make sure conda environment is activated conda activate vista3d-nv python -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json']" + --config_file="['configs/train.json','configs/train_continual.json']" ``` Multi-GPU: + ```bash # Activate conda environment first (CRITICAL!) conda activate vista3d-nv # Change --nproc_per_node to match your number of GPUs torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json','configs/multi_gpu_train.json']" + --config_file="['configs/train.json','configs/train_continual.json','configs/multi_gpu_train.json']" ``` -#### MLFlow Visualization +### MLFlow Visualization MLFlow is enabled by default (defined in train.json, use_mlflow) and the data is stored in the `mlruns/` folder under the bundle's root directory. To launch the MLflow UI and track your experiment data, follow these steps: @@ -239,78 +267,87 @@ MLFlow is enabled by default (defined in train.json, use_mlflow) and the data is 2. Execute the following command to start the MLflow server. This will make the MLflow UI accessible. -```Bash +```bash mlflow ui ``` -# Evaluation +## Evaluation + Evaluation can be used to calculate dice scores for the model or a finetuned model. Change the `ckpt_path` to the checkpoint you wish to evaluate. The dice score is calculated on the original image spacing using `invertd`, while the dice score during finetuning is calculated on resampled space. -``` +```text NOTE: Evaluation does not support point evaluation.`"validate#evaluator#hyper_kwargs#val_head` is always set to `auto`. ``` Single-GPU: + ```bash # Make sure conda environment is activated conda activate vista3d-nv python -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json']" + --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json']" ``` Multi-GPU: + ```bash # Activate conda environment first (CRITICAL!) conda activate vista3d-nv # Change --nproc_per_node to match your number of GPUs torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json','configs/mgpu_evaluate.json']" + --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json','configs/mgpu_evaluate.json']" ``` -#### Other explanatory items + +### Other explanatory items + The `label_mapping` in `evaluation.json` does not include `0` because the postprocessing step performs argmax (`VistaPostTransformd`), and a `0` prediction would negatively impact performance. In continuous learning, however, `0` is included for validation because no argmax is performed, and validation is done channel-wise (include_background=False). Additionally, `Relabeld` in `postprocessing` is required to map `label` and `pred` back to sequential indexes like `0, 1, 2, 3, 4` for dice calculation, as they are not in one-hot format. Evaluation does not support `point`, but finetuning does, as it does not perform argmax. +## FAQ + +### TroubleShoot for Out-of-Memory -# FAQ -## TroubleShoot for Out-of-Memory - Changing `patch_size` to a smaller value such as `"patch_size": [96, 96, 96]` would reduce the training/inference memory footprint. - Changing `train_dataset_cache_rate` and `val_dataset_cache_rate` to a smaller value like `0.1` can solve the out-of-cpu memory issue when using huge finetuning dataset. - Set `"postprocessing#transforms#0#_disabled_": false` to move the postprocessing to cpu to reduce the GPU memory footprint. -## Multi-channel input +### Multi-channel input + - Change `input_channels` in `train.json` to your desired channel number - Data split json can be a single multi-channel image or can be a list of single channeled images. Those images must have the same spatial shape and aligned/registered. -``` + +```json { "image": ["modality1.nii.gz", "modality2.nii.gz", "modality3.nii.gz"] "label": "label.nii.gz" }, ``` -## Wrong inference results from finetuned checkpoint + +### Wrong inference results from finetuned checkpoint + - Make sure you removed the `subclass` dictionary from inference.json if you ever mapped local index to [2,20,21] - Make sure `0` is not included in your inference prompt for automatic segmentation. +## References -# References -- Antonelli, M., Reinke, A., Bakas, S. et al. The Medical Segmentation Decathlon. Nat Commun 13, 4128 (2022). https://doi.org/10.1038/s41467-022-30695-9 - -- VISTA3D: Versatile Imaging SegmenTation and Annotation model for 3D Computed Tomography. arxiv (2024) https://arxiv.org/abs/2406.05285 +- Antonelli, M., Reinke, A., Bakas, S. et al. The Medical Segmentation Decathlon. Nat Commun 13, 4128 (2022). +- VISTA3D: Versatile Imaging SegmenTation and Annotation model for 3D Computed Tomography. arxiv (2024) -# License +## License -## Code License +### Code License This project includes code licensed under the Apache License 2.0. You may obtain a copy of the License at - http://www.apache.org/licenses/LICENSE-2.0 + -## Model Weights License +### Model Weights License The model weights included in this project are licensed under the NCLS v1 License. Both licenses' full texts have been combined into a single `LICENSE` file. Please refer to this `LICENSE` file for more details about the terms and conditions of both licenses. -For MRI CT joint model. The license is non-commercial and needs future discussion. \ No newline at end of file +For MRI CT joint model. The license is non-commercial and needs future discussion. diff --git a/NV-Segment-CTMR/docs/data.md b/NV-Segment-CTMR/docs/data.md index 45595a8..031ece4 100644 --- a/NV-Segment-CTMR/docs/data.md +++ b/NV-Segment-CTMR/docs/data.md @@ -1,5 +1,9 @@ -### Best practice to generate data list +# Data + +## Best practice to generate data list + User can use monai to generate the 5-fold data lists. Full exampls can be found in VISTA3D open source [codebase](https://github.com/Project-MONAI/VISTA/blob/main/vista3d/data/make_datalists.py) + ```python from monai.data.utils import partition_dataset from monai.bundle import ConfigParser diff --git a/NV-Segment-CTMR/docs/finetune.md b/NV-Segment-CTMR/docs/finetune.md index 855350a..7625d6e 100644 --- a/NV-Segment-CTMR/docs/finetune.md +++ b/NV-Segment-CTMR/docs/finetune.md @@ -1,38 +1,46 @@ -### Configurations +# Finetune configurations +## Configurations - -#### Best practice to set label_mapping +### Best practice to set label_mapping For a class that represent the same or similar class as the global index, directly map it to the global index. For example, "mouse left lung" (e.g. index 2 in the mouse dataset) can be mapped to the 28 "left lung upper lobe"(or 29 "left lung lower lobe") with [[2,28]]. After finetuning, 28 now represents "mouse left lung" and will be used for segmentation. If you want to segment 4 substructures of aorta, you can map one of the substructuress to 6 aorta and the rest to any value, [[1,6],[2,133],[3,134],[4,135]]. -``` + +```text NOTE: Do not map to global index value >= 255. `num_classes=255` in the config only represent the maximum mapping index, while the actual output class number only depends on your label_mapping definition. The 255 value in the inference output is also used to represent 'NaN' value. ``` -#### `val_at_start` +### `val_at_start` + Default `true`, VISTA3D will perform out-of-the-box segmentation before the training. Users can disable if the validation takes too long. +### `n_train_samples` and `n_val_samples` -#### `n_train_samples` and `n_val_samples` In `train_continual.json`, only `n_train_samples` and `n_val_samples` are used for training and validation. -#### `patch_size` +### `patch_size` + The patch size parameter is defined in `configs/train_continual.json`: `"patch_size": [128, 128, 128]`. For finetuning purposes, this value needs to be changed acccording to user's task and GPU memory. Usually a larger patch_size will give better final results. `[192,192,128]` is a good value for larger memory GPU. -#### `resample_to_spacing` +### `resample_to_spacing` + The resample_to_spacing parameter is defined in `configs/train_continual.json` and it represents the resolution the model will be trained on. The `1.5,1.5,1.5` mm default is suitable for large CT organs, but for other tasks, this value should be changed to achive the optimal performance. -#### Advanced user: `drop_label_prob` and `drop_point_prob` (in train.json) +### Advanced user: `drop_label_prob` and `drop_point_prob` (in train.json) + VISTA3D is trained to perform both automatic (class prompts) and interactive point segmentation. `drop_label_prob` and `drop_point_prob` means percentage to remove class prompts and point prompts during training respectively. If `drop_point_prob=1`, the model is only finetuning for automatic segmentation, while `drop_label_prob=1` means only finetuning for interactive segmentation. The VISTA3D foundation model is trained with interactive only (drop_label_prob=1) and then froze the point branch and trained with fully automatic segmentation (`drop_point_prob=1`). In this bundle, the training is simplified by jointly training with class prompts and point prompts and both of the drop ratio is set to 0.25. -``` + +```text NOTE: If user doesn't use interactive segmentation, set `drop_point_prob=1` and `drop_label_prob=0` in train.json might provide a faster and easier finetuning process. ``` -#### Other explanatory items + +### Other explanatory items + In `train.json`, `validate[evaluator][val_head]` can be `auto` and `point`. If `auto`, the validation results will be automatic segmentation. If `point`, the validation results will be sampling one positive point per object per patch. The validation scheme of combining auto and point is deprecated due to speed issue. @@ -44,7 +52,8 @@ to the VISTA3D global class index. `label_set` is used to identify the VISTA model classes for providing training prompts. `val_label_set` is used to identify the original training label classes for computing foreground/background mask during validation. The default configs for both variables are derived from the `label_mappings` config and include `[0]`: -``` + +```json "label_set": "$[0] + list(x[1] for x in @label_mappings#default)" "val_label_set": "$[0] + list(x[0] for x in @label_mappings#default)" ``` diff --git a/NV-Segment-CTMR/docs/inference.md b/NV-Segment-CTMR/docs/inference.md index 220a621..653f70a 100644 --- a/NV-Segment-CTMR/docs/inference.md +++ b/NV-Segment-CTMR/docs/inference.md @@ -1,10 +1,16 @@ +# Inference configurations + All the configurations for inference is stored in inference.json, change those parameters: -### `input_dict` + +## `input_dict` + `input_dict` defines the image to segment and the prompt for segmentation. -``` + +```json "input_dict": "$[{'image': '/data/Task09_Spleen/imagesTs/spleen_15.nii.gz', 'label_prompt':[1]}]", "input_dict": "$[{'image': '/data/Task09_Spleen/imagesTs/spleen_15.nii.gz', 'points':[[138,245,18], [271,343,27]], 'point_labels':[1,0]}]" ``` + - The input_dict must include the key `image` which contain the absolute path to the nii image file, and includes prompt keys of `label_prompt`, `points` and `point_labels`. - The `label_prompt` is a list of length `B`, which can perform `B` foreground objects segmentation, e.g. `[2,3,4,5]`. If `B>1`, Point prompts must NOT be provided. - The `points` is of shape `[N, 3]` like `[[x1,y1,z1],[x2,y2,z2],...[xN,yN,zN]]`, representing `N` point coordinates **IN THE ORIGINAL IMAGE SPACE** of a single foreground object. `point_labels` is a list of length [N] like [1,1,0,-1,...], which @@ -12,7 +18,7 @@ matches the `points`. 0 means background, 1 means foreground, -1 means ignoring - **B must be 1 if label_prompt and points are provided together**. The inferer only supports SINGLE OBJECT point click segmentatation. - If no prompt is provided, the model will use `everything_labels` to segment 117 classes: -```Python +```python list(set([i+1 for i in range(132)]) - set([2,16,18,20,21,23,24,25,26,27,128,129,130,131,132])) ``` @@ -20,10 +26,11 @@ list(set([i+1 for i in range(132)]) - set([2,16,18,20,21,23,24,25,26,27,128,129, - To specify a new class for zero-shot segmentation, set the `label_prompt` to a value between 133 and 254. Ensure that `points` and `point_labels` are also provided; otherwise, the inference result will be a tensor of zeros. ### `label_prompt` and `label_dict` + The `label_dict` defined in `configs/metadata.json` has in total 132 classes. However, there are 5 we do not support and we keep them due to legacy issue. So in total VISTA3D support 127 classes. -``` +```text "16, # prostate or uterus" since we already have "prostate" class, "18, # rectum", insufficient data or dataset excluded. "130, # liver tumor" already have hepatic tumor. @@ -34,34 +41,36 @@ VISTA3D support 127 classes. These 5 are excluded in the `everything_labels`. Another 7 tumor and vessel classes are also removed since they will overlap with other organs and make the output messy. To segment those 7 classes, we recommend users to directly set `label_prompt` to those indexes and avoid using them in `everything_labels`. For "Kidney", "Lung", "Bone" (class index [2, 20, 21]), VISTA3D did not directly use the class index for segmentation, but instead convert them to their subclass indexes as defined by `subclass` dict. For example, "2-Kidney" is converted to "14-Left Kidney" + "5-Right Kidney" since "2" is defined in `subclasss` dict. ### `resample_spacing` + The optimal inference resample spacing should be changed according to the task. For monkey data, a high resolution of [1,1,1] showed better automatic inference results. This spacing applies to both automatic and interactive segmentation. For zero-shot interactive segmentation for non-human CTs e.g. mouse CT or even rock/stone CT, using original resolution (set `resample_spacing` to [-1,-1,-1]) may give better interactive results. ### `use_point_window` + When user click a point, there is no need to perform whole image sliding window inference. Set "use_point_window" to true in the inference.json to enable this function. A window centered at the clicked points will be used for inference. All values outside of the window will set to be "NaN" unless "prev_mask" is passed to the inferer (255 is used to represent NaN). If no point click exists, this function will not be used. Notice if "use_point_window" is true and user provided point clicks, there will be obvious cut-off box artefacts. - ### Inference GPU benchmarks + Benchmarks on a 16GB V100 GPU with 400G system cpu memory. + | Volume size at 1.5x1.5x1.5 mm | 333x333x603 | 512x512x512 | 512x512x768 | 1024x1024x512 | 1024x1024x768 | | :---: | :---: | :---: | :---: | :---: | :---: | |RunTime| 1m07s | 2m09s | 3m25s| 9m20s| killed | +### Execute inference with the TensorRT model - -### Execute inference with the TensorRT model: - -``` +```bash python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']" ``` By default, the argument `head_trt_enabled` is set to `false` in `configs/inference_trt.json`. This means that the `class_head` module of the network will not be converted into a TensorRT model. Setting this to `true` may accelerate the process, but there are some limitations: -Since the `label_prompt` will be converted into a tensor and input into the `class_head` module, the batch size of this input tensor will equal the length of the original `label_prompt` list (if no prompt is provided, the length is 117). To make the TensorRT model work on the `class_head` module, you should set a suitable dynamic batch size range. The maximum dynamic batch size can be configured using the argument `max_prompt_size` in `configs/inference_trt.json`. If the length of the `label_prompt` list exceeds `max_prompt_size`, the engine will fall back to using the normal PyTorch model for inference. Setting a larger `max_prompt_size` can cover more input cases but may require more GPU memory (the default value is 4, which requires 16 GB of GPU memory). Therefore, please set it to a reasonable value according to your actual requirements. - +Since the `label_prompt` will be converted into a tensor and input into the `class_head` module, the batch size of this input tensor will equal the length of the original `label_prompt` list (if no prompt is provided, the length is 117). To make the TensorRT model work on the `class_head` module, you should set a suitable dynamic batch size range. The maximum dynamic batch size can be configured using the argument `max_prompt_size` in `configs/inference_trt.json`. If the length of the `label_prompt` list exceeds `max_prompt_size`, the engine will fall back to using the normal PyTorch model for inference. +Setting a larger `max_prompt_size` can cover more input cases but may require more GPU memory (the default value is 4, which requires 16 GB of GPU memory). Therefore, please set it to a reasonable value according to your actual requirements. ### TensorRT speedup + The `vista3d` bundle supports acceleration with TensorRT. The table below displays the speedup ratios observed on an A100 80G GPU. Please note for 32bit precision models, they are benchmarked with tf32 weight format. | method | torch_tf32(ms) | torch_amp(ms) | trt_tf32(ms) | trt_fp16(ms) | speedup amp | speedup tf32 | speedup fp16 | amp vs fp16| @@ -70,6 +79,7 @@ The `vista3d` bundle supports acceleration with TensorRT. The table below displa | end2end | 6740 | 5166 | 5242 | 3386 | 1.30 | 1.29 | 1.99 | 1.53 | Where: + - `model computation` means the speedup ratio of model's inference with a random input without preprocessing and postprocessing - `end2end` means run the bundle end-to-end with the TensorRT based model. - `torch_tf32` and `torch_amp` are for the PyTorch models with or without `amp` mode. @@ -78,10 +88,11 @@ Where: - `amp vs fp16` is the speedup ratio between the PyTorch amp model and the TensorRT float16 based model. This result is benchmarked under: - - TensorRT: 10.3.0+cuda12.6 - - Torch-TensorRT Version: 2.4.0 - - CPU Architecture: x86-64 - - OS: ubuntu 20.04 - - Python version:3.10.12 - - CUDA version: 12.6 - - GPU models and configuration: A100 80G + +- TensorRT: 10.3.0+cuda12.6 +- Torch-TensorRT Version: 2.4.0 +- CPU Architecture: x86-64 +- OS: ubuntu 20.04 +- Python version:3.10.12 +- CUDA version: 12.6 +- GPU models and configuration: A100 80G diff --git a/README.md b/README.md index 8c5da51..9a1bca7 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ This repository contains two NVIDIA medical segmentation foundation models for 3 ## Model Comparison -Both models follow the MONAI bundle architecture. +Both models follow the MONAI bundle architecture. | Feature | NV-Segment-CT | NV-Segment-CTMR | |---------|---------------|-----------------| @@ -13,19 +13,20 @@ Both models follow the MONAI bundle architecture. | **Segmentation Type** | Automatic + Interactive (point-click) | Automatic only | | **Model Weights** | [NV-Segment-CT on HuggingFace](https://huggingface.co/nvidia/NV-Segment-CT) | [NV-Segment-CTMR on HuggingFace](https://huggingface.co/nvidia/NV-Segment-CTMR) | - **NV-Segment-CT** ([`Paper`](https://arxiv.org/pdf/2406.05285)) is a foundation model trained systematically on 11,454 volumes encompassing 127 types of human anatomical structures and various lesions. The model provides State-of-the-art performances on: + - out-of-the-box automatic segmentation on 3D CT scans - zero-shot interactive segmentation in 3D CT scans - automatic segemntation + interactive refinement -**NV-Segment-CTMR** starts from NV-Segment-CT checkpoint and finetuned on over 30K CT and MRI scans, supporting over 300 classes. +**NV-Segment-CTMR** starts from NV-Segment-CT checkpoint and finetuned on over 30K CT and MRI scans, supporting over 300 classes. + - out-of-the-box automatic segmentation on 3D CT scans - share the same architecture with VISTA3D-CT model but we only trained the automatic segmentation branch with larger CT and MRI datasets. ## Performance on held-out test set -
+![Benchmark CT](./NV-Segment-CTMR/docs/benchmarkct.png) ![Benchmark MR](./NV-Segment-CTMR/docs/benchmarkmr.png) ## Resources From 93038192bf064a283064abe88708fd6758c3f3d2 Mon Sep 17 00:00:00 2001 From: Julien Jomier Date: Wed, 25 Feb 2026 13:19:28 -0500 Subject: [PATCH 3/6] Fixing json --- NV-Segment-CT/configs/metadata.json | 2 +- NV-Segment-CT/configs/mgpu_inference.json | 2 +- .../run_brain_segmentation.sh | 49 +++++++++---------- NV-Segment-CTMR/configs/label_dict.json | 2 +- NV-Segment-CTMR/configs/label_mappings.json | 2 +- NV-Segment-CTMR/configs/metadata.json | 2 +- NV-Segment-CTMR/configs/mgpu_inference.json | 2 +- 7 files changed, 30 insertions(+), 31 deletions(-) diff --git a/NV-Segment-CT/configs/metadata.json b/NV-Segment-CT/configs/metadata.json index 01768f2..253aa3d 100644 --- a/NV-Segment-CT/configs/metadata.json +++ b/NV-Segment-CT/configs/metadata.json @@ -767,4 +767,4 @@ } } } -} \ No newline at end of file +} diff --git a/NV-Segment-CT/configs/mgpu_inference.json b/NV-Segment-CT/configs/mgpu_inference.json index 3004868..bc0e9c4 100644 --- a/NV-Segment-CT/configs/mgpu_inference.json +++ b/NV-Segment-CT/configs/mgpu_inference.json @@ -25,4 +25,4 @@ "finalize": [ "$dist.is_initialized() and dist.destroy_process_group()" ] -} \ No newline at end of file +} diff --git a/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh b/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh index 7e52b24..22f60bd 100755 --- a/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh +++ b/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh @@ -55,19 +55,19 @@ check_conda_env() { echo -e "${RED}Error: conda command not found. Please install conda or activate your environment manually.${NC}" >&2 exit 1 fi - + # Check if environment exists if ! conda env list | grep -q "^${CONDA_ENV} "; then echo -e "${YELLOW}Warning: Conda environment '${CONDA_ENV}' not found. Attempting to activate anyway...${NC}" >&2 fi - + # Activate conda environment eval "$(conda shell.bash hook)" conda activate "${CONDA_ENV}" || { echo -e "${RED}Error: Failed to activate conda environment '${CONDA_ENV}'.${NC}" >&2 exit 1 } - + echo -e "${GREEN}Activated conda environment: ${CONDA_ENV}${NC}" } @@ -75,36 +75,36 @@ check_conda_env() { process_single_file() { local input_file="$1" local output_dir="${OUTPUT_DIR:-./eval}" - + if [[ ! -f "$input_file" ]]; then echo -e "${RED}Error: Input file not found: $input_file${NC}" >&2 exit 1 fi - + # Get absolute paths input_file=$(realpath "$input_file") output_dir=$(realpath -m "$output_dir") mkdir -p "$output_dir" - + # Extract filename without extension local file_basename=$(basename "$input_file" .nii.gz) file_basename=$(basename "$file_basename" .nii) local file_dir=$(dirname "$input_file") - + # Create temporary directory for this file local temp_dir="${file_dir}/${file_basename}_temp" mkdir -p "$temp_dir" - + # Temporary file paths local skull_stripped="${temp_dir}/${file_basename}_skull_stripped.nii.gz" local preprocess_tmp="${temp_dir}/${file_basename}_preprocessed.nii.gz" local preprocess_meta="${temp_dir}/${file_basename}_preprocessed.meta.json" local preprocess_tmp_seg="${output_dir}/${file_basename}_preprocessed/${file_basename}_preprocessed_trans.nii.gz" local final_output="${output_dir}/${file_basename}_trans.nii.gz" - + echo -e "${GREEN}Processing: $input_file${NC}" echo -e "${GREEN}Output will be saved to: $final_output${NC}" - + # Step 1: Skull stripping with SynthStrip echo -e "${YELLOW}Step 1/4: Skull stripping...${NC}" if [[ ! -f "$skull_stripped" ]]; then @@ -117,7 +117,7 @@ process_single_file() { else echo -e "${YELLOW} Skull-stripped file already exists, skipping...${NC}" fi - + # Step 2: Affine align to the LUMIR template echo -e "${YELLOW}Step 2/4: Affine alignment to LUMIR template...${NC}" cd "$BUNDLE_ROOT" @@ -130,7 +130,7 @@ process_single_file() { [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" exit 1 } - + # Step 3: Segment the brain echo -e "${YELLOW}Step 3/4: Running segmentation...${NC}" cd "$BUNDLE_ROOT" @@ -142,7 +142,7 @@ process_single_file() { [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" exit 1 } - + # Step 4: Revert the segmentation back to original space echo -e "${YELLOW}Step 4/4: Reverting to original space...${NC}" if [[ ! -f "$preprocess_tmp_seg" ]]; then @@ -150,7 +150,7 @@ process_single_file() { [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" exit 1 fi - + cd "$BUNDLE_ROOT" python brain_t1_preprocess/revert_preprocess.py \ "$preprocess_tmp" \ @@ -162,7 +162,7 @@ process_single_file() { [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" exit 1 } - + # Clean up temporary files if not keeping them if [[ "$KEEP_TEMP" == "false" ]]; then echo -e "${YELLOW}Cleaning up temporary files...${NC}" @@ -175,7 +175,7 @@ process_single_file() { else echo -e "${GREEN}Temporary files kept in: $temp_dir${NC}" fi - + echo -e "${GREEN}✓ Successfully processed: $input_file${NC}" echo -e "${GREEN} Output saved to: $final_output${NC}" } @@ -184,34 +184,34 @@ process_single_file() { process_folder() { local input_folder="$1" local output_dir="${OUTPUT_DIR:-./eval}" - + if [[ ! -d "$input_folder" ]]; then echo -e "${RED}Error: Input folder not found: $input_folder${NC}" >&2 exit 1 fi - + # Get absolute paths input_folder=$(realpath "$input_folder") output_dir=$(realpath -m "$output_dir") mkdir -p "$output_dir" - + # Find all NIfTI files local files=() while IFS= read -r -d '' file; do files+=("$file") done < <(find "$input_folder" -maxdepth 1 -type f \( -name "*.nii.gz" -o -name "*.nii" \) -print0) - + if [[ ${#files[@]} -eq 0 ]]; then echo -e "${YELLOW}Warning: No NIfTI files found in $input_folder${NC}" >&2 exit 1 fi - + echo -e "${GREEN}Found ${#files[@]} file(s) to process${NC}" - + # Process each file local success_count=0 local fail_count=0 - + for file in "${files[@]}"; do echo "" echo -e "${GREEN}========================================${NC}" @@ -222,7 +222,7 @@ process_folder() { echo -e "${RED}Failed to process: $file${NC}" >&2 fi done - + echo "" echo -e "${GREEN}========================================${NC}" echo -e "${GREEN}Batch processing complete!${NC}" @@ -303,4 +303,3 @@ else fi echo -e "${GREEN}All done!${NC}" - diff --git a/NV-Segment-CTMR/configs/label_dict.json b/NV-Segment-CTMR/configs/label_dict.json index 0c359ea..b69852f 100644 --- a/NV-Segment-CTMR/configs/label_dict.json +++ b/NV-Segment-CTMR/configs/label_dict.json @@ -3472,4 +3472,4 @@ "LUMIR" ] } -} \ No newline at end of file +} diff --git a/NV-Segment-CTMR/configs/label_mappings.json b/NV-Segment-CTMR/configs/label_mappings.json index b9c5d03..5d72f18 100644 --- a/NV-Segment-CTMR/configs/label_mappings.json +++ b/NV-Segment-CTMR/configs/label_mappings.json @@ -2614,4 +2614,4 @@ 12 ] ] -} \ No newline at end of file +} diff --git a/NV-Segment-CTMR/configs/metadata.json b/NV-Segment-CTMR/configs/metadata.json index 88e3deb..7cfb169 100644 --- a/NV-Segment-CTMR/configs/metadata.json +++ b/NV-Segment-CTMR/configs/metadata.json @@ -755,4 +755,4 @@ } } } -} \ No newline at end of file +} diff --git a/NV-Segment-CTMR/configs/mgpu_inference.json b/NV-Segment-CTMR/configs/mgpu_inference.json index 3004868..bc0e9c4 100644 --- a/NV-Segment-CTMR/configs/mgpu_inference.json +++ b/NV-Segment-CTMR/configs/mgpu_inference.json @@ -25,4 +25,4 @@ "finalize": [ "$dist.is_initialized() and dist.destroy_process_group()" ] -} \ No newline at end of file +} From 4b79d67bfa2bef8ad534806486d234060ca265d4 Mon Sep 17 00:00:00 2001 From: Julien Jomier Date: Wed, 25 Feb 2026 13:19:38 -0500 Subject: [PATCH 4/6] Fixing README --- NV-Segment-CT/docs/README.md | 114 ++++++++++++++++++++++------------- 1 file changed, 73 insertions(+), 41 deletions(-) diff --git a/NV-Segment-CT/docs/README.md b/NV-Segment-CT/docs/README.md index b85cce7..1d4dc17 100644 --- a/NV-Segment-CT/docs/README.md +++ b/NV-Segment-CT/docs/README.md @@ -1,9 +1,12 @@ # Model Overview + NV-Segment-CT is a copy from the VISTA3D monai model zoo. This is the Vista3D model fintuning/evaluation/inference pipeline. VISTA3D is trained using over 20 partial datasets with more complicated pipeline. To avoid confusion, we will only provide finetuning/continual learning APIs for users to finetune on their -own datasets. To reproduce the paper results, please refer to https://github.com/Project-MONAI/VISTA/tree/main/vista3d +own datasets. To reproduce the paper results, please refer to [VISTA3D repo](https://github.com/Project-MONAI/VISTA/tree/main/vista3d). + +## Quick Start + +### Installation -### Quick Start -#### Installation ```bash # use the same conda env as this repo conda create -y -n vista3d-nv python=3.9 @@ -17,10 +20,12 @@ mkdir NV-Segment-CT/models; wget -O NV-Segment-CT/models/model.pt https://huggingface.co/nvidia/NV-Segment-CT/resolve/main/vista3d_pretrained_model/model.pt ``` -## 1.1 **VISTA3D-CT**[[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CT)[[Huggingface]](https://huggingface.co/nvidia/NV-Segment-CT) +## 1.1 **VISTA3D-CT** [[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CT) [[Huggingface]](https://huggingface.co/nvidia/NV-Segment-CT) + +### Automatic Segmentation (support multi-gpu batch processing) -#### Automatic Segmentation (support multi-gpu batch processing) [class definition](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/blob/main/NV-Segment-CTMR/configs/label_dict.json) + ```bash # CT sementation cd NV-Segment-CT @@ -33,32 +38,36 @@ python -m monai.bundle run --config_file="['configs/inference.json', 'configs/ba # Automatic Batch segmentation for the whole folder with multi-gpu support. mgpu_inference.json is below. change nproc_per_node to your GPU number. torchrun --nproc_per_node=2 --nnodes=1 -m monai.bundle run --config_file="['configs/inference.json', 'configs/batch_inference.json', 'configs/mgpu_inference.json']" --input_dir="example/" --output_dir="example/" ``` -#### Interactive segmentation + +### Interactive segmentation + ```bash # Points must be three dimensional (x,y,z) in the shape of [[x,y,z],...,[x,y,z]]. Point labels can only be -1(ignore), 0(negative), 1(positive) and 2(negative for special overlaped class like tumor), 3(positive for special class). Only supporting 1 class per inference. The output 255 represents NaN value which means not processed region. If you provide label_prompt at the same time, the results will be auto + interactive refinement. cd NV-Segment-CT python -m monai.bundle run --config_file configs/inference.json --input_dict "{'image':'example/spleen_03.nii.gz','points':[[128,128,16], [100,100,16]],'point_labels':[1, 0]}" ``` -**NOTE** MONAI bundle accepts multiple json config files and input arguments. The latter configs/arguments will overide the previous configs/arguments if they have overlapping keys. +**NOTE** MONAI bundle accepts multiple json config files and input arguments. The latter configs/arguments will overide the previous configs/arguments if they have overlapping keys. ## Configuration details and interactive segmentation For inference, VISTA3d bundle requires at least one prompt for segmentation. It supports label prompt, which is the index of the class for automatic segmentation. It also supports point click prompts for binary interactive segmentation. User can provide both prompts at the same time. Please refer to [this](inference.md). -## Execute inference with the TensorRT model: +## Execute inference with the TensorRT model -``` +```bash python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']" ``` + For more details, please refer to [this](inference.md). +## Continual learning / Finetuning -# Continual learning / Finetuning +### Step1: Generate Data json file -## Step1: Generate Data json file Users need to provide a json data split for continuous learning (`configs/msd_task09_spleen_folds.json` from the [MSD](http://medicaldecathlon.com/) is provided as an example). The data split should meet the following format ('testing' labels are optional): + ```json { "training": [ @@ -73,14 +82,20 @@ Users need to provide a json data split for continuous learning (`configs/msd_ta ] } ``` + Example code for 5 fold cross-validation generation can be found [here](data.md) -``` + +```text Note the data is not the absolute path to the image and label file. The actual image file will be `os.path.join(dataset_dir, data["training"][item]["image"])`, where `dataset_dir` is defined in `configs/train_continual.json`. Also 5-fold cross-validation is not required! `fold=0` is defined in train.json, which means any data item with fold==0 will be used as validation and other fold will be used for training. So if you only have train/val split, you can manually set validation data with "fold": 0 in its datalist and the other to be training by setting "fold" to any number other than 0. ``` -## Step2: Changing hyperparameters + +### Step2: Changing hyperparameters + For continual learning, user can change `configs/train_continual.json`. More advanced users can change configurations in `configs/train.json`. Most hyperparameters are straighforward and user can tell based on their names. The users must manually change the following keys in `configs/train_continual.json`. + #### 1. `label_mappings` -``` + +```json "label_mappings": { "default": [ [ @@ -98,27 +113,35 @@ For continual learning, user can change `configs/train_continual.json`. More adv ] }, ``` -`index_1_in_user_data`,...,`index_N_in_user_data` is the class index value in the groundtruth that user tries to segment. `mapped_index_1`,...,`mapped_index_N` is the mapped index value that the bundle will output. You can make these two the same for finetuning, but we suggest finding the semantic relevant mappings from our unified [global label index](../configs/metadata.json). For example, "Spleen" in MSD09 groundtruth label is represented by 1, but "Spleen" is 3 in `docs/labels.json`. So by defining label mapping `[[1, 3]]`, VISTA3D can segment "Spleen" using its pretrained weights out-of-the-box, and can speed up the finetuning convergence speed. If you cannot find a relevant semantic label for your class, just use any value < `num_classes` defined in train_continue.json. For more details about this label_mapping, please read [this](finetune.md). + +`index_1_in_user_data`,...,`index_N_in_user_data` is the class index value in the groundtruth that user tries to segment. `mapped_index_1`,...,`mapped_index_N` is the mapped index value that the bundle will output. You can make these two the same for finetuning, but we suggest finding the semantic relevant mappings from our unified [global label index](../configs/metadata.json). For example, "Spleen" in MSD09 groundtruth label is represented by 1, but "Spleen" is 3 in `docs/labels.json`. So by defining label mapping `[[1, 3]]`, VISTA3D can segment "Spleen" using its pretrained weights out-of-the-box, and can speed up the finetuning convergence speed. +If you cannot find a relevant semantic label for your class, just use any value < `num_classes` defined in train_continue.json. +For more details about this label_mapping, please read [this](finetune.md). #### 2. `data_list_file_path` and `dataset_dir` + Change `data_list_file_path` to the absolute path of your data json split. Change `dataset_dir` to the root folder that combines with the relative path in the data json split. -#### 3. Optional hyperparameters and details are [here](finetune.md). +#### 3. Optional hyperparameters and details are [here](finetune.md) + Hyperparameteers finetuning is important and varies from task to task. -## Step3: Run finetuning +### Step3: Run finetuning + The hyperparameters in `configs/train_continual.json` will overwrite ones in `configs/train.json`. Configs in the back will overide the previous ones if they have the same key. Single-GPU: + ```bash python -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json']" + --config_file="['configs/train.json','configs/train_continual.json']" ``` Multi-GPU: + ```bash torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json','configs/multi_gpu_train.json']" + --config_file="['configs/train.json','configs/train_continual.json','configs/multi_gpu_train.json']" ``` #### MLFlow Visualization @@ -129,71 +152,80 @@ MLFlow is enabled by default (defined in train.json, use_mlflow) and the data is 2. Execute the following command to start the MLflow server. This will make the MLflow UI accessible. -```Bash +```bash mlflow ui ``` -# Evaluation +## Evaluation + Evaluation can be used to calculate dice scores for the model or a finetuned model. Change the `ckpt_path` to the checkpoint you wish to evaluate. The dice score is calculated on the original image spacing using `invertd`, while the dice score during finetuning is calculated on resampled space. -``` +```text NOTE: Evaluation does not support point evaluation.`"validate#evaluator#hyper_kwargs#val_head` is always set to `auto`. ``` Single-GPU: -``` + +```bash python -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json']" + --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json']" ``` Multi-GPU: -``` + +```bash torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run \ - --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json','configs/mgpu_evaluate.json']" + --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json','configs/mgpu_evaluate.json']" ``` -#### Other explanatory items + +### Other explanatory items + The `label_mapping` in `evaluation.json` does not include `0` because the postprocessing step performs argmax (`VistaPostTransformd`), and a `0` prediction would negatively impact performance. In continuous learning, however, `0` is included for validation because no argmax is performed, and validation is done channel-wise (include_background=False). Additionally, `Relabeld` in `postprocessing` is required to map `label` and `pred` back to sequential indexes like `0, 1, 2, 3, 4` for dice calculation, as they are not in one-hot format. Evaluation does not support `point`, but finetuning does, as it does not perform argmax. +## FAQ + +### TroubleShoot for Out-of-Memory -# FAQ -## TroubleShoot for Out-of-Memory - Changing `patch_size` to a smaller value such as `"patch_size": [96, 96, 96]` would reduce the training/inference memory footprint. - Changing `train_dataset_cache_rate` and `val_dataset_cache_rate` to a smaller value like `0.1` can solve the out-of-cpu memory issue when using huge finetuning dataset. - Set `"postprocessing#transforms#0#_disabled_": false` to move the postprocessing to cpu to reduce the GPU memory footprint. -## Multi-channel input +### Multi-channel input + - Change `input_channels` in `train.json` to your desired channel number - Data split json can be a single multi-channel image or can be a list of single channeled images. Those images must have the same spatial shape and aligned/registered. -``` + +```json { "image": ["modality1.nii.gz", "modality2.nii.gz", "modality3.nii.gz"] "label": "label.nii.gz" }, ``` -## Wrong inference results from finetuned checkpoint + +### Wrong inference results from finetuned checkpoint + - Make sure you removed the `subclass` dictionary from inference.json if you ever mapped local index to [2,20,21] - Make sure `0` is not included in your inference prompt for automatic segmentation. +## References -# References -- Antonelli, M., Reinke, A., Bakas, S. et al. The Medical Segmentation Decathlon. Nat Commun 13, 4128 (2022). https://doi.org/10.1038/s41467-022-30695-9 - -- VISTA3D: Versatile Imaging SegmenTation and Annotation model for 3D Computed Tomography. arxiv (2024) https://arxiv.org/abs/2406.05285 +- Antonelli, M., Reinke, A., Bakas, S. et al. The Medical Segmentation Decathlon. Nat Commun 13, 4128 (2022). +- VISTA3D: Versatile Imaging SegmenTation and Annotation model for 3D Computed Tomography. arxiv (2024) -# License +## License -## Code License +### Code License This project includes code licensed under the Apache License 2.0. You may obtain a copy of the License at - http://www.apache.org/licenses/LICENSE-2.0 + -## Model Weights License +### Model Weights License The model weights included in this project are licensed under the NCLS v1 License. Both licenses' full texts have been combined into a single `LICENSE` file. Please refer to this `LICENSE` file for more details about the terms and conditions of both licenses. -For MRI CT joint model. The license is non-commercial and needs furture discussion. \ No newline at end of file +For MRI CT joint model. The license is non-commercial and needs furture discussion. From 8e95a1b454175fc63d125ccf2f639ef10bb9f472 Mon Sep 17 00:00:00 2001 From: Julien Jomier Date: Wed, 25 Feb 2026 13:19:49 -0500 Subject: [PATCH 5/6] Fixing Python --- NV-Segment-CT/scripts/evaluator.py | 15 +-- NV-Segment-CT/scripts/inferer.py | 3 +- NV-Segment-CT/scripts/trainer.py | 7 +- .../intensity_normalization/base_cli.py | 21 ++-- .../intensity_normalization/normalize/base.py | 59 ++++------- .../intensity_normalization/normalize/fcm.py | 18 ++-- .../intensity_normalization/normalize/lsq.py | 39 +++----- .../intensity_normalization/normalize/nyul.py | 14 +-- .../normalize/ravel.py | 27 ++--- .../normalize/whitestripe.py | 9 +- .../intensity_normalization/plot/histogram.py | 17 +--- .../intensity_normalization/typing.py | 99 +++++++------------ .../util/coregister.py | 12 +-- .../util/histogram_tools.py | 13 +-- .../intensity_normalization/util/io.py | 18 ++-- .../util/preprocess.py | 15 +-- .../util/tissue_membership.py | 7 +- .../brain_t1_preprocess/preprocess.py | 85 +++++++++------- .../brain_t1_preprocess/revert_preprocess.py | 89 ++++++++++------- .../brain_t1_preprocess/run_synthstrip.py | 14 ++- .../brain_t1_preprocess/synthstrip-docker | 41 ++++---- NV-Segment-CTMR/scripts/evaluator.py | 15 +-- NV-Segment-CTMR/scripts/inferer.py | 3 +- NV-Segment-CTMR/scripts/trainer.py | 7 +- 24 files changed, 272 insertions(+), 375 deletions(-) diff --git a/NV-Segment-CT/scripts/evaluator.py b/NV-Segment-CT/scripts/evaluator.py index f20261b..c2629e5 100644 --- a/NV-Segment-CT/scripts/evaluator.py +++ b/NV-Segment-CT/scripts/evaluator.py @@ -11,7 +11,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -225,23 +226,17 @@ def _iteration(self, engine: SupervisedEvaluator, batchdata: dict[str, torch.Ten label_prompt, points, point_labels = self.check_prompts_format(label_prompt, points, point_labels) inputs = inputs.to(engine.device) # For N foreground object, label_prompt is [1, N], but the batch number 1 needs to be removed. Convert to [N, 1] - label_prompt = ( - torch.as_tensor([label_prompt]).to(inputs.device)[0].unsqueeze(-1) if label_prompt is not None else None - ) + label_prompt = torch.as_tensor([label_prompt]).to(inputs.device)[0].unsqueeze(-1) if label_prompt is not None else None # For points, the size can only be [1, K, 3], where K is the number of points for this single foreground object. if points is not None: points = torch.as_tensor([points]) - points = self.transform_points( - points, np.linalg.inv(inputs.affine[0]) @ inputs.meta["original_affine"][0].numpy() - ) + points = self.transform_points(points, np.linalg.inv(inputs.affine[0]) @ inputs.meta["original_affine"][0].numpy()) points = torch.from_numpy(points).to(inputs.device) point_labels = torch.as_tensor([point_labels]).to(inputs.device) if point_labels is not None else None # If validation with ground truth label available. else: - inputs, labels = engine.prepare_batch( - batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs - ) + inputs, labels = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) # create label prompt, this should be consistent with the label prompt used for training. if label_set is None: output_classes = engine.hyper_kwargs["output_classes"] diff --git a/NV-Segment-CT/scripts/inferer.py b/NV-Segment-CT/scripts/inferer.py index 18d2290..c396e9f 100644 --- a/NV-Segment-CT/scripts/inferer.py +++ b/NV-Segment-CT/scripts/inferer.py @@ -10,7 +10,6 @@ # limitations under the License. import copy -from typing import List, Union import torch from monai.apps.vista3d.inferer import point_based_window_inferer @@ -36,7 +35,7 @@ def __init__(self, roi_size, overlap, use_point_window=False, sw_batch_size=1) - def __call__( self, - inputs: Union[List[Tensor], Tensor], + inputs: list[Tensor] | Tensor, network, point_coords, point_labels, diff --git a/NV-Segment-CT/scripts/trainer.py b/NV-Segment-CT/scripts/trainer.py index 43d06dc..52cb52d 100644 --- a/NV-Segment-CT/scripts/trainer.py +++ b/NV-Segment-CT/scripts/trainer.py @@ -11,7 +11,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -177,9 +178,7 @@ def _iteration(self, engine, batchdata: dict[str, torch.Tensor]): ) def _compute_pred_loss(): - outputs = engine.network( - input_images=inputs, point_coords=point, point_labels=point_label, class_vector=label_prompt - ) + outputs = engine.network(input_images=inputs, point_coords=point, point_labels=point_label, class_vector=label_prompt) # engine.state.output[Keys.PRED] = outputs engine.fire_event(IterationEvents.FORWARD_COMPLETED) loss, loss_n = torch.tensor(0.0, device=engine.state.device), torch.tensor(0.0, device=engine.state.device) diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/base_cli.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/base_cli.py index 29821e3..17f7567 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/base_cli.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/base_cli.py @@ -14,11 +14,10 @@ import sys import typing -import pymedio.image as mioi - import intensity_normalization as intnorm import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio +import pymedio.image as mioi from intensity_normalization import __version__ as int_norm_version logger = logging.getLogger(__name__) @@ -93,9 +92,7 @@ def parser(cls) -> argparse.ArgumentParser: return parser @classmethod - def main( - cls, parser: argparse.ArgumentParser - ) -> typing.Callable[[intnormt.ArgType], int]: + def main(cls, parser: argparse.ArgumentParser) -> typing.Callable[[intnormt.ArgType], int]: def _main(args: intnormt.ArgType = None) -> int: if args is None: if len(sys.argv) == 2 and sys.argv[1] == "--version": @@ -117,13 +114,11 @@ def _main(args: intnormt.ArgType = None) -> int: @classmethod @abc.abstractmethod - def from_argparse_args(cls: typing.Type[T], args: argparse.Namespace) -> T: + def from_argparse_args(cls: type[T], args: argparse.Namespace) -> T: raise NotImplementedError @abc.abstractmethod - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: raise NotImplementedError @staticmethod @@ -196,9 +191,7 @@ def get_parent_parser( ) return parser - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: image = self.load_image(args.image) mask: intnormt.ImageLike | None if hasattr(args, "mask") and args.mask is not None: @@ -278,7 +271,5 @@ def get_parent_parser( return parser @abc.abstractmethod - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: raise NotImplementedError diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/base.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/base.py index 0a4de33..b8d379c 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/base.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/base.py @@ -21,23 +21,24 @@ import typing import warnings -import pymedio.image as mioi - import intensity_normalization as intnorm import intensity_normalization.base_cli as intnormcli import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio +import pymedio.image as mioi logger = logging.getLogger(__name__) T = typing.TypeVar("T") -#ImageSeq = collections.abc.Sequence[intnormt.ImageLike] -#MaskSeqOrNone = typing.Union[ImageSeq, None] -from typing import Sequence, Union +# ImageSeq = collections.abc.Sequence[intnormt.ImageLike] +# MaskSeqOrNone = typing.Union[ImageSeq, None] +from collections.abc import Sequence +from typing import Union ImageSeq = Sequence[intnormt.ImageLike] MaskSeqOrNone = Union[ImageSeq, None] + class NormalizeMixin(metaclass=abc.ABCMeta): def __call__( self, @@ -80,9 +81,7 @@ def estimate_foreground(image: intnormt.ImageLike, /) -> intnormt.ImageLike: return foreground @staticmethod - def skull_stripped_foreground( - image: intnormt.ImageLike, /, *, background_threshold: float = 1e-6 - ) -> intnormt.ImageLike: + def skull_stripped_foreground(image: intnormt.ImageLike, /, *, background_threshold: float = 1e-6) -> intnormt.ImageLike: if image.min() < 0.0: msg = "Data contains negative values; " msg += "skull-stripped functionality assumes " @@ -102,9 +101,7 @@ def _get_mask( background_threshold: float = 1e-6, ) -> intnormt.ImageLike: if mask is None: - mask = self.skull_stripped_foreground( - image, background_threshold=background_threshold - ) + mask = self.skull_stripped_foreground(image, background_threshold=background_threshold) out: intnormt.ImageLike = mask > 0.0 return out @@ -174,14 +171,12 @@ def normalize_from_filename( modality: intnormt.Modality = intnormt.Modality.T1, ) -> tuple[mioi.Image, mioi.Image | None]: image: mioi.Image = mioi.Image.from_path(image_path) - mask: typing.Optional[mioi.Image] + mask: mioi.Image | None mask = None if mask_path is None else mioi.Image.from_path(mask_path) if out_path is None: out_path = self.append_name_to_file(image_path) logger.info(f"Normalizing image: {image_path}") - normalized = typing.cast( - mioi.Image, self.normalize_image(image, mask, modality=modality) - ) + normalized = typing.cast(mioi.Image, self.normalize_image(image, mask, modality=modality)) logger.info(f"Saving normalized image: {out_path}") normalized.to_filename(out_path) return normalized, mask @@ -193,9 +188,7 @@ def get_parent_parser( valid_modalities: frozenset[str] = intnorm.VALID_MODALITIES, **kwargs: typing.Any, ) -> argparse.ArgumentParser: - parser = super().get_parent_parser( - desc, valid_modalities=valid_modalities, **kwargs - ) + parser = super().get_parent_parser(desc, valid_modalities=valid_modalities, **kwargs) parser.add_argument( "-p", "--plot-histogram", @@ -205,14 +198,12 @@ def get_parent_parser( return parser @abc.abstractmethod - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: raise NotImplementedError @classmethod @abc.abstractmethod - def from_argparse_args(cls: typing.Type[T], args: argparse.Namespace, /) -> T: + def from_argparse_args(cls: type[T], args: argparse.Namespace, /) -> T: raise NotImplementedError def save_additional_info( @@ -231,9 +222,7 @@ def get_parent_parser( valid_modalities: frozenset[str] = intnorm.VALID_MODALITIES, **kwargs: typing.Any, ) -> argparse.ArgumentParser: - parser = super().get_parent_parser( - desc, valid_modalities=valid_modalities, **kwargs - ) + parser = super().get_parent_parser(desc, valid_modalities=valid_modalities, **kwargs) parser.add_argument( "-n", "--norm-value", @@ -244,7 +233,7 @@ def get_parent_parser( return parser @classmethod - def from_argparse_args(cls: typing.Type[T], args: argparse.Namespace, /) -> T: + def from_argparse_args(cls: type[T], args: argparse.Namespace, /) -> T: return cls(norm_value=args.norm_value) # type: ignore[call-arg] @@ -256,9 +245,8 @@ def plot_histogram_from_args( normalized: intnormt.ImageLike, mask: intnormt.ImageLike | None = None, ) -> None: - import matplotlib.pyplot as plt - import intensity_normalization.plot.histogram as intnormhist + import matplotlib.pyplot as plt if args.output is None: output = pathlib.Path(args.image).parent / "hist.pdf" @@ -268,9 +256,7 @@ def plot_histogram_from_args( ax.set_title(self.fullname()) plt.savefig(output) - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: normalized, mask = self.normalize_from_filename( args.image, args.mask, @@ -325,9 +311,8 @@ def plot_histogram_from_args( normalized: ImageSeq, masks: MaskSeqOrNone = None, ) -> None: - import matplotlib.pyplot as plt - import intensity_normalization.plot.histogram as intnormhist + import matplotlib.pyplot as plt if args.output_dir is None: output = pathlib.Path(args.image_dir) / "hist.pdf" @@ -356,9 +341,7 @@ def call_from_argparse_args( normalized, masks = out assert isinstance(normalized, list) image_filenames = intnormio.glob_ext(args.image_dir, ext=args.extension) - output_filenames = [ - self.append_name_to_file(fn, args.output_dir) for fn in image_filenames - ] + output_filenames = [self.append_name_to_file(fn, args.output_dir) for fn in image_filenames] n_images = len(normalized) assert n_images == len(output_filenames) for i, (norm_image, fn) in enumerate(zip(normalized, output_filenames), 1): @@ -375,9 +358,7 @@ def call_from_argparse_args( self.plot_histogram_from_args(args, normalized, _masks) -class DirectoryNormalizeCLI( - SampleNormalizeCLIMixin, intnormcli.DirectoryCLI, metaclass=abc.ABCMeta -): +class DirectoryNormalizeCLI(SampleNormalizeCLIMixin, intnormcli.DirectoryCLI, metaclass=abc.ABCMeta): def fit( self, images: ImageSeq, diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/fcm.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/fcm.py index ba04ea7..0cebfd1 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/fcm.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/fcm.py @@ -12,15 +12,14 @@ import pathlib import typing -import numpy as np -import numpy.typing as npt -import pymedio.image as mioi - import intensity_normalization as intnorm import intensity_normalization.normalize.base as intnormb import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio import intensity_normalization.util.tissue_membership as intnormtm +import numpy as np +import numpy.typing as npt +import pymedio.image as mioi logger = logging.getLogger(__name__) @@ -167,9 +166,7 @@ def add_method_specific_arguments( choices=("wm", "gm", "csf"), help="Reference tissue to use for normalization.", ) - exclusive = parent_parser.add_argument_group( - "mutually exclusive optional arguments" - ) + exclusive = parent_parser.add_argument_group("mutually exclusive optional arguments") group = exclusive.add_mutually_exclusive_group(required=False) group.add_argument( "-m", @@ -183,8 +180,7 @@ def add_method_specific_arguments( "-tm", "--tissue-mask", type=intnormt.file_path(), - help="Path to a mask of a target tissue (usually found through FCM). " - "Provide this if not providing the foreground mask.", + help="Path to a mask of a target tissue (usually found through FCM). " "Provide this if not providing the foreground mask.", ) return parent_parser @@ -193,9 +189,7 @@ def from_argparse_args(cls, args: argparse.Namespace, /) -> FCMNormalize: tt = intnormt.TissueType.from_string(args.tissue_type) return cls(norm_value=args.norm_value, tissue_type=tt) - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: if args.mask is not None: if args.modality is not None: if args.modality.lower() != "t1": diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/lsq.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/lsq.py index 4fa092e..6143212 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/lsq.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/lsq.py @@ -13,25 +13,22 @@ import pathlib import typing -import numpy as np -import numpy.typing as npt -import pymedio.image as mioi - import intensity_normalization as intnorm import intensity_normalization.errors as intnorme import intensity_normalization.normalize.base as intnormb import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio import intensity_normalization.util.tissue_membership as intnormtm +import numpy as np +import numpy.typing as npt +import pymedio.image as mioi logger = logging.getLogger(__name__) S = typing.TypeVar("S", bound=intnormt.ImageLike) -class LeastSquaresNormalize( - intnormb.LocationScaleCLIMixin, intnormb.DirectoryNormalizeCLI -): +class LeastSquaresNormalize(intnormb.LocationScaleCLIMixin, intnormb.DirectoryNormalizeCLI): def __init__(self, *, norm_value: float = 1.0, **kwargs: typing.Any): """Minimize the distance tissue means in a set of images via least-squares""" super().__init__(norm_value=norm_value, **kwargs) @@ -100,9 +97,7 @@ def _fit( tissue_membership, ) - def _fix_tissue_membership( - self, image: intnormt.ImageLike, tissue_membership: S - ) -> S: + def _fix_tissue_membership(self, image: intnormt.ImageLike, tissue_membership: S) -> S: image_ndim = int(image.ndim) tm_ndim = int(tissue_membership.ndim) if tissue_membership.shape[:image_ndim] != image.shape and tm_ndim == 4: @@ -114,14 +109,9 @@ def _fix_tissue_membership( return tissue_membership @staticmethod - def tissue_means( - image: intnormt.ImageLike, /, tissue_membership: intnormt.ImageLike - ) -> npt.NDArray: + def tissue_means(image: intnormt.ImageLike, /, tissue_membership: intnormt.ImageLike) -> npt.NDArray: n_tissues = tissue_membership.shape[-1] - weighted_avgs = [ - np.average(image, weights=tissue_membership[..., i]) - for i in range(n_tissues) - ] + weighted_avgs = [np.average(image, weights=tissue_membership[..., i]) for i in range(n_tissues)] return np.asarray([weighted_avgs]).T def scaling_factor(self, tissue_means: npt.NDArray) -> float: @@ -197,9 +187,7 @@ def from_argparse_args(cls, args: argparse.Namespace, /) -> LeastSquaresNormaliz out = cls(norm_value=args.norm_value) return out - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: if args.load_standard_tissue_means is not None: self.load_standard_tissue_means(args.load_standard_tissue_means) self.fit = lambda *args, **kwargs: None # type: ignore[method-assign] @@ -299,24 +287,19 @@ def add_method_specific_arguments( type=intnormt.file_path(), help="Load a standard tissue means previously fit by the method.", ) - exclusive = parent_parser.add_argument_group( - "mutually exclusive optional arguments" - ) + exclusive = parent_parser.add_argument_group("mutually exclusive optional arguments") group = exclusive.add_mutually_exclusive_group(required=False) group.add_argument( "-m", "--mask-dir", type=intnormt.dir_path(), default=None, - help="Path to a foreground mask for the image. " - "Provide this if not providing a tissue mask " - "(if image is not skull-stripped).", + help="Path to a foreground mask for the image. " "Provide this if not providing a tissue mask " "(if image is not skull-stripped).", ) group.add_argument( "-tm", "--tissue-membership-dir", type=intnormt.dir_path(), - help="Path to a mask of a tissue memberships. " - "Provide this if not providing the foreground mask.", + help="Path to a mask of a tissue memberships. " "Provide this if not providing the foreground mask.", ) return parent_parser diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/nyul.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/nyul.py index 2423d11..392ae56 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/nyul.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/nyul.py @@ -11,14 +11,13 @@ import collections.abc import typing -import numpy as np -import numpy.typing as npt -from scipy.interpolate import interp1d - import intensity_normalization.errors as intnorme import intensity_normalization.normalize.base as intnormb import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio +import numpy as np +import numpy.typing as npt +from scipy.interpolate import interp1d class NyulNormalize(intnormb.DirectoryNormalizeCLI): @@ -221,14 +220,11 @@ def add_method_specific_arguments( "--percentile-step", type=float, default=10.0, - help="Percentile steps between 'percentile-after-min' and " - "'prev-percentile-before-max' for finding standard histogram", + help="Percentile steps between 'percentile-after-min' and " "'prev-percentile-before-max' for finding standard histogram", ) return parent_parser - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: if args.load_standard_histogram is not None: self.load_standard_histogram(args.load_standard_histogram) self.fit = lambda *args, **kwargs: None # type: ignore[method-assign] diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/ravel.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/ravel.py index af68fd3..5c1e9f8 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/ravel.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/ravel.py @@ -15,18 +15,17 @@ import pathlib import typing -import numpy as np -import numpy.typing as npt -import pymedio.image as mioi -import scipy.sparse -import scipy.sparse.linalg - import intensity_normalization.errors as intnorme import intensity_normalization.normalize.base as intnormb import intensity_normalization.normalize.whitestripe as intnormws import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio import intensity_normalization.util.tissue_membership as intnormtm +import numpy as np +import numpy.typing as npt +import pymedio.image as mioi +import scipy.sparse +import scipy.sparse.linalg logger = logging.getLogger(__name__) @@ -151,9 +150,7 @@ def _find_csf_mask( return csf_mask @staticmethod - def _ravel_correction( - control_voxels: npt.NDArray, unwanted_factors: npt.NDArray - ) -> npt.NDArray: + def _ravel_correction(control_voxels: npt.NDArray, unwanted_factors: npt.NDArray) -> npt.NDArray: """Correct control voxels by removing trend from unwanted factors Args: @@ -260,11 +257,7 @@ def create_image_matrix_and_control_voxels( for i, registered in enumerate(registered_images): ctrl_vox = registered[intersection] control_voxels[:, i] = ctrl_vox - logger.debug( - f"Image {i+1} control voxels - " - f"mean: {ctrl_vox.mean():.3f}; " - f"std: {ctrl_vox.std():.3f}" - ) + logger.debug(f"Image {i+1} control voxels - " f"mean: {ctrl_vox.mean():.3f}; " f"std: {ctrl_vox.std():.3f}") else: control_voxels = image_matrix[intersection.flatten(), :] @@ -363,8 +356,7 @@ def add_method_specific_arguments( action="store_false", dest="register", default=True, - help="Do not do deformable registration to find control mask. " - "(Assumes images are deformably co-registered).", + help="Do not do deformable registration to find control mask. " "(Assumes images are deformably co-registered).", ) parser.add_argument( "--sparse-svd", @@ -376,8 +368,7 @@ def add_method_specific_arguments( "--masks-are-csf", action="store_true", default=False, - help="Use this flag if mask directory corresponds to CSF masks " - "instead of brain masks. (Assumes images are deformably co-registered).", + help="Use this flag if mask directory corresponds to CSF masks " "instead of brain masks. (Assumes images are deformably co-registered).", ) parser.add_argument( "--quantile-to-label-csf", diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/whitestripe.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/whitestripe.py index ac69f7d..173d100 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/whitestripe.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/normalize/whitestripe.py @@ -10,17 +10,14 @@ import argparse import typing -import numpy as np -import numpy.typing as npt - import intensity_normalization.normalize.base as intnormb import intensity_normalization.typing as intnormt import intensity_normalization.util.histogram_tools as intnormhisttool +import numpy as np +import numpy.typing as npt -class WhiteStripeNormalize( - intnormb.LocationScaleCLIMixin, intnormb.SingleImageNormalizeCLI -): +class WhiteStripeNormalize(intnormb.LocationScaleCLIMixin, intnormb.SingleImageNormalizeCLI): def __init__( self, *, diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/plot/histogram.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/plot/histogram.py index d62ce7f..96e1de2 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/plot/histogram.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/plot/histogram.py @@ -14,13 +14,12 @@ import typing import warnings -import matplotlib.pyplot as plt -import numpy as np - import intensity_normalization as intnorm import intensity_normalization.base_cli as intnormcli import intensity_normalization.typing as intnormt import intensity_normalization.util.io as intnormio +import matplotlib.pyplot as plt +import numpy as np logger = logging.getLogger(__name__) @@ -96,9 +95,7 @@ def from_directories( exclude: collections.abc.Sequence[str] = ("membership",), **kwargs: typing.Any, ) -> plt.Axes: - images, masks = intnormio.gather_images_and_masks( - image_dir, mask_dir, ext=ext, exclude=exclude - ) + images, masks = intnormio.gather_images_and_masks(image_dir, mask_dir, ext=ext, exclude=exclude) return self(images, masks, **kwargs) @staticmethod @@ -197,12 +194,8 @@ def get_parent_parser( def from_argparse_args(cls, args: argparse.Namespace) -> HistogramPlotter: return cls(figsize=args.figsize, alpha=args.alpha, title=args.title) - def call_from_argparse_args( - self, args: argparse.Namespace, /, **kwargs: typing.Any - ) -> None: - _ = self.from_directories( - args.image_dir, args.mask_dir, ext=args.extension, exclude=args.exclude - ) + def call_from_argparse_args(self, args: argparse.Namespace, /, **kwargs: typing.Any) -> None: + _ = self.from_directories(args.image_dir, args.mask_dir, ext=args.extension, exclude=args.exclude) if args.output is None: args.output = pathlib.Path(args.image_dir).resolve() / "hist.pdf" logger.info(f"Saving histogram: {args.output}.") diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/typing.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/typing.py index 0886366..553411c 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/typing.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/typing.py @@ -32,18 +32,18 @@ ] import argparse -import collections.abc import enum import os import pathlib import typing +from collections.abc import Sequence +from typing import SupportsIndex, Union + +import intensity_normalization as intnorm import numpy as np import numpy.typing as npt -from typing import Union, List, Optional -from typing import Union, Sequence, SupportsIndex -import intensity_normalization as intnorm -ArgType = Union[argparse.Namespace, List[str], Optional[None]] +ArgType = Union[argparse.Namespace, list[str], None | None] PathLike = Union[str, os.PathLike] ShapeLike = Union[SupportsIndex, Sequence[SupportsIndex]] @@ -59,7 +59,7 @@ class Modality(enum.Enum): T2: str = "t2" @classmethod - def from_string(cls: typing.Type, string: str | Modality) -> Modality: + def from_string(cls: type, string: str | Modality) -> Modality: if isinstance(string, cls): modality: Modality = string return modality @@ -285,9 +285,7 @@ class SplitFilename(typing.NamedTuple): ) -def return_none( - func: typing.Callable[[typing.Any, typing.Any], typing.Any] -) -> typing.Callable[[typing.Any, typing.Any], typing.Any]: +def return_none(func: typing.Callable[[typing.Any, typing.Any], typing.Any]) -> typing.Callable[[typing.Any, typing.Any], typing.Any]: def new_func(self: object, string: typing.Any) -> typing.Any: if string is None: return None @@ -416,9 +414,7 @@ def __call__(self, val: typing.Any) -> typing.Any: return self.func(val) -def new_parse_type( - func: typing.Callable[[typing.Any], typing.Any], name: str -) -> NewParseType: +def new_parse_type(func: typing.Callable[[typing.Any], typing.Any], name: str) -> NewParseType: return NewParseType(func, name) @@ -427,95 +423,70 @@ def new_parse_type( U_co = typing.TypeVar("U_co", bound="ImageLike", covariant=True) NBit = typing.TypeVar("NBit", bound=npt.NBitBase) -#Float = typing.Union[np.floating[NBit], float] +# Float = typing.Union[np.floating[NBit], float] Float = Union[np.floating, float] -#Int = typing.Union[np.integer[NBit], int] +# Int = typing.Union[np.integer[NBit], int] Int = Union[int, int] class ImageLike(typing.Protocol[S_co, T_co, U_co]): """support anything that implements the methods here""" - def __gt__(self: T_co, other: typing.Any) -> U_co: - ... + def __gt__(self: T_co, other: typing.Any) -> U_co: ... - def __ge__(self: T_co, other: typing.Any) -> U_co: - ... + def __ge__(self: T_co, other: typing.Any) -> U_co: ... - def __lt__(self: T_co, other: typing.Any) -> U_co: - ... + def __lt__(self: T_co, other: typing.Any) -> U_co: ... - def __le__(self: T_co, other: typing.Any) -> U_co: - ... + def __le__(self: T_co, other: typing.Any) -> U_co: ... - def __and__(self: T_co, other: typing.Any) -> U_co: - ... + def __and__(self: T_co, other: typing.Any) -> U_co: ... - def __or__(self: T_co, other: typing.Any) -> U_co: - ... + def __or__(self: T_co, other: typing.Any) -> U_co: ... - def __add__(self: T_co, other: typing.Any) -> S_co: - ... + def __add__(self: T_co, other: typing.Any) -> S_co: ... - def __sub__(self: T_co, other: typing.Any) -> S_co: - ... + def __sub__(self: T_co, other: typing.Any) -> S_co: ... - def __mul__(self: T_co, other: typing.Any) -> S_co: - ... + def __mul__(self: T_co, other: typing.Any) -> S_co: ... - def __truediv__(self: T_co, other: typing.Any) -> S_co: - ... + def __truediv__(self: T_co, other: typing.Any) -> S_co: ... - def __getitem__(self: T_co, item: typing.Any) -> typing.Any: - ... + def __getitem__(self: T_co, item: typing.Any) -> typing.Any: ... - def __iter__(self: T_co) -> T_co: - ... + def __iter__(self: T_co) -> T_co: ... - def __array__(self) -> npt.NDArray: - ... + def __array__(self) -> npt.NDArray: ... - def sum(self) -> Float | Int: - ... + def sum(self) -> Float | Int: ... @property - def ndim(self) -> Int: - ... + def ndim(self) -> Int: ... def any( self, axis: int | tuple[int, ...] | None = None, - ) -> typing.Any: - ... + ) -> typing.Any: ... - def nonzero(self) -> typing.Any: - ... + def nonzero(self) -> typing.Any: ... - def squeeze(self) -> typing.Any: - ... + def squeeze(self) -> typing.Any: ... @property - def shape(self) -> tuple[int, ...]: - ... + def shape(self) -> tuple[int, ...]: ... - def mean(self) -> float: - ... + def mean(self) -> float: ... - def std(self) -> float: - ... + def std(self) -> float: ... - def min(self) -> float: - ... + def min(self) -> float: ... - def flatten(self: T_co) -> T_co: - ... + def flatten(self: T_co) -> T_co: ... def reshape( self: T_co, *shape: typing.SupportsIndex, order: typing.Literal["A", "C", "F"] | None = ..., - ) -> T_co: - ... + ) -> T_co: ... - def transpose(self: T_co, *axes: int) -> T_co: - ... + def transpose(self: T_co, *axes: int) -> T_co: ... diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/coregister.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/coregister.py index d6080c4..f5727eb 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/coregister.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/coregister.py @@ -12,12 +12,11 @@ import logging import typing -import nibabel as nib -import numpy as np - import intensity_normalization as intnorm import intensity_normalization.base_cli as intnormcli import intensity_normalization.typing as intnormt +import nibabel as nib +import numpy as np logger = logging.getLogger(__name__) @@ -47,13 +46,13 @@ def to_ants(image: ValidImage, /) -> ants.ANTsImage: def register( image: ValidImage, /, - template: typing.Optional[ValidImage] = None, + template: ValidImage | None = None, *, type_of_transform: str = "Affine", interpolator: str = "bSpline", metric: str = "mattes", initial_rigid: bool = True, - template_mask: typing.Optional[ValidImage] = None, + template_mask: ValidImage | None = None, ) -> nib.nifti1.Nifti1Image | ants.ANTsImage: if template is None: standard_mni = ants.get_ants_data("mni") @@ -229,8 +228,7 @@ def get_parent_parser( "-ir", "--initial-rigid", action="store_true", - help="Do a rigid registration before doing " - "the `type_of_transform` registration.", + help="Do a rigid registration before doing " "the `type_of_transform` registration.", ) parser.add_argument( "-v", diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/histogram_tools.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/histogram_tools.py index eebf50a..c40bf38 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/histogram_tools.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/histogram_tools.py @@ -13,17 +13,14 @@ "smooth_histogram", ] +import intensity_normalization as intnorm +import intensity_normalization.typing as intnormt import numpy as np import scipy.signal import statsmodels.api as sm -import intensity_normalization as intnorm -import intensity_normalization.typing as intnormt - -def smooth_histogram( - image: intnormt.ImageLike, / -) -> tuple[intnormt.ImageLike, intnormt.ImageLike]: +def smooth_histogram(image: intnormt.ImageLike, /) -> tuple[intnormt.ImageLike, intnormt.ImageLike]: """Use kernel density estimate to get smooth histogram Args: @@ -118,9 +115,7 @@ def get_first_tissue_mode( return first_tissue_mode -def get_tissue_mode( - image: intnormt.ImageLike, /, *, modality: intnormt.Modality -) -> float: +def get_tissue_mode(image: intnormt.ImageLike, /, *, modality: intnormt.Modality) -> float: """Find the appropriate tissue mode given a modality""" modality_ = modality.value if modality_ in intnorm.PEAK["last"]: diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/io.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/io.py index a648b78..07e0266 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/io.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/io.py @@ -16,12 +16,11 @@ import collections.abc import pathlib import typing -from typing import List -import pymedio.image as mioi import intensity_normalization.typing as intnormt +import pymedio.image as mioi -PymedioImageList = List[mioi.Image] +PymedioImageList = list[mioi.Image] PymedioMaskListOrNone = typing.Union[PymedioImageList, None] @@ -70,11 +69,7 @@ def glob_ext( dirpath = pathlib.Path(dirpath) if not dirpath.is_dir(): raise ValueError("'dirpath' must be a directory.") - filenames = sorted( - dp - for dp in dirpath.resolve().glob(f"*.{ext}") - if all(exc not in str(dp) for exc in exclude) - ) + filenames = sorted(dp for dp in dirpath.resolve().glob(f"*.{ext}") if all(exc not in str(dp) for exc in exclude)) return filenames @@ -106,10 +101,11 @@ def split_filename( return intnormt.SplitFilename(pathlib.Path(path), base, ext) -from typing import Generator, Any, Tuple +from collections.abc import Generator +from typing import Any -Zipped = Generator[Tuple[Any, ...], None, None] -#Zipped = typing.Generator[tuple[typing.Any, ...], None, None] +Zipped = Generator[tuple[Any, ...], None, None] +# Zipped = typing.Generator[tuple[typing.Any, ...], None, None] def zip_with_nones(*args: typing.Sequence[typing.Any] | None) -> Zipped: diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/preprocess.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/preprocess.py index 2dc0647..1d1580f 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/preprocess.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/preprocess.py @@ -17,13 +17,12 @@ import logging import typing -import nibabel as nib -import numpy as np -import pymedio.image as mioi - import intensity_normalization as intnorm import intensity_normalization.base_cli as intnormcli import intensity_normalization.typing as intnormt +import nibabel as nib +import numpy as np +import pymedio.image as mioi logger = logging.getLogger(__name__) @@ -79,9 +78,7 @@ def preprocess( ants_mask = ants_image.get_mask() logger.debug("Starting bias field correction.") - ants_image = ants.n4_bias_field_correction( - ants_image, convergence=n4_convergence_options - ) + ants_image = ants.n4_bias_field_correction(ants_image, convergence=n4_convergence_options) if second_n4_with_smoothed_mask: smoothed_mask = ants.smooth_image(ants_mask, 1.0) logger.debug("Starting 2nd bias field correction.") @@ -260,9 +257,7 @@ def _to_ants(image: typing.Any) -> ants.ANTsImage: if isinstance(image, nib.nifti1.Nifti1Image): ants_image = ants.from_nibabel(image) elif isinstance(image, mioi.Image): - ants_image = ants.from_numpy( - image, origin=image.origin, spacing=image.spacing, direction=image.direction - ) + ants_image = ants.from_numpy(image, origin=image.origin, spacing=image.spacing, direction=image.direction) elif isinstance(image, np.ndarray): ants_image = ants.from_numpy(image) elif isinstance(image, ants.ANTsImage): diff --git a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/tissue_membership.py b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/tissue_membership.py index b808d54..a427850 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/tissue_membership.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/intensity_normalization/util/tissue_membership.py @@ -11,15 +11,14 @@ import operator import typing +import intensity_normalization as intnorm +import intensity_normalization.base_cli as intnormcli +import intensity_normalization.typing as intnormt import numpy as np import numpy.typing as npt import pymedio.image as mioi from skfuzzy import cmeans -import intensity_normalization as intnorm -import intensity_normalization.base_cli as intnormcli -import intensity_normalization.typing as intnormt - def find_tissue_memberships( image: intnormt.ImageLike, diff --git a/NV-Segment-CTMR/brain_t1_preprocess/preprocess.py b/NV-Segment-CTMR/brain_t1_preprocess/preprocess.py index 6cb6e4e..6bc8fc4 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/preprocess.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/preprocess.py @@ -1,35 +1,43 @@ -from intensity_normalization.normalize.kde import KDENormalize -from intensity_normalization.typing import Modality, TissueType -import nibabel as nib -import numpy as np -from scipy.ndimage import zoom -import ants import argparse +import base64 import json import os import shutil -import base64 + +import ants +import nibabel as nib +import numpy as np +from intensity_normalization.normalize.kde import KDENormalize +from intensity_normalization.typing import Modality +from scipy.ndimage import zoom + def reorient_image_to_match(reference_nii, target_nii): reference_ornt = nib.aff2axcodes(reference_nii.affine) target_reoriented = nib.as_closest_canonical(target_nii, enforce_diag=False) target_ornt = nib.aff2axcodes(target_reoriented.affine) - + # If orientations don't match, perform reorientation if target_ornt != reference_ornt: # Calculate the transformation matrix to match the reference orientation - ornt_trans = nib.orientations.ornt_transform(nib.io_orientation(target_reoriented.affine), - nib.io_orientation(reference_nii.affine)) + ornt_trans = nib.orientations.ornt_transform(nib.io_orientation(target_reoriented.affine), nib.io_orientation(reference_nii.affine)) target_reoriented = target_reoriented.as_reoriented(ornt_trans) return target_reoriented -def resampling(img_npy, img_pixdim, tar_pixdim, order, mode='constant'): +def resampling(img_npy, img_pixdim, tar_pixdim, order, mode="constant"): if order == 0: img_npy = img_npy.astype(np.uint16) - img_npy = zoom(img_npy, ((img_pixdim[0] / tar_pixdim[0]), (img_pixdim[1] / tar_pixdim[1]), (img_pixdim[2] / tar_pixdim[2])), order=order, prefilter=False, mode=mode) + img_npy = zoom( + img_npy, + ((img_pixdim[0] / tar_pixdim[0]), (img_pixdim[1] / tar_pixdim[1]), (img_pixdim[2] / tar_pixdim[2])), + order=order, + prefilter=False, + mode=mode, + ) return img_npy + def intensity_norm(img_npy: np.ndarray, mod: Modality) -> np.ndarray: kde_norm = KDENormalize(norm_value=110) out = kde_norm(img_npy.astype(np.float32), modality=mod) @@ -38,6 +46,7 @@ def intensity_norm(img_npy: np.ndarray, mod: Modality) -> np.ndarray: out = np.clip(out / max(vmax, 1e-6), 0, 1).astype(np.float32) return out + def make_affine_from_pixdim(pixdim): # Create a 4x4 affine with spacing along the diagonal affine = np.eye(4) @@ -46,15 +55,20 @@ def make_affine_from_pixdim(pixdim): affine[2, 2] = pixdim[2] return affine + def main(): parser = argparse.ArgumentParser(description="Affine register and normalize image to template space.") parser.add_argument("img_path", help="Path to moving image (NIfTI)") parser.add_argument("template_path", help="Path to template image (NIfTI)") parser.add_argument("output_path", help="Path to save output image (NIfTI)") - parser.add_argument("-m", "--mask", dest="mask_path", default=None, - help="Path to brain mask of the moving image (NIfTI)") - parser.add_argument("-s", "--save-preprocess", dest="meta_path", default=None, - help="Path to save preprocessing metadata (JSON). Defaults to output_path + '.preprocess.json'") + parser.add_argument("-m", "--mask", dest="mask_path", default=None, help="Path to brain mask of the moving image (NIfTI)") + parser.add_argument( + "-s", + "--save-preprocess", + dest="meta_path", + default=None, + help="Path to save preprocessing metadata (JSON). Defaults to output_path + '.preprocess.json'", + ) args = parser.parse_args() img_modality = Modality.T1 @@ -62,22 +76,22 @@ def main(): template_nib = nib.load(args.template_path) original_affine = img_nib.affine.copy() original_shape = img_nib.shape - original_pixdim = img_nib.header.structarr['pixdim'][1:-4] + original_pixdim = img_nib.header.structarr["pixdim"][1:-4] if args.mask_path is not None: mask_nib = nib.load(args.mask_path) mask_nib = reorient_image_to_match(template_nib, mask_nib) img_nib = reorient_image_to_match(template_nib, img_nib) - affine_type = 'Affine' - affine_metric = 'meanSquares' + affine_type = "Affine" + affine_metric = "meanSquares" tar_pixdim = [1.0, 1.0, 1.0] # Target pixel dimensions - img_pixdim = img_nib.header.structarr['pixdim'][1:-4] + img_pixdim = img_nib.header.structarr["pixdim"][1:-4] img_npy = img_nib.get_fdata() if args.mask_path is not None: mask_npy = mask_nib.get_fdata() print(img_npy.shape, mask_npy.shape) img_npy = img_npy * (mask_npy > 0) - + # N4 bias field correction img_ants = ants.from_numpy(img_npy) img_ants = ants.n4_bias_field_correction(img_ants) @@ -85,29 +99,33 @@ def main(): # Intensity normalization img_npy = intensity_norm(img_npy, img_modality) img_npy = resampling(img_npy, img_pixdim, tar_pixdim, order=2) - + # Affine registration tmp_npy = template_nib.get_fdata() tmp_npy = intensity_norm(tmp_npy, img_modality) - + tmp_ants = ants.from_numpy(tmp_npy) img_ants = ants.from_numpy(img_npy) regMovTmp = ants.registration(fixed=tmp_ants, moving=img_ants, type_of_transform=affine_type, aff_metric=affine_metric) - img_ants = ants.apply_transforms(fixed=tmp_ants, moving=img_ants, transformlist=regMovTmp['fwdtransforms'],) - + img_ants = ants.apply_transforms( + fixed=tmp_ants, + moving=img_ants, + transformlist=regMovTmp["fwdtransforms"], + ) + img_npy = img_ants.numpy() nib_img = nib.Nifti1Image(img_npy, template_nib.affine, header=template_nib.header) nib.save(nib_img, args.output_path) output_root = args.output_path - if output_root.endswith('.nii.gz'): + if output_root.endswith(".nii.gz"): output_root = output_root[:-7] else: output_root = os.path.splitext(output_root)[0] - + meta_path = args.meta_path or f"{output_root}.preprocess.json" saved_transforms = [] - for idx, tfm_path in enumerate(regMovTmp['fwdtransforms']): + for idx, tfm_path in enumerate(regMovTmp["fwdtransforms"]): tfm_name = os.path.basename(tfm_path) tfm_out = f"{output_root}.{idx}.{tfm_name}" shutil.copy(tfm_path, tfm_out) @@ -125,13 +143,14 @@ def main(): "original_affine": original_affine.tolist(), "original_shape": list(original_shape), "original_pixdim": original_pixdim.tolist(), - "original_header": base64.b64encode(header_bytes).decode('ascii'), + "original_header": base64.b64encode(header_bytes).decode("ascii"), "target_pixdim": tar_pixdim, "transformlist": saved_transforms, "output_path": os.path.abspath(args.output_path), } - with open(meta_path, 'w', encoding='utf-8') as f: + with open(meta_path, "w", encoding="utf-8") as f: json.dump(meta, f, indent=2) - -if __name__ == '__main__': - main() \ No newline at end of file + + +if __name__ == "__main__": + main() diff --git a/NV-Segment-CTMR/brain_t1_preprocess/revert_preprocess.py b/NV-Segment-CTMR/brain_t1_preprocess/revert_preprocess.py index f709b2e..84a0da1 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/revert_preprocess.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/revert_preprocess.py @@ -1,12 +1,14 @@ -import nibabel as nib -import numpy as np -from scipy.ndimage import zoom -import ants import argparse -import json -import os import base64 import io +import json +import os + +import ants +import nibabel as nib +import numpy as np +from scipy.ndimage import zoom + def reorient_image_to_match(reference_nii, target_nii): reference_ornt = nib.aff2axcodes(reference_nii.affine) @@ -16,45 +18,56 @@ def reorient_image_to_match(reference_nii, target_nii): # If orientations don't match, perform reorientation if target_ornt != reference_ornt: # Calculate the transformation matrix to match the reference orientation - ornt_trans = nib.orientations.ornt_transform(nib.io_orientation(target_reoriented.affine), - nib.io_orientation(reference_nii.affine)) + ornt_trans = nib.orientations.ornt_transform(nib.io_orientation(target_reoriented.affine), nib.io_orientation(reference_nii.affine)) target_reoriented = target_reoriented.as_reoriented(ornt_trans) return target_reoriented -def resampling(img_npy, img_pixdim, tar_pixdim, order, mode='constant'): +def resampling(img_npy, img_pixdim, tar_pixdim, order, mode="constant"): if order == 0: img_npy = img_npy.astype(np.uint16) - img_npy = zoom(img_npy, ((img_pixdim[0] / tar_pixdim[0]), (img_pixdim[1] / tar_pixdim[1]), (img_pixdim[2] / tar_pixdim[2])), order=order, prefilter=False, mode=mode) + img_npy = zoom( + img_npy, + ((img_pixdim[0] / tar_pixdim[0]), (img_pixdim[1] / tar_pixdim[1]), (img_pixdim[2] / tar_pixdim[2])), + order=order, + prefilter=False, + mode=mode, + ) return img_npy def main(): parser = argparse.ArgumentParser(description="Revert preprocessing back to original image space (except intensity normalization).") parser.add_argument("processed_path", help="Path to processed image in template space (NIfTI)") - parser.add_argument("--out", dest="output_path", nargs='?', default=None, - help="Path to save reverted image (NIfTI)") - parser.add_argument("--mask", dest="mask_path", nargs='?', default=None, - help="Path to mask associated with processed image (NIfTI)") - parser.add_argument("--mask-out", dest="mask_output_path", nargs='?', default=None, - help="Path to save reverted mask (NIfTI)") - parser.add_argument("-t", "--transform", dest="transformlist", nargs='+', default=None, - help="Forward transforms from preprocessing (moving->template). Will be inverted.") - parser.add_argument("--meta", dest="meta_path", default=None, - help="Path to preprocessing metadata JSON saved by preprocess.py") - parser.add_argument("--interp", default="linear", choices=["linear", "nearestNeighbor", "bspline"], - help="Interpolator for inverse transform") - parser.add_argument("--resample-order", type=int, default=2, - help="Interpolation order for resampling back to original spacing") - parser.add_argument("--org", dest="original_path", nargs='?', default=None, - help="Path to original image (NIfTI). Optional if --meta includes original geometry") - parser.add_argument("--template", dest="template_path", nargs='?', default=None, - help="Path to template image used in preprocessing (NIfTI). Optional if --meta includes template geometry") + parser.add_argument("--out", dest="output_path", nargs="?", default=None, help="Path to save reverted image (NIfTI)") + parser.add_argument("--mask", dest="mask_path", nargs="?", default=None, help="Path to mask associated with processed image (NIfTI)") + parser.add_argument("--mask-out", dest="mask_output_path", nargs="?", default=None, help="Path to save reverted mask (NIfTI)") + parser.add_argument( + "-t", + "--transform", + dest="transformlist", + nargs="+", + default=None, + help="Forward transforms from preprocessing (moving->template). Will be inverted.", + ) + parser.add_argument("--meta", dest="meta_path", default=None, help="Path to preprocessing metadata JSON saved by preprocess.py") + parser.add_argument("--interp", default="linear", choices=["linear", "nearestNeighbor", "bspline"], help="Interpolator for inverse transform") + parser.add_argument("--resample-order", type=int, default=2, help="Interpolation order for resampling back to original spacing") + parser.add_argument( + "--org", dest="original_path", nargs="?", default=None, help="Path to original image (NIfTI). Optional if --meta includes original geometry" + ) + parser.add_argument( + "--template", + dest="template_path", + nargs="?", + default=None, + help="Path to template image used in preprocessing (NIfTI). Optional if --meta includes template geometry", + ) args = parser.parse_args() meta = None if args.meta_path is not None: - with open(args.meta_path, 'r', encoding='utf-8') as f: + with open(args.meta_path, encoding="utf-8") as f: meta = json.load(f) if args.template_path is None: args.template_path = meta.get("template_path") @@ -104,7 +117,7 @@ def main(): tar_pixdim = [1.0, 1.0, 1.0] if meta is not None and meta.get("target_pixdim") is not None: tar_pixdim = meta.get("target_pixdim") - original_pixdim = original_reoriented.header.structarr['pixdim'][1:-4] + original_pixdim = original_reoriented.header.structarr["pixdim"][1:-4] original_reoriented_npy = original_reoriented.get_fdata() ref_resampled_npy = resampling(original_reoriented_npy, original_pixdim, tar_pixdim, order=2) @@ -128,18 +141,17 @@ def main(): inv_npy = inv_ants.numpy() reverted_npy = resampling(inv_npy, tar_pixdim, original_pixdim, order=args.resample_order) - # Reorient back to original orientation reverted_nib = nib.Nifti1Image(reverted_npy, original_reoriented.affine, header=original_reoriented.header) reverted_nib = reorient_image_to_match(original_nib, reverted_nib) # Save with original header/affine out_header = original_header if original_header is not None else original_nib.header.copy() - qform_code = out_header['qform_code'] - sform_code = out_header['sform_code'] - if hasattr(qform_code, '__len__'): + qform_code = out_header["qform_code"] + sform_code = out_header["sform_code"] + if hasattr(qform_code, "__len__"): qform_code = int(np.array(qform_code).ravel()[0]) - if hasattr(sform_code, '__len__'): + if hasattr(sform_code, "__len__"): sform_code = int(np.array(sform_code).ravel()[0]) out_header.set_qform(original_affine, code=int(qform_code)) out_header.set_sform(original_affine, code=int(sform_code)) @@ -159,7 +171,7 @@ def main(): ) inv_mask_npy = inv_mask_ants.numpy() reverted_mask_npy = resampling(inv_mask_npy, tar_pixdim, original_pixdim, order=0) - + reverted_mask_nib = nib.Nifti1Image(reverted_mask_npy, original_reoriented.affine, header=original_reoriented.header) reverted_mask_nib = reorient_image_to_match(original_nib, reverted_mask_nib) @@ -171,12 +183,13 @@ def main(): if mask_output_path is None: if args.output_path is None: raise ValueError("mask_out is required when output_path is not provided") - mask_output_path = args.output_path.replace('.nii.gz', '.mask.nii.gz') + mask_output_path = args.output_path.replace(".nii.gz", ".mask.nii.gz") if mask_output_path == args.output_path: mask_output_path = f"{args.output_path}.mask.nii.gz" out_mask_nib = nib.Nifti1Image(reverted_mask_nib.get_fdata(), original_affine, header=mask_out_header) nib.save(out_mask_nib, mask_output_path) -if __name__ == '__main__': + +if __name__ == "__main__": main() diff --git a/NV-Segment-CTMR/brain_t1_preprocess/run_synthstrip.py b/NV-Segment-CTMR/brain_t1_preprocess/run_synthstrip.py index aa01658..15e3599 100644 --- a/NV-Segment-CTMR/brain_t1_preprocess/run_synthstrip.py +++ b/NV-Segment-CTMR/brain_t1_preprocess/run_synthstrip.py @@ -1,10 +1,14 @@ import glob import subprocess +img_dir = "MRI/Image/Directory/" -img_dir = 'MRI/Image/Directory/' - -for img_i in glob.glob(img_dir+"*.nii.gz"): - img_name = img_i.split('/')[-1].split('_')[0] +for img_i in glob.glob(img_dir + "*.nii.gz"): + img_name = img_i.split("/")[-1].split("_")[0] print(img_name) - subprocess.call('python3.8 ./synthstrip-docker -i {} -o {} -m {}'.format(img_i, img_dir+'t1w_stripped/'+img_name+'_t1w.nii.gz', img_dir+'t1w_mask/'+img_name+'_t1w_mask.nii.gz'), shell=True) \ No newline at end of file + subprocess.call( + "python3.8 ./synthstrip-docker -i {} -o {} -m {}".format( + img_i, img_dir + "t1w_stripped/" + img_name + "_t1w.nii.gz", img_dir + "t1w_mask/" + img_name + "_t1w_mask.nii.gz" + ), + shell=True, + ) diff --git a/NV-Segment-CTMR/brain_t1_preprocess/synthstrip-docker b/NV-Segment-CTMR/brain_t1_preprocess/synthstrip-docker index 738c950..8e81fb7 100755 --- a/NV-Segment-CTMR/brain_t1_preprocess/synthstrip-docker +++ b/NV-Segment-CTMR/brain_t1_preprocess/synthstrip-docker @@ -1,6 +1,5 @@ #!/usr/bin/env python -from __future__ import print_function # -------------------------------- SynthStrip -------------------------------- @@ -11,23 +10,23 @@ from __future__ import print_function # default FreeSurfer `mri_synthstrip` command (use the --help flag for more info). # Upon first use, the relevant docker image will be automatically pulled from # DockerHub. To use a different SynthStrip version, update the variable below. -version = '1.8' +version = "1.8" # ---------------------------------------------------------------------------- import os -import sys -import subprocess import shutil +import subprocess +import sys # Sanity check on env -if shutil.which('docker') is None: - print('Cannot find docker in PATH. Make sure it is installed.') +if shutil.which("docker") is None: + print("Cannot find docker in PATH. Make sure it is installed.") exit(1) # Since we're wrapping a Docker image, we want to get the full paths of all input and output # files so that we can mount their corresponding paths. Tedious, but a fine option for now... -flags = ['-i', '--input', '-o', '--output', '-m', '--mask', '-d', '--sdt', '--model'] +flags = ["-i", "--input", "-o", "--output", "-m", "--mask", "-d", "--sdt", "--model"] # Loop through the arguments and expand any necessary paths idx = 1 @@ -42,47 +41,43 @@ while idx < len(sys.argv): args.append(path) paths.append(path) idx += 1 -args = ' '.join(args) +args = " ".join(args) # Get the unique mount points mounts = list(set([os.path.dirname(p) for p in paths])) -mounts = ' '.join(['-v %s:%s' % (p, p) for p in mounts]) +mounts = " ".join(["-v %s:%s" % (p, p) for p in mounts]) # Set UID and GID to avoid output files owned by root -user = '-u %s:%s' % (os.getuid(), os.getgid()) +user = "-u %s:%s" % (os.getuid(), os.getgid()) -print('Running SynthStrip version %s from Docker' % version) +print("Running SynthStrip version %s from Docker" % version) # Get image tag -image = 'freesurfer/synthstrip:' + version +image = "freesurfer/synthstrip:" + version # Let's check to see if we have this container on the system -proc = subprocess.Popen('docker images -q %s' % image, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - shell=True, - universal_newlines=True) +proc = subprocess.Popen("docker images -q %s" % image, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, universal_newlines=True) stdout, stderr = proc.communicate() if proc.returncode != 0: print(stderr) - print('Error running docker command. Make sure Docker is installed.') + print("Error running docker command. Make sure Docker is installed.") exit(proc.returncode) # If not, let's download it. Normally, docker run will do this automatically, # but we're trying to be transparent here... if not stdout: - print('Docker image %s is not installed. Downloading now. This only needs to be done once.' % image) - proc = subprocess.Popen('docker pull %s' % image, shell=True) + print("Docker image %s is not installed. Downloading now. This only needs to be done once." % image) + proc = subprocess.Popen("docker pull %s" % image, shell=True) proc.communicate() if proc.returncode != 0: - print('Error running docker pull.') + print("Error running docker pull.") exit(proc.returncode) # Go ahead and run the entry point -command = 'docker run %s %s %s %s' % (user, mounts, image, args) +command = "docker run %s %s %s %s" % (user, mounts, image, args) proc = subprocess.Popen(command, shell=True) proc.communicate() if proc.returncode == 137: - print('Container ran out of memory, try increasing RAM in Docker preferences.') + print("Container ran out of memory, try increasing RAM in Docker preferences.") exit(proc.returncode) diff --git a/NV-Segment-CTMR/scripts/evaluator.py b/NV-Segment-CTMR/scripts/evaluator.py index b20e87d..d95603d 100644 --- a/NV-Segment-CTMR/scripts/evaluator.py +++ b/NV-Segment-CTMR/scripts/evaluator.py @@ -11,7 +11,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -225,23 +226,17 @@ def _iteration(self, engine: SupervisedEvaluator, batchdata: dict[str, torch.Ten label_prompt, points, point_labels = self.check_prompts_format(label_prompt, points, point_labels) inputs = inputs.to(engine.device) # For N foreground object, label_prompt is [1, N], but the batch number 1 needs to be removed. Convert to [N, 1] - label_prompt = ( - torch.as_tensor([label_prompt]).to(inputs.device)[0].unsqueeze(-1) if label_prompt is not None else None - ) + label_prompt = torch.as_tensor([label_prompt]).to(inputs.device)[0].unsqueeze(-1) if label_prompt is not None else None # For points, the size can only be [1, K, 3], where K is the number of points for this single foreground object. if points is not None: points = torch.as_tensor([points]) - points = self.transform_points( - points, np.linalg.inv(inputs.affine[0]) @ inputs.meta["original_affine"][0].numpy() - ) + points = self.transform_points(points, np.linalg.inv(inputs.affine[0]) @ inputs.meta["original_affine"][0].numpy()) points = torch.from_numpy(points).to(inputs.device) point_labels = torch.as_tensor([point_labels]).to(inputs.device) if point_labels is not None else None # If validation with ground truth label available. else: - inputs, labels = engine.prepare_batch( - batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs - ) + inputs, labels = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) # create label prompt, this should be consistent with the label prompt used for training. if label_set is None: output_classes = engine.hyper_kwargs["output_classes"] diff --git a/NV-Segment-CTMR/scripts/inferer.py b/NV-Segment-CTMR/scripts/inferer.py index 18d2290..c396e9f 100644 --- a/NV-Segment-CTMR/scripts/inferer.py +++ b/NV-Segment-CTMR/scripts/inferer.py @@ -10,7 +10,6 @@ # limitations under the License. import copy -from typing import List, Union import torch from monai.apps.vista3d.inferer import point_based_window_inferer @@ -36,7 +35,7 @@ def __init__(self, roi_size, overlap, use_point_window=False, sw_batch_size=1) - def __call__( self, - inputs: Union[List[Tensor], Tensor], + inputs: list[Tensor] | Tensor, network, point_coords, point_labels, diff --git a/NV-Segment-CTMR/scripts/trainer.py b/NV-Segment-CTMR/scripts/trainer.py index 43d06dc..52cb52d 100644 --- a/NV-Segment-CTMR/scripts/trainer.py +++ b/NV-Segment-CTMR/scripts/trainer.py @@ -11,7 +11,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -177,9 +178,7 @@ def _iteration(self, engine, batchdata: dict[str, torch.Tensor]): ) def _compute_pred_loss(): - outputs = engine.network( - input_images=inputs, point_coords=point, point_labels=point_label, class_vector=label_prompt - ) + outputs = engine.network(input_images=inputs, point_coords=point, point_labels=point_label, class_vector=label_prompt) # engine.state.output[Keys.PRED] = outputs engine.fire_event(IterationEvents.FORWARD_COMPLETED) loss, loss_n = torch.tensor(0.0, device=engine.state.device), torch.tensor(0.0, device=engine.state.device) From 5938a580baa8ba6903f517846f12198c4611f268 Mon Sep 17 00:00:00 2001 From: Julien Jomier Date: Wed, 25 Feb 2026 14:30:45 -0500 Subject: [PATCH 6/6] Small fix --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3bbbd10..004b27b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,4 +1,4 @@ -# CI for NV-Generate-CTMR: lint and format checks via pre-commit +# CI for NV-Segment-CTMR: lint and format checks via pre-commit name: CI on: