Nv Segment Ctmr

scripts/run_ctmr.py

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""NVIDIA-Medtech NV-Segment-CTMR skill wrapper.

Thin wrapper around the upstream MONAI bundle command documented by
NVIDIA-Medtech/NV-Segment-CTMR. The wrapper does not implement inference; it
launches `python -m monai.bundle run`, captures logs, and summarizes the
resulting NIfTI label map as JSON.

Engineering verification only. Output is NOT clinically meaningful.
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

import nibabel as nib
import numpy as np
import typer

app = typer.Typer(add_completion=False)

SKILL_NAME = "nv_segment_ctmr"
DEFAULT_UPSTREAM_COMMIT = "cb921f5c58837c0f42a713855d68b32af88e1cdd"
MODEL_REPO = (
    "https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/"
    f"{DEFAULT_UPSTREAM_COMMIT}/NV-Segment-CTMR"
)
SUPPORTED_MODALITIES = ("CT_BODY", "MRI_BODY", "MRI_BRAIN")
GEOMETRY_TOLERANCE = float("1e-4")
REPO_ROOT = Path(__file__).resolve().parents[int("3")]
_CHILD_ENV_KEYS = (
    "CUDA_DEVICE_ORDER",
    "CUDA_HOME",
    "CUDA_PATH",
    "CUDA_VISIBLE_DEVICES",
    "CUBLAS_WORKSPACE_CONFIG",
    "CURL_CA_BUNDLE",
    "HF_HOME",
    "HF_HUB_CACHE",
    "HF_HUB_OFFLINE",
    "HOME",
    "HUGGINGFACE_HUB_CACHE",
    "LANG",
    "LC_ALL",
    "LD_LIBRARY_PATH",
    "MONAI_DATA_DIRECTORY",
    "NVIDIA_DRIVER_CAPABILITIES",
    "NVIDIA_VISIBLE_DEVICES",
    "OMP_NUM_THREADS",
    "PATH",
    "PYTHONNOUSERSITE",
    "PYTHONPATH",
    "PYTORCH_CUDA_ALLOC_CONF",
    "REQUESTS_CA_BUNDLE",
    "SSL_CERT_FILE",
    "TMPDIR",
    "TORCH_EXTENSIONS_DIR",
    "TORCH_HOME",
    "TRANSFORMERS_CACHE",
    "XDG_CACHE_HOME",
)


def _child_process_env(overrides: dict[str, str] | None = None) -> dict[str, str]:
    """Build the minimal environment needed by the trusted upstream bundle."""
    env: dict[str, str] = {}
    for name in _CHILD_ENV_KEYS:
        value = os.environ.get(name)
        if value is not None:
            env[name] = value
    env.setdefault("PATH", os.defpath)
    if overrides:
        env.update(overrides)
    return env


def emit(payload: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(payload, indent=2))
    sys.stdout.flush()


def tail(s: str, n_chars: int = int("4000")) -> str:
    if len(s) <= n_chars:
        return s
    return "..." + s[-n_chars:]


def git_commit(root: Path) -> str:
    try:
        proc = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=str(root),
            env=_child_process_env(),
            capture_output=True,
            text=True,
            timeout=int("10"),
            check=False,
        )
    except Exception:
        return ""
    if proc.returncode == 0:
        return proc.stdout.strip()
    return ""


def _strip_nifti_suffix(path: Path) -> str:
    name = path.name
    for suffix in (".nii.gz", ".nii"):
        if name.endswith(suffix):
            return name[: -len(suffix)]
    return path.stem


def _is_nifti_path(path: Path) -> bool:
    return path.name.endswith(".nii.gz") or path.suffix == ".nii"


def _round_floats(values, ndigits: int = int("6")) -> list[float]:
    return [round(float(v), ndigits) for v in values]


def _spacing3(img: nib.spatialimages.SpatialImage) -> list[float]:
    zooms = list(img.header.get_zooms())
    while len(zooms) < int("3"):
        zooms.append(1.0)
    return _round_floats(zooms[: int("3")])


def _input_summary(img: nib.spatialimages.SpatialImage) -> dict[str, Any]:
    return {
        "shape": [int(v) for v in img.shape],
        "ndim": len(img.shape),
        "spacing": _spacing3(img),
    }


