scripts/generate_region_text.py

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd

from pet_vlm_dataset import load_suvr_vector


def readable_region_name(name: str) -> str:
    side = ""
    stem = name
    if name.endswith("_L"):
        side = "left "
        stem = name[:-2]
    elif name.endswith("_R"):
        side = "right "
        stem = name[:-2]
    stem = stem.replace("_", " ").lower()
    return f"{side}{stem}".strip()


def make_region_text(labels: list[str], values: list[float], top_k: int) -> tuple[str, list[str], list[str]]:
    pairs = sorted(zip(labels, values), key=lambda item: item[1])
    low = pairs[:top_k]
    high = pairs[-top_k:][::-1]

    low_text = ", ".join(readable_region_name(name) for name, _ in low)
    high_text = ", ".join(readable_region_name(name) for name, _ in high)
    text = (
        "FDG-PET regional metabolic summary. "
        f"Relatively low metabolism is observed in {low_text}. "
        f"Relatively high uptake is observed in {high_text}."
    )
    return text, [name for name, _ in low], [name for name, _ in high]


def generate(manifest_path: Path, out_path: Path, top_k: int) -> None:
    manifest = pd.read_csv(manifest_path)
    rows = []
    for row in manifest.to_dict(orient="records"):
        labels, suvr = load_suvr_vector(row["suvr_csv_path"], include_background=False)
        text, low_regions, high_regions = make_region_text(labels, suvr.tolist(), top_k)
        enriched = dict(row)
        enriched["region_text"] = text
        enriched["low_regions"] = "|".join(low_regions)
        enriched["high_regions"] = "|".join(high_regions)
        enriched["top_k"] = top_k
        rows.append(enriched)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    pd.DataFrame(rows).to_csv(out_path, index=False)
    print(f"wrote {len(rows)} rows to {out_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate controlled region-text summaries from SUVR profiles.")
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--top-k", type=int, default=5)
    args = parser.parse_args()
    generate(args.manifest, args.out, args.top_k)


if __name__ == "__main__":
    main()