diff --git a/README.md b/README.md index d0885de..8484522 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,11 @@ In the paper setting, MedSegAgent integrates 23 datasets and supports 343 segmentation targets across CT, MRI, PET/CT, and ultrasound-related scenarios. +## Skill +This repository also includes a reusable skill at [`skills/medsegagent-nnunet-runner/`](skills/medsegagent-nnunet-runner) for inspecting `nnUNet_results`, selecting a deployed task, and running or preparing nnUNet inference. + +The skill is intended for environments that already have nnUNet v2 available. To run inference, the target environment must provide a configured `nnUNet_results` path and usable nnUNet CLI or Python API support. + ## Supported Datasets The repository currently includes metadata for the following datasets and targets. The summary below is adapted from the dataset table in the paper. diff --git a/skills/medsegagent-nnunet-runner/SKILL.md b/skills/medsegagent-nnunet-runner/SKILL.md new file mode 100644 index 0000000..a258cc5 --- /dev/null +++ b/skills/medsegagent-nnunet-runner/SKILL.md @@ -0,0 +1,84 @@ +--- +name: medsegagent-nnunet-runner +description: Inspect `nnUNet_results`, choose a deployed nnUNet task, and run or prepare segmentation inference. Use when a user wants to match a free-form target to deployed labels, select the right task, or run prediction and postprocessing. +--- + +# nnUNet Segmentation Runner + +## Overview + +Inspect what is deployed, choose the right nnUNet task, and run or prepare inference. Do not claim inference succeeded unless it was actually run and outputs were produced. + +## Prerequisites + +- An nnUNet v2 environment is available. +- `nnUNet_results` is configured. `nnUNet_raw` and `nnUNet_preprocessed` may also be needed. +- To actually run inference, the nnUNet CLI or Python API must be installed and usable. +- Custom trainers must be importable in the target environment. + +## Workflow + +1. Inspect the nnUNet environment first. + - Prefer `nnUNet_results`. + - Use `RESULTS_FOLDER` only as a legacy fallback. + - If the environment is unclear, run `scripts/nnunetv2-doctor.sh` first. + - Run `scripts/list_nnunet_results.py --json` to enumerate real datasets, model folders, modalities, labels, folds, and checkpoints. + +2. Parse the request into: + - modality + - target + - qualifiers such as left/right, tumor/cyst, whole organ vs substructure + +3. Read metadata in this order: + - `dataset.json` under nnUNet result folders for what is actually available + - `dataset/*.json` in this repo for supplemental routing metadata + - `dataset/README.md` only when adding or fixing dataset entries + +4. Filter by modality, then match labels conservatively. + - Preserve laterality and substructure specificity. + - Keep tumor, cyst, lesion, vessel, and organ concepts separate. + - Treat repo metadata as secondary when it disagrees with deployed nnUNet results. + +5. Choose how to run inference. + - Prefer `inference_instructions.txt` when it exists. + - For default model folders, use `nnUNetv2_predict`, plus ensembling or postprocessing if needed. + - For custom trainers, prefer exact model-folder inference through the Python predictor API. + - Use all available trained folds by default unless the environment only supports a specific fold such as `all`. + - Keep input naming and file endings consistent with the selected task `dataset.json`. + +6. Return either: + - a compact dataset-to-label mapping + - or the selected task plus the exact inference command + +```json +{ + "DatasetNameA": [1, 2], + "DatasetNameB": [7] +} +``` + +Return `{}` if no suitable modality or label is available, and say why. + +## Rules + +- Prefer real nnUNet availability over paper claims or static repo metadata. +- If a dataset exists in repo metadata but not in `nnUNet_results`, describe it as known but not deployed. +- Group multiple trainer folders under one dataset unless the user explicitly asks about checkpoints or trainer variants. +- Prefer officially generated `inference_instructions.txt` when present. +- If input files do not match the trained task modality, naming scheme, or file ending, stop and say the input is not ready. +- Do not silently guess a configuration when multiple materially different deployed options exist. +- If repo metadata and nnUNet metadata disagree, report the mismatch explicitly. +- Do not guess custom trainer or plans settings from memory; read them from the selected model folder. + +## Extending Coverage + +When adding or updating routing metadata: + +1. Follow `dataset/README.md`. +2. Keep names explicit, such as `kidney_left`. +3. Do not invent deployed availability if no trained output exists in `nnUNet_results`. + +## References + +- Read `references/nnunetv2.md` only when you need the official nnUNet v2 path conventions, result layout, prediction commands, or `dataset.json` details. +- Read `references/examples.md` only when you need compact routing or inference examples. diff --git a/skills/medsegagent-nnunet-runner/references/examples.md b/skills/medsegagent-nnunet-runner/references/examples.md new file mode 100644 index 0000000..896746e --- /dev/null +++ b/skills/medsegagent-nnunet-runner/references/examples.md @@ -0,0 +1,70 @@ +# Examples + +## Example 1: Routing + +Request: + +```text +segment the spleen in CT +``` + +Answer style: + +```json +{ + "Dataset009_Spleen": [1], + "Dataset022_FLARE22": [3], + "Dataset051_AMOS_CT": [1] +} +``` + +Note: + +```text +Found CT tasks in nnUNet results. Labels come from task dataset.json files. +``` + +## Example 2: No Match + +Request: + +```text +segment the left kidney in MRI +``` + +Answer style: + +```json +{} +``` + +Note: + +```text +No deployed MRI task with a confirmed left-kidney label was found. +``` + +## Example 3: Prediction Command + +Request: + +```text +run liver segmentation on a CT case with a deployed nnUNet task +``` + +Answer style: + +```text +Selected task: Dataset051_AMOS_CT +Selected labels: [6] +Selected model folder: nnUNetTrainer__nnUNetPlans__3d_fullres +Selected configuration: 3d_fullres +Command: +nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d 51 -c 3d_fullres +``` + +Note: + +```text +Replace dataset ID, model folder, and paths with the deployed task you actually select. +``` diff --git a/skills/medsegagent-nnunet-runner/references/nnunetv2.md b/skills/medsegagent-nnunet-runner/references/nnunetv2.md new file mode 100644 index 0000000..f901eed --- /dev/null +++ b/skills/medsegagent-nnunet-runner/references/nnunetv2.md @@ -0,0 +1,102 @@ +# nnUNet v2 Notes + +Use this reference when you need to verify official nnUNet v2 conventions. + +## Environment Variables + +- Official nnUNet v2 variables are `nnUNet_raw`, `nnUNet_preprocessed`, and `nnUNet_results`. +- `nnUNet_results` is the primary source for trained or downloaded model availability. +- `RESULTS_FOLDER` is not the primary nnUNet v2 convention; use it only as a legacy fallback. + +## Result Layout + +The official training output layout is: + +```text +nnUNet_results/DatasetXXX_NAME/TRAINER__PLANS__CONFIGURATION/fold_x +``` + +Typical trainer-level files: + +- `dataset.json` +- `dataset_fingerprint.json` +- `plans.json` + +Typical fold-level files: + +- `checkpoint_final.pth` +- `checkpoint_best.pth` +- `debug.json` + +`checkpoint_final.pth` is the standard inference checkpoint. + +If `nnUNetv2_find_best_configuration` was run, nnUNet may also create: + +- `inference_instructions.txt` +- `inference_information.json` + +Prefer `inference_instructions.txt` when it exists. + +## dataset.json + +When reading nnUNet v2 metadata: + +- modality or channel information usually comes from `channel_names` +- classes or regions come from `labels` +- `background` should be `0` + +Label values may appear as integers, strings, or region definitions depending on dataset style. Normalize simple numeric strings when summarizing, but do not change semantic meaning. + +## Practical Guidance + +- Use nnUNet result-folder `dataset.json` as the source of truth for deployed labels and modalities. +- Use repo-local `dataset/*.json` as a supplemental routing library. +- If the two disagree, surface the mismatch instead of silently choosing one. + +## Inference Commands + +Core prediction pattern: + +```bash +nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -c CONFIGURATION +``` + +Notes: + +- use `--save_probabilities` only if you intend to ensemble predictions later +- each configuration should write to its own output folder +- by default nnUNet uses all 5 folds as an ensemble +- if predicting with a single model trained on all cases, specify `-f all` + +Ensembling: + +```bash +nnUNetv2_ensemble -i FOLDER1 FOLDER2 -o OUTPUT_FOLDER -np NUM_PROCESSES +``` + +Postprocessing: + +```bash +nnUNetv2_apply_postprocessing -i FOLDER_WITH_PREDICTIONS -o OUTPUT_FOLDER --pp_pkl_file POSTPROCESSING_FILE -plans_json PLANS_FILE -dataset_json DATASET_JSON_FILE +``` + +## Python Predictor + +For custom trainer folders, the safer path is to initialize from the trained model folder directly. + +```python +from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor +``` + +Useful methods: + +- `initialize_from_trained_model_folder(...)` +- `predict_from_files(...)` + +This avoids reconstructing trainer or plans choices from memory when the deployed model folder is non-default. + +## Official Sources + +- `documentation/setting_up_paths.md` +- `documentation/dataset_format.md` +- `documentation/how_to_use_nnunet.md` diff --git a/skills/medsegagent-nnunet-runner/scripts/list_nnunet_results.py b/skills/medsegagent-nnunet-runner/scripts/list_nnunet_results.py new file mode 100644 index 0000000..37b909a --- /dev/null +++ b/skills/medsegagent-nnunet-runner/scripts/list_nnunet_results.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +from pathlib import Path + + +def resolve_results_dir(): + for key in ("nnUNet_results", "RESULTS_FOLDER"): + value = os.environ.get(key) + if value: + return key, Path(value) + return None, None + + +def load_json(path): + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + + +def normalize_label_value(value): + if isinstance(value, str) and value.isdigit(): + return int(value) + if isinstance(value, list): + return [normalize_label_value(v) for v in value] + return value + + +def first_dataset_json(dataset_dir): + candidates = sorted(dataset_dir.glob("*/dataset.json")) + if candidates: + return candidates[0] + direct = dataset_dir / "dataset.json" + if direct.is_file(): + return direct + return None + + +def checkpoint_count(trainer_dir): + count = 0 + for pattern in ("**/checkpoint_final.pth", "**/checkpoint_best.pth"): + count += len(list(trainer_dir.glob(pattern))) + return count + + +def parse_model_folder_name(name): + parts = name.split("__") + return { + "trainer": parts[0] if len(parts) > 0 else None, + "plans": parts[1] if len(parts) > 1 else None, + "configuration": parts[2] if len(parts) > 2 else None, + } + + +def list_available_folds(trainer_dir): + folds = [] + for fold_dir in sorted(p for p in trainer_dir.iterdir() if p.is_dir() and p.name.startswith("fold_")): + fold_name = fold_dir.name.removeprefix("fold_") + fold_info = { + "fold": fold_name, + "has_checkpoint_final": (fold_dir / "checkpoint_final.pth").is_file(), + "has_checkpoint_best": (fold_dir / "checkpoint_best.pth").is_file(), + } + folds.append(fold_info) + return folds + + +def summarize_dataset(dataset_dir): + trainer_dirs = sorted([p for p in dataset_dir.iterdir() if p.is_dir()]) + dataset_json_path = first_dataset_json(dataset_dir) + dataset_json = load_json(dataset_json_path) if dataset_json_path else {} + + labels = (dataset_json or {}).get("labels", {}) + channel_names = (dataset_json or {}).get("channel_names", {}) + input_modalities = (dataset_json or {}).get("input_modalities", {}) + modalities = list(channel_names.values()) or list(input_modalities.values()) + + trainers = [] + for trainer_dir in trainer_dirs: + parsed = parse_model_folder_name(trainer_dir.name) + folds = list_available_folds(trainer_dir) + trainer_info = { + "name": trainer_dir.name, + "trainer": parsed["trainer"], + "plans": parsed["plans"], + "configuration": parsed["configuration"], + "has_dataset_json": (trainer_dir / "dataset.json").is_file(), + "has_plans_json": (trainer_dir / "plans.json").is_file(), + "checkpoint_count": checkpoint_count(trainer_dir), + "folds": folds, + } + trainers.append(trainer_info) + + normalized_labels = { + key: normalize_label_value(value) + for key, value in labels.items() + if key != "background" + } + + return { + "folder": dataset_dir.name, + "dataset_name": (dataset_json or {}).get("name", dataset_dir.name), + "modalities": modalities, + "labels": normalized_labels, + "trainer_count": len(trainers), + "trainers": trainers, + } + + +def collect_results(results_dir): + dataset_dirs = sorted( + p for p in results_dir.iterdir() if p.is_dir() and p.name.startswith("Dataset") + ) + return [summarize_dataset(dataset_dir) for dataset_dir in dataset_dirs] + + +def main(): + parser = argparse.ArgumentParser( + description="List datasets, trainers, modalities, and labels from an nnUNet results directory." + ) + parser.add_argument("--json", action="store_true", help="Emit JSON output") + args = parser.parse_args() + + env_key, results_dir = resolve_results_dir() + if not results_dir: + raise SystemExit("No nnUNet results directory found in nnUNet_results or RESULTS_FOLDER.") + if not results_dir.exists(): + raise SystemExit(f"{env_key} points to a missing path: {results_dir}") + + datasets = collect_results(results_dir) + + if args.json: + print( + json.dumps( + { + "env_var": env_key, + "results_dir": str(results_dir), + "dataset_count": len(datasets), + "datasets": datasets, + }, + ensure_ascii=False, + indent=2, + ) + ) + return + + print(f"{env_key}={results_dir}") + print(f"datasets={len(datasets)}") + for dataset in datasets: + print(f"- {dataset['folder']} | {dataset['dataset_name']}") + if dataset["modalities"]: + print(f" modalities: {', '.join(dataset['modalities'])}") + print(f" labels: {len(dataset['labels'])}") + print(f" trainers: {dataset['trainer_count']}") + for trainer in dataset["trainers"]: + print( + f" - {trainer['trainer']} | {trainer['configuration']} | folds={len(trainer['folds'])} | checkpoints={trainer['checkpoint_count']}" + ) + + +if __name__ == "__main__": + main() diff --git a/skills/medsegagent-nnunet-runner/scripts/nnunetv2-doctor.sh b/skills/medsegagent-nnunet-runner/scripts/nnunetv2-doctor.sh new file mode 100644 index 0000000..2b90d44 --- /dev/null +++ b/skills/medsegagent-nnunet-runner/scripts/nnunetv2-doctor.sh @@ -0,0 +1,189 @@ +#!/usr/bin/env bash +set -euo pipefail + +CONDA_ROOT="${CONDA_ROOT:-}" +CONDA_ENV_NAME="${CONDA_ENV_NAME:-}" +NNUNET_RAW="${NNUNET_RAW:-${nnUNet_raw:-}}" +NNUNET_PREPROCESSED="${NNUNET_PREPROCESSED:-${nnUNet_preprocessed:-}}" +NNUNET_RESULTS="${NNUNET_RESULTS:-${nnUNet_results:-}}" +DATASET="" + +pick_first_existing_dir() { + local candidate + for candidate in "$@"; do + if [[ -n "${candidate}" && -d "${candidate}" ]]; then + printf '%s\n' "${candidate}" + return 0 + fi + done + return 1 +} + +resolve_conda_root() { + if [[ -n "${CONDA_ROOT}" ]]; then + printf '%s\n' "${CONDA_ROOT}" + return 0 + fi + + if [[ -n "${CONDA_EXE:-}" ]]; then + printf '%s\n' "$(cd "$(dirname "${CONDA_EXE}")/.." && pwd)" + return 0 + fi + + pick_first_existing_dir \ + "${HOME}/miniconda3" \ + "${HOME}/anaconda3" +} + +resolve_conda_env_name() { + if [[ -n "${CONDA_ENV_NAME}" ]]; then + printf '%s\n' "${CONDA_ENV_NAME}" + return 0 + fi + + if [[ -n "${CONDA_DEFAULT_ENV:-}" ]]; then + printf '%s\n' "${CONDA_DEFAULT_ENV}" + return 0 + fi + + printf '%s\n' "nnunetv2" +} + +resolve_nnunet_path() { + local current="$1" + shift + if [[ -n "${current}" ]]; then + printf '%s\n' "${current}" + return 0 + fi + pick_first_existing_dir "$@" +} + +usage() { + cat <<'EOF' +Usage: + nnunetv2-doctor.sh [--dataset 22|Dataset022_FLARE22] + +Environment overrides: + CONDA_ROOT + CONDA_ENV_NAME + NNUNET_RAW + NNUNET_PREPROCESSED + NNUNET_RESULTS +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --dataset|-d) + DATASET="${2:-}" + shift 2 + ;; + --help|-h) + usage + exit 0 + ;; + *) + echo "Unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +CONDA_ROOT="$(resolve_conda_root || true)" +CONDA_ENV_NAME="$(resolve_conda_env_name)" +NNUNET_RAW="$(resolve_nnunet_path "${NNUNET_RAW}" "${HOME}/nnUNet_raw" || true)" +NNUNET_PREPROCESSED="$(resolve_nnunet_path "${NNUNET_PREPROCESSED}" "${HOME}/nnUNet_preprocessed" || true)" +NNUNET_RESULTS="$(resolve_nnunet_path "${NNUNET_RESULTS}" "${HOME}/nnUNet_results" || true)" + +if [[ -z "${CONDA_ROOT}" || ! -f "${CONDA_ROOT}/etc/profile.d/conda.sh" ]]; then + echo "Missing conda activation script: ${CONDA_ROOT}/etc/profile.d/conda.sh" >&2 + exit 1 +fi + +source "${CONDA_ROOT}/etc/profile.d/conda.sh" +conda activate "${CONDA_ENV_NAME}" + +if [[ -n "${NNUNET_RAW}" ]]; then + export nnUNet_raw="${NNUNET_RAW}" +fi +if [[ -n "${NNUNET_PREPROCESSED}" ]]; then + export nnUNet_preprocessed="${NNUNET_PREPROCESSED}" +fi +if [[ -n "${NNUNET_RESULTS}" ]]; then + export nnUNet_results="${NNUNET_RESULTS}" +fi + +echo "=== Conda ===" +echo "CONDA_PREFIX=${CONDA_PREFIX}" +echo "python=$(command -v python)" +echo "nnUNetv2_predict=$(command -v nnUNetv2_predict || true)" +echo + +echo "=== nnUNet Paths ===" +echo "nnUNet_raw=${nnUNet_raw:-}" +echo "nnUNet_preprocessed=${nnUNet_preprocessed:-}" +echo "nnUNet_results=${nnUNet_results:-}" +echo + +echo "=== GPU ===" +if command -v nvidia-smi >/dev/null 2>&1; then + nvidia-smi +else + echo "nvidia-smi not found" +fi +echo + +echo "=== Torch CUDA ===" +python - <<'PY' +import os +import torch + +print(f"torch={torch.__version__}") +print(f"torch.version.cuda={torch.version.cuda}") +print(f"cuda_available={torch.cuda.is_available()}") +try: + print(f"device_count={torch.cuda.device_count()}") +except Exception as exc: + print(f"device_count_error={exc!r}") +try: + if torch.cuda.device_count() > 0: + print(f"device0_name={torch.cuda.get_device_name(0)}") +except Exception as exc: + print(f"device0_name_error={exc!r}") +for key in ("nnUNet_raw", "nnUNet_preprocessed", "nnUNet_results"): + print(f"{key}={os.environ.get(key)}") +PY +echo + +if [[ -n "${DATASET}" ]]; then + echo "=== Dataset Lookup ===" + python - "${DATASET}" <<'PY' +import os +import sys +from pathlib import Path + +from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name + +dataset_arg = sys.argv[1] +dataset_name = maybe_convert_to_dataset_name(dataset_arg) +print(f"dataset={dataset_name}") + +roots = { + "raw": Path(os.environ["nnUNet_raw"]), + "preprocessed": Path(os.environ["nnUNet_preprocessed"]), + "results": Path(os.environ["nnUNet_results"]), +} +for label, root in roots.items(): + path = root / dataset_name + print(f"{label}_path={path} exists={path.exists()}") + +results_dir = roots["results"] / dataset_name +if results_dir.exists(): + print("model_folders:") + for child in sorted(results_dir.iterdir()): + if child.is_dir(): + print(f" {child.name}") +PY +fi