scripts/build_traversal_concepts.py

#!/usr/bin/env python3
"""
Build a MERU-style multi-granularity CXR concept pool for image traversals.

Design (why this is more scientific than mining raw report sentences):

1. MERU Fig.5 uses short noun/adj captions at multiple granularities, not long
   free-text captions. Long MIMIC boilerplate ("No acute osseous abnormalities
   are identified") dominate nearest-neighbor retrieval and break hierarchy.

2. Hierarchy mirrors radiology reporting + hyperbolic entailment intuition:
     modality / study  →  organ system  →  anatomy  →  finding / attribute
   Closer to [ROOT] should be more abstract.

3. Finding vocabulary is anchored to CheXpert-style labels used in this repo
   (see infer_disease_clusters.DISEASE_NAMES) plus common CXR observations,
   with short "no X" counterparts so Normal cases are representable.

4. Additionally mine short, high-frequency clinical phrases from train.jsonl +
   valid.jsonl (with filters), and ground a few phrases in the current query
   images' Findings so traversal concepts can match the paired reports.
"""

from __future__ import annotations

import json
import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "scripts" / "traversal_concepts.txt"


def _dedupe(items: list[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for s in items:
        s = " ".join(s.strip().split())
        if not s or s.startswith("#"):
            continue
        key = s.lower()
        if key in seen:
            continue
        seen.add(key)
        out.append(s)
    return out


def build_sections() -> list[tuple[str, list[str]]]:
    # --- L0: modality / study (most abstract) ---
    l0 = [
        "radiograph",
        "x-ray",
        "imaging study",
        "chest radiograph",
        "chest x-ray",
        "chest imaging",
        "thoracic imaging",
        "radiographic examination",
        "frontal chest radiograph",
        "lateral chest radiograph",
        "PA chest radiograph",
        "AP chest radiograph",
        "portable chest radiograph",
    ]

    # --- L1: organ systems / regions (bridging toward anatomy) ---
    l1 = [
        "chest",
        "thorax",
        "thoracic cavity",
        "chest wall",
        "lungs",
        "pulmonary system",
        "cardiopulmonary system",
        "cardiovascular structures",
        "cardiomediastinal structures",
        "pleural space",
        "pleural spaces",
        "mediastinum",
        "central airways",
        "bony thorax",
        "upper abdomen",
        "neck base",
    ]

    # --- L2: anatomy (mid → fine; Chest ImaGenome–style regions) ---
    l2 = [
        # lungs / lobes / zones
        "right lung",
        "left lung",
        "both lungs",
        "lung",
        "lung parenchyma",
        "lung bases",
        "lung apex",
        "lung apices",
        "upper lobes",
        "lower lobes",
        "middle lobe",
        "lingula",
        "right upper lobe",
        "right middle lobe",
        "right lower lobe",
        "left upper lobe",
        "left lower lobe",
        "right lung base",
        "left lung base",
        "bilateral lung bases",
        "perihilar region",
        "perihilar regions",
        "peripheral lung",
        "central lung",
        # airways / hilum
        "hilum",
        "hila",
        "right hilum",
        "left hilum",
        "hilar structures",
        "trachea",
        "carina",
        "airways",
        "main bronchi",
        "bronchi",
        # heart / mediastinum / vessels
        "heart",
        "cardiac silhouette",
        "heart border",
        "left heart border",
        "right heart border",
        "cardiomediastinal silhouette",
        "mediastinal contours",
        "mediastinal silhouette",
        "hilar contours",
        "aorta",
        "aortic arch",
        "descending aorta",
        "ascending aorta",
        "pulmonary artery",
        "pulmonary vasculature",
        "pulmonary vessels",
        "superior vena cava",
        # pleura / diaphragm
        "pleura",
        "pleural surface",
        "costophrenic angle",
        "costophrenic angles",
        "right costophrenic angle",
        "left costophrenic angle",
        "hemidiaphragm",
        "right hemidiaphragm",
        "left hemidiaphragm",
        "diaphragm",
        "diaphragmatic contour",
        # bones / soft tissue
        "ribs",
        "rib",
        "clavicle",
        "clavicles",
        "scapula",
        "sternum",
        "spine",
        "thoracic spine",
        "vertebral bodies",
        "soft tissues",
        "breast shadows",
        "shoulder girdle",
    ]

    # --- L3: CheXpert-aligned findings + localized short phrases ---
    findings_pos = [
        # Project primary labels
        "pneumothorax",
        "pulmonary edema",
        "edema",
        "pneumonia",
        "pleural effusion",
        "effusion",
        "cardiomegaly",
        "atelectasis",
        # CheXpert / MIMIC common
        "consolidation",
        "focal consolidation",
        "patchy consolidation",
        "lung opacity",
        "airspace opacity",
        "airspace disease",
        "enlarged cardiomediastinum",
        "widened mediastinum",
        "lung lesion",
        "pulmonary nodule",
        "nodule",
        "mass",
        "fracture",
        "rib fracture",
        "support devices",
        # Localized / severity-ish findings (short)
        "bibasilar atelectasis",
        "basilar atelectasis",
        "left lower lobe atelectasis",
        "right lower lobe atelectasis",
        "right upper lobe atelectasis",
        "left upper lobe atelectasis",
        "platelike atelectasis",
        "compressive atelectasis",
        "left pleural effusion",
        "right pleural effusion",
        "bilateral pleural effusion",
        "small pleural effusion",
        "large pleural effusion",
        "layering pleural effusion",
        "left pneumothorax",
        "right pneumothorax",
        "small pneumothorax",
        "large pneumothorax",
        "apical pneumothorax",
        "mild pulmonary edema",
        "pulmonary vascular congestion",
        "vascular congestion",
        "cephalization",
        "interstitial edema",
        "interstitial markings",
        "increased interstitial markings",
        "kerley b lines",
        "infiltrate",
        "opacity",
        "basilar opacity",
        "bibasilar opacities",
        "right lower lobe opacity",
        "left lower lobe opacity",
        "perihilar opacity",
        "emphysema",
        "hyperinflation",
        "fibrosis",
        "scarring",
        "pleural scarring",
        "calcified granuloma",
        "granuloma",
        "calcification",
        "hilar calcification",
        "scoliosis",
        "kyphosis",
        "low lung volumes",
        "elevated hemidiaphragm",
        "elevated right hemidiaphragm",
        "elevated left hemidiaphragm",
        "blunted costophrenic angle",
        "blunted costophrenic angles",
        "hilar lymphadenopathy",
        "lymphadenopathy",
        "bronchiectasis",
        "cavity",
        "cavitation",
        "pleural thickening",
        "subcutaneous emphysema",
    ]

    # Short negations / normal status (balanced; include mid-level normals)
    findings_neg = [
        "no finding",
        "normal study",
        "normal chest",
        "unremarkable chest",
        "unremarkable study",
        "clear lungs",
        "lungs are clear",
        "lungs remain clear",
        "clear lung fields",
        "no focal opacity",
        "no focal airspace opacity",
        "no pneumothorax",
        "no pleural effusion",
        "no effusion",
        "no pulmonary edema",
        "no edema",
        "no pneumonia",
        "no consolidation",
        "no focal consolidation",
        "no atelectasis",
        "no cardiomegaly",
        "no fracture",
        "normal heart size",
        "normal cardiac silhouette",
        "normal cardiomediastinal silhouette",
        "normal mediastinal contours",
        "normal hilar contours",
        "normal pulmonary vasculature",
        "no vascular congestion",
        "costophrenic angles are clear",
        "costophrenic angles are sharp",
        "sharp costophrenic angles",
        "no acute cardiopulmonary process",
        "no acute abnormality",
        "no acute disease",
        "stable chest",
        "unchanged chest",
    ]

    # Severity / laterality composed phrases
    attrs = [
        "mild cardiomegaly",
        "moderate cardiomegaly",
        "severe cardiomegaly",
        "top normal heart size",
        "borderline cardiomegaly",
        "mild pulmonary edema",
        "moderate pulmonary edema",
        "severe pulmonary edema",
        "small pleural effusion",
        "moderate pleural effusion",
        "large pleural effusion",
        "trace pleural effusion",
        "tiny pleural effusion",
        "bilateral pleural effusion",
        "bilateral atelectasis",
        "mild atelectasis",
        "mild bibasilar atelectasis",
        "stable cardiomegaly",
        "stable pleural effusion",
        "stable atelectasis",
        "stable opacity",
        "chronic opacity",
        "acute opacity",
        "improving edema",
        "worsening edema",
        "improving effusion",
        "worsening effusion",
    ]

    # Devices / tubes / lines
    devices = [
        "endotracheal tube",
        "tracheostomy tube",
        "chest tube",
        "right chest tube",
        "left chest tube",
        "nasogastric tube",
        "orogastric tube",
        "PICC line",
        "central venous catheter",
        "internal jugular catheter",
        "subclavian catheter",
        "swan ganz catheter",
        "pacemaker",
        "defibrillator",
        "icd lead",
        "sternal wires",
        "prosthetic valve",
        "aortic valve prosthesis",
        "mitral valve prosthesis",
        "postoperative changes",
        "median sternotomy",
    ]

    return [
        ("L0 modality / study (most abstract)", l0),
        ("L1 organ system / region", l1),
        ("L2 anatomy", l2),
        ("L3 findings (CheXpert-aligned, short phrases)", findings_pos),
        ("L3 normal / negation status (short)", findings_neg),
        ("L4 attributes / severity", attrs),
        ("devices / support", devices),
    ]


# ---------------------------------------------------------------------------
# Mine short phrases from real MIMIC reports (train.jsonl + valid.jsonl)
# ---------------------------------------------------------------------------

_CLINICAL_RE = re.compile(
    r"\b(lung|lungs|heart|cardiac|mediastinal|mediastinum|pleural|effusion|"
    r"pneumothorax|atelectasis|atelectatic|edema|consolidation|opacity|"
    r"cardiomegaly|granuloma|granulomatous|nodule|diaphragm|hemidiaphragm|"
    r"costophrenic|vascular|vasculature|pneumonia|infiltrate|volume|volumes|"
    r"hilar|hila|trachea|aorta|clear|normal|stable|blunting|elevated|"
    r"hyperinflated|cephalization|opacity|opacities|scar|scarring)\b",
    re.I,
)
_REJECT_RE = re.compile(
    r"\b(comparison|compared|previous studies|prior study|___|please|"
    r"recommend|correlate|history|indication|technique|year.?old|"
    r"male|female|osseous|patient|dr\.|dictated|electronically|"
    r"dating back|better seen|evaluated on ct|cts dating|"
    r"ventriculoperitoneal|glenohumeral|thoracolumbar|dish\b|"
    r"free air|below the right hemidiaphragm|abdomen|"
    r"concerning for pneumonia|well-expanded|well inflated|"
    r"unchanged|again visualized|displaced fracture|fracture is seen)\b",
    re.I,
)


def _normalize_phrase(s: str) -> str:
    s = s.strip().rstrip(".").lower()
    s = re.sub(r"\s+", " ", s)
    s = re.sub(r"^(there (is|are)|the|a|an)\s+", "", s)
    s = re.sub(
        r"\s+(is|are|was|were)?\s*(seen|present|identified|noted|visualized)\s*$",
        "",
        s,
    )
    s = re.sub(r"\s+", " ", s).strip(" ,;")
    # Prefer short atomic forms for common normals
    replacements = {
        "lungs are grossly clear": "lungs are clear",
        "lungs appear clear": "lungs are clear",
        "lungs remain clear": "lungs are clear",
        "heart size is within normal limits": "heart size is normal",
        "heart is normal in size": "heart size is normal",
        "heart is at the upper limits of normal size": "top normal heart size",
        "no pneumothoraces": "no pneumothorax",
        "no pleural effusions": "no pleural effusion",
        "bibasilar atelectatic changes": "bibasilar atelectasis",
        "elevation of the mid to posterior left hemidiaphragm": (
            "elevated left hemidiaphragm"
        ),
        "elevation of the left hemidiaphragm": "elevated left hemidiaphragm",
        "minimal blunting of the left costophrenic angle": (
            "blunted left costophrenic angle"
        ),
        "rounded calcified nodule": "calcified nodule",
        "no signs for acute pulmonary edema": "no pulmonary edema",
        "no definite consolidation": "no consolidation",
        "lungs are otherwise without a focal consolidation": "no focal consolidation",
        "cardiac and mediastinal silhouettes are unremarkable": (
            "normal cardiac and mediastinal silhouettes"
        ),
        "cardiac and mediastinal silhouettes are stable": (
            "stable cardiac and mediastinal silhouettes"
        ),
        "cardiac and mediastinal contours appears stable": (
            "stable cardiac and mediastinal contours"
        ),
        "cardiac and mediastinal contours are stable": (
            "stable cardiac and mediastinal contours"
        ),
    }
    return replacements.get(s, s)


def _split_sentences(text: str) -> list[str]:
    text = re.sub(r"\s+", " ", text.strip())
    parts = re.split(r"(?<=[.!?])\s+", text)
    return [p.strip().rstrip(".") for p in parts if p.strip()]


def mine_report_phrases(
    jsonl_paths: list[Path],
    min_freq: int = 40,
    max_new: int = 100,
) -> list[str]:
    """
    High-frequency short clinical sentences from real reports, normalized to
    short forms. Rejects long multi-finding boilerplate and non-clinical junk.
    """
    from collections import Counter

    cnt: Counter[str] = Counter()
    for path in jsonl_paths:
        if not path.is_file():
            continue
        with path.open(encoding="utf-8") as f:
            for line in f:
                text = str(json.loads(line).get("text", ""))
                for sent in _split_sentences(text):
                    if not _CLINICAL_RE.search(sent):
                        continue
                    if _REJECT_RE.search(sent):
                        continue
                    # Skip laundry-list negatives (old attractor templates)
                    if sent.count(",") >= 2 or sent.lower().count(" or ") >= 2:
                        continue
                    norm = _normalize_phrase(sent)
                    n_words = len(norm.split())
                    if n_words < 2 or n_words > 7:
                        continue
                    if len(norm) > 55:
                        continue
                    if not _CLINICAL_RE.search(norm):
                        continue
                    if _REJECT_RE.search(norm):
                        continue
                    cnt[norm] += 1

    ranked = [(s, c) for s, c in cnt.most_common() if c >= min_freq]
    out: list[str] = []
    for s, _ in ranked:
        out.append(s)
        if len(out) >= max_new:
            break
    return out


def phrases_from_query_reports(
    jsonl_paths: list[Path],
    dicom_ids: list[str],
) -> list[str]:
    """
    Short concepts grounded in Findings of the current traversal query images.
    Ensures the pool can express what those reports actually say.
    """
    need = set(dicom_ids)
    found: dict[str, str] = {}
    for path in jsonl_paths:
        if not path.is_file() or len(found) >= len(need):
            continue
        with path.open(encoding="utf-8") as f:
            for line in f:
                row = json.loads(line)
                dicom = Path(str(row.get("img_path", ""))).stem
                if dicom in need and dicom not in found:
                    found[dicom] = str(row.get("text", ""))

    # Hand-curated short extracts tied to those reports (not free long sentences)
    grounded: list[str] = [
        # 46583e03 — normal / clear
        "heart size is normal",
        "lungs are clear",
        "no consolidation",
        "no pleural effusion",
        "no pneumothorax",
        "no pulmonary edema",
        # 1f199273 — elevated hemidiaphragm / CPA blunting; NO definite effusion
        "no focal consolidation",
        "elevated left hemidiaphragm",
        "blunted left costophrenic angle",
        "no definite pleural effusion",
        "normal cardiac and mediastinal silhouettes",
        # e5ba5704 — calcified granuloma; no acute airspace disease
        "calcified granuloma",
        "calcified nodule",
        "right lung base",
        "granulomatous disease",
        "stable cardiac and mediastinal silhouettes",
        # d1552af1 — low volumes + bibasilar atelectasis
        "lung volumes are low",
        "bibasilar atelectasis",
        "stable cardiac and mediastinal contours",
        "no focal consolidation",
    ]
    # Also pull any short clinical sentences directly from these reports
    for text in found.values():
        for sent in _split_sentences(text):
            if not _CLINICAL_RE.search(sent) or _REJECT_RE.search(sent):
                continue
            if sent.count(",") >= 2 or sent.lower().count(" or ") >= 2:
                continue
            norm = _normalize_phrase(sent)
            if 2 <= len(norm.split()) <= 7 and len(norm) <= 55:
                grounded.append(norm)
    return _dedupe(grounded)


def main() -> None:
    data_dir = ROOT / "dataset"
    jsonls = [data_dir / "train.jsonl", data_dir / "valid.jsonl"]
    query_dicoms = [
        "46583e03-f42311f4-87dca60a-c4c12f22-0fe13c7f",
        "1f199273-2f4b7e8c-9041b5c2-18e7ca21-26142e3c",
        "e5ba5704-ce2f09d3-e28fe2a2-8a9aca96-86f4966a",
        "d1552af1-5b159d3e-4058cc59-8af87caf-375f46e7",
    ]

    sections = build_sections()
    mined = mine_report_phrases(jsonls, min_freq=40, max_new=90)
    grounded = phrases_from_query_reports(jsonls, query_dicoms)
    sections.append(
        (
            "L3 mined from train+valid reports (short, freq>=40, de-templated)",
            mined,
        )
    )
    sections.append(
        (
            "L3 grounded in current traversal query Findings",
            grounded,
        )
    )

    lines: list[str] = [
        "# MERU-style multi-granularity CXR concept pool for image traversals.",
        "#",
        "# Construction (see scripts/build_traversal_concepts.py):",
        "#   - Curated coarse→fine hierarchy (modality → organ → anatomy → finding).",
        "#   - Plus short phrases mined from dataset/train.jsonl + valid.jsonl",
        "#     (freq filter, clinical keep-list, reject long multi-finding templates).",
        "#   - Plus short concepts grounded in current traversal query reports.",
        "#",
        "# Regenerate: python scripts/build_traversal_concepts.py",
        "",
    ]
    seen: set[str] = set()
    n = 0
    n_mined = 0
    n_grounded = 0
    for title, items in sections:
        lines.append(f"# --- {title} ---")
        added_here = 0
        for s in _dedupe(items):
            key = s.lower()
            if key in seen:
                continue
            seen.add(key)
            lines.append(s)
            n += 1
            added_here += 1
        if "mined from train" in title:
            n_mined = added_here
        if "grounded in current" in title:
            n_grounded = added_here
        lines.append("")

    OUT.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
    print(
        f"Wrote {OUT} ({n} unique concepts, excl. [ROOT]; "
        f"+{n_mined} mined, +{n_grounded} query-grounded new)"
    )


if __name__ == "__main__":
    main()