def _geometry_summary(
    input_img: nib.spatialimages.SpatialImage,
    output_img: nib.spatialimages.SpatialImage,
) -> dict[str, Any]:
    input_shape = [int(v) for v in input_img.shape]
    output_shape = [int(v) for v in output_img.shape]
    input_spacing = _spacing3(input_img)
    output_spacing = _spacing3(output_img)
    affine_max_abs_diff = float(np.max(np.abs(input_img.affine - output_img.affine)))
    return {
        "input_shape": input_shape,
        "output_shape": output_shape,
        "shape_match": input_shape == output_shape,
        "input_spacing": input_spacing,
        "output_spacing": output_spacing,
        "spacing_match": input_spacing == output_spacing,
        "affine_max_abs_diff": round(affine_max_abs_diff, int("8")),
        "affine_match": affine_max_abs_diff <= GEOMETRY_TOLERANCE,
    }


def _coerce_label_id(value: Any) -> int | None:
    if isinstance(value, bool):
        return None
    if isinstance(value, int):
        return int(value)
    if isinstance(value, str):
        try:
            return int(value)
        except ValueError:
            return None
    return None


def _walk_label_records(raw: Any) -> list[tuple[int, str]]:
    records: list[tuple[int, str]] = []
    if isinstance(raw, list):
        for item in raw:
            records.extend(_walk_label_records(item))
        return records
    if not isinstance(raw, dict):
        return records

    for key, value in raw.items():
        key_id = _coerce_label_id(key)
        if key_id is not None:
            if isinstance(value, str):
                records.append((key_id, value))
            elif isinstance(value, dict):
                name = (
                    value.get("name")
                    or value.get("label")
                    or value.get("organ")
                    or value.get("class")
                    or f"label_id_{key_id}"
                )
                records.append((key_id, str(name)))
                records.extend(_walk_label_records(value))
            else:
                records.append((key_id, f"label_id_{key_id}"))
            continue

        value_id = _coerce_label_id(value)
        if value_id is not None:
            records.append((value_id, str(key)))
            continue

        if isinstance(value, dict):
            id_value = None
            for id_key in ("id", "index", "label_id", "label_index", "value"):
                if id_key in value:
                    id_value = _coerce_label_id(value[id_key])
                    if id_value is not None:
                        break
            if id_value is not None:
                name = (
                    value.get("name")
                    or value.get("label")
                    or value.get("organ")
                    or value.get("class")
                    or key
                )
                records.append((id_value, str(name)))
            records.extend(_walk_label_records(value))
        elif isinstance(value, list):
            records.extend(_walk_label_records(value))
    return records


def _load_label_map(upstream_root: Path) -> tuple[dict[int, str], Path | None]:
    candidates = [
        upstream_root / "configs" / "label_dict.json",
        upstream_root / "configs" / "metadata.json",
        upstream_root / "label_dict.json",
        upstream_root / "metadata.json",
    ]
    for path in candidates:
        if not path.is_file():
            continue
        try:
            raw = json.loads(path.read_text())
        except Exception:
            continue
        if path.name == "metadata.json" and isinstance(raw, dict):
            channel_def = (
                raw.get("network_data_format", {})
                .get("outputs", {})
                .get("pred", {})
                .get("channel_def")
            )
            if channel_def:
                raw = channel_def
        label_by_id: dict[int, str] = {}
        for label_id, name in _walk_label_records(raw):
            if label_id >= 0 and label_id not in label_by_id:
                label_by_id[label_id] = name
        if label_by_id:
            return label_by_id, path
    return {}, None


def _parse_label_prompts(raw: str | None) -> list[int] | None:
    if raw is None or not raw.strip():
        return None
    values: list[int] = []
    for part in raw.split(","):
        item = part.strip()
        if not item:
            continue
        values.append(int(item))
    return values


def _build_command(
    input_path: Path,
    output_dir: Path,
    modality: str,
    label_prompts: list[int] | None,
) -> list[str]:
    input_dict: dict[str, Any] = {"image": str(input_path)}
    if label_prompts is not None:
        input_dict["label_prompt"] = label_prompts
    return [
        sys.executable,
        "-m",
        "monai.bundle",
        "run",
        "--config_file",
        "configs/inference.json",
        "--input_dict",
        repr(input_dict),
        "--output_dir",
        str(output_dir),
        "--modality",
        modality,
    ]


