diff --git a/NV-Segment-CT/configs/batch_inference.json b/NV-Segment-CT/configs/batch_inference.json index 605e66d..2118459 100644 --- a/NV-Segment-CT/configs/batch_inference.json +++ b/NV-Segment-CT/configs/batch_inference.json @@ -1,7 +1,14 @@ { "input_dir": "@bundle_root", "input_suffix": "*.nii.gz", - "input_list": "$sorted(glob.glob(os.path.join(@input_dir, @input_suffix)))", + "input_root_abs": "$os.path.abspath(@input_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_dicts": "$[{'image': x, 'label_prompt': @everything_labels} for x in @input_list]", - "dataset#data": "@input_dicts" + "dataset#data": "@input_dicts", + "postprocessing#transforms#4#data_root_dir": "@input_root_abs" } diff --git a/NV-Segment-CT/configs/mgpu_inference.json b/NV-Segment-CT/configs/mgpu_inference.json index bc0e9c4..eecf12b 100644 --- a/NV-Segment-CT/configs/mgpu_inference.json +++ b/NV-Segment-CT/configs/mgpu_inference.json @@ -17,7 +17,9 @@ "initialize": [ "$import torch.distributed as dist", "$dist.is_initialized() or dist.init_process_group(backend='nccl')", - "$torch.cuda.set_device(@device)" + "$torch.cuda.set_device(@device)", + "$monai.utils.set_determinism(seed=123)", + "$@checkpointloader(@evaluator)" ], "run": [ "$@evaluator.run()" diff --git a/NV-Segment-CT/docs/README.md b/NV-Segment-CT/docs/README.md index f1d1290..6df05bc 100644 --- a/NV-Segment-CT/docs/README.md +++ b/NV-Segment-CT/docs/README.md @@ -22,7 +22,7 @@ mv NV-Segment-CT/models/vista3d_pretrained_model/model.pt NV-Segment-CT/models/m rmdir NV-Segment-CT/models/vista3d_pretrained_model ``` -## 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 **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 fc5f109..5b58bf6 100644 --- a/NV-Segment-CT/scripts/__init__.py +++ b/NV-Segment-CT/scripts/__init__.py @@ -13,3 +13,6 @@ # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer from .early_stop_score_function import score_function + +# 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/batch_inference_utils.py b/NV-Segment-CT/scripts/batch_inference_utils.py new file mode 100644 index 0000000..16e12f3 --- /dev/null +++ b/NV-Segment-CT/scripts/batch_inference_utils.py @@ -0,0 +1,262 @@ +""" +Cohort batch inference: discover NIfTI inputs, match MONAI SaveImaged paths, optional resume. + +Behavior is controlled from ``configs/batch_inference.json`` (skip dir lists, resume, cache). +To change which classes are segmented, edit ``everything_labels`` in ``configs/inference.json``. +To customize path filtering further, edit :func:`should_skip_path_by_parent_rules`. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +import time +from pathlib import Path +from typing import Any + + +def should_skip_path_by_parent_rules( + path: Path, + *, + skip_dir_names: list[str] | None = None, + skip_dir_prefixes: list[str] | None = None, +) -> bool: + """ + Return True if this path should be excluded from batch discovery. + + Default (empty lists): do not skip any path based on directory names. + + - ``skip_dir_names``: any **parent** directory component that equals a name (case-insensitive). + - ``skip_dir_prefixes``: any **parent** directory component whose name **starts with** a prefix + (case-insensitive). + """ + names = {n.strip().lower() for n in (skip_dir_names or []) if n and str(n).strip()} + prefixes = tuple(p.strip().lower() for p in (skip_dir_prefixes or []) if p and str(p).strip()) + for part in path.parts[:-1]: + pl = part.lower() + if pl in names: + return True + for prefix in prefixes: + if prefix and pl.startswith(prefix): + return True + return False + + +def expected_output_path( + input_path: Path, + input_root: Path, + output_dir: Path, + postfix: str, + ext: str, +) -> Path: + """Match MONAI FolderLayout + separate_folder + data_root_dir.""" + input_path = input_path.resolve() + input_root = input_root.resolve() + output_dir = output_dir.resolve() + rel = os.path.relpath(input_path, input_root) + rel_dir = os.path.dirname(rel) + stem = input_path.name + if stem.endswith(".nii.gz"): + stem = stem[: -len(".nii.gz")] + elif stem.endswith(".nii"): + stem = stem[: -len(".nii")] + sub = Path(rel_dir) if rel_dir else Path() + return output_dir / sub / stem / f"{stem}_{postfix}{ext}" + + +def collect_input_paths( + input_root: Path, + pattern: str = "**/*.nii.gz", + *, + skip_dir_names: list[str] | None = None, + skip_dir_prefixes: list[str] | None = None, +) -> list[Path]: + input_root = input_root.resolve() + out: list[Path] = [] + 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 + ): + continue + out.append(p) + return out + + +def _cache_path( + input_dir: str, + output_dir: str, + postfix: str, + ext: str, + skip: bool, + skip_dir_names: list[str] | None, + skip_dir_prefixes: list[str] | None, +) -> Path: + sig = ( + f"{os.path.abspath(input_dir)}|{os.path.abspath(output_dir)}|{postfix}|{ext}|{skip}|" + f"{sorted(skip_dir_names or [])}|{sorted(skip_dir_prefixes or [])}" + ) + h = hashlib.sha256(sig.encode()).hexdigest()[:24] + base = Path(tempfile.gettempdir()) + return base / f"nvseg_batch_input_{h}.json" + + +def _compute_input_list( + input_dir: str, + output_dir: str, + postfix: str, + ext: str, + *, + skip_existing: bool, + skip_dir_names: list[str] | None, + skip_dir_prefixes: list[str] | None, +) -> tuple[list[str], int]: + """ + Returns ``(paths_to_run, n_discovered)`` where ``n_discovered`` is the number of + ``*.nii.gz`` paths after directory filters (before resume skip). + """ + root = Path(input_dir) + out_root = Path(output_dir) + all_paths = collect_input_paths( + root, + skip_dir_names=skip_dir_names, + skip_dir_prefixes=skip_dir_prefixes, + ) + n_discovered = len(all_paths) + if not skip_existing: + return [str(p) for p in all_paths], n_discovered + + missing: list[str] = [] + for inp in all_paths: + exp = expected_output_path(inp, root, out_root, postfix, ext) + if not exp.is_file() or exp.stat().st_size == 0: + missing.append(str(inp)) + return missing, n_discovered + + +def _parse_cache_payload(raw: Any) -> tuple[list[str], int]: + """Load cache written by rank 0. Supports legacy JSON list for backward compatibility.""" + if isinstance(raw, list): + # Legacy: empty [] is ambiguous (stale file or old format) -> signal recompute on workers. + return [str(p) for p in raw], len(raw) if raw else -1 + if isinstance(raw, dict) and "paths" in raw: + paths = raw["paths"] + if not isinstance(paths, list): + raise RuntimeError("[nvseg] batch: bad cache (paths); rm /tmp/nvseg_batch_input_*.json") + n_raw = raw.get("n_discovered") + if n_raw is None: + n_discovered = len(paths) if paths else -1 + else: + n_discovered = int(n_raw) + return [str(p) for p in paths], n_discovered + raise RuntimeError("[nvseg] batch: bad cache format; rm /tmp/nvseg_batch_input_*.json") + + +def build_input_list( + input_dir: str, + output_dir: str, + output_postfix: str, + output_ext: str, + batch_skip_dir_names: list | None = None, + batch_skip_dir_prefixes: list | None = None, + batch_resume_skip_existing: bool = True, + batch_use_input_list_cache: bool = True, + batch_cache_wait_sec: float = 120.0, +) -> list[str]: + """ + Called from ``configs/batch_inference.json``. + + ``batch_resume_skip_existing``: if True, only queue inputs whose output file is missing or empty. + ``batch_skip_dir_names`` / ``batch_skip_dir_prefixes``: filter discovery (see + :func:`should_skip_path_by_parent_rules`). + + ``LOCAL_RANK`` (set by ``torchrun``) is still read from the environment for multi-GPU cache. + + If resume mode leaves nothing to run (outputs already exist), prints a message and raises + ``SystemExit(0)`` so the process exits before MONAI builds a zero-length DataLoader (which + would fail under ``DistributedSampler``). If no ``*.nii.gz`` files are discovered, raises + ``RuntimeError``. + """ + names = list(batch_skip_dir_names) if batch_skip_dir_names is not None else [] + prefixes = list(batch_skip_dir_prefixes) if batch_skip_dir_prefixes is not None else [] + + skip = bool(batch_resume_skip_existing) + use_cache = bool(batch_use_input_list_cache) + wait_sec = float(batch_cache_wait_sec) + + local_rank = os.environ.get("LOCAL_RANK", "0") + + if not use_cache or local_rank == "0": + paths, n_discovered = _compute_input_list( + input_dir, + output_dir, + output_postfix, + output_ext, + skip_existing=skip, + skip_dir_names=names, + skip_dir_prefixes=prefixes, + ) + 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.parent.mkdir(parents=True, exist_ok=True) + tmp = cache.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload)) + tmp.replace(cache) + + if local_rank == "0": + mode = "resume (skip existing outputs)" if skip else "full pass (all inputs)" + print( + f"[nvseg] batch {mode}: {len(paths)} volume(s) " + f"(input_dir={os.path.abspath(input_dir)}, output_dir={os.path.abspath(output_dir)})", + flush=True, + ) + else: + 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: + if cache.is_file(): + payload = json.loads(cache.read_text()) + 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)" + ) + paths, n_discovered = _parse_cache_payload(payload) + # Stale legacy `[]` or missing n_discovered: recompute so workers agree with rank 0. + if n_discovered < 0: + paths, n_discovered = _compute_input_list( + input_dir, + output_dir, + output_postfix, + output_ext, + skip_existing=skip, + skip_dir_names=names, + skip_dir_prefixes=prefixes, + ) + + if not paths: + if n_discovered == 0: + raise RuntimeError("[nvseg] batch: no *.nii.gz under input_dir (check paths and skip rules)") + # Resume: all outputs exist — exit before an empty DistributedSampler / dataloader. + print("[nvseg] batch: nothing to run (resume); ok", flush=True) + raise SystemExit(0) + + return paths + + +__all__ = [ + "build_input_list", + "collect_input_paths", + "expected_output_path", + "should_skip_path_by_parent_rules", +] 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 c9b5e09..9bd8e60 100755 --- a/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh +++ b/NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh @@ -1,22 +1,43 @@ #!/bin/bash # Brain MRI Segmentation Script -# Supports single file or folder batch processing with optional temporary file retention +# Supports single file, folder batch processing, or file list processing +# Supports optional skull stripping and temporary file retention -set -e # Exit on error +# Print startup message +echo "Starting brain segmentation script..." >&2 # Default values KEEP_TEMP=false OUTPUT_DIR="" MODALITY="MRI_BRAIN" -CONDA_ENV="vista3d-nv" -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -BUNDLE_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" +SKIP_SKULLSTRIP=false +SKIP_EXISTING=true +NUM_PARTITIONS=1 +PARTITION_NUM=1 + +# Get script directory - handle both direct execution and sourced execution +if [[ "${BASH_SOURCE[0]}" != "" ]]; then + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" || { + echo "Error: Failed to determine script directory" >&2 + exit 1 + } +else + SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" || { + echo "Error: Failed to determine script directory" >&2 + exit 1 + } +fi +BUNDLE_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" || { + echo "Error: Failed to determine bundle root directory" >&2 + exit 1 +} # Colors for output RED='\033[0;31m' GREEN='\033[0;32m' YELLOW='\033[1;33m' +BLUE='\033[0;34m' NC='\033[0m' # No Color # Function to print usage @@ -25,114 +46,146 @@ usage() { Usage: $0 [OPTIONS] Brain MRI Segmentation Script for NV-Segment-CTMR +Supports single file, folder batch processing, or file list processing OPTIONS: --input FILE Single NIfTI file to segment --input_folder FOLDER Folder containing NIfTI files to process (batch mode) + --file_list FILE Text file with one file path per line (relative to root_path) + --root_path PATH Root path where file paths in file_list are relative to --output_dir DIR Output directory (default: ./eval) --keep-temp Keep temporary preprocessing files (default: false) + --no-skullstrip Skip skull stripping step (default: false, performs skull stripping) --modality MODALITY Segmentation modality: MRI_BRAIN (default), MRI_BODY, CT_BODY - --conda-env ENV Conda environment name (default: vista3d-nv) + --no-skip Don't skip existing output files (default: skip if exists, only for file_list mode) + --num_partitions N Split file list into N partitions (default: 1, no split, only for file_list mode) + --partition M Process partition M (1-indexed, default: 1). Requires --num_partitions (only for file_list mode) -h, --help Show this help message EXAMPLES: # Single file processing $0 --input example/brain_t1.nii.gz - # Batch processing + # Single file without skull stripping + $0 --input example/brain_t1.nii.gz --no-skullstrip + + # Batch processing from folder $0 --input_folder example/ --output_dir results/ - # Keep temporary files for debugging - $0 --input example/brain_t1.nii.gz --keep-temp + # Process files from a list + $0 --file_list file_list.txt --root_path /path/to/root --output_dir /path/to/output -EOF - exit 1 -} + # Process files from a list without skull stripping + $0 --file_list file_list.txt --root_path /path/to/root --output_dir /path/to/output --no-skullstrip -# Function to check if conda environment is activated -check_conda_env() { - if ! command -v conda &> /dev/null; then - echo -e "${RED}Error: conda command not found. Please install conda or activate your environment manually.${NC}" >&2 - exit 1 - fi + # Split into 10 partitions and process partition 3 + $0 --file_list file_list.txt --root_path /path/to/root --output_dir /path/to/output --num_partitions 10 --partition 3 - # 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 +NOTES: + - When using --file_list, output files maintain the same directory structure as input, with _seg suffix added + - Existing output files are skipped by default to support job resubmission (file_list mode only) - # 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}" +EOF + exit 1 } # Function to process a single file 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") + 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") + else + if [[ ! "$input_file" =~ ^/ ]]; then + 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) - local file_dir=$(dirname "$input_file") - - # Create temporary directory for this file - local temp_dir="${file_dir}/${file_basename}_temp" + + # 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" local preprocess_meta="${temp_dir}/${file_basename}_preprocessed.meta.json" + # Monai bundle saves to output_dir from config, which defaults to bundle_root/eval + # 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" - 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 - cd "$BUNDLE_ROOT" - ./brain_t1_preprocess/synthstrip-docker -i "$input_file" -o "$skull_stripped" || { - echo -e "${RED}Error: Skull stripping failed${NC}" >&2 - [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" - exit 1 - } + + # Cleanup function + cleanup_on_error() { + if [[ "$KEEP_TEMP" == "false" ]]; then + rm -rf "$temp_dir" + 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" + echo -e "${RED}$reason: $input_file${NC}" >&2 + if [[ -n "$failed_log" ]]; then + echo "$input_file" >> "$failed_log" + fi + 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 + echo -e "${YELLOW}Step 1/4: Skull stripping...${NC}" >&2 + if [[ ! -f "$skull_stripped" ]]; then + cd "$BUNDLE_ROOT" + ./brain_t1_preprocess/synthstrip-docker -i "$input_file" -o "$skull_stripped" || { + log_failure "Error: Skull stripping failed" + return 1 + } + else + echo -e "${YELLOW} Skull-stripped file already exists, skipping...${NC}" >&2 + fi + preprocess_input="$skull_stripped" + step_num=2 else - echo -e "${YELLOW} Skull-stripped file already exists, skipping...${NC}" + echo -e "${YELLOW}Note: Skull stripping step is skipped${NC}" >&2 fi - + # Step 2: Affine align to the LUMIR template - echo -e "${YELLOW}Step 2/4: Affine alignment to LUMIR template...${NC}" + echo -e "${YELLOW}Step ${step_num}/${total_steps}: Affine alignment to LUMIR template...${NC}" >&2 cd "$BUNDLE_ROOT" python brain_t1_preprocess/preprocess.py \ - "$skull_stripped" \ + "$preprocess_input" \ brain_t1_preprocess/LUMIR_template.nii.gz \ "$preprocess_tmp" \ --save-preprocess "$preprocess_meta" || { - echo -e "${RED}Error: Preprocessing failed${NC}" >&2 - [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" - exit 1 + log_failure "Error: Preprocessing failed" + return 1 } - + # Step 3: Segment the brain - echo -e "${YELLOW}Step 3/4: Running segmentation...${NC}" + ((step_num++)) + echo -e "${YELLOW}Step ${step_num}/${total_steps}: Running segmentation...${NC}" >&2 cd "$BUNDLE_ROOT" # Override output_dir in config to use our output_dir so segmentation saves to the right place python -m monai.bundle run \ @@ -140,34 +193,32 @@ process_single_file() { --input_dict "{'image':'$preprocess_tmp'}" \ --output_dir "$output_dir" \ --modality "$MODALITY" || { - echo -e "${RED}Error: Segmentation failed${NC}" >&2 - [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" - exit 1 + log_failure "Error: Segmentation failed" + return 1 } - + # Step 4: Revert the segmentation back to original space - echo -e "${YELLOW}Step 4/4: Reverting to original space...${NC}" + ((step_num++)) + echo -e "${YELLOW}Step ${step_num}/${total_steps}: Reverting to original space...${NC}" >&2 if [[ ! -f "$preprocess_tmp_seg" ]]; then - echo -e "${RED}Error: Segmentation output not found: $preprocess_tmp_seg${NC}" >&2 - [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" - exit 1 + log_failure "Error: Segmentation output not found" + return 1 fi - + cd "$BUNDLE_ROOT" python brain_t1_preprocess/revert_preprocess.py \ "$preprocess_tmp" \ --out "${temp_dir}/${file_basename}_revert.nii.gz" \ --mask "$preprocess_tmp_seg" \ - --mask-out "$final_output" \ + --mask-out "$output_file" \ --meta "$preprocess_meta" || { - echo -e "${RED}Error: Reversion failed${NC}" >&2 - [[ "$KEEP_TEMP" == "false" ]] && rm -rf "$temp_dir" - exit 1 + 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}" + echo -e "${YELLOW}Cleaning up temporary files...${NC}" >&2 rm -rf "$temp_dir" # Also clean up the preprocessed output directory if it only contains temp files local preprocess_output_dir="${output_dir}/${file_basename}_preprocessed" @@ -175,11 +226,14 @@ process_single_file() { rm -rf "$preprocess_output_dir" fi else - echo -e "${GREEN}Temporary files kept in: $temp_dir${NC}" + echo -e "${GREEN}Temporary files kept in: $temp_dir${NC}" >&2 fi - - echo -e "${GREEN}✓ Successfully processed: $input_file${NC}" - echo -e "${GREEN} Output saved to: $final_output${NC}" + + trap - ERR + + echo -e "${GREEN}✓ Successfully processed: $input_file${NC}" >&2 + echo -e "${GREEN} Output saved to: $output_file${NC}" >&2 + return 0 } # Function to process a folder (batch mode) @@ -193,8 +247,17 @@ process_folder() { fi # Get absolute paths - input_folder=$(realpath "$input_folder") - output_dir=$(realpath -m "$output_dir") + if command -v realpath &> /dev/null; then + input_folder=$(realpath "$input_folder" 2>/dev/null || echo "$input_folder") + output_dir=$(realpath -m "$output_dir" 2>/dev/null || echo "$output_dir") + else + if [[ ! "$input_folder" =~ ^/ ]]; then + input_folder="$(cd "$input_folder" && pwd)" + fi + if [[ ! "$output_dir" =~ ^/ ]]; then + output_dir="$(cd "$(dirname "$output_dir")" && pwd)/$(basename "$output_dir")" + fi + fi mkdir -p "$output_dir" # Find all NIfTI files @@ -208,16 +271,28 @@ process_folder() { exit 1 fi - echo -e "${GREEN}Found ${#files[@]} file(s) to process${NC}" + echo -e "${GREEN}Found ${#files[@]} file(s) to process${NC}" >&2 # Process each file local success_count=0 local fail_count=0 for file in "${files[@]}"; do - echo "" - echo -e "${GREEN}========================================${NC}" - if process_single_file "$file"; then + echo "" >&2 + echo -e "${GREEN}========================================${NC}" >&2 + + # Determine output file path + local file_basename=$(basename "$file" .nii.gz) + file_basename=$(basename "$file_basename" .nii) + local rel_ext="" + if [[ "$file" == *.nii.gz ]]; then + rel_ext=".nii.gz" + elif [[ "$file" == *.nii ]]; then + 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 ((fail_count++)) @@ -225,18 +300,259 @@ process_folder() { fi done - echo "" - echo -e "${GREEN}========================================${NC}" - echo -e "${GREEN}Batch processing complete!${NC}" - echo -e "${GREEN} Successful: $success_count${NC}" + echo "" >&2 + echo -e "${GREEN}========================================${NC}" >&2 + echo -e "${GREEN}Batch processing complete!${NC}" >&2 + echo -e "${GREEN} Successful: $success_count${NC}" >&2 if [[ $fail_count -gt 0 ]]; then - echo -e "${RED} Failed: $fail_count${NC}" + echo -e "${RED} Failed: $fail_count${NC}" >&2 + fi +} + +# Function to process files from txt file list +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") + root_path=$(realpath "$root_path" 2>/dev/null || echo "$root_path") + output_dir=$(realpath -m "$output_dir" 2>/dev/null || echo "$output_dir") + else + # Fallback if realpath is not available + if [[ ! "$file_list" =~ ^/ ]]; then + file_list="$(cd "$(dirname "$file_list")" && pwd)/$(basename "$file_list")" + fi + if [[ ! "$root_path" =~ ^/ ]]; then + root_path="$(cd "$root_path" && pwd)" + fi + if [[ ! "$output_dir" =~ ^/ ]]; then + output_dir="$(cd "$(dirname "$output_dir")" && pwd)/$(basename "$output_dir")" + 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 + while IFS= read -r line || [[ -n "$line" ]]; do + ((line_num++)) + # Skip empty lines and comments + line=$(echo "$line" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//') + 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") + local rel_filename=$(basename "$rel_path") + local rel_basename=$(basename "$rel_filename" .nii.gz) + rel_basename=$(basename "$rel_basename" .nii) + local rel_ext="" + if [[ "$rel_filename" == *.nii.gz ]]; then + rel_ext=".nii.gz" + 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 + echo -e "${BLUE} Output: $output_file${NC}" >&2 + ((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 + echo "$input_file" >> "$failed_log" + ((fail_count++)) + process_result=1 + elif [[ $exit_code -eq 0 ]]; then + # Process completed, check if output was created + if [[ -f "$output_file" ]]; then + ((success_count++)) + process_result=0 + else + # Output not created despite success exit code + echo -e "${RED}Error: Output file not created: $input_file${NC}" >&2 + echo "$input_file" >> "$failed_log" + ((fail_count++)) + process_result=1 + fi + else + # Process failed (error already logged in process_single_file) + ((fail_count++)) + process_result=1 + fi + else + # Fallback: run without timeout wrapper + # Note: Without timeout command, we can't enforce the 5-minute limit + # but we can still catch failures + if process_single_file "$input_file" "$output_file" "$failed_log"; then + ((success_count++)) + process_result=0 + else + ((fail_count++)) + process_result=1 + fi + fi + done + + echo "" >&2 + echo -e "${GREEN}========================================${NC}" >&2 + echo -e "${GREEN}Batch processing complete!${NC}" >&2 + echo -e "${GREEN} Successful: $success_count${NC}" >&2 + if [[ $skip_count -gt 0 ]]; then + echo -e "${BLUE} Skipped (existing): $skip_count${NC}" >&2 + fi + if [[ $fail_count -gt 0 ]]; then + echo -e "${RED} Failed/Timeout: $fail_count${NC}" >&2 + echo -e "${YELLOW} Failed files logged to: $failed_log${NC}" >&2 fi } # Parse command line arguments INPUT_FILE="" INPUT_FOLDER="" +FILE_LIST="" +ROOT_PATH="" while [[ $# -gt 0 ]]; do case $1 in @@ -248,6 +564,14 @@ while [[ $# -gt 0 ]]; do INPUT_FOLDER="$2" shift 2 ;; + --file_list) + FILE_LIST="$2" + shift 2 + ;; + --root_path) + ROOT_PATH="$2" + shift 2 + ;; --output_dir) OUTPUT_DIR="$2" shift 2 @@ -256,12 +580,24 @@ while [[ $# -gt 0 ]]; do KEEP_TEMP=true shift ;; + --no-skullstrip) + SKIP_SKULLSTRIP=true + shift + ;; --modality) MODALITY="$2" shift 2 ;; - --conda-env) - CONDA_ENV="$2" + --no-skip) + SKIP_EXISTING=false + shift + ;; + --num_partitions) + NUM_PARTITIONS="$2" + shift 2 + ;; + --partition) + PARTITION_NUM="$2" shift 2 ;; -h|--help) @@ -275,13 +611,34 @@ while [[ $# -gt 0 ]]; do done # Validate arguments -if [[ -z "$INPUT_FILE" && -z "$INPUT_FOLDER" ]]; then - echo -e "${RED}Error: Either --input or --input_folder must be specified${NC}" >&2 +input_count=0 +[[ -n "$INPUT_FILE" ]] && ((input_count++)) +[[ -n "$INPUT_FOLDER" ]] && ((input_count++)) +[[ -n "$FILE_LIST" ]] && ((input_count++)) + +if [[ $input_count -eq 0 ]]; then + echo -e "${RED}Error: One of --input, --input_folder, or --file_list must be specified${NC}" >&2 + usage +fi + +if [[ $input_count -gt 1 ]]; then + echo -e "${RED}Error: Cannot specify multiple input options (--input, --input_folder, --file_list)${NC}" >&2 + usage +fi + +if [[ -n "$FILE_LIST" ]] && [[ -z "$ROOT_PATH" ]]; then + echo -e "${RED}Error: --root_path must be specified when using --file_list${NC}" >&2 + usage +fi + +# Validate partition arguments +if [[ $NUM_PARTITIONS -lt 1 ]]; then + echo -e "${RED}Error: --num_partitions must be at least 1${NC}" >&2 usage fi -if [[ -n "$INPUT_FILE" && -n "$INPUT_FOLDER" ]]; then - echo -e "${RED}Error: Cannot specify both --input and --input_folder${NC}" >&2 +if [[ $PARTITION_NUM -lt 1 ]] || [[ $PARTITION_NUM -gt $NUM_PARTITIONS ]]; then + echo -e "${RED}Error: --partition must be between 1 and $NUM_PARTITIONS${NC}" >&2 usage fi @@ -291,17 +648,36 @@ if [[ ! "$MODALITY" =~ ^(MRI_BRAIN|MRI_BODY|CT_BODY)$ ]]; then MODALITY="MRI_BRAIN" fi -# Check and activate conda environment -check_conda_env - # Change to bundle root directory -cd "$BUNDLE_ROOT" +if [[ ! -d "$BUNDLE_ROOT" ]]; then + echo -e "${RED}Error: Bundle root directory not found: $BUNDLE_ROOT${NC}" >&2 + exit 1 +fi + +cd "$BUNDLE_ROOT" || { + echo -e "${RED}Error: Failed to change to bundle root directory: $BUNDLE_ROOT${NC}" >&2 + exit 1 +} # Process based on input type if [[ -n "$INPUT_FILE" ]]; then - process_single_file "$INPUT_FILE" -else + # Single file mode + file_basename=$(basename "$INPUT_FILE" .nii.gz) + file_basename=$(basename "$file_basename" .nii) + rel_ext="" + if [[ "$INPUT_FILE" == *.nii.gz ]]; then + rel_ext=".nii.gz" + elif [[ "$INPUT_FILE" == *.nii ]]; then + rel_ext=".nii" + fi + output_file="${OUTPUT_DIR:-./eval}/${file_basename}_trans${rel_ext}" + process_single_file "$INPUT_FILE" "$output_file" "" +elif [[ -n "$INPUT_FOLDER" ]]; then + # Folder batch mode process_folder "$INPUT_FOLDER" +else + # File list mode + process_file_list "$FILE_LIST" "$ROOT_PATH" fi -echo -e "${GREEN}All done!${NC}" +echo -e "${GREEN}All done!${NC}" >&2 diff --git a/NV-Segment-CTMR/configs/batch_inference.json b/NV-Segment-CTMR/configs/batch_inference.json index 605e66d..2118459 100644 --- a/NV-Segment-CTMR/configs/batch_inference.json +++ b/NV-Segment-CTMR/configs/batch_inference.json @@ -1,7 +1,14 @@ { "input_dir": "@bundle_root", "input_suffix": "*.nii.gz", - "input_list": "$sorted(glob.glob(os.path.join(@input_dir, @input_suffix)))", + "input_root_abs": "$os.path.abspath(@input_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_dicts": "$[{'image': x, 'label_prompt': @everything_labels} for x in @input_list]", - "dataset#data": "@input_dicts" + "dataset#data": "@input_dicts", + "postprocessing#transforms#4#data_root_dir": "@input_root_abs" } diff --git a/NV-Segment-CTMR/configs/mgpu_inference.json b/NV-Segment-CTMR/configs/mgpu_inference.json index bc0e9c4..eecf12b 100644 --- a/NV-Segment-CTMR/configs/mgpu_inference.json +++ b/NV-Segment-CTMR/configs/mgpu_inference.json @@ -17,7 +17,9 @@ "initialize": [ "$import torch.distributed as dist", "$dist.is_initialized() or dist.init_process_group(backend='nccl')", - "$torch.cuda.set_device(@device)" + "$torch.cuda.set_device(@device)", + "$monai.utils.set_determinism(seed=123)", + "$@checkpointloader(@evaluator)" ], "run": [ "$@evaluator.run()" diff --git a/NV-Segment-CTMR/docs/README.md b/NV-Segment-CTMR/docs/README.md index 3e89a6b..600b9b1 100644 --- a/NV-Segment-CTMR/docs/README.md +++ b/NV-Segment-CTMR/docs/README.md @@ -2,9 +2,21 @@ 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 +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). -![Benchmark CT](./benchmarkct.png) ![Benchmark MR](./benchmarkmr.png) + +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 @@ -33,12 +45,12 @@ rmdir NV-Segment-CTMR/models/vista3d_pretrained_model ## 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). + + ## 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. +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 # Make sure conda environment is activated @@ -57,8 +69,28 @@ The detailed automatic segmentation class index can be found [here](../configs/l python -m monai.bundle run --config_file configs/inference.json --input_dict "{'image':'example/s0289.nii.gz','label_prompt':[3]}" ``` -## Batch inference with multiGPU support for segmenting everything (automatic) +## Batch inference with multiGPU support (automatic) +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. + - **Discovery filters:** edit `batch_skip_dir_names` (exact parent folder name) and/or `batch_skip_dir_prefixes` (parent folder name starts with…) in `configs/batch_inference.json` (JSON arrays of strings). + - **Optional keys** in `configs/batch_inference.json` (defaults in the file): `batch_skip_dir_names`, `batch_skip_dir_prefixes`, `batch_resume_skip_existing`, `batch_use_input_list_cache`, `batch_cache_wait_sec`. + - **`batch_use_input_list_cache`:** With `torchrun` (multi-process), only rank 0 walks the tree to build `input_list` and writes a small JSON cache under the system temp directory; other ranks read that file so you do not repeat a huge filesystem scan on every GPU. Set to `false` if you want every rank to compute the list itself (simpler, slower on large cohorts). Single-process runs are unaffected in practice. + - **`batch_cache_wait_sec`:** When `batch_use_input_list_cache` is `true`, non-zero ranks wait up to this many seconds for rank 0’s cache file. Increase if rank 0’s scan is slow; decrease only if the list is always built quickly. + - **Which classes to segment:** edit **`everything_labels`** in **`configs/inference.json`** (and `modality` / `--modality` as needed). See `configs/label_dict.json` and `docs/inference.md`. + - 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. + - 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` +``` + "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 @@ -69,68 +101,82 @@ conda activate vista3d-nv python -m monai.bundle run --config_file="['configs/inference.json', 'configs/batch_inference.json']" --input_dir="example/" --output_dir="example/" --modality MRI_BODY ``` -### Multi-GPU Batch Inference +### Multi-GPU batch inference (cohorts, resume, optional folder filters) -**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 -# Activate conda environment first (CRITICAL!) 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/" +# Example: multi-GPU batch (same command for first run or resume) +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/" \ + --modality MRI_BODY ``` -`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 +## Brain MRI segmentation (any MRI sequence) -For brain MRI segmentation, we only support T1 and require preprocessing. We provide a convenient bash script that handles all preprocessing steps automatically. +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 `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). +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). + +#### Single file -#### Single File Processing +Output path: `{output_dir}/{basename}_trans.nii.gz` (default `output_dir` is `./eval`). ```bash -# Process a single brain MRI file -./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz +conda activate vista3d-nv # or your env with MONAI + deps -# Specify custom output directory +./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 temporary files for debugging ./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz --keep-temp +./brain_t1_preprocess/run_brain_segmentation.sh --input example/brain_t1.nii.gz --no-skullstrip ``` -#### Batch Processing +#### Folder batch (`--input_folder`) + +Only `*.nii` / `*.nii.gz` files **directly inside** the given folder are processed (`find` with `-maxdepth 1`; no subfolders). Each output is written as `{output_dir}/{basename}_trans.nii.gz` (flat layout). ```bash -# Process all NIfTI files in a folder ./brain_t1_preprocess/run_brain_segmentation.sh --input_folder example/ --output_dir results/ - -# Batch processing with temporary files kept -./brain_t1_preprocess/run_brain_segmentation.sh --input_folder example/ --keep-temp +./brain_t1_preprocess/run_brain_segmentation.sh --input_folder example/ --output_dir results/ --keep-temp ``` -#### Script Options +#### File list (`--file_list` + `--root_path`) -- `--input FILE`: Single NIfTI file to segment -- `--input_folder FOLDER`: Folder containing NIfTI files (batch mode) -- `--output_dir DIR`: Output directory (default: `./eval`) -- `--keep-temp`: Keep temporary preprocessing files (default: false, files are cleaned up automatically) -- `--modality MODALITY`: Segmentation modality: `MRI_BRAIN` (default), `MRI_BODY`, `CT_BODY` -- `--conda-env ENV`: Conda environment name (default: `vista3d-nv`) -- `-h, --help`: Show help message +Lines in the list are paths relative to `root_path` (comments and empty lines allowed). Outputs mirror that relative layout under `output_dir`, with `_seg` before the extension (e.g. `root_path/sub/scan.nii.gz` → `output_dir/sub/scan_seg.nii.gz`). By default, existing outputs are skipped (resume-friendly). Optional: split the sorted list across jobs with `--num_partitions N` and `--partition M` (1-based). -**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. +```bash +./brain_t1_preprocess/run_brain_segmentation.sh \ + --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 +./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 +``` + +#### Script options + +- `--input FILE`: Single NIfTI to segment (mutually exclusive with `--input_folder` / `--file_list`). +- `--input_folder FOLDER`: Batch only the NIfTI files in that folder (not recursive). +- `--file_list FILE` and `--root_path PATH`: Batch from a text list (requires both). +- `--output_dir DIR`: Output directory (default `./eval`). +- `--keep-temp`: Keep per-case temporary preprocessing files (default: delete after success). +- `--no-skullstrip`: Skip SynthStrip; preprocessing starts from the original image. +- `--modality MODALITY`: `MRI_BRAIN` (default), `MRI_BODY`, or `CT_BODY`. +- `--no-skip`: Recompute even when the expected output already exists (**file list mode only**). +- `--num_partitions N`, `--partition M`: Deterministic shards of the sorted file list (**file list mode only**). +- `-h`, `--help`: Print usage. + +**Note:** Skull stripping calls `brain_t1_preprocess/synthstrip-docker` (Docker must be available unless you use `--no-skullstrip`). In file list mode, failed or timed-out cases are appended to a timestamped log under the output directory. ### Manual Processing (Advanced) @@ -174,169 +220,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) - -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']" -``` - -Multi-GPU: +See [details](docs/inference.md) -```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']" -``` - -### 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 -# 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']" -``` - -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']" -``` - -### 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 -- 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) +- 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) ## License @@ -345,11 +235,10 @@ The `label_mapping` in `evaluation.json` does not include `0` because the postpr 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](http://www.apache.org/licenses/LICENSE-2.0) ### Model Weights License 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) - +[NCLS v1 License](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf?t=eyJscyI6InJlZiIsImxzZCI6IlJFRi1naXRodWIuY29tL252aWRpYS1ob2xvc2NhbiJ9) \ No newline at end of file diff --git a/NV-Segment-CTMR/docs/ctmr2.png b/NV-Segment-CTMR/docs/ctmr2.png new file mode 100644 index 0000000..029a2b7 Binary files /dev/null and b/NV-Segment-CTMR/docs/ctmr2.png differ diff --git a/NV-Segment-CTMR/docs/finetune.md b/NV-Segment-CTMR/docs/finetune.md index 7625d6e..9154533 100644 --- a/NV-Segment-CTMR/docs/finetune.md +++ b/NV-Segment-CTMR/docs/finetune.md @@ -1,4 +1,102 @@ # 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) + +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']" +``` + +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']" +``` + +### 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 +``` ## Configurations @@ -59,3 +157,62 @@ 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. + +```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']" +``` + +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']" +``` + +### 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. \ No newline at end of file diff --git a/NV-Segment-CTMR/docs/inference.md b/NV-Segment-CTMR/docs/inference.md deleted file mode 100644 index 653f70a..0000000 --- a/NV-Segment-CTMR/docs/inference.md +++ /dev/null @@ -1,98 +0,0 @@ -# Inference configurations - -All the configurations for inference is stored in inference.json, change those parameters: - -## `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 -matches the `points`. 0 means background, 1 means foreground, -1 means ignoring this point. `points` and `point_labels` must pe provided together and match length. -- **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 -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])) -``` - -- The `points` together with `label_prompts` for "Kidney", "Lung", "Bone" (class index [2, 20, 21]) are not allowed since those prompts will be divided into sub-categories (e.g. left kidney and right kidney). Use `points` for the sub-categories as defined in the `inference.json`. -- 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. -"129, # kidney mass" insufficient data or dataset excluded. -"131, # vertebrae L6", insufficient data or dataset excluded. -``` - -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 - -```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. - -### 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| -| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | -| model computation | 108.53| 91.9 | 106.84 | 60.02 | 1.18 | 1.02 | 1.81 | 1.53 | -| 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. -- `trt_tf32` and `trt_fp16` are for the TensorRT based models converted in corresponding precision. -- `speedup amp`, `speedup tf32` and `speedup fp16` are the speedup ratios of corresponding models versus the PyTorch float32 model -- `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 diff --git a/NV-Segment-CTMR/scripts/__init__.py b/NV-Segment-CTMR/scripts/__init__.py index fc5f109..5b58bf6 100644 --- a/NV-Segment-CTMR/scripts/__init__.py +++ b/NV-Segment-CTMR/scripts/__init__.py @@ -13,3 +13,6 @@ # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer from .early_stop_score_function import score_function + +# 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/batch_inference_utils.py b/NV-Segment-CTMR/scripts/batch_inference_utils.py new file mode 100644 index 0000000..16e12f3 --- /dev/null +++ b/NV-Segment-CTMR/scripts/batch_inference_utils.py @@ -0,0 +1,262 @@ +""" +Cohort batch inference: discover NIfTI inputs, match MONAI SaveImaged paths, optional resume. + +Behavior is controlled from ``configs/batch_inference.json`` (skip dir lists, resume, cache). +To change which classes are segmented, edit ``everything_labels`` in ``configs/inference.json``. +To customize path filtering further, edit :func:`should_skip_path_by_parent_rules`. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +import time +from pathlib import Path +from typing import Any + + +def should_skip_path_by_parent_rules( + path: Path, + *, + skip_dir_names: list[str] | None = None, + skip_dir_prefixes: list[str] | None = None, +) -> bool: + """ + Return True if this path should be excluded from batch discovery. + + Default (empty lists): do not skip any path based on directory names. + + - ``skip_dir_names``: any **parent** directory component that equals a name (case-insensitive). + - ``skip_dir_prefixes``: any **parent** directory component whose name **starts with** a prefix + (case-insensitive). + """ + names = {n.strip().lower() for n in (skip_dir_names or []) if n and str(n).strip()} + prefixes = tuple(p.strip().lower() for p in (skip_dir_prefixes or []) if p and str(p).strip()) + for part in path.parts[:-1]: + pl = part.lower() + if pl in names: + return True + for prefix in prefixes: + if prefix and pl.startswith(prefix): + return True + return False + + +def expected_output_path( + input_path: Path, + input_root: Path, + output_dir: Path, + postfix: str, + ext: str, +) -> Path: + """Match MONAI FolderLayout + separate_folder + data_root_dir.""" + input_path = input_path.resolve() + input_root = input_root.resolve() + output_dir = output_dir.resolve() + rel = os.path.relpath(input_path, input_root) + rel_dir = os.path.dirname(rel) + stem = input_path.name + if stem.endswith(".nii.gz"): + stem = stem[: -len(".nii.gz")] + elif stem.endswith(".nii"): + stem = stem[: -len(".nii")] + sub = Path(rel_dir) if rel_dir else Path() + return output_dir / sub / stem / f"{stem}_{postfix}{ext}" + + +def collect_input_paths( + input_root: Path, + pattern: str = "**/*.nii.gz", + *, + skip_dir_names: list[str] | None = None, + skip_dir_prefixes: list[str] | None = None, +) -> list[Path]: + input_root = input_root.resolve() + out: list[Path] = [] + 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 + ): + continue + out.append(p) + return out + + +def _cache_path( + input_dir: str, + output_dir: str, + postfix: str, + ext: str, + skip: bool, + skip_dir_names: list[str] | None, + skip_dir_prefixes: list[str] | None, +) -> Path: + sig = ( + f"{os.path.abspath(input_dir)}|{os.path.abspath(output_dir)}|{postfix}|{ext}|{skip}|" + f"{sorted(skip_dir_names or [])}|{sorted(skip_dir_prefixes or [])}" + ) + h = hashlib.sha256(sig.encode()).hexdigest()[:24] + base = Path(tempfile.gettempdir()) + return base / f"nvseg_batch_input_{h}.json" + + +def _compute_input_list( + input_dir: str, + output_dir: str, + postfix: str, + ext: str, + *, + skip_existing: bool, + skip_dir_names: list[str] | None, + skip_dir_prefixes: list[str] | None, +) -> tuple[list[str], int]: + """ + Returns ``(paths_to_run, n_discovered)`` where ``n_discovered`` is the number of + ``*.nii.gz`` paths after directory filters (before resume skip). + """ + root = Path(input_dir) + out_root = Path(output_dir) + all_paths = collect_input_paths( + root, + skip_dir_names=skip_dir_names, + skip_dir_prefixes=skip_dir_prefixes, + ) + n_discovered = len(all_paths) + if not skip_existing: + return [str(p) for p in all_paths], n_discovered + + missing: list[str] = [] + for inp in all_paths: + exp = expected_output_path(inp, root, out_root, postfix, ext) + if not exp.is_file() or exp.stat().st_size == 0: + missing.append(str(inp)) + return missing, n_discovered + + +def _parse_cache_payload(raw: Any) -> tuple[list[str], int]: + """Load cache written by rank 0. Supports legacy JSON list for backward compatibility.""" + if isinstance(raw, list): + # Legacy: empty [] is ambiguous (stale file or old format) -> signal recompute on workers. + return [str(p) for p in raw], len(raw) if raw else -1 + if isinstance(raw, dict) and "paths" in raw: + paths = raw["paths"] + if not isinstance(paths, list): + raise RuntimeError("[nvseg] batch: bad cache (paths); rm /tmp/nvseg_batch_input_*.json") + n_raw = raw.get("n_discovered") + if n_raw is None: + n_discovered = len(paths) if paths else -1 + else: + n_discovered = int(n_raw) + return [str(p) for p in paths], n_discovered + raise RuntimeError("[nvseg] batch: bad cache format; rm /tmp/nvseg_batch_input_*.json") + + +def build_input_list( + input_dir: str, + output_dir: str, + output_postfix: str, + output_ext: str, + batch_skip_dir_names: list | None = None, + batch_skip_dir_prefixes: list | None = None, + batch_resume_skip_existing: bool = True, + batch_use_input_list_cache: bool = True, + batch_cache_wait_sec: float = 120.0, +) -> list[str]: + """ + Called from ``configs/batch_inference.json``. + + ``batch_resume_skip_existing``: if True, only queue inputs whose output file is missing or empty. + ``batch_skip_dir_names`` / ``batch_skip_dir_prefixes``: filter discovery (see + :func:`should_skip_path_by_parent_rules`). + + ``LOCAL_RANK`` (set by ``torchrun``) is still read from the environment for multi-GPU cache. + + If resume mode leaves nothing to run (outputs already exist), prints a message and raises + ``SystemExit(0)`` so the process exits before MONAI builds a zero-length DataLoader (which + would fail under ``DistributedSampler``). If no ``*.nii.gz`` files are discovered, raises + ``RuntimeError``. + """ + names = list(batch_skip_dir_names) if batch_skip_dir_names is not None else [] + prefixes = list(batch_skip_dir_prefixes) if batch_skip_dir_prefixes is not None else [] + + skip = bool(batch_resume_skip_existing) + use_cache = bool(batch_use_input_list_cache) + wait_sec = float(batch_cache_wait_sec) + + local_rank = os.environ.get("LOCAL_RANK", "0") + + if not use_cache or local_rank == "0": + paths, n_discovered = _compute_input_list( + input_dir, + output_dir, + output_postfix, + output_ext, + skip_existing=skip, + skip_dir_names=names, + skip_dir_prefixes=prefixes, + ) + 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.parent.mkdir(parents=True, exist_ok=True) + tmp = cache.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload)) + tmp.replace(cache) + + if local_rank == "0": + mode = "resume (skip existing outputs)" if skip else "full pass (all inputs)" + print( + f"[nvseg] batch {mode}: {len(paths)} volume(s) " + f"(input_dir={os.path.abspath(input_dir)}, output_dir={os.path.abspath(output_dir)})", + flush=True, + ) + else: + 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: + if cache.is_file(): + payload = json.loads(cache.read_text()) + 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)" + ) + paths, n_discovered = _parse_cache_payload(payload) + # Stale legacy `[]` or missing n_discovered: recompute so workers agree with rank 0. + if n_discovered < 0: + paths, n_discovered = _compute_input_list( + input_dir, + output_dir, + output_postfix, + output_ext, + skip_existing=skip, + skip_dir_names=names, + skip_dir_prefixes=prefixes, + ) + + if not paths: + if n_discovered == 0: + raise RuntimeError("[nvseg] batch: no *.nii.gz under input_dir (check paths and skip rules)") + # Resume: all outputs exist — exit before an empty DistributedSampler / dataloader. + print("[nvseg] batch: nothing to run (resume); ok", flush=True) + raise SystemExit(0) + + return paths + + +__all__ = [ + "build_input_list", + "collect_input_paths", + "expected_output_path", + "should_skip_path_by_parent_rules", +]