run_three_models_parallel.py

#!/usr/bin/env python3
"""Run the three RECISTto3D model backends concurrently from Python.

Primary use:

    from run_three_models_parallel import load_all_models, infer_with_loaded_models

    models = load_all_models(device="cuda:0")
    results = infer_with_loaded_models(
        models=models,
        image="image.nii.gz",
        recist="recist.nii.gz",
        intensity="window",
        window="-175,275",
    )

The implementation imports and calls ``recist_infer.py`` functions directly in
three Python worker processes. It does not shell out to ``recist_infer.py``.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import multiprocessing as mp
import os
import time
from argparse import Namespace
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable, Sequence


ROOT = Path(__file__).resolve().parent
MODELS = ("eff-medsam2", "medsam2", "nninteractive")


@dataclass(frozen=True)
class ModelResult:
    model: str
    output_nifti: str
    duration_s: float
    metadata: dict


@dataclass
class LoadedModel:
    model: str
    handle: object
    device: str
    metadata: dict


@dataclass
class LoadedModels:
    eff_medsam2: LoadedModel
    medsam2: LoadedModel
    nninteractive: LoadedModel
    metadata: dict[str, dict]


@dataclass(frozen=True)
class _ModelConfig:
    model: str
    image: str
    output_nifti: str
    recist: str | None
    recist_lines: tuple[str, ...]
    image_key: str | None
    recist_key: str | None
    recist_space: str
    spacing: str | None
    shift: int
    seed: int
    intensity: str
    window: str | None
    device: str | None
    checkpoint: str | None
    nninteractive_model_dir: str | None
    nninteractive_model_name: str
    nninteractive_fold: str | None
    nninteractive_checkpoint: str
    nninteractive_prompt: str
    nninteractive_compile: bool
    no_autozoom: bool
    torch_threads: int | None
    verbose: bool


# --- Self-contained RECIST-from-GT extraction -------------------------------
# Turns a GT label mask into the app's RECIST text format, one line per lesion:
#   "z,x1,y1,x2,y2,label"  (axial voxel indices, label == connected-component id)
# Same canonical algorithm as the rest of the pipeline (cc3d-26 instances, drop
# tiny lesions, key slice = max axial area, longest external-contour pair), but
# intentionally inlined here so this helper depends ONLY on third-party libs and
# never imports another module in this repo.

_MIN_LESION_VOXELS = 10        # drop lesions smaller than this (benchmark convention)
_RECIST_MAX_CONTOUR_PTS = 500  # subsample contour to at most this many points before pdist


def recist_lines_from_gt_mask(
    gt_mask: "str | Path | object",
    *,
    min_voxels: int = _MIN_LESION_VOXELS,
    expected_shape: "tuple[int, int, int] | None" = None,
) -> list[str]:
    """Extract one RECIST diameter line per GT lesion as app-ready text.

    Returns a list of ``"z,x1,y1,x2,y2,label"`` strings (axial voxel indices,
    ``label`` == connected-component instance id), a drop-in for the app's
    ``recist-line-box`` textbox and ``_validate_recist_lines``. Returns ``[]``
    when no lesion qualifies. Fully self-contained: depends only on numpy /
    SimpleITK / scipy / (cv2 or skimage), never on another module in this repo.

    ``gt_mask`` may be a NIfTI/.npy/.npz path or a 3D ``(z, y, x)`` ndarray.
    ``expected_shape`` optionally guards that the mask matches the loaded image.
    """
    import numpy as np

    gt = _as_mask_array(gt_mask)
    if expected_shape is not None and gt.shape != tuple(expected_shape):
        raise ValueError(f"GT mask shape {gt.shape} != image shape {tuple(expected_shape)}")

    instance = _connected_components_26(gt > 0)
    lines: list[str] = []
    for lid in (int(v) for v in np.unique(instance) if v != 0):
        lesion = instance == lid
        if int(lesion.sum()) < min_voxels:
            continue
        z = int(np.argmax(lesion.sum(axis=(1, 2))))  # key slice = max axial area
        endpoints = _longest_diameter_xy(lesion[z].astype(np.uint8))
        if endpoints is None:
            continue
        (x1, y1), (x2, y2) = endpoints
        lines.append(f"{z},{int(x1)},{int(y1)},{int(x2)},{int(y2)},{lid}")
    return lines


def _as_mask_array(gt_mask: "str | Path | object"):
    """Load a GT mask into a 3D ``(z, y, x)`` integer ndarray (self-contained)."""
    import numpy as np

    if isinstance(gt_mask, np.ndarray):
        arr = gt_mask
    else:
        p = str(gt_mask)
        lower = p.lower()
        if lower.endswith((".nii", ".nii.gz")):
            import SimpleITK as sitk

            arr = sitk.GetArrayFromImage(sitk.ReadImage(p))  # -> (z, y, x)
        elif lower.endswith(".npz"):
            with np.load(p, allow_pickle=True) as data:
                key = next((k for k in ("gts", "mask", "recist", "arr_0") if k in data), None)
                if key is None:
                    raise ValueError(f"{p} has no gts/mask/recist/arr_0 array")
                arr = data[key]
        elif lower.endswith(".npy"):
            arr = np.load(p, allow_pickle=True)
        else:
            raise ValueError(f"Unsupported GT mask format: {p}")
    arr = np.asarray(arr)
    if arr.ndim != 3:
        raise ValueError(f"GT mask must be 3D (z, y, x); got shape {arr.shape}")
    return arr


def _connected_components_26(binary):
    """26-connectivity instance labeling; cc3d if available, else scipy."""
    import numpy as np

    try:
        import cc3d

        return cc3d.connected_components(binary.astype(np.uint8), connectivity=26)
    except ImportError:
        from scipy import ndimage

        structure = np.ones((3, 3, 3), dtype=int)  # full 26-connectivity
        instance, _ = ndimage.label(binary.astype(np.uint8), structure=structure)
        return instance


def _longest_diameter_xy(mask_2d):
    """Farthest external-contour point pair as ``((x1, y1), (x2, y2))`` or None."""
    import numpy as np
    from scipy.spatial.distance import pdist, squareform

    try:
        import cv2

        contours, _ = cv2.findContours(mask_2d, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if not contours:
            return None
        pts = np.vstack(contours).squeeze()  # (x, y)
    except ImportError:
        from skimage import measure

        contours_yx = measure.find_contours(mask_2d, 0.5)
        if not contours_yx:
            return None
        yx = np.concatenate(contours_yx, axis=0)
        pts = np.stack([yx[:, 1], yx[:, 0]], axis=1)  # (row, col) -> (x, y)

    if pts.ndim != 2 or len(pts) < 2:
        return None
    if len(pts) > _RECIST_MAX_CONTOUR_PTS:
        pts = pts[np.linspace(0, len(pts) - 1, _RECIST_MAX_CONTOUR_PTS, dtype=int)]
    dist_matrix = squareform(pdist(pts))
    i, j = np.unravel_index(np.argmax(dist_matrix), dist_matrix.shape)
    return np.rint(pts[i]).astype(int), np.rint(pts[j]).astype(int)


def _strip_known_suffixes(path: Path) -> str:
    name = path.name
    for suffix in (".nii.gz", ".npz", ".npy", ".png", ".jpg", ".jpeg", ".tif", ".tiff"):
        if name.lower().endswith(suffix):
            return name[: -len(suffix)]
    return path.stem


def _as_path_string(value: str | Path | None) -> str | None:
    if value is None:
        return None
    return str(Path(value).expanduser())


def _resolve_output_dir(output_dir: str | Path) -> Path:
    path = Path(output_dir).expanduser()
    if path.is_absolute():
        return path
    return (Path.cwd() / path).resolve()


def _model_device(
    model: str,
    common_device: str | None,
    eff_medsam2_device: str | None,
    medsam2_device: str | None,
    nninteractive_device: str | None,
) -> str | None:
    if model == "eff-medsam2":
        return eff_medsam2_device or common_device
    if model == "medsam2":
        return medsam2_device or common_device
    if model == "nninteractive":
        return nninteractive_device or common_device
    raise ValueError(f"unknown model: {model}")


def _model_checkpoint(
    model: str,
    eff_medsam2_checkpoint: str | Path | None,
    medsam2_checkpoint: str | Path | None,
) -> str | None:
    if model == "eff-medsam2":
        return _as_path_string(eff_medsam2_checkpoint)
    if model == "medsam2":
        return _as_path_string(medsam2_checkpoint)
    if model == "nninteractive":
        return None
    raise ValueError(f"unknown model: {model}")


def _namespace_from_config(config: _ModelConfig) -> Namespace:
    return Namespace(
        image=config.image,
        model=config.model,
        recist=config.recist,
        recist_line=list(config.recist_lines),
        image_key=config.image_key,
        recist_key=config.recist_key,
        recist_space=config.recist_space,
        spacing=config.spacing,
        output_nifti=config.output_nifti,
        shift=config.shift,
        seed=config.seed,
        device=config.device,
        checkpoint=config.checkpoint,
        intensity=config.intensity,
        window=config.window,
        nninteractive_model_dir=config.nninteractive_model_dir,
        nninteractive_model_name=config.nninteractive_model_name,
        nninteractive_fold=config.nninteractive_fold,
        nninteractive_checkpoint=config.nninteractive_checkpoint,
        nninteractive_prompt=config.nninteractive_prompt,
        nninteractive_compile=config.nninteractive_compile,
        no_autozoom=config.no_autozoom,
        torch_threads=config.torch_threads,
        verbose=config.verbose,
    )


def _normalize_device(device: str | None) -> str | None:
    if device is None:
        return None
    if device == "gpu":
        return "cuda:0"
    return device


def count_parameters(handle) -> int | None:
    """Total parameter count of a loaded model handle, or None if not countable.

    The SAM-based predictors are themselves nn.Modules; the nnInteractive handle
    is a session object that holds its net on `.network`. Measured rather than
    hardcoded so the number tracks whatever checkpoint was actually loaded.
    """
    import torch

    if isinstance(handle, torch.nn.Module):
        return sum(p.numel() for p in handle.parameters())
    network = getattr(handle, "network", None)
    if isinstance(network, torch.nn.Module):
        return sum(p.numel() for p in network.parameters())
    return None


def format_param_count(n: int | None) -> str:
    """Render a parameter count as e.g. "34.1M"; empty string when unknown.

    Falls back to "K" below 100k so small nets don't all render as "0.0M".
    """
    if not n:
        return ""
    if n < 100_000:
        return f"{n / 1e3:.1f}K"
    return f"{n / 1e6:.1f}M"


def _model_loader_args(
    *,
    model: str,
    device: str | None,
    checkpoint: str | Path | None = None,
    nninteractive_model_dir: str | Path | None = None,
    nninteractive_model_name: str = "nnInteractive_v1.0",
    nninteractive_fold: str | None = None,
    nninteractive_checkpoint: str = "checkpoint_final.pth",
    nninteractive_compile: bool = False,
    no_autozoom: bool = False,
    torch_threads: int | None = None,
    verbose: bool = False,
) -> Namespace:
    return Namespace(
        model=model,
        device=_normalize_device(device),
        checkpoint=_as_path_string(checkpoint),
        nninteractive_model_dir=_as_path_string(nninteractive_model_dir),
        nninteractive_model_name=nninteractive_model_name,
        nninteractive_fold=nninteractive_fold,
        nninteractive_checkpoint=nninteractive_checkpoint,
        nninteractive_compile=nninteractive_compile,
        no_autozoom=no_autozoom,
        torch_threads=torch_threads,
        verbose=verbose,
    )


def _load_medsam2_model(args: Namespace) -> LoadedModel:
    from recist_infer import MEDSAM2_ROOT, load_medsam2_predictor, pushd

    with pushd(MEDSAM2_ROOT):
        predictor, checkpoint, model_name, device = load_medsam2_predictor(args)

    metadata = {
        "model": args.model,
        "model_name": model_name,
        "checkpoint": checkpoint,
        "device": device,
    }
    return LoadedModel(model=args.model, handle=predictor, device=str(device), metadata=metadata)


def _load_nninteractive_model(args: Namespace) -> LoadedModel:
    import torch
    from recist_infer import NNINTERACTIVE_ROOT, add_to_syspath, resolve_nninteractive_model_dir

    add_to_syspath(NNINTERACTIVE_ROOT)
    from nnInteractive.inference.inference_session import nnInteractiveInferenceSession

    model_dir = resolve_nninteractive_model_dir(args)
    device = args.device or "cuda:0"
    session = nnInteractiveInferenceSession(
        device=torch.device(device),
        use_torch_compile=args.nninteractive_compile,
        verbose=args.verbose,
        torch_n_threads=args.torch_threads or os.cpu_count(),
        do_autozoom=not args.no_autozoom,
    )
    session.initialize_from_trained_model_folder(
        model_dir,
        use_fold=args.nninteractive_fold,
        checkpoint_name=args.nninteractive_checkpoint,
    )

    metadata = {
        "model": "nninteractive",
        "model_dir": model_dir,
        "device": device,
        "nninteractive_checkpoint": args.nninteractive_checkpoint,
        "nninteractive_fold": args.nninteractive_fold,
    }
    return LoadedModel(model="nninteractive", handle=session, device=str(device), metadata=metadata)


def load_all_models(
    *,
    device: str | None = "cuda:0",
    eff_medsam2_device: str | None = None,
    medsam2_device: str | None = None,
    nninteractive_device: str | None = None,
    eff_medsam2_checkpoint: str | Path | None = None,
    medsam2_checkpoint: str | Path | None = None,
    nninteractive_model_dir: str | Path | None = None,
    nninteractive_model_name: str = "nnInteractive_v1.0",
    nninteractive_fold: str | None = None,
    nninteractive_checkpoint: str = "checkpoint_final.pth",
    nninteractive_compile: bool = False,
    no_autozoom: bool = False,
    torch_threads: int | None = None,
    verbose: bool = False,
) -> LoadedModels:
    """Load eff-medsam2, medsam2, and nninteractive into this Python process.

    Pass ``device="cpu"`` to keep all models on CPU, ``device="cuda:0"`` to
    place them on one GPU, or use the per-model device arguments to split models
    across devices. ``device="gpu"`` is accepted as an alias for ``"cuda:0"``.
    """

    eff_args = _model_loader_args(
        model="eff-medsam2",
        device=eff_medsam2_device or device,
        checkpoint=eff_medsam2_checkpoint,
    )
    medsam_args = _model_loader_args(
        model="medsam2",
        device=medsam2_device or device,
        checkpoint=medsam2_checkpoint,
    )
    nninteractive_args = _model_loader_args(
        model="nninteractive",
        device=nninteractive_device or device,
        nninteractive_model_dir=nninteractive_model_dir,
        nninteractive_model_name=nninteractive_model_name,
        nninteractive_fold=nninteractive_fold,
        nninteractive_checkpoint=nninteractive_checkpoint,
        nninteractive_compile=nninteractive_compile,
        no_autozoom=no_autozoom,
        torch_threads=torch_threads,
        verbose=verbose,
    )

    def _timed_load(loader, loader_args, key: str) -> LoadedModel:
        start = time.time()
        loaded = loader(loader_args)
        load_s = time.time() - start
        loaded.metadata["load_duration_s"] = load_s
        n_params = count_parameters(loaded.handle)
        loaded.metadata["n_parameters"] = n_params
        size_note = f", {format_param_count(n_params)} params" if n_params else ""
        print(f"[load] {key}: {load_s:.2f}s{size_note}", flush=True)
        return loaded

    eff_medsam2 = _timed_load(_load_medsam2_model, eff_args, "eff_medsam2")
    medsam2 = _timed_load(_load_medsam2_model, medsam_args, "medsam2")
    nninteractive = _timed_load(_load_nninteractive_model, nninteractive_args, "nninteractive")
    total_load_duration_s = (
        eff_medsam2.metadata["load_duration_s"]
        + medsam2.metadata["load_duration_s"]
        + nninteractive.metadata["load_duration_s"]
    )
    metadata = {
        "eff_medsam2": eff_medsam2.metadata,
        "medsam2": medsam2.metadata,
        "nninteractive": nninteractive.metadata,
        "total_load_duration_s": total_load_duration_s,
    }
    return LoadedModels(
        eff_medsam2=eff_medsam2,
        medsam2=medsam2,
        nninteractive=nninteractive,
        metadata=metadata,
    )



# Gradio insertion point:
# After ``infer_with_loaded_models(...)`` returns three model-specific NIfTI
# masks, call ``make_bitmask_labels_image(...)`` to create one display mask and
# return ``_file_url(combined_mask_path)`` through app.py's existing
# ``mask_url_state`` output. The current frontend can keep calling
# ``window.recistTo3DViewer.loadMask(maskUrl)``, because ``loadMask`` already
# uses NiiVue's ``nv.loadDrawingFromUrl(maskUrl, false)``.
#
# The combined NIfTI is a single 3D label image with bit-encoded labels:
#   0 = background
#   1 = eff-medsam2 only
#   2 = medsam2 only
#   4 = nninteractive only
#   3 = eff-medsam2 + medsam2 overlap
#   5 = eff-medsam2 + nninteractive overlap
#   6 = medsam2 + nninteractive overlap
#   7 = all three overlap
#
# In app.py, add a Gradio CheckboxGroup for model visibility and a JavaScript
# method that toggles NiiVue colormap alpha for labels whose bits match the
# selected masks. This preserves overlap information while still loading just
# one NIfTI into the existing viewer.
def make_bitmask_labels_image(
    eff_medsam2_mask: str | Path,
    medsam2_mask: str | Path,
    nninteractive_mask: str | Path,
    output_nifti: str | Path,
    *,
    reference_nifti: str | Path | None = None,
    require_same_geometry: bool = True,
) -> str:
    """Write one bit-encoded display NIfTI from three binary/model masks."""

    import numpy as np
    import SimpleITK as sitk

    mask_paths = {
        "eff-medsam2": Path(eff_medsam2_mask),
        "medsam2": Path(medsam2_mask),
        "nninteractive": Path(nninteractive_mask),
    }
    images = {name: sitk.ReadImage(str(path)) for name, path in mask_paths.items()}
    reference_image = sitk.ReadImage(str(reference_nifti)) if reference_nifti is not None else images["eff-medsam2"]

    def geometry_tuple(image) -> tuple:
        return (image.GetSize(), image.GetSpacing(), image.GetOrigin(), image.GetDirection())

    if require_same_geometry:
        reference_geometry = geometry_tuple(reference_image)
        mismatches = [
            name
            for name, image in images.items()
            if geometry_tuple(image) != reference_geometry
        ]
        if mismatches:
            raise ValueError(
                "All masks must share the same NIfTI geometry as the reference image. "
                f"Mismatched masks: {mismatches}"
            )

    arrays = {name: sitk.GetArrayFromImage(image) > 0 for name, image in images.items()}
    shape = arrays["eff-medsam2"].shape
    shape_mismatches = [name for name, array in arrays.items() if array.shape != shape]
    if shape_mismatches:
        raise ValueError(f"All masks must have the same array shape. Mismatched masks: {shape_mismatches}")

    combined = np.zeros(shape, dtype=np.uint8)
    combined[arrays["eff-medsam2"]] |= 1
    combined[arrays["medsam2"]] |= 2
    combined[arrays["nninteractive"]] |= 4

    output_path = Path(output_nifti)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_image = sitk.GetImageFromArray(combined)
    output_image.CopyInformation(reference_image)
    sitk.WriteImage(output_image, str(output_path))
    return str(output_path)




def _loaded_model_for_name(models: LoadedModels, model: str) -> LoadedModel:
    if model == "eff-medsam2":
        return models.eff_medsam2
    if model == "medsam2":
        return models.medsam2
    if model == "nninteractive":
        return models.nninteractive
    raise ValueError(f"unknown model: {model}")


def _infer_medsam2_with_loaded_model(image, recist, spacing, args: Namespace, loaded_model: LoadedModel):
    import numpy as np
    from recist_infer import InferenceResult, run_medsam2_loop

    rng = np.random.RandomState(args.seed)
    predictor = loaded_model.handle
    segs, boxes_array, labels = run_medsam2_loop(
        image, recist, spacing, args, predictor, loaded_model.device, rng=rng
    )

    metadata = dict(loaded_model.metadata)
    metadata.update(
        {
            "model": args.model,
            "prompt": "recist-box",
            "preprocessing": "medsam2_uint8_512_rgb_imagenet_norm",
            "intensity": args.intensity,
            "window": args.window,
            "spacing_zyx": spacing.tolist(),
            "labels": labels,
            "loaded_model_reused": True,
        }
    )
    return InferenceResult(segs=segs, boxes_xyzxyz=boxes_array, metadata=metadata)


def _infer_nninteractive_with_loaded_model(image, recist, spacing, args: Namespace, loaded_model: LoadedModel):
    import math as _math
    import numpy as np
    from recist_infer import InferenceResult, prompt_specs_from_recist

    rng = np.random.RandomState(args.seed)
    prompt_specs = prompt_specs_from_recist(recist, spacing, args, rng, target=args.nninteractive_prompt)
    if not prompt_specs:
        raise ValueError("No prompts were provided")

    session = loaded_model.handle
    if hasattr(session, "reset_interactions"):
        session.reset_interactions()

    raw_image = np.asarray(image)
    session.set_image(raw_image[None])
    target = np.zeros(raw_image.shape, dtype=np.uint8)
    session.set_target_buffer(target)

    segs = np.zeros(raw_image.shape, dtype=np.uint16)
    boxes = []
    labels = [prompt.label for prompt in prompt_specs]

    for i, prompt in enumerate(prompt_specs):
        if i > 0:
            session.reset_interactions()

        prediction_center = None
        prediction_zoom_out_factor = None

        if prompt.kind == "box":
            x1, y1, x2, y2 = prompt.box_xyxy.astype(int)
            spacing_z = float(spacing[0])
            spacing_y = float(spacing[1])
            spacing_x = float(spacing[2])
            diameter_mm = _math.hypot((int(y2) - int(y1)) * spacing_y, (int(x2) - int(x1)) * spacing_x)
            half_dz = max(10, int(round(diameter_mm / max(spacing_z, 0.1) / 2)))
            z_lo = max(0, int(prompt.z) - half_dz)
            z_hi = min(raw_image.shape[0], int(prompt.z) + half_dz + 1)
            bbox_dhw = [[z_lo, z_hi], [int(y1), int(y2) + 1], [int(x1), int(x2) + 1]]
            boxes.append([x1, y1, prompt.z, x2, y2, prompt.z])
            session.add_bbox_interaction(bbox_dhw, include_interaction=True, run_prediction=False)
            prediction_center = session.new_interaction_centers[-1]
            prediction_zoom_out_factor = session.new_interaction_zoom_out_factors[-1]
        elif prompt.kind == "points":
            if prompt.points_xy is not None and len(prompt.points_xy):
                cx, cy = np.mean(prompt.points_xy, axis=0)
                prediction_center = (int(prompt.z), int(round(float(cy))), int(round(float(cx))))
            for x, y in prompt.points_xy:
                session.add_point_interaction(
                    (prompt.z, int(round(y)), int(round(x))),
                    include_interaction=True,
                    run_prediction=False,
                )
            if prompt.negative_points_xy is not None:
                for x, y in prompt.negative_points_xy:
                    session.add_point_interaction(
                        (prompt.z, int(round(y)), int(round(x))),
                        include_interaction=False,
                        run_prediction=False,
                    )
        else:
            raise ValueError(f"Unsupported prompt kind: {prompt.kind}")

        if prediction_center is not None:
            session.new_interaction_centers = [prediction_center]
        else:
            session.new_interaction_centers = [session.new_interaction_centers[-1]]
        if prediction_zoom_out_factor is not None:
            session.new_interaction_zoom_out_factors = [prediction_zoom_out_factor]
        else:
            session.new_interaction_zoom_out_factors = [session.new_interaction_zoom_out_factors[-1]]
        session._predict()
        segs[target > 0] = prompt.label

    boxes_array = np.asarray(boxes, dtype=np.float32).reshape((-1, 6)) if boxes else np.zeros((0, 6), dtype=np.float32)
    metadata = dict(loaded_model.metadata)
    metadata.update(
        {
            "model": args.model,
            "prompt": f"recist-{args.nninteractive_prompt.replace('_', '-')}",
            "preprocessing": "nninteractive_internal_raw_set_image",
            "spacing_zyx": spacing.tolist(),
            "labels": labels,
            "loaded_model_reused": True,
        }
    )
    return InferenceResult(segs=segs, boxes_xyzxyz=boxes_array, metadata=metadata)


def _run_one_loaded_model(config: _ModelConfig, loaded_model: LoadedModel, loaded_image, recist, recist_source: str) -> ModelResult:
    from recist_infer import write_nifti

    args = _namespace_from_config(config)
    output_nifti_path = Path(args.output_nifti)

    start = time.time()
    if args.model in {"eff-medsam2", "medsam2"}:
        result = _infer_medsam2_with_loaded_model(loaded_image.array, recist, loaded_image.spacing, args, loaded_model)
    elif args.model == "nninteractive":
        result = _infer_nninteractive_with_loaded_model(loaded_image.array, recist, loaded_image.spacing, args, loaded_model)
    else:
        raise ValueError(f"unknown model: {args.model}")

    duration_s = time.time() - start
    metadata = dict(result.metadata)
    metadata.update(
        {
            "image_path": str(Path(args.image).resolve()),
            "image_key": loaded_image.source_key,
            "recist_source": recist_source,
            "image_shape_dhw": list(loaded_image.array.shape),
            "duration_s": duration_s,
            "output_format": "nifti",
        }
    )
    result.metadata = metadata
    write_nifti(output_nifti_path, result.segs, loaded_image)
    return ModelResult(
        model=args.model,
        output_nifti=str(output_nifti_path),
        duration_s=duration_s,
        metadata=metadata,
    )


def infer_with_loaded_models(
    *,
    models: LoadedModels,
    image: str | Path,
    recist: str | Path | None = None,
    recist_lines: Sequence[str] = (),
    output_dir: str | Path = "outputs/three_models",
    output_prefix: str | None = None,
    image_key: str | None = None,
    recist_key: str | None = None,
    recist_space: str = "strict",
    spacing: str | None = None,
    shift: int = 0,
    seed: int = 2024,
    intensity: str = "preserve",
    window: str | None = None,
    nninteractive_prompt: str = "5_points",
    verbose: bool = False,
    run_concurrent: bool = False,
) -> list[ModelResult]:
    """Run inference with models already returned by ``load_all_models``.

    Loaded model objects stay in this process. The default is sequential
    execution because repeated threaded calls can retain per-thread runtime
    caches and raise process RSS in Gradio-style long-running services.
    """

    configs = _build_configs(
        image=image,
        recist=recist,
        recist_lines=recist_lines,
        output_dir=output_dir,
        output_prefix=output_prefix,
        image_key=image_key,
        recist_key=recist_key,
        recist_space=recist_space,
        spacing=spacing,
        shift=shift,
        seed=seed,
        intensity=intensity,
        window=window,
        device=None,
        eff_medsam2_device=None,
        medsam2_device=None,
        nninteractive_device=None,
        eff_medsam2_checkpoint=None,
        medsam2_checkpoint=None,
        nninteractive_model_dir=None,
        nninteractive_model_name="nnInteractive_v1.0",
        nninteractive_fold=None,
        nninteractive_checkpoint="checkpoint_final.pth",
        nninteractive_prompt=nninteractive_prompt,
        nninteractive_compile=False,
        no_autozoom=False,
        torch_threads=None,
        verbose=verbose,
    )

    from recist_infer import load_image, load_recist

    image_path = Path(_as_path_string(image) or "")
    loaded_image = load_image(image_path, image_key, spacing)
    recist_array, recist_source = load_recist(
        image_path=image_path,
        recist_path=Path(recist) if recist is not None else None,
        recist_key=recist_key,
        recist_lines=list(recist_lines),
        shape=loaded_image.array.shape,
        reference_sitk_image=loaded_image.sitk_image,
        recist_space=recist_space,
    )

    if not run_concurrent:
        return [
            _run_one_loaded_model(config, _loaded_model_for_name(models, config.model), loaded_image, recist_array, recist_source)
            for config in configs
        ]

    results_by_model: dict[str, ModelResult] = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(configs)) as executor:
        future_to_model = {
            executor.submit(
                _run_one_loaded_model,
                config,
                _loaded_model_for_name(models, config.model),
                loaded_image,
                recist_array,
                recist_source,
            ): config.model
            for config in configs
        }
        for future in concurrent.futures.as_completed(future_to_model):
            model = future_to_model[future]
            results_by_model[model] = future.result()

    return [results_by_model[model] for model in MODELS]




def _run_one_model(config: _ModelConfig) -> ModelResult:
    from recist_infer import infer_medsam2, infer_nninteractive, load_image, load_recist, write_nifti

    args = _namespace_from_config(config)
    image_path = Path(args.image)
    output_nifti_path = Path(args.output_nifti)

    start = time.time()
    loaded = load_image(image_path, args.image_key, args.spacing)
    recist, recist_source = load_recist(
        image_path=image_path,
        recist_path=Path(args.recist) if args.recist else None,
        recist_key=args.recist_key,
        recist_lines=args.recist_line,
        shape=loaded.array.shape,
        reference_sitk_image=loaded.sitk_image,
        recist_space=args.recist_space,
    )

    if args.model in {"eff-medsam2", "medsam2"}:
        result = infer_medsam2(loaded.array, recist, loaded.spacing, args)
    elif args.model == "nninteractive":
        result = infer_nninteractive(loaded.array, recist, loaded.spacing, args)
    else:
        raise ValueError(f"unknown model: {args.model}")

    duration_s = time.time() - start
    metadata = dict(result.metadata)
    metadata.update(
        {
            "image_path": str(Path(args.image).resolve()),
            "image_key": loaded.source_key,
            "recist_source": recist_source,
            "image_shape_dhw": list(loaded.array.shape),
            "duration_s": duration_s,
            "output_format": "nifti",
        }
    )
    result.metadata = metadata
    write_nifti(output_nifti_path, result.segs, loaded)
    return ModelResult(
        model=args.model,
        output_nifti=str(output_nifti_path),
        duration_s=duration_s,
        metadata=metadata,
    )


def _build_configs(
    *,
    image: str | Path,
    recist: str | Path | None,
    recist_lines: Sequence[str],
    output_dir: str | Path,
    output_prefix: str | None,
    image_key: str | None,
    recist_key: str | None,
    recist_space: str,
    spacing: str | None,
    shift: int,
    seed: int,
    intensity: str,
    window: str | None,
    device: str | None,
    eff_medsam2_device: str | None,
    medsam2_device: str | None,
    nninteractive_device: str | None,
    eff_medsam2_checkpoint: str | Path | None,
    medsam2_checkpoint: str | Path | None,
    nninteractive_model_dir: str | Path | None,
    nninteractive_model_name: str,
    nninteractive_fold: str | None,
    nninteractive_checkpoint: str,
    nninteractive_prompt: str,
    nninteractive_compile: bool,
    no_autozoom: bool,
    torch_threads: int | None,
    verbose: bool,
) -> list[_ModelConfig]:
    output_root = _resolve_output_dir(output_dir)
    output_root.mkdir(parents=True, exist_ok=True)

    image_str = _as_path_string(image)
    if image_str is None:
        raise ValueError("image is required")
    recist_str = _as_path_string(recist)
    prefix = output_prefix or _strip_known_suffixes(Path(image_str))

    configs = []
    for model in MODELS:
        safe_model = model.replace("-", "_")
        output_nii = output_root / f"{prefix}_{safe_model}.nii.gz"
        configs.append(
            _ModelConfig(
                model=model,
                image=image_str,
                output_nifti=str(output_nii),
                recist=recist_str,
                recist_lines=tuple(recist_lines),
                image_key=image_key,
                recist_key=recist_key,
                recist_space=recist_space,
                spacing=spacing,
                shift=shift,
                seed=seed,
                intensity=intensity,
                window=window,
                device=_model_device(model, device, eff_medsam2_device, medsam2_device, nninteractive_device),
                checkpoint=_model_checkpoint(model, eff_medsam2_checkpoint, medsam2_checkpoint),
                nninteractive_model_dir=_as_path_string(nninteractive_model_dir),
                nninteractive_model_name=nninteractive_model_name,
                nninteractive_fold=nninteractive_fold,
                nninteractive_checkpoint=nninteractive_checkpoint,
                nninteractive_prompt=nninteractive_prompt,
                nninteractive_compile=nninteractive_compile,
                no_autozoom=no_autozoom,
                torch_threads=torch_threads,
                verbose=verbose,
            )
        )
    return configs


def run_three_models(
    *,
    image: str | Path,
    loaded_models: LoadedModels | None = None,
    recist: str | Path | None = None,
    recist_lines: Sequence[str] = (),
    output_dir: str | Path = "outputs/three_models",
    output_prefix: str | None = None,
    image_key: str | None = None,
    recist_key: str | None = None,
    recist_space: str = "strict",
    spacing: str | None = None,
    shift: int = 0,
    seed: int = 2024,
    intensity: str = "preserve",
    window: str | None = None,
    device: str | None = None,
    eff_medsam2_device: str | None = None,
    medsam2_device: str | None = None,
    nninteractive_device: str | None = None,
    eff_medsam2_checkpoint: str | Path | None = None,
    medsam2_checkpoint: str | Path | None = None,
    nninteractive_model_dir: str | Path | None = None,
    nninteractive_model_name: str = "nnInteractive_v1.0",
    nninteractive_fold: str | None = None,
    nninteractive_checkpoint: str = "checkpoint_final.pth",
    nninteractive_prompt: str = "5_points",
    nninteractive_compile: bool = False,
    no_autozoom: bool = False,
    torch_threads: int | None = None,
    verbose: bool = False,
    start_method: str | None = None,
) -> list[ModelResult]:
    """Run eff-medsam2, medsam2, and nninteractive concurrently.

    If ``loaded_models`` is provided, the already-loaded model handles are reused
    in this process. Otherwise, this function keeps the older process-based path
    where each worker loads its own model.
    """

    if loaded_models is not None:
        return infer_with_loaded_models(
            models=loaded_models,
            image=image,
            recist=recist,
            recist_lines=recist_lines,
            output_dir=output_dir,
            output_prefix=output_prefix,
                image_key=image_key,
            recist_key=recist_key,
            recist_space=recist_space,
            spacing=spacing,
            shift=shift,
            seed=seed,
            intensity=intensity,
            window=window,
            nninteractive_prompt=nninteractive_prompt,
            verbose=verbose,
        )

    configs = _build_configs(
        image=image,
        recist=recist,
        recist_lines=recist_lines,
        output_dir=output_dir,
        output_prefix=output_prefix,
        image_key=image_key,
        recist_key=recist_key,
        recist_space=recist_space,
        spacing=spacing,
        shift=shift,
        seed=seed,
        intensity=intensity,
        window=window,
        device=device,
        eff_medsam2_device=eff_medsam2_device,
        medsam2_device=medsam2_device,
        nninteractive_device=nninteractive_device,
        eff_medsam2_checkpoint=eff_medsam2_checkpoint,
        medsam2_checkpoint=medsam2_checkpoint,
        nninteractive_model_dir=nninteractive_model_dir,
        nninteractive_model_name=nninteractive_model_name,
        nninteractive_fold=nninteractive_fold,
        nninteractive_checkpoint=nninteractive_checkpoint,
        nninteractive_prompt=nninteractive_prompt,
        nninteractive_compile=nninteractive_compile,
        no_autozoom=no_autozoom,
        torch_threads=torch_threads,
        verbose=verbose,
    )

    context = mp.get_context(start_method) if start_method is not None else None
    executor_kwargs = {"max_workers": len(configs)}
    if context is not None:
        executor_kwargs["mp_context"] = context

    results_by_model: dict[str, ModelResult] = {}
    with concurrent.futures.ProcessPoolExecutor(**executor_kwargs) as executor:
        future_to_model = {executor.submit(_run_one_model, config): config.model for config in configs}
        for future in concurrent.futures.as_completed(future_to_model):
            model = future_to_model[future]
            results_by_model[model] = future.result()

    return [results_by_model[model] for model in MODELS]


def results_to_dicts(results: Iterable[ModelResult]) -> list[dict]:
    return [asdict(result) for result in results]


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Run eff-medsam2, medsam2, and nninteractive concurrently via Python function calls.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--image", required=True)
    parser.add_argument("--recist")
    parser.add_argument("--recist-line", action="append", default=[])
    parser.add_argument("--output-dir", default="outputs/three_models")
    parser.add_argument("--output-prefix")
    parser.add_argument("--image-key")
    parser.add_argument("--recist-key")
    parser.add_argument("--recist-space", choices=("strict", "index"), default="strict")
    parser.add_argument("--spacing")
    parser.add_argument("--shift", type=int, default=0)
    parser.add_argument("--seed", type=int, default=2024)
    parser.add_argument("--intensity", choices=("preserve", "minmax", "window"), default="preserve")
    parser.add_argument("--window")
    parser.add_argument("--device")
    parser.add_argument("--eff-medsam2-device")
    parser.add_argument("--medsam2-device")
    parser.add_argument("--nninteractive-device")
    parser.add_argument("--eff-medsam2-checkpoint")
    parser.add_argument("--medsam2-checkpoint")
    parser.add_argument("--nninteractive-model-dir")
    parser.add_argument("--nninteractive-model-name", default="nnInteractive_v1.0")
    parser.add_argument("--nninteractive-fold")
    parser.add_argument("--nninteractive-checkpoint", default="checkpoint_final.pth")
    parser.add_argument(
        "--nninteractive-prompt",
        choices=("3_points", "5_points", "5pos_4neg", "5pos_6neg"),
        default="5_points",
    )
    parser.add_argument("--nninteractive-compile", action="store_true")
    parser.add_argument("--no-autozoom", action="store_true")
    parser.add_argument("--torch-threads", type=int)
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--start-method", choices=tuple(mp.get_all_start_methods()))
    return parser


def main() -> int:
    args = build_parser().parse_args()
    results = run_three_models(
        image=args.image,
        recist=args.recist,
        recist_lines=args.recist_line,
        output_dir=args.output_dir,
        output_prefix=args.output_prefix,
        image_key=args.image_key,
        recist_key=args.recist_key,
        recist_space=args.recist_space,
        spacing=args.spacing,
        shift=args.shift,
        seed=args.seed,
        intensity=args.intensity,
        window=args.window,
        device=args.device,
        eff_medsam2_device=args.eff_medsam2_device,
        medsam2_device=args.medsam2_device,
        nninteractive_device=args.nninteractive_device,
        eff_medsam2_checkpoint=args.eff_medsam2_checkpoint,
        medsam2_checkpoint=args.medsam2_checkpoint,
        nninteractive_model_dir=args.nninteractive_model_dir,
        nninteractive_model_name=args.nninteractive_model_name,
        nninteractive_fold=args.nninteractive_fold,
        nninteractive_checkpoint=args.nninteractive_checkpoint,
        nninteractive_prompt=args.nninteractive_prompt,
        nninteractive_compile=args.nninteractive_compile,
        no_autozoom=args.no_autozoom,
        torch_threads=args.torch_threads,
        verbose=args.verbose,
        start_method=args.start_method,
    )
    print(json.dumps(results_to_dicts(results), indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())