pipelines/base.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal, Optional

ModelKey = Literal["ct", "mr", "mr_brain"]


@dataclass
class GenerationRequest:
    model: ModelKey
    output_size: tuple[int, int, int]
    spacing: tuple[float, float, float]
    seed: int = 0
    num_steps: int = 30
    cfg_guidance_scale: float = 0.0

    # CT-only
    body_region: Optional[list[str]] = None
    anatomy_list: Optional[list[str]] = None
    generate_masks: bool = True

    # MR-only (modality int from configs/modality_mapping.json: 8/9/10/11/...)
    modality_class: Optional[int] = None

    # MR-Brain-only (mapped to a modality int)
    contrast: Optional[str] = None  # "T1" / "T2" / "FLAIR" / "SWI" / "T1_skull_stripped" / ...


@dataclass
class GenerationResult:
    volume_path: str
    mask_path: Optional[str] = None
    used_anatomy_labels: dict[int, str] = field(default_factory=dict)
    runtime_seconds: float = 0.0
    seed: int = 0
    modality: Optional[int] = None