From b851d0c714c91c3ead55793c9dba86c1aa08d400 Mon Sep 17 00:00:00 2001 From: heyufan Date: Mon, 18 May 2026 13:24:48 -0400 Subject: [PATCH 1/2] Trigger Hugging Face access during inference --- NV-Segment-CT/configs/batch_inference.json | 6 +- NV-Segment-CT/configs/inference.json | 7 +- NV-Segment-CT/configs/mgpu_inference.json | 2 + NV-Segment-CT/docs/README.md | 7 +- NV-Segment-CT/scripts/__init__.py | 1 + NV-Segment-CT/scripts/huggingface_download.py | 88 +++++++++++++++++++ NV-Segment-CT/scripts/test_batch_inference.sh | 41 +++++++++ NV-Segment-CTMR/configs/batch_inference.json | 6 +- NV-Segment-CTMR/configs/inference.json | 7 +- NV-Segment-CTMR/configs/mgpu_inference.json | 2 + NV-Segment-CTMR/docs/README.md | 12 +-- NV-Segment-CTMR/scripts/__init__.py | 1 + .../scripts/huggingface_download.py | 88 +++++++++++++++++++ .../scripts/test_batch_inference.sh | 44 ++++++++++ 14 files changed, 293 insertions(+), 19 deletions(-) create mode 100644 NV-Segment-CT/scripts/huggingface_download.py create mode 100755 NV-Segment-CT/scripts/test_batch_inference.sh create mode 100644 NV-Segment-CTMR/scripts/huggingface_download.py create mode 100755 NV-Segment-CTMR/scripts/test_batch_inference.sh diff --git a/NV-Segment-CT/configs/batch_inference.json b/NV-Segment-CT/configs/batch_inference.json index 2118459..f94160b 100644 --- a/NV-Segment-CT/configs/batch_inference.json +++ b/NV-Segment-CT/configs/batch_inference.json @@ -1,14 +1,16 @@ { "input_dir": "@bundle_root", "input_suffix": "*.nii.gz", - "input_root_abs": "$os.path.abspath(@input_dir)", + "input_root_abs": "$os.path.realpath(@input_dir)", + "output_dir_abs": "$os.path.realpath(@output_dir)", "batch_skip_dir_names": [], "batch_skip_dir_prefixes": [], "batch_resume_skip_existing": true, "batch_use_input_list_cache": true, "batch_cache_wait_sec": 120, - "input_list": "$scripts.batch_inference_utils.build_input_list(os.path.abspath(@input_dir), os.path.abspath(@output_dir), @output_postfix, @output_ext, @batch_skip_dir_names, @batch_skip_dir_prefixes, @batch_resume_skip_existing, @batch_use_input_list_cache, @batch_cache_wait_sec)", + "input_list": "$scripts.batch_inference_utils.build_input_list(@input_root_abs, @output_dir_abs, @output_postfix, @output_ext, @batch_skip_dir_names, @batch_skip_dir_prefixes, @batch_resume_skip_existing, @batch_use_input_list_cache, @batch_cache_wait_sec)", "input_dicts": "$[{'image': x, 'label_prompt': @everything_labels} for x in @input_list]", "dataset#data": "@input_dicts", + "postprocessing#transforms#4#output_dir": "@output_dir_abs", "postprocessing#transforms#4#data_root_dir": "@input_root_abs" } diff --git a/NV-Segment-CT/configs/inference.json b/NV-Segment-CT/configs/inference.json index 92f2c0d..84faef6 100644 --- a/NV-Segment-CT/configs/inference.json +++ b/NV-Segment-CT/configs/inference.json @@ -9,6 +9,10 @@ "$import pathlib" ], "bundle_root": "./", + "huggingface_repo_id": "nvidia/NV-Segment-CT", + "huggingface_checkpoint_file": "vista3d_pretrained_model/model.pt", + "huggingface_download_counter_file": "config.json", + "checkpoint_path": "$@bundle_root + '/models/model.pt'", "image_key": "image", "output_dir": "$@bundle_root + '/eval'", "output_ext": ".nii.gz", @@ -174,7 +178,7 @@ ], "checkpointloader": { "_target_": "CheckpointLoader", - "load_path": "$@bundle_root + '/models/model.pt'", + "load_path": "@checkpoint_path", "load_dict": { "model": "@network" }, @@ -195,6 +199,7 @@ } }, "initialize": [ + "$scripts.prepare_huggingface_checkpoint(@huggingface_repo_id, @huggingface_checkpoint_file, @checkpoint_path, @huggingface_download_counter_file)", "$monai.utils.set_determinism(seed=123)", "$@checkpointloader(@evaluator)" ], diff --git a/NV-Segment-CT/configs/mgpu_inference.json b/NV-Segment-CT/configs/mgpu_inference.json index eecf12b..a377e98 100644 --- a/NV-Segment-CT/configs/mgpu_inference.json +++ b/NV-Segment-CT/configs/mgpu_inference.json @@ -18,6 +18,8 @@ "$import torch.distributed as dist", "$dist.is_initialized() or dist.init_process_group(backend='nccl')", "$torch.cuda.set_device(@device)", + "$scripts.prepare_huggingface_checkpoint(@huggingface_repo_id, @huggingface_checkpoint_file, @checkpoint_path, @huggingface_download_counter_file)", + "$dist.barrier()", "$monai.utils.set_determinism(seed=123)", "$@checkpointloader(@evaluator)" ], diff --git a/NV-Segment-CT/docs/README.md b/NV-Segment-CT/docs/README.md index 0903791..9daa3a0 100644 --- a/NV-Segment-CT/docs/README.md +++ b/NV-Segment-CT/docs/README.md @@ -14,13 +14,10 @@ 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; - -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 ``` +Model weights are prepared automatically during inference. The first run downloads the checkpoint from Hugging Face into the local Hugging Face cache and links it at `models/model.pt`; later runs reuse the cached weights while still touching Hugging Face download stats for each inference. + ## 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) ### Automatic Segmentation (support multi-gpu batch processing) diff --git a/NV-Segment-CT/scripts/__init__.py b/NV-Segment-CT/scripts/__init__.py index 5b58bf6..6378b53 100644 --- a/NV-Segment-CT/scripts/__init__.py +++ b/NV-Segment-CT/scripts/__init__.py @@ -13,6 +13,7 @@ # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer from .early_stop_score_function import score_function +from .huggingface_download import prepare_huggingface_checkpoint, touch_huggingface_download_counter # Ensures bundle expressions like ``scripts.batch_inference_utils.build_input_list`` resolve. from . import batch_inference_utils # noqa: F401 diff --git a/NV-Segment-CT/scripts/huggingface_download.py b/NV-Segment-CT/scripts/huggingface_download.py new file mode 100644 index 0000000..6c4a7b9 --- /dev/null +++ b/NV-Segment-CT/scripts/huggingface_download.py @@ -0,0 +1,88 @@ +import os +import shutil +from pathlib import Path +from typing import Optional + + +def _is_rank_zero() -> bool: + for name in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): + value = os.environ.get(name) + if value not in (None, "", "0"): + return False + return True + + +def touch_huggingface_download_counter( + repo_id: str, + filename: str = "config.json", + revision: str = "main", + rank_zero_only: bool = True, +) -> Optional[str]: + """Force a tiny Hugging Face file request without re-downloading weights.""" + + if rank_zero_only and not _is_rank_zero(): + return None + + try: + from huggingface_hub import hf_hub_download + except ImportError: + print("[nvseg] warning: huggingface_hub is not installed; skipping Hugging Face download counter touch.") + return None + + try: + path = hf_hub_download( + repo_id=repo_id, + filename=filename, + repo_type="model", + revision=revision, + force_download=True, + ) + except Exception as exc: # noqa: BLE001 + print(f"[nvseg] warning: could not touch Hugging Face download counter for {repo_id}/{filename}: {exc}") + return None + + print(f"[nvseg] touched Hugging Face download counter for {repo_id}/{filename}") + return path + + +def prepare_huggingface_checkpoint( + repo_id: str, + checkpoint_filename: str, + local_checkpoint_path: str, + counter_filename: str = "config.json", + revision: str = "main", + rank_zero_only: bool = True, +) -> str: + """Ensure the local MONAI checkpoint path exists and touch HF stats for this inference.""" + + local_path = Path(local_checkpoint_path) + + if rank_zero_only and not _is_rank_zero(): + return str(local_path) + + touch_huggingface_download_counter(repo_id, counter_filename, revision, rank_zero_only=False) + if local_path.exists(): + return str(local_path) + + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise RuntimeError( + f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}." + ) from exc + + checkpoint_path = hf_hub_download( + repo_id=repo_id, + filename=checkpoint_filename, + repo_type="model", + revision=revision, + ) + + local_path.parent.mkdir(parents=True, exist_ok=True) + try: + local_path.symlink_to(checkpoint_path) + except OSError: + shutil.copy2(checkpoint_path, local_path) + + print(f"[nvseg] prepared checkpoint at {local_path}") + return str(local_path) diff --git a/NV-Segment-CT/scripts/test_batch_inference.sh b/NV-Segment-CT/scripts/test_batch_inference.sh new file mode 100755 index 0000000..ed233a8 --- /dev/null +++ b/NV-Segment-CT/scripts/test_batch_inference.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +BUNDLE_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd) +PYTHON_BIN=${PYTHON_BIN:-python} +WORK_DIR=${WORK_DIR:-$(mktemp -d)} +KEEP_TEST_WORKDIR=${KEEP_TEST_WORKDIR:-0} + +if [[ "${KEEP_TEST_WORKDIR}" != "1" ]]; then + trap 'rm -rf "${WORK_DIR}"' EXIT +fi + +INPUT_DIR="${WORK_DIR}/input" +OUTPUT_DIR="${WORK_DIR}/output" +mkdir -p "${INPUT_DIR}" "${OUTPUT_DIR}" + +cp "${BUNDLE_ROOT}/example/spleen_03.nii.gz" "${INPUT_DIR}/" + +cd "${BUNDLE_ROOT}" + +"${PYTHON_BIN}" -m monai.bundle run \ + --config_file="['configs/inference.json', 'configs/batch_inference.json']" \ + --input_dir="${INPUT_DIR}" \ + --output_dir="${OUTPUT_DIR}" \ + 2>&1 | tee "${WORK_DIR}/batch_inference.log" + +EXPECTED_OUTPUT="${OUTPUT_DIR}/spleen_03/spleen_03_trans.nii.gz" +test -s "${EXPECTED_OUTPUT}" +grep -q "\[nvseg\] batch resume (skip existing outputs): 1 volume" "${WORK_DIR}/batch_inference.log" +grep -q "\[nvseg\] touched Hugging Face download counter for nvidia/NV-Segment-CT/config.json" "${WORK_DIR}/batch_inference.log" + +"${PYTHON_BIN}" -m monai.bundle run \ + --config_file="['configs/inference.json', 'configs/batch_inference.json']" \ + --input_dir="${INPUT_DIR}" \ + --output_dir="${OUTPUT_DIR}" \ + 2>&1 | tee "${WORK_DIR}/batch_resume.log" + +grep -q "\[nvseg\] batch: nothing to run (resume); ok" "${WORK_DIR}/batch_resume.log" + +echo "[nvseg-test] CT batch inference smoke test passed: ${EXPECTED_OUTPUT}" diff --git a/NV-Segment-CTMR/configs/batch_inference.json b/NV-Segment-CTMR/configs/batch_inference.json index 2118459..f94160b 100644 --- a/NV-Segment-CTMR/configs/batch_inference.json +++ b/NV-Segment-CTMR/configs/batch_inference.json @@ -1,14 +1,16 @@ { "input_dir": "@bundle_root", "input_suffix": "*.nii.gz", - "input_root_abs": "$os.path.abspath(@input_dir)", + "input_root_abs": "$os.path.realpath(@input_dir)", + "output_dir_abs": "$os.path.realpath(@output_dir)", "batch_skip_dir_names": [], "batch_skip_dir_prefixes": [], "batch_resume_skip_existing": true, "batch_use_input_list_cache": true, "batch_cache_wait_sec": 120, - "input_list": "$scripts.batch_inference_utils.build_input_list(os.path.abspath(@input_dir), os.path.abspath(@output_dir), @output_postfix, @output_ext, @batch_skip_dir_names, @batch_skip_dir_prefixes, @batch_resume_skip_existing, @batch_use_input_list_cache, @batch_cache_wait_sec)", + "input_list": "$scripts.batch_inference_utils.build_input_list(@input_root_abs, @output_dir_abs, @output_postfix, @output_ext, @batch_skip_dir_names, @batch_skip_dir_prefixes, @batch_resume_skip_existing, @batch_use_input_list_cache, @batch_cache_wait_sec)", "input_dicts": "$[{'image': x, 'label_prompt': @everything_labels} for x in @input_list]", "dataset#data": "@input_dicts", + "postprocessing#transforms#4#output_dir": "@output_dir_abs", "postprocessing#transforms#4#data_root_dir": "@input_root_abs" } diff --git a/NV-Segment-CTMR/configs/inference.json b/NV-Segment-CTMR/configs/inference.json index 5463bea..32b7372 100644 --- a/NV-Segment-CTMR/configs/inference.json +++ b/NV-Segment-CTMR/configs/inference.json @@ -9,6 +9,10 @@ "$import pathlib" ], "bundle_root": "./", + "huggingface_repo_id": "nvidia/NV-Segment-CTMR", + "huggingface_checkpoint_file": "vista3d_pretrained_model/model.pt", + "huggingface_download_counter_file": "config.json", + "checkpoint_path": "$@bundle_root + '/models/model.pt'", "image_key": "image", "output_dir": "$@bundle_root + '/eval'", "output_ext": ".nii.gz", @@ -179,7 +183,7 @@ ], "checkpointloader": { "_target_": "CheckpointLoader", - "load_path": "$@bundle_root + '/models/model.pt'", + "load_path": "@checkpoint_path", "load_dict": { "model": "@network" }, @@ -200,6 +204,7 @@ } }, "initialize": [ + "$scripts.prepare_huggingface_checkpoint(@huggingface_repo_id, @huggingface_checkpoint_file, @checkpoint_path, @huggingface_download_counter_file)", "$monai.utils.set_determinism(seed=123)", "$@checkpointloader(@evaluator)" ], diff --git a/NV-Segment-CTMR/configs/mgpu_inference.json b/NV-Segment-CTMR/configs/mgpu_inference.json index eecf12b..a377e98 100644 --- a/NV-Segment-CTMR/configs/mgpu_inference.json +++ b/NV-Segment-CTMR/configs/mgpu_inference.json @@ -18,6 +18,8 @@ "$import torch.distributed as dist", "$dist.is_initialized() or dist.init_process_group(backend='nccl')", "$torch.cuda.set_device(@device)", + "$scripts.prepare_huggingface_checkpoint(@huggingface_repo_id, @huggingface_checkpoint_file, @checkpoint_path, @huggingface_download_counter_file)", + "$dist.barrier()", "$monai.utils.set_determinism(seed=123)", "$@checkpointloader(@evaluator)" ], diff --git a/NV-Segment-CTMR/docs/README.md b/NV-Segment-CTMR/docs/README.md index 2c32722..489e4cf 100644 --- a/NV-Segment-CTMR/docs/README.md +++ b/NV-Segment-CTMR/docs/README.md @@ -33,14 +33,10 @@ cd NV-Segment-CTMR/NV-Segment-CTMR # Install dependencies pip install -r requirements.txt - -# Create models directory and download pretrained model -mkdir -p models -# Option 1: Download using hf and move to expected location -hf download nvidia/NV-Segment-CTMR --local-dir models/ && \ -mv models/vista3d_pretrained_model/model.pt models/model.pt ``` +Model weights are prepared automatically during inference. The first run downloads the checkpoint from Hugging Face into the local Hugging Face cache and links it at `models/model.pt`; later runs reuse the cached weights while still touching Hugging Face download stats for each inference. + ## Automatic Segmentation (support multi-gpu batch processing) @@ -82,7 +78,7 @@ The `configs/batch_inference.json` defines the batch inference, you can - If **resume** leaves nothing to run (all outputs already present), the run **exits successfully** with `[nvseg] batch: nothing to run (resume); ok` (avoids a zero-length dataloader / `DistributedSampler` failure). If **no** `*.nii.gz` files are discovered under `input_dir`, you get a short `[nvseg] batch: no *.nii.gz…` error. - Rank 0 logs: `[nvseg] batch resume (skip existing outputs): N volume(s) (...)`. - **Multi-GPU:** `--nproc_per_node` must be ≤ the number of volumes in `input_list` after filtering. - - **Outputs:** With `data_root_dir` and `separate_folder: true`, `input_dir/patient1/mri/scan.nii.gz` → `output_dir/patient1/mri/scan/scan_trans.nii.gz`. Ensure `models/model.pt` exists. + - **Outputs:** With `data_root_dir` and `separate_folder: true`, `input_dir/patient1/mri/scan.nii.gz` → `output_dir/patient1/mri/scan/scan_trans.nii.gz`. If `models/model.pt` is missing, inference prepares it automatically from Hugging Face. - Advanced: edit `should_skip_path_by_parent_rules()` in `scripts/batch_inference_utils.py` for custom path rules. 2. Segment based on a filelist.txt file, you can change the `input_list` in `configs/batch_inference.json` @@ -241,4 +237,4 @@ You may obtain a copy of the License at The model weights included in this project are licensed under the Non-Commercial -[NCLS v1 License](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf?t=eyJscyI6InJlZiIsImxzZCI6IlJFRi1naXRodWIuY29tL252aWRpYS1ob2xvc2NhbiJ9) \ No newline at end of file +[NCLS v1 License](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf?t=eyJscyI6InJlZiIsImxzZCI6IlJFRi1naXRodWIuY29tL252aWRpYS1ob2xvc2NhbiJ9) diff --git a/NV-Segment-CTMR/scripts/__init__.py b/NV-Segment-CTMR/scripts/__init__.py index 5b58bf6..6378b53 100644 --- a/NV-Segment-CTMR/scripts/__init__.py +++ b/NV-Segment-CTMR/scripts/__init__.py @@ -13,6 +13,7 @@ # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer from .early_stop_score_function import score_function +from .huggingface_download import prepare_huggingface_checkpoint, touch_huggingface_download_counter # Ensures bundle expressions like ``scripts.batch_inference_utils.build_input_list`` resolve. from . import batch_inference_utils # noqa: F401 diff --git a/NV-Segment-CTMR/scripts/huggingface_download.py b/NV-Segment-CTMR/scripts/huggingface_download.py new file mode 100644 index 0000000..6c4a7b9 --- /dev/null +++ b/NV-Segment-CTMR/scripts/huggingface_download.py @@ -0,0 +1,88 @@ +import os +import shutil +from pathlib import Path +from typing import Optional + + +def _is_rank_zero() -> bool: + for name in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): + value = os.environ.get(name) + if value not in (None, "", "0"): + return False + return True + + +def touch_huggingface_download_counter( + repo_id: str, + filename: str = "config.json", + revision: str = "main", + rank_zero_only: bool = True, +) -> Optional[str]: + """Force a tiny Hugging Face file request without re-downloading weights.""" + + if rank_zero_only and not _is_rank_zero(): + return None + + try: + from huggingface_hub import hf_hub_download + except ImportError: + print("[nvseg] warning: huggingface_hub is not installed; skipping Hugging Face download counter touch.") + return None + + try: + path = hf_hub_download( + repo_id=repo_id, + filename=filename, + repo_type="model", + revision=revision, + force_download=True, + ) + except Exception as exc: # noqa: BLE001 + print(f"[nvseg] warning: could not touch Hugging Face download counter for {repo_id}/{filename}: {exc}") + return None + + print(f"[nvseg] touched Hugging Face download counter for {repo_id}/{filename}") + return path + + +def prepare_huggingface_checkpoint( + repo_id: str, + checkpoint_filename: str, + local_checkpoint_path: str, + counter_filename: str = "config.json", + revision: str = "main", + rank_zero_only: bool = True, +) -> str: + """Ensure the local MONAI checkpoint path exists and touch HF stats for this inference.""" + + local_path = Path(local_checkpoint_path) + + if rank_zero_only and not _is_rank_zero(): + return str(local_path) + + touch_huggingface_download_counter(repo_id, counter_filename, revision, rank_zero_only=False) + if local_path.exists(): + return str(local_path) + + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise RuntimeError( + f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}." + ) from exc + + checkpoint_path = hf_hub_download( + repo_id=repo_id, + filename=checkpoint_filename, + repo_type="model", + revision=revision, + ) + + local_path.parent.mkdir(parents=True, exist_ok=True) + try: + local_path.symlink_to(checkpoint_path) + except OSError: + shutil.copy2(checkpoint_path, local_path) + + print(f"[nvseg] prepared checkpoint at {local_path}") + return str(local_path) diff --git a/NV-Segment-CTMR/scripts/test_batch_inference.sh b/NV-Segment-CTMR/scripts/test_batch_inference.sh new file mode 100755 index 0000000..6cbe11f --- /dev/null +++ b/NV-Segment-CTMR/scripts/test_batch_inference.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +BUNDLE_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd) +PYTHON_BIN=${PYTHON_BIN:-python} +WORK_DIR=${WORK_DIR:-$(mktemp -d)} +KEEP_TEST_WORKDIR=${KEEP_TEST_WORKDIR:-0} +MODALITY=${MODALITY:-MRI_BODY} + +if [[ "${KEEP_TEST_WORKDIR}" != "1" ]]; then + trap 'rm -rf "${WORK_DIR}"' EXIT +fi + +INPUT_DIR="${WORK_DIR}/input" +OUTPUT_DIR="${WORK_DIR}/output" +mkdir -p "${INPUT_DIR}" "${OUTPUT_DIR}" + +cp "${BUNDLE_ROOT}/example/s0289.nii.gz" "${INPUT_DIR}/" + +cd "${BUNDLE_ROOT}" + +"${PYTHON_BIN}" -m monai.bundle run \ + --config_file="['configs/inference.json', 'configs/batch_inference.json']" \ + --input_dir="${INPUT_DIR}" \ + --output_dir="${OUTPUT_DIR}" \ + --modality="${MODALITY}" \ + 2>&1 | tee "${WORK_DIR}/batch_inference.log" + +EXPECTED_OUTPUT="${OUTPUT_DIR}/s0289/s0289_trans.nii.gz" +test -s "${EXPECTED_OUTPUT}" +grep -q "\[nvseg\] batch resume (skip existing outputs): 1 volume" "${WORK_DIR}/batch_inference.log" +grep -q "\[nvseg\] touched Hugging Face download counter for nvidia/NV-Segment-CTMR/config.json" "${WORK_DIR}/batch_inference.log" + +"${PYTHON_BIN}" -m monai.bundle run \ + --config_file="['configs/inference.json', 'configs/batch_inference.json']" \ + --input_dir="${INPUT_DIR}" \ + --output_dir="${OUTPUT_DIR}" \ + --modality="${MODALITY}" \ + 2>&1 | tee "${WORK_DIR}/batch_resume.log" + +grep -q "\[nvseg\] batch: nothing to run (resume); ok" "${WORK_DIR}/batch_resume.log" + +echo "[nvseg-test] CTMR batch inference smoke test passed: ${EXPECTED_OUTPUT}" From 0fd2b98557bd1662667c5735f67182ea521c004a Mon Sep 17 00:00:00 2001 From: heyufan Date: Mon, 18 May 2026 13:44:32 -0400 Subject: [PATCH 2/2] Fix pre-commit lint for inference PR --- NV-Segment-CT/docs/README.md | 6 +- NV-Segment-CT/docs/finetune.md | 19 ++-- NV-Segment-CT/requirements.txt | 2 +- NV-Segment-CT/scripts/__init__.py | 12 +- .../scripts/batch_inference_utils.py | 16 +-- NV-Segment-CT/scripts/huggingface_download.py | 7 +- .../run_brain_segmentation.sh | 104 +++++++++--------- NV-Segment-CTMR/docs/README.md | 33 +++--- NV-Segment-CTMR/docs/finetune.md | 18 +-- NV-Segment-CTMR/requirements.txt | 2 +- NV-Segment-CTMR/scripts/__init__.py | 12 +- .../scripts/batch_inference_utils.py | 16 +-- .../scripts/huggingface_download.py | 7 +- README.md | 2 +- pyproject.toml | 10 +- 15 files changed, 130 insertions(+), 136 deletions(-) diff --git a/NV-Segment-CT/docs/README.md b/NV-Segment-CT/docs/README.md index 9daa3a0..d323313 100644 --- a/NV-Segment-CT/docs/README.md +++ b/NV-Segment-CT/docs/README.md @@ -36,9 +36,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,12 +66,14 @@ For more details, please refer to [this](inference.md). 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 + - 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 diff --git a/NV-Segment-CT/docs/finetune.md b/NV-Segment-CT/docs/finetune.md index d9dd7a9..3f74e82 100644 --- a/NV-Segment-CT/docs/finetune.md +++ b/NV-Segment-CT/docs/finetune.md @@ -1,6 +1,6 @@ # Finetune configurations -### 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): @@ -25,11 +25,11 @@ Example code for 5 fold cross-validation generation can be found [here](data.md) 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` +### 1. `label_mappings` ```json "label_mappings": { @@ -54,15 +54,15 @@ For continual learning, user can change `configs/train_continual.json`. More adv 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` +### 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. @@ -80,7 +80,7 @@ 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 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: @@ -114,7 +114,7 @@ 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 +### Evaluation 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. @@ -143,7 +143,6 @@ The `label_mapping` in `evaluation.json` does not include `0` because the postpr - 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 @@ -183,7 +182,7 @@ In this bundle, the training is simplified by jointly training with class prompt 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 +### Training 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 diff --git a/NV-Segment-CT/requirements.txt b/NV-Segment-CT/requirements.txt index 04d60a5..360f690 100644 --- a/NV-Segment-CT/requirements.txt +++ b/NV-Segment-CT/requirements.txt @@ -16,4 +16,4 @@ timm pytorch-ignite tensorboardX mlflow -huggingface_hub \ No newline at end of file +huggingface_hub diff --git a/NV-Segment-CT/scripts/__init__.py b/NV-Segment-CT/scripts/__init__.py index 6378b53..870bec8 100644 --- a/NV-Segment-CT/scripts/__init__.py +++ b/NV-Segment-CT/scripts/__init__.py @@ -12,8 +12,12 @@ # from .evaluator import EnsembleEvaluator, Evaluator, SupervisedEvaluator # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer -from .early_stop_score_function import score_function -from .huggingface_download import prepare_huggingface_checkpoint, touch_huggingface_download_counter - # Ensures bundle expressions like ``scripts.batch_inference_utils.build_input_list`` resolve. -from . import batch_inference_utils # noqa: F401 +from . import batch_inference_utils as batch_inference_utils +from .early_stop_score_function import score_function as score_function +from .huggingface_download import ( + prepare_huggingface_checkpoint as prepare_huggingface_checkpoint, +) +from .huggingface_download import ( + touch_huggingface_download_counter as touch_huggingface_download_counter, +) diff --git a/NV-Segment-CT/scripts/batch_inference_utils.py b/NV-Segment-CT/scripts/batch_inference_utils.py index 16e12f3..2446c8e 100644 --- a/NV-Segment-CT/scripts/batch_inference_utils.py +++ b/NV-Segment-CT/scripts/batch_inference_utils.py @@ -78,9 +78,7 @@ def collect_input_paths( for p in sorted(input_root.glob(pattern)): if not p.is_file(): continue - if should_skip_path_by_parent_rules( - p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes - ): + if should_skip_path_by_parent_rules(p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes): continue out.append(p) return out @@ -201,9 +199,7 @@ def build_input_list( ) payload = {"paths": paths, "n_discovered": n_discovered} if use_cache and local_rank == "0": - cache = _cache_path( - input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes - ) + cache = _cache_path(input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes) cache.parent.mkdir(parents=True, exist_ok=True) tmp = cache.with_suffix(".json.tmp") tmp.write_text(json.dumps(payload)) @@ -217,9 +213,7 @@ def build_input_list( flush=True, ) else: - cache = _cache_path( - input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes - ) + cache = _cache_path(input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes) deadline = time.time() + wait_sec payload = None while time.time() < deadline: @@ -228,9 +222,7 @@ def build_input_list( break time.sleep(0.05) else: - raise RuntimeError( - "[nvseg] batch: cache timeout (raise batch_cache_wait_sec or set batch_use_input_list_cache false)" - ) + raise RuntimeError("[nvseg] batch: cache timeout (raise batch_cache_wait_sec or set batch_use_input_list_cache false)") paths, n_discovered = _parse_cache_payload(payload) # Stale legacy `[]` or missing n_discovered: recompute so workers agree with rank 0. if n_discovered < 0: diff --git a/NV-Segment-CT/scripts/huggingface_download.py b/NV-Segment-CT/scripts/huggingface_download.py index 6c4a7b9..5ce5e6b 100644 --- a/NV-Segment-CT/scripts/huggingface_download.py +++ b/NV-Segment-CT/scripts/huggingface_download.py @@ -1,7 +1,6 @@ import os import shutil from pathlib import Path -from typing import Optional def _is_rank_zero() -> bool: @@ -17,7 +16,7 @@ def touch_huggingface_download_counter( filename: str = "config.json", revision: str = "main", rank_zero_only: bool = True, -) -> Optional[str]: +) -> str | None: """Force a tiny Hugging Face file request without re-downloading weights.""" if rank_zero_only and not _is_rank_zero(): @@ -67,9 +66,7 @@ def prepare_huggingface_checkpoint( try: from huggingface_hub import hf_hub_download except ImportError as exc: - raise RuntimeError( - f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}." - ) from exc + raise RuntimeError(f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}.") from exc checkpoint_path = hf_hub_download( repo_id=repo_id, 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 9bd8e60..9c0b54d 100755 --- a/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh +++ b/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh @@ -95,10 +95,10 @@ process_single_file() { local output_file="$2" local failed_log="$3" local output_dir=$(dirname "$output_file") - + # Ensure output directory exists mkdir -p "$output_dir" - + # Get absolute paths if command -v realpath &> /dev/null; then input_file=$(realpath "$input_file" 2>/dev/null || echo "$input_file") @@ -107,15 +107,15 @@ process_single_file() { input_file="$(cd "$(dirname "$input_file")" && pwd)/$(basename "$input_file")" fi fi - + # Extract filename without extension local file_basename=$(basename "$input_file" .nii.gz) file_basename=$(basename "$file_basename" .nii) - + # Create temporary directory in the output directory (use a unique name to avoid conflicts) local temp_dir="${output_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" @@ -124,7 +124,7 @@ process_single_file() { # We'll override it to use our output_dir, and it will save as: # {output_dir}/{basename}_preprocessed/{basename}_preprocessed_trans.nii.gz local preprocess_tmp_seg="${output_dir}/${file_basename}_preprocessed/${file_basename}_preprocessed_trans.nii.gz" - + # Cleanup function cleanup_on_error() { if [[ "$KEEP_TEMP" == "false" ]]; then @@ -132,10 +132,10 @@ process_single_file() { fi } trap cleanup_on_error ERR - + echo -e "${GREEN}Processing: $input_file${NC}" >&2 echo -e "${GREEN}Output will be saved to: $output_file${NC}" >&2 - + # Function to log failure and return log_failure() { local reason="$1" @@ -146,12 +146,12 @@ process_single_file() { cleanup_on_error return 1 } - + # Determine which file to use for preprocessing local preprocess_input="$input_file" local step_num=1 local total_steps=3 - + # Step 1: Skull stripping with SynthStrip (if not skipped) if [[ "$SKIP_SKULLSTRIP" == "false" ]]; then total_steps=4 @@ -170,7 +170,7 @@ process_single_file() { else echo -e "${YELLOW}Note: Skull stripping step is skipped${NC}" >&2 fi - + # Step 2: Affine align to the LUMIR template echo -e "${YELLOW}Step ${step_num}/${total_steps}: Affine alignment to LUMIR template...${NC}" >&2 cd "$BUNDLE_ROOT" @@ -182,7 +182,7 @@ process_single_file() { log_failure "Error: Preprocessing failed" return 1 } - + # Step 3: Segment the brain ((step_num++)) echo -e "${YELLOW}Step ${step_num}/${total_steps}: Running segmentation...${NC}" >&2 @@ -196,7 +196,7 @@ process_single_file() { log_failure "Error: Segmentation failed" return 1 } - + # Step 4: Revert the segmentation back to original space ((step_num++)) echo -e "${YELLOW}Step ${step_num}/${total_steps}: Reverting to original space...${NC}" >&2 @@ -204,7 +204,7 @@ process_single_file() { log_failure "Error: Segmentation output not found" return 1 fi - + cd "$BUNDLE_ROOT" python brain_t1_preprocess/revert_preprocess.py \ "$preprocess_tmp" \ @@ -215,7 +215,7 @@ process_single_file() { log_failure "Error: Reversion failed" return 1 } - + # Clean up temporary files if not keeping them if [[ "$KEEP_TEMP" == "false" ]]; then echo -e "${YELLOW}Cleaning up temporary files...${NC}" >&2 @@ -228,9 +228,9 @@ process_single_file() { else echo -e "${GREEN}Temporary files kept in: $temp_dir${NC}" >&2 fi - + trap - ERR - + echo -e "${GREEN}✓ Successfully processed: $input_file${NC}" >&2 echo -e "${GREEN} Output saved to: $output_file${NC}" >&2 return 0 @@ -280,7 +280,7 @@ process_folder() { for file in "${files[@]}"; do echo "" >&2 echo -e "${GREEN}========================================${NC}" >&2 - + # Determine output file path local file_basename=$(basename "$file" .nii.gz) file_basename=$(basename "$file_basename" .nii) @@ -291,7 +291,7 @@ process_folder() { rel_ext=".nii" fi local output_file="${output_dir}/${file_basename}_trans${rel_ext}" - + if process_single_file "$file" "$output_file" ""; then ((success_count++)) else @@ -314,17 +314,17 @@ process_file_list() { local file_list="$1" local root_path="$2" local output_dir="${OUTPUT_DIR:-./eval}" - + if [[ ! -f "$file_list" ]]; then echo -e "${RED}Error: File list not found: $file_list${NC}" >&2 exit 1 fi - + if [[ ! -d "$root_path" ]]; then echo -e "${RED}Error: Root path not found: $root_path${NC}" >&2 exit 1 fi - + # Get absolute paths if command -v realpath &> /dev/null; then file_list=$(realpath "$file_list" 2>/dev/null || echo "$file_list") @@ -343,12 +343,12 @@ process_file_list() { fi fi mkdir -p "$output_dir" - + # Create log file for failed/timeout files (after directory is created) local failed_log="${output_dir}/failed_files_$(date +%Y%m%d_%H%M%S).txt" touch "$failed_log" echo -e "${YELLOW}Failed/timeout files will be logged to: $failed_log${NC}" >&2 - + # Read file paths from the list local files=() local line_num=0 @@ -359,47 +359,47 @@ process_file_list() { if [[ -z "$line" ]] || [[ "$line" =~ ^# ]]; then continue fi - + # Remove leading ./ if present line="${line#./}" - + # Construct full path local full_path="${root_path}/${line}" - + if [[ ! -f "$full_path" ]]; then echo -e "${YELLOW}Warning: File not found (line $line_num): $full_path${NC}" >&2 continue fi - + files+=("$full_path") done < "$file_list" - + if [[ ${#files[@]} -eq 0 ]]; then echo -e "${YELLOW}Warning: No valid files found in $file_list${NC}" >&2 exit 1 fi - + # Sort files deterministically for consistent partitioning local sorted_files=() while IFS= read -r line; do sorted_files+=("$line") done < <(printf '%s\n' "${files[@]}" | sort) files=("${sorted_files[@]}") - + # Apply partitioning if requested local total_files=${#files[@]} local partition_files=() - + if [[ $NUM_PARTITIONS -gt 1 ]]; then if [[ $PARTITION_NUM -lt 1 ]] || [[ $PARTITION_NUM -gt $NUM_PARTITIONS ]]; then echo -e "${RED}Error: Partition number must be between 1 and $NUM_PARTITIONS${NC}" >&2 exit 1 fi - + # Calculate partition boundaries (deterministic split) local files_per_partition=$((total_files / NUM_PARTITIONS)) local remainder=$((total_files % NUM_PARTITIONS)) - + # Calculate start and end indices for this partition (0-indexed) local start_idx=0 for ((i=1; i&2 echo -e "${GREEN}Partition $PARTITION_NUM of $NUM_PARTITIONS: ${#partition_files[@]} file(s)${NC}" >&2 files=("${partition_files[@]}") else echo -e "${GREEN}Found ${#files[@]} file(s) to process${NC}" >&2 fi - + echo -e "${BLUE}Root path: $root_path${NC}" >&2 echo -e "${BLUE}Output directory: $output_dir${NC}" >&2 - + # Process each file local success_count=0 local fail_count=0 local skip_count=0 local total_in_partition=${#files[@]} local processed_count=0 - + for input_file in "${files[@]}"; do ((processed_count++)) local remaining=$((total_in_partition - processed_count)) - + echo "" >&2 echo -e "${GREEN}========================================${NC}" >&2 echo -e "${BLUE}Progress: [$((processed_count-1))/$total_in_partition] completed, $((remaining+1)) remaining${NC}" >&2 - + # Get relative path from root local rel_path="${input_file#$root_path/}" - + # Construct output path maintaining directory structure # Change filename to add _seg before extension local rel_dir=$(dirname "$rel_path") @@ -462,9 +462,9 @@ process_file_list() { elif [[ "$rel_filename" == *.nii ]]; then rel_ext=".nii" fi - + local output_file="${output_dir}/${rel_dir}/${rel_basename}_seg${rel_ext}" - + # Check if output already exists (before processing) if [[ "$SKIP_EXISTING" == "true" ]] && [[ -f "$output_file" ]]; then echo -e "${BLUE}Skipping (output exists): $input_file${NC}" >&2 @@ -472,32 +472,32 @@ process_file_list() { ((skip_count++)) continue fi - + # Process the file with overall timeout of 5 minutes local timeout_seconds=300 # 5 minutes total per scan local process_result=0 - + if command -v timeout &> /dev/null; then # Use timeout command to limit total processing time per scan # Export necessary variables for the function export BUNDLE_ROOT KEEP_TEMP MODALITY SKIP_SKULLSTRIP - + # Export the function so it's available in subshell # If export -f fails, we'll declare it inline in bash -c export -f process_single_file 2>/dev/null - + # Run with timeout - use bash -c to ensure function is available # Escape file paths safely using printf %q (bash-recommended method) printf -v escaped_input_file %q "$input_file" printf -v escaped_output_file %q "$output_file" printf -v escaped_failed_log %q "$failed_log" - + timeout $timeout_seconds bash -c " $(declare -f process_single_file) process_single_file $escaped_input_file $escaped_output_file $escaped_failed_log " 2>&1 local exit_code=$? - + if [[ $exit_code -eq 124 ]]; then # Timeout occurred (exit code 124 is timeout) echo -e "${RED}Error: Processing timed out after ${timeout_seconds}s: $input_file${NC}" >&2 @@ -534,7 +534,7 @@ process_file_list() { fi fi done - + echo "" >&2 echo -e "${GREEN}========================================${NC}" >&2 echo -e "${GREEN}Batch processing complete!${NC}" >&2 diff --git a/NV-Segment-CTMR/docs/README.md b/NV-Segment-CTMR/docs/README.md index 489e4cf..9658921 100644 --- a/NV-Segment-CTMR/docs/README.md +++ b/NV-Segment-CTMR/docs/README.md @@ -2,21 +2,19 @@ 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. -We defined 345 classes as in [metadata.json](../configs/metadata.json) and their corresponding dataset 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 [metadata.json](../configs/metadata.json) and their corresponding dataset 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). Example segmentations for **CT_BODY** (CT whole-body), **MRI_BRAIN**, and **MRI_BODY** (MRI torso): ![CT_BODY, MRI_BRAIN, and MRI_BODY segmentation examples](ctmr.png) -``` + Note: The predefined segment everything does not cover all labels, user can select more classes as output. Below is a segmentation using the label list from AutoPetAtals. User can extract the label list from each dataset defined in configs/label_mappings.json -``` + ![CT_BODY, MRI_BRAIN, and MRI_BODY segmentation examples](ctmr2.png) -``` Note: For Brain MRI segmentation, the model is able to segment 133 classes across diverse MRI sequences including T1, T2, Flair e.t.c. -``` ## Quick Start @@ -39,9 +37,6 @@ Model weights are prepared automatically during inference. The first run downloa ## Automatic Segmentation (support multi-gpu batch processing) - - - ## 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. @@ -65,8 +60,10 @@ python -m monai.bundle run --config_file configs/inference.json --input_dict "{' ## Batch inference with multiGPU support (automatic) -The `configs/batch_inference.json` defines the batch inference, you can +The `configs/batch_inference.json` defines the batch inference, you can: + 1. Segment all NIfTI files within a folder and subfolders + - `configs/batch_inference.json` builds `input_list` with `scripts/batch_inference_utils.build_input_list()`: - Recursively discovers `**/*.nii.gz` under `--input_dir`. - **Resume (default):** with `batch_resume_skip_existing: true` in `batch_inference.json`, only volumes whose expected output is **missing or empty** under `--output_dir` are queued (same layout as `SaveImaged`). Re-run the **same** command to finish leftovers. Set `batch_resume_skip_existing` to false to segment every discovered file again. @@ -82,9 +79,11 @@ The `configs/batch_inference.json` defines the batch inference, you can - Advanced: edit `should_skip_path_by_parent_rules()` in `scripts/batch_inference_utils.py` for custom path rules. 2. Segment based on a filelist.txt file, you can change the `input_list` in `configs/batch_inference.json` -``` + +```json "input_list": "$sorted([os.path.abspath(line.strip()) for line in open('/absolute/path/to/filelist.txt') if line.strip() and not line.strip().startswith('#')])", ``` + ### Single-GPU Batch Inference ```bash @@ -97,7 +96,6 @@ python -m monai.bundle run --config_file="['configs/inference.json', 'configs/ba ### Multi-GPU batch inference (cohorts, resume, optional folder filters) - ```bash conda activate vista3d-nv @@ -124,7 +122,7 @@ 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 +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/ @@ -211,17 +209,18 @@ conda activate vista3d-nv python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']" ``` +## Continual learning / Finetuning -# 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 + - 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 @@ -235,6 +234,6 @@ You may obtain a copy of the License at ### Model Weights License -The model weights included in this project are licensed under the Non-Commercial +The model weights included in this project are licensed under the Non-Commercial [NCLS v1 License](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf?t=eyJscyI6InJlZiIsImxzZCI6IlJFRi1naXRodWIuY29tL252aWRpYS1ob2xvc2NhbiJ9) diff --git a/NV-Segment-CTMR/docs/finetune.md b/NV-Segment-CTMR/docs/finetune.md index 9154533..10f8efb 100644 --- a/NV-Segment-CTMR/docs/finetune.md +++ b/NV-Segment-CTMR/docs/finetune.md @@ -1,5 +1,6 @@ # Finetune configurations -### 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): @@ -24,11 +25,11 @@ Example code for 5 fold cross-validation generation can be found [here](data.md) 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` +### 1. `label_mappings` ```json "label_mappings": { @@ -54,11 +55,11 @@ 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` +### 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. @@ -137,7 +138,7 @@ In this bundle, the training is simplified by jointly training with class prompt 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 +### Training 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 @@ -158,7 +159,6 @@ The default configs for both variables are derived from the `label_mappings` con Note: Please ensure the input data header is correct. The output file will use the same header as the input data, but if the input data is missing header information, MONAI will automatically provide some default values for missing values (e.g. `np.eye(4)` will be used if affine information is absent). This may cause a visualization misalignment depending on the visualization tool. - ## 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. @@ -188,7 +188,7 @@ 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 +### Evaluation 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. @@ -215,4 +215,4 @@ The `label_mapping` in `evaluation.json` does not include `0` because the postpr ### 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. \ No newline at end of file +- Make sure `0` is not included in your inference prompt for automatic segmentation. diff --git a/NV-Segment-CTMR/requirements.txt b/NV-Segment-CTMR/requirements.txt index 906ac85..6c15baa 100644 --- a/NV-Segment-CTMR/requirements.txt +++ b/NV-Segment-CTMR/requirements.txt @@ -22,4 +22,4 @@ SimpleITK pydicom statsmodels antspyx -huggingface_hub \ No newline at end of file +huggingface_hub diff --git a/NV-Segment-CTMR/scripts/__init__.py b/NV-Segment-CTMR/scripts/__init__.py index 6378b53..870bec8 100644 --- a/NV-Segment-CTMR/scripts/__init__.py +++ b/NV-Segment-CTMR/scripts/__init__.py @@ -12,8 +12,12 @@ # from .evaluator import EnsembleEvaluator, Evaluator, SupervisedEvaluator # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer -from .early_stop_score_function import score_function -from .huggingface_download import prepare_huggingface_checkpoint, touch_huggingface_download_counter - # Ensures bundle expressions like ``scripts.batch_inference_utils.build_input_list`` resolve. -from . import batch_inference_utils # noqa: F401 +from . import batch_inference_utils as batch_inference_utils +from .early_stop_score_function import score_function as score_function +from .huggingface_download import ( + prepare_huggingface_checkpoint as prepare_huggingface_checkpoint, +) +from .huggingface_download import ( + touch_huggingface_download_counter as touch_huggingface_download_counter, +) diff --git a/NV-Segment-CTMR/scripts/batch_inference_utils.py b/NV-Segment-CTMR/scripts/batch_inference_utils.py index 16e12f3..2446c8e 100644 --- a/NV-Segment-CTMR/scripts/batch_inference_utils.py +++ b/NV-Segment-CTMR/scripts/batch_inference_utils.py @@ -78,9 +78,7 @@ def collect_input_paths( for p in sorted(input_root.glob(pattern)): if not p.is_file(): continue - if should_skip_path_by_parent_rules( - p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes - ): + if should_skip_path_by_parent_rules(p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes): continue out.append(p) return out @@ -201,9 +199,7 @@ def build_input_list( ) payload = {"paths": paths, "n_discovered": n_discovered} if use_cache and local_rank == "0": - cache = _cache_path( - input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes - ) + cache = _cache_path(input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes) cache.parent.mkdir(parents=True, exist_ok=True) tmp = cache.with_suffix(".json.tmp") tmp.write_text(json.dumps(payload)) @@ -217,9 +213,7 @@ def build_input_list( flush=True, ) else: - cache = _cache_path( - input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes - ) + cache = _cache_path(input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes) deadline = time.time() + wait_sec payload = None while time.time() < deadline: @@ -228,9 +222,7 @@ def build_input_list( break time.sleep(0.05) else: - raise RuntimeError( - "[nvseg] batch: cache timeout (raise batch_cache_wait_sec or set batch_use_input_list_cache false)" - ) + raise RuntimeError("[nvseg] batch: cache timeout (raise batch_cache_wait_sec or set batch_use_input_list_cache false)") paths, n_discovered = _parse_cache_payload(payload) # Stale legacy `[]` or missing n_discovered: recompute so workers agree with rank 0. if n_discovered < 0: diff --git a/NV-Segment-CTMR/scripts/huggingface_download.py b/NV-Segment-CTMR/scripts/huggingface_download.py index 6c4a7b9..5ce5e6b 100644 --- a/NV-Segment-CTMR/scripts/huggingface_download.py +++ b/NV-Segment-CTMR/scripts/huggingface_download.py @@ -1,7 +1,6 @@ import os import shutil from pathlib import Path -from typing import Optional def _is_rank_zero() -> bool: @@ -17,7 +16,7 @@ def touch_huggingface_download_counter( filename: str = "config.json", revision: str = "main", rank_zero_only: bool = True, -) -> Optional[str]: +) -> str | None: """Force a tiny Hugging Face file request without re-downloading weights.""" if rank_zero_only and not _is_rank_zero(): @@ -67,9 +66,7 @@ def prepare_huggingface_checkpoint( try: from huggingface_hub import hf_hub_download except ImportError as exc: - raise RuntimeError( - f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}." - ) from exc + raise RuntimeError(f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}.") from exc checkpoint_path = hf_hub_download( repo_id=repo_id, diff --git a/README.md b/README.md index b615fc7..80c7b29 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Both models follow the MONAI bundle architecture. - 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. -![CTMR](./NV-Segment-CTMR/docs/ctmr.png) +![CTMR](./NV-Segment-CTMR/docs/ctmr.png) ## Performance on held-out test set diff --git a/pyproject.toml b/pyproject.toml index a0254ab..b3505f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,15 @@ [tool.ruff] target-version = "py311" line-length = 150 -exclude = [".git", "__pycache__", "data", "figures", "assets", "*.ipynb"] +exclude = [ + ".git", + "__pycache__", + "data", + "figures", + "assets", + "*.ipynb", + "NV-Segment-CTMR/brain_t1_preprocess", +] [tool.ruff.lint] select = ["E", "F", "I", "N", "W", "UP"]