def _expected_output_candidates(output_dir: Path, input_path: Path) -> list[Path]:
    stem = _strip_nifti_suffix(input_path)
    folder = output_dir / stem
    suffixes = ("_trans.nii.gz", "_seg.nii.gz", ".nii.gz", ".nii")
    return [folder / f"{stem}{suffix}" for suffix in suffixes] + [
        output_dir / f"{stem}{suffix}" for suffix in suffixes
    ]


def _find_output_mask(output_dir: Path, input_path: Path, run_started: float) -> Path | None:
    for candidate in _expected_output_candidates(output_dir, input_path):
        if candidate.is_file() and candidate.stat().st_size > 0:
            return candidate

    candidates: list[Path] = []
    if output_dir.is_dir():
        for path in output_dir.rglob("*"):
            if not path.is_file() or not _is_nifti_path(path):
                continue
            try:
                if path.resolve() == input_path.resolve():
                    continue
            except OSError:
                pass
            try:
                if path.stat().st_size > 0 and path.stat().st_mtime >= run_started - 1:
                    candidates.append(path)
            except OSError:
                continue
    if not candidates:
        return None
    return sorted(candidates, key=lambda p: p.stat().st_mtime, reverse=True)[0]


def _empty_geometry(input_summary: dict[str, Any]) -> dict[str, Any]:
    return {
        "input_shape": input_summary["shape"],
        "output_shape": [],
        "shape_match": False,
        "input_spacing": input_summary["spacing"],
        "output_spacing": [],
        "spacing_match": False,
        "affine_max_abs_diff": None,
        "affine_match": False,
    }


def _empty_output_summary(
    input_summary: dict[str, Any],
    label_prompts: list[int] | None,
    label_map: dict[int, str],
    label_map_source: Path | None,
) -> dict[str, Any]:
    return {
        "path": None,
        "shape": [],
        "label_prompts_requested": label_prompts,
        "label_ids_present": [],
        "unexpected_label_ids": [],
        "label_set_valid": False,
        "label_map_loaded": bool(label_map),
        "label_map_source": str(label_map_source) if label_map_source is not None else None,
        "class_counts": {},
        "voxel_volume_ml": None,
        "class_volumes_ml": {},
        "any_label_present": False,
        "geometry": _empty_geometry(input_summary),
    }


def _mask_summary(
    mask_path: Path,
    input_img: nib.spatialimages.SpatialImage,
    label_prompts: list[int] | None,
    label_map: dict[int, str],
    label_map_source: Path | None,
) -> dict[str, Any]:
    mask_img = nib.load(str(mask_path))
    arr = np.asarray(mask_img.get_fdata()).astype(np.int64)
    voxel_volume_ml = float(np.prod(_spacing3(mask_img))) / float("1000.0")
    unique, counts = np.unique(arr, return_counts=True)
    class_counts: dict[str, int] = {}
    class_volumes_ml: dict[str, float] = {}
    label_ids_present: list[int] = []
    unexpected: list[int] = []
    valid_ids = set(label_map)

    for value, count in zip(unique.tolist(), counts.tolist()):
        label_id = int(value)
        if label_id == 0:
            continue
        label_ids_present.append(label_id)
        if label_map:
            if label_id not in valid_ids:
                unexpected.append(label_id)
        elif label_id < 0:
            unexpected.append(label_id)
        name = label_map.get(label_id, f"label_id_{label_id}")
        class_counts[name] = int(count)
        class_volumes_ml[name] = round(int(count) * voxel_volume_ml, int("4"))

    return {
        "shape": [int(v) for v in arr.shape],
        "label_prompts_requested": label_prompts,
        "label_ids_present": sorted(label_ids_present),
        "unexpected_label_ids": sorted(unexpected),
        "label_set_valid": len(unexpected) == 0,
        "label_map_loaded": bool(label_map),
        "label_map_source": str(label_map_source) if label_map_source is not None else None,
        "class_counts": class_counts,
        "voxel_volume_ml": round(voxel_volume_ml, int("8")),
        "class_volumes_ml": class_volumes_ml,
        "any_label_present": len(class_counts) > 0,
        "geometry": _geometry_summary(input_img, mask_img),
    }


