diff --git a/NV-Segment-CT/docs/README.md b/NV-Segment-CT/docs/README.md index 6df05bc..7a7013f 100644 --- a/NV-Segment-CT/docs/README.md +++ b/NV-Segment-CT/docs/README.md @@ -14,12 +14,11 @@ conda activate vista3d-nv git clone https://github.com/NVIDIA-Medtech/NV-Segment-CTMR.git cd NV-Segment-CTMR/NV-Segment-CT; pip install -r requirements.txt; -cd ..; -mkdir NV-Segment-CT/models; -# download from huggingface link -hf download nvidia/NV-Segment-CT vista3d_pretrained_model/model.pt --local-dir NV-Segment-CT/models/ && \ -mv NV-Segment-CT/models/vista3d_pretrained_model/model.pt NV-Segment-CT/models/model.pt && \ -rmdir NV-Segment-CT/models/vista3d_pretrained_model + +mkdir -p models +# Option 1: Download using hf and move to expected location +hf download nvidia/NV-Segment-CT --local-dir models/ && \ +mv models/vista3d_pretrained_model/model.pt models/model.pt ``` ## 1.1 **NV-Segment-CT** [[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CT) [[Huggingface]](https://huggingface.co/nvidia/NV-Segment-CT) @@ -40,7 +39,9 @@ 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/" ``` - +``` +Note: For more details about batch processing, please refer to NV-Segment-CTMR readme.md +``` ### Interactive segmentation ```bash @@ -66,149 +67,13 @@ For more details, please refer to [this](inference.md). ## Continual learning / Finetuning -### 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": [ - {"image": "img0001.nii.gz", "label": "label0001.nii.gz", "fold": 0}, - {"image": "img0002.nii.gz", "label": "label0002.nii.gz", "fold": 2}, - ... - ], - "testing": [ - {"image": "img0003.nii.gz", "label": "label0003.nii.gz"}, - {"image": "img0004.nii.gz", "label": "label0004.nii.gz"}, - ... - ] -} -``` - -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 - -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": [ - [ - index_1_in_user_data, # e.g. 1 - mapped_index_1, # e.g. 1 - ], - [ - index_2_in_user_data, # e.g. 2 - mapped_index_2, # e.g. 2 - ], ..., - [ - index_last_in_user_data, # e.g. N - mapped_index_N, # e.g. N - ] - ] - }, -``` - -`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) - -Hyperparameteers 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 -python -m monai.bundle run \ - --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']" -``` - -#### 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: - -1. Open a terminal and navigate to the root directory of your bundle where the `mlruns/` folder is located. - -2. Execute the following command to start the MLflow server. This will make the MLflow UI accessible. - -```bash -mlflow ui -``` - -## 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`. +We provide predefined finetuning tutorial in [details](inference.md). +For complicated finetuning, we suggest users to do vibe coding to generate finetuning pipelines by simply reuse the model and checkpoint +```python +from monai.networks.nets.vista3d import vista3d132 +vista3d132.load_state_dict(pretrained_ckpt, strict=True) ``` -Single-GPU: - -```bash -python -m monai.bundle run \ - --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']" -``` - -### 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 - -- 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 - -- 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 - -- 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 - He, Yufan, et al. "VISTA3D: A unified segmentation foundation model for 3D medical imaging." Proceedings of the Computer Vision and Pattern Recognition Conference. 2025. diff --git a/NV-Segment-CT/docs/finetune.md b/NV-Segment-CT/docs/finetune.md index 7625d6e..d9dd7a9 100644 --- a/NV-Segment-CT/docs/finetune.md +++ b/NV-Segment-CT/docs/finetune.md @@ -1,5 +1,149 @@ # Finetune configurations +### 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": [ + {"image": "img0001.nii.gz", "label": "label0001.nii.gz", "fold": 0}, + {"image": "img0002.nii.gz", "label": "label0002.nii.gz", "fold": 2}, + ... + ], + "testing": [ + {"image": "img0003.nii.gz", "label": "label0003.nii.gz"}, + {"image": "img0004.nii.gz", "label": "label0004.nii.gz"}, + ... + ] +} +``` + +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 + +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": [ + [ + index_1_in_user_data, # e.g. 1 + mapped_index_1, # e.g. 1 + ], + [ + index_2_in_user_data, # e.g. 2 + mapped_index_2, # e.g. 2 + ], ..., + [ + index_last_in_user_data, # e.g. N + mapped_index_N, # e.g. N + ] + ] + }, +``` + +`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) + +Hyperparameteers 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 +python -m monai.bundle run \ + --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']" +``` + +#### 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: + +1. Open a terminal and navigate to the root directory of your bundle where the `mlruns/` folder is located. + +2. Execute the following command to start the MLflow server. This will make the MLflow UI accessible. + +```bash +mlflow ui +``` + +## 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']" +``` + +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']" +``` + +### 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 + +- 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 + +- 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 + +- 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. + + ## Configurations ### Best practice to set label_mapping diff --git a/NV-Segment-CTMR/docs/README.md b/NV-Segment-CTMR/docs/README.md index 600b9b1..4147d17 100644 --- a/NV-Segment-CTMR/docs/README.md +++ b/NV-Segment-CTMR/docs/README.md @@ -35,12 +35,10 @@ cd NV-Segment-CTMR/NV-Segment-CTMR pip install -r requirements.txt # Create models directory and download pretrained model -cd .. -mkdir -p NV-Segment-CTMR/models +mkdir -p models # Option 1: Download using hf and move to expected location -hf download nvidia/NV-Segment-CTMR vista3d_pretrained_model/model.pt --local-dir NV-Segment-CTMR/models/ && \ -mv NV-Segment-CTMR/models/vista3d_pretrained_model/model.pt NV-Segment-CTMR/models/model.pt && \ -rmdir NV-Segment-CTMR/models/vista3d_pretrained_model +hf download nvidia/NV-Segment-CTMR --local-dir models/ && \ +mv models/vista3d_pretrained_model/model.pt models/model.pt ``` ## Automatic Segmentation (support multi-gpu batch processing) @@ -121,8 +119,6 @@ Note: if using the finetuned checkpoint and the finetuning label_mapping mapped ## Brain MRI segmentation (any MRI sequence) -For brain MRI segmentation, we require preprocessing. The script `brain_t1_preprocess/run_brain_segmentation.sh` runs from the bundle root: it changes into the NV-Segment-CTMR directory next to `brain_t1_preprocess/`, so invoke it from the repo (paths like `example/...` are relative to that root). Activate your Python/conda environment (for example `vista3d-nv`) **before** running the script; the script does not run `conda activate` for you. - ### Using the Brain Segmentation Script The script automates skull stripping (SynthStrip via `brain_t1_preprocess/synthstrip-docker`), affine alignment to the LUMIR template, MONAI bundle inference, and reverting the mask to the original image space. Temporary files are removed after each case unless you pass `--keep-temp`. It is modified from [MIR tutorials](https://github.com/junyuchen245/MIR/tree/main/tutorials/brain_MRI_preprocessing). @@ -132,11 +128,13 @@ The script automates skull stripping (SynthStrip via `brain_t1_preprocess/synths Output path: `{output_dir}/{basename}_trans.nii.gz` (default `output_dir` is `./eval`). ```bash -conda activate vista3d-nv # or your env with MONAI + deps +conda activate vista3d-nv ./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz ./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz --output_dir results/ +# keep temperary files for skull stripping and registration ./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz --keep-temp +# if you cannot perform skull stripping with synthstrip-docker, e.g. on a remote cluster in docker env, you can skull strip the data first then run segmentation with skull strip skipped. ./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz --no-skullstrip ``` @@ -158,6 +156,7 @@ Lines in the list are paths relative to `root_path` (comments and empty lines al --file_list file_list.txt --root_path /path/to/root --output_dir /path/to/output ./brain_t1_preprocess/run_brain_segmentation.sh \ --file_list file_list.txt --root_path /path/to/root --output_dir /path/to/output --no-skullstrip +# split into multiple partitions to submit multiple jobs ./brain_t1_preprocess/run_brain_segmentation.sh \ --file_list file_list.txt --root_path /path/to/root --output_dir /path/to/output \ --num_partitions 10 --partition 3 @@ -216,17 +215,18 @@ conda activate vista3d-nv 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 -See [details](docs/inference.md) +# Continual learning / Finetuning +We provide predefined finetuning tutorial in [details](inference.md). +For complicated finetuning, we suggest users to do vibe coding to generate finetuning pipelines by simply reuse the model and checkpoint +```python +from monai.networks.nets.vista3d import vista3d132 +vista3d132.load_state_dict(pretrained_ckpt, strict=True) +``` ## 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](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](https://arxiv.org/abs/2406.05285) +- He, Yufan, et al. "VISTA3D: A unified segmentation foundation model for 3D medical imaging." Proceedings of the Computer Vision and Pattern Recognition Conference. 2025. ## License