pipelines/mr_brain.py

from __future__ import annotations

import time

from .base import GenerationRequest, GenerationResult
from . import _runner

CONTRAST_TO_MODALITY: dict[str, int] = {
    "T1": 9,
    "T2": 10,
    "FLAIR": 11,
    "SWI": 20,
    "T1_skull_stripped": 29,
    "T2_skull_stripped": 30,
    "FLAIR_skull_stripped": 31,
    "SWI_skull_stripped": 32,
}


def generate(req: GenerationRequest) -> GenerationResult:
    """MR Brain (rflow-mr-brain) — image-only multi-contrast brain MRI."""
    contrast = req.contrast or "T1"
    if contrast not in CONTRAST_TO_MODALITY:
        raise ValueError(f"Unsupported contrast: {contrast}. Choose from {list(CONTRAST_TO_MODALITY)}")
    modality = CONTRAST_TO_MODALITY[contrast]

    t0 = time.time()
    path = _runner.run_image_only(
        version="rflow-mr-brain",
        output_size=req.output_size,
        spacing=req.spacing,
        modality=modality,
        seed=req.seed,
        num_inference_steps=req.num_steps,
        cfg_guidance_scale=req.cfg_guidance_scale or 10.0,
    )
    return GenerationResult(
        volume_path=str(path),
        runtime_seconds=time.time() - t0,
        seed=req.seed,
        modality=modality,
    )