def _model_inventory(upstream_root: Path, label_map_source: Path | None) -> dict[str, Any]:
    model_pt = upstream_root / "models" / "model.pt"
    return {
        "model_pt_present": model_pt.is_file(),
        "model_pt_path": str(model_pt),
        "label_map_present": label_map_source is not None,
        "label_map_path": str(label_map_source) if label_map_source is not None else None,
    }


def _resolve_device(requested: str) -> str:
    if requested != "auto":
        return requested
    try:
        import torch  # noqa: PLC0415

        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "unknown"


def _error_payload(message: str, detail: str) -> dict[str, Any]:
    return {
        "skill": SKILL_NAME,
        "error": message,
        "detail": detail,
        "model_repo": MODEL_REPO,
    }


def _valid_upstream_root(path: Path) -> bool:
    return (path / "configs" / "inference.json").is_file()


def _candidate_upstream_roots(env_value: str) -> list[Path]:
    candidates: list[Path] = []
    if env_value:
        candidates.append(Path(env_value).expanduser())
    candidates.extend(
        [
            REPO_ROOT / ".workbench_data/upstreams/NV-Segment-CTMR/NV-Segment-CTMR",
            Path.home() / "NV-Segment-CTMR/NV-Segment-CTMR",
            Path.home() / "NV-Segment-CTMR",
        ]
    )
    deduped: list[Path] = []
    seen: set[str] = set()
    for candidate in candidates:
        key = str(candidate)
        if key not in seen:
            seen.add(key)
            deduped.append(candidate)
    return deduped


def _resolve_upstream_root(
    explicit_root: Path | None,
    env_value: str,
) -> tuple[Path | None, list[str]]:
    if explicit_root is not None:
        resolved = explicit_root.expanduser().resolve()
        return (resolved if _valid_upstream_root(resolved) else None), [str(resolved)]
    checked: list[str] = []
    for candidate in _candidate_upstream_roots(env_value):
        resolved = candidate.resolve()
        checked.append(str(resolved))
        if _valid_upstream_root(resolved):
            return resolved, checked
    return None, checked


@app.command()
def main(
    nifti_path: Path = typer.Argument(..., exists=True, dir_okay=False),
    output_dir: Path | None = typer.Option(
        None, "--output-dir", "-o", help="dir for produced masks"
    ),
    modality: str = typer.Option("CT_BODY", "--modality", help="CT_BODY | MRI_BODY | MRI_BRAIN"),
    label_prompts: str | None = typer.Option(
        None,
        "--label-prompts",
        help="Optional comma-separated upstream class IDs, e.g. '3,14'.",
    ),
    device: str = typer.Option("auto", "--device", help="Recorded device hint: auto | cuda | cpu"),
    upstream_root: Path | None = typer.Option(
        None,
        "--upstream-root",
        help="Path to NV-Segment-CTMR/NV-Segment-CTMR; defaults to $NV_SEGMENT_CTMR_ROOT.",
    ),
    timeout_seconds: float = typer.Option(float("3600.0"), "--timeout-seconds"),
    ground_truth: Path | None = typer.Option(
        None,
        "--ground-truth",
        exists=True,
        dir_okay=False,
        help=(
            "Optional reference label map. Recorded under input.ground_truth_path "
            "for downstream verifiers. The skill does not compute GT metrics."
        ),
    ),
) -> None:
    """Run NV-Segment-CTMR on a CT or MRI NIfTI volume."""
    if modality not in SUPPORTED_MODALITIES:
        raise typer.BadParameter(f"--modality must be one of {SUPPORTED_MODALITIES}")

    env_root = os.environ.get("NV_SEGMENT_CTMR_ROOT", "").strip()
    resolved_root, checked_roots = _resolve_upstream_root(upstream_root, env_root)
    if resolved_root is None and upstream_root is None and not env_root:
        emit(
            _error_payload(
                "NV_SEGMENT_CTMR_ROOT is unset",
                "Clone https://github.com/NVIDIA-Medtech/NV-Segment-CTMR and export "
                "NV_SEGMENT_CTMR_ROOT to the nested NV-Segment-CTMR directory, or place the clone at "
                ".workbench_data/upstreams/NV-Segment-CTMR/NV-Segment-CTMR.",
            )
            | {"checked_roots": checked_roots}
        )
        raise typer.Exit(2)
    if resolved_root is None:
        emit(
            _error_payload(
                "NV_SEGMENT_CTMR_ROOT layout invalid",
                "configs/inference.json not found in any checked root",
            )
            | {"checked_roots": checked_roots}
        )
        raise typer.Exit(2)
    config_file = resolved_root / "configs" / "inference.json"

    nifti_path = nifti_path.expanduser().resolve()
    if output_dir is None:
        output_dir = nifti_path.parent / f"{_strip_nifti_suffix(nifti_path)}_nv_segment_ctmr_out"
    output_dir = output_dir.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    parsed_label_prompts = _parse_label_prompts(label_prompts)
    label_map, label_map_source = _load_label_map(resolved_root)
    inventory = _model_inventory(resolved_root, label_map_source)
    resolved_device = _resolve_device(device)

    input_img = nib.load(str(nifti_path))
    input_summary = _input_summary(input_img)
    output_summary = _empty_output_summary(
        input_summary,
        parsed_label_prompts,
        label_map,
        label_map_source,
    )

    cmd = _build_command(nifti_path, output_dir, modality, parsed_label_prompts)
    run_env = _child_process_env()
    run_env.setdefault("MONAI_DATA_DIRECTORY", str(output_dir / "_monai_data"))
    run_env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128,expandable_segments:True")

    run_started = time.time()
    t0 = time.monotonic()
    try:
        proc = subprocess.run(
            cmd,
            cwd=str(resolved_root),
            env=run_env,
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
            check=False,
        )
        rc = proc.returncode
        stdout = proc.stdout
        stderr = proc.stderr
    except subprocess.TimeoutExpired as e:
        rc = int("124")
        stdout = e.stdout.decode() if isinstance(e.stdout, bytes) else (e.stdout or "")
        stderr_raw = e.stderr.decode() if isinstance(e.stderr, bytes) else (e.stderr or "")
        stderr = stderr_raw + f"\n[TIMEOUT after {timeout_seconds}s]"
    elapsed = time.monotonic() - t0

    mask_path = _find_output_mask(output_dir, nifti_path, run_started)
    if mask_path is not None:
        output_summary = _mask_summary(
            mask_path,
            input_img,
            parsed_label_prompts,
            label_map,
            label_map_source,
        )
        output_summary["path"] = str(mask_path)

    payload: dict[str, Any] = {
        "skill": SKILL_NAME,
        "model": "NVIDIA-Medtech/NV-Segment-CTMR (VISTA3D CT/MRI)",
        "model_repo": MODEL_REPO,
        "license": "Wrapper Apache-2.0; upstream model and repository licenses apply.",
        "input": {
            "path": str(nifti_path),
            **input_summary,
            "modality": modality,
            "ground_truth_path": str(ground_truth) if ground_truth is not None else None,
        },
        "output": output_summary,
        "invocation": {
            "official_entrypoint": "python -m monai.bundle run",
            "upstream_root": str(resolved_root),
            "upstream_commit": git_commit(resolved_root),
            "config_file": str(config_file),
            "output_dir": str(output_dir),
            "modality": modality,
            "label_prompts": parsed_label_prompts,
            "command": cmd,
            "exit_code": rc,
            "model_inventory": inventory,
        },
        "runtime": {
            "subprocess_seconds": round(elapsed, int("3")),
            "device": resolved_device,
        },
        "logs": {
            "stdout_tail": tail(stdout),
            "stderr_tail": tail(stderr),
        },
        "intended_use_disclaimer": (
            "Engineering verification only. Output is NOT clinically meaningful. "
            "This wrapper invokes the upstream MONAI bundle entry point from the "
            "NVIDIA-Medtech/NV-Segment-CTMR README; it does not modify inference, "
            "preprocessing, or postprocessing."
        ),
    }
    emit(payload)
    raise typer.Exit(0)


if __name__ == "__main__":
    app()