diff --git a/.gitignore b/.gitignore index d1fe96f..622672f 100644 --- a/.gitignore +++ b/.gitignore @@ -41,4 +41,6 @@ backend/*.pyc backend/application_default_credentials.json dataset_cxr_primary.json -ingest_cxr.py \ No newline at end of file +ingest_cxr.py + +venv/ \ No newline at end of file diff --git a/data_pipeline/build_cxr_dataset.py b/data_pipeline/build_cxr_dataset.py new file mode 100644 index 0000000..198b9e5 --- /dev/null +++ b/data_pipeline/build_cxr_dataset.py @@ -0,0 +1,200 @@ +""" +Build a chest X-ray dataset from the whole_multicare_dataset. + +Inclusion rule: any PMCID that has at least one chest X-ray image. +For those PMCIDs, ALL images (CXR, CT, pathology, etc.) are included. +Each image has an "is_chest_xray" boolean flag. + +Output structure: + cxr_full_dataset/ + dataset.json <- array of JSON objects, one per PMCID + images/ <- copies of ALL images for qualifying PMCIDs +""" + +import os +import json +import shutil +import pandas as pd +from pathlib import Path + +# ── Paths ────────────────────────────────────────────────────────────────── +BASE = Path("medical_datasets/whole_multicare_dataset") +OUTPUT_DIR = Path("cxr_full_dataset") +IMAGES_DIR = OUTPUT_DIR / "images" +DATASET_JSON = OUTPUT_DIR / "dataset.json" + +os.makedirs(IMAGES_DIR, exist_ok=True) +print(f"Output folder: {OUTPUT_DIR.resolve()}") + +# ── Load all source files ────────────────────────────────────────────────── +print("Loading source files...") +captions_df = pd.read_csv(BASE / "captions_and_labels.csv") +metadata_df = pd.read_parquet(BASE / "metadata.parquet") +abstracts_df = pd.read_parquet(BASE / "abstracts.parquet") +cases_df = pd.read_parquet(BASE / "cases.parquet") +ci_df = pd.read_parquet(BASE / "case_images.parquet") + +# ── Filter for chest X-rays only ────────────────────────────────────────── +cxr_mask = ( + (captions_df["image_subtype"] == "x_ray") & + (captions_df["image_type"] == "radiology") & + (captions_df["caption"].str.contains(r"chest\s*x-?ray|cxr", case=False, na=False)) +) +cxr_df = captions_df[cxr_mask].copy() +print(f"Chest X-ray image rows found: {len(cxr_df)}") + +pmc_ids = sorted(cxr_df["patient_id"].str.extract(r"(PMC\d+)")[0].unique()) +print(f"Unique PMCIDs with chest X-rays: {len(pmc_ids)}") + +# ── Build image file index (filename -> full path) ───────────────────────── +# Use the naming convention from data_dictionary: PMC1/PMC10/filename +def image_path_from_filename(fname: str) -> Path: + return BASE / fname[:4] / fname[:5] / fname + +# ── Build per-image text_references lookup from case_images.parquet ──────── +# Flatten: image_id -> text_references list +text_ref_lookup: dict[str, list[str]] = {} +for _, ci_row in ci_df.iterrows(): + for img in ci_row["case_images"]: + image_id = img.get("image_id", "") + refs = img.get("text_references", []) + text_ref_lookup[image_id] = list(refs) if refs is not None else [] + +# ── Helpers to safely convert numpy/pandas types to plain Python ─────────── +def to_list(val): + if val is None: + return [] + try: + return list(val) + except TypeError: + return [] + +def to_python(val): + """Recursively convert numpy scalars / NaN to plain Python types.""" + import numpy as np + if isinstance(val, float) and pd.isna(val): + return None + if isinstance(val, (np.integer,)): + return int(val) + if isinstance(val, (np.floating,)): + return float(val) + if isinstance(val, (np.ndarray,)): + return [to_python(v) for v in val] + if isinstance(val, list): + return [to_python(v) for v in val] + if isinstance(val, dict): + return {k: to_python(v) for k, v in val.items()} + return val + +# ── Main loop ───────────────────────────────────────────────────────────── +dataset = [] +copied_images = 0 +missing_images = 0 + +for pmc_id in pmc_ids: + + # -- Metadata -- + meta_row = metadata_df[metadata_df["article_id"] == pmc_id] + if len(meta_row): + m = meta_row.iloc[0]["article_metadata"] + meta = { + "title": m.get("title", ""), + "authors": to_list(m.get("authors")), + "journal": m.get("journal", ""), + "journal_detail": m.get("journal_detail", ""), + "year": m.get("year", ""), + "doi": m.get("doi", ""), + "pmid": m.get("pmid", ""), + "license": m.get("license", ""), + "keywords": to_list(m.get("keywords")), + "mesh_terms": to_list(m.get("mesh_terms")), + "major_mesh_terms":to_list(m.get("major_mesh_terms")), + "link": m.get("link", ""), + "case_amount": int(m.get("case_amount", 0)), + } + else: + meta = {} + + # -- Abstract -- + abs_row = abstracts_df[abstracts_df["article_id"] == pmc_id] + abstract = abs_row.iloc[0]["abstract"] if len(abs_row) else "" + + # -- Cases -- + cases_row = cases_df[cases_df["article_id"] == pmc_id] + cases_out = [] + if len(cases_row): + for case in cases_row.iloc[0]["cases"]: + cases_out.append({ + "case_id": case.get("case_id", ""), + "age": to_python(case.get("age")), + "gender": case.get("gender", ""), + "case_text": case.get("case_text", ""), + }) + + # -- All images for this PMCID (CXR cases include all modalities) -- + # Build a set of CXR file_ids for flagging + cxr_file_ids = set(cxr_df[cxr_df["patient_id"].str.startswith(pmc_id)]["file_id"].values) + all_pmc_images = captions_df[captions_df["patient_id"].str.startswith(pmc_id)] + images_out = [] + + for _, img_row in all_pmc_images.iterrows(): + filename = img_row["file"] + src_path = image_path_from_filename(filename) + dst_fname = filename + dst_path = IMAGES_DIR / dst_fname + local_rel = f"images/{dst_fname}" + + # Copy image + if src_path.exists(): + if not dst_path.exists(): + shutil.copy2(src_path, dst_path) + copied_images += 1 + else: + missing_images += 1 + + # text_references via main_image id + main_image_id = img_row.get("main_image", "") + text_refs = text_ref_lookup.get(main_image_id, []) + + images_out.append({ + "file_id": img_row.get("file_id", ""), + "file": filename, + "main_image": main_image_id, + "image_component": img_row.get("image_component", ""), + "patient_id": img_row.get("patient_id", ""), + "license": img_row.get("license", ""), + "file_size": int(img_row.get("file_size", 0)), + "image_type": img_row.get("image_type", ""), + "image_subtype": img_row.get("image_subtype", ""), + "is_chest_xray": img_row.get("file_id", "") in cxr_file_ids, + "caption": img_row.get("caption", ""), + "case_substring": img_row.get("case_substring", ""), + "radiology_region": to_python(img_row.get("radiology_region")), + "radiology_region_granular": to_python(img_row.get("radiology_region_granular")), + "radiology_view": to_python(img_row.get("radiology_view")), + "ml_labels": img_row.get("ml_labels_for_supervised_classification", ""), + "gt_labels": img_row.get("gt_labels_for_semisupervised_classification", ""), + "text_references": text_refs, + "local_image_path": local_rel, + }) + + # -- Assemble PMCID record -- + record = { + "pmc_id": pmc_id, + **meta, + "abstract": abstract, + "cases": cases_out, + "images": images_out, + } + dataset.append(record) + +# ── Write JSON ───────────────────────────────────────────────────────────── +print(f"\nWriting dataset.json ({len(dataset)} PMCIDs)...") +with open(DATASET_JSON, "w") as f: + json.dump(dataset, f, indent=2, default=str) + +print(f"\nDone.") +print(f" PMCIDs: {len(dataset)}") +print(f" Images copied: {copied_images}") +print(f" Images missing: {missing_images}") +print(f" Output: {OUTPUT_DIR.resolve()}") diff --git a/data_pipeline/build_schema_dataset.py b/data_pipeline/build_schema_dataset.py new file mode 100644 index 0000000..0b70236 --- /dev/null +++ b/data_pipeline/build_schema_dataset.py @@ -0,0 +1,312 @@ +""" +Extract structured fields from multicare CXR dataset using Gemini API +and produce a final dataset.json conforming to the target schema. + +Fields covered: + - Directly mapped : source, study, images, patient age/sex, abstract→summary + - Gemini-extracted : chief_complaint, symptom_duration, comorbidities, + medications, immunocompromised, findings_structured, + radiology_labels, outcomes, clinical_text structured note + - Removed (N/A) : dicom, facility, routing, embeddings, audit.models_used, + pregnancy, smoking_status, family_history, portable, + laterality_marker_present, transfer_memo +""" + +import os, json, uuid, time, re +from pathlib import Path +from dotenv import load_dotenv +from google import genai +from tqdm import tqdm + +# ── Config ───────────────────────────────────────────────────────────────── +load_dotenv(dotenv_path=Path("/blue/gtyson.fsu/tp22o.fsu/medgemma/.env"), override=True) +client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) +MODEL_NAME = "gemini-2.0-flash" + +INPUT_JSON = Path("cxr_full_dataset/dataset.json") +OUTPUT_DIR = Path("cxr_schema_dataset") +OUTPUT_JSON = OUTPUT_DIR / "dataset.json" +CACHE_DIR = OUTPUT_DIR / ".cache" # per-record cache to allow resume + +OUTPUT_DIR.mkdir(exist_ok=True) +CACHE_DIR.mkdir(exist_ok=True) + +# ── Gemini extraction prompt ─────────────────────────────────────────────── +EXTRACT_PROMPT = """ +You are a medical NLP assistant. Given a clinical case report text and abstract, extract the following fields and return ONLY valid JSON — no markdown, no explanation. + +Return exactly this structure (use null for unknown/missing values, use the exact enum strings shown): + +{{ + "chief_complaint": "", + "symptom_duration": "", + "comorbidities": [""], + "medications_used": "", + "immunocompromised": "unknown | yes | no", + "clinical_note_hpi": "<1-3 sentence history of present illness or null>", + "clinical_note_pmh": "", + "clinical_note_meds": "", + "clinical_note_allergies": "", + "primary_suspected": [""], + "differential": [""], + "infectious_concern": "unknown | yes | no", + "icu_candidate": "unknown | yes | no", + "lungs_consolidation_present": "unknown | yes | no", + "lungs_consolidation_location": [""], + "lungs_consolidation_extent": "mild | moderate | severe | unknown", + "lungs_atelectasis_present": "unknown | yes | no", + "lungs_atelectasis_location": [], + "lungs_edema_present": "unknown | yes | no", + "lungs_edema_pattern": "interstitial | alveolar | mixed | unknown", + "pleura_effusion_present": "unknown | yes | no", + "pleura_effusion_side": "left | right | bilateral | unknown", + "pleura_effusion_size": "small | moderate | large | unknown", + "pleura_pneumothorax_present": "unknown | yes | no", + "pleura_pneumothorax_side": "left | right | bilateral | unknown", + "cardiomegaly": "unknown | yes | no", + "mediastinal_widening": "unknown | yes | no", + "lines_tubes_present": "unknown | yes | no", + "device_list": [""], + "summary_1_2_lines": "<1-2 sentence case summary>", + "bullets": ["", "", ""], + "red_flags": [""], + "uncertainties": [""], + "urgency": "routine | urgent | emergent", + "outcome_success": "unknown | yes | no", + "outcome_detail": "", + "ground_truth_diagnosis": "", + "ground_truth_source": "case_text | abstract | unknown" +}} + +ABSTRACT: +{abstract} + +CASE TEXT: +{case_text} +""" + +def call_gemini(abstract: str, case_text: str) -> dict: + prompt = EXTRACT_PROMPT.format( + abstract=abstract or "Not available", + case_text=case_text or "Not available" + ) + for attempt in range(4): + try: + resp = client.models.generate_content(model=MODEL_NAME, contents=prompt) + raw = resp.text.strip() + # Strip markdown code fences if present + raw = re.sub(r"^```(?:json)?\s*", "", raw) + raw = re.sub(r"\s*```$", "", raw) + return json.loads(raw) + except json.JSONDecodeError: + time.sleep(2 ** attempt) + except Exception as e: + if "429" in str(e) or "quota" in str(e).lower(): + time.sleep(30) + else: + time.sleep(2 ** attempt) + return {} + +# ── View position normalisation ──────────────────────────────────────────── +VIEW_MAP = { + "frontal": "PA", + "sagittal": "LATERAL", + "axial": "UNKNOWN", + "oblique": "UNKNOWN", + None: "UNKNOWN", +} + +def normalise_view(v): + return VIEW_MAP.get(v, "UNKNOWN") + +# ── Build one schema record ──────────────────────────────────────────────── +def build_record(rec: dict, extracted: dict) -> dict: + pmc_id = rec["pmc_id"] + cases = rec.get("cases", []) + case0 = cases[0] if cases else {} + images = rec.get("images", []) + cxr_imgs = [img for img in images if img.get("is_chest_xray")] + first_cxr = cxr_imgs[0] if cxr_imgs else (images[0] if images else {}) + e = extracted # shorthand + + return { + "case_id": str(uuid.uuid5(uuid.NAMESPACE_URL, pmc_id)), + "source": { + "dataset_name": "multicare", + "dataset_record_id": pmc_id, + "pmc_link": rec.get("link", ""), + "doi": rec.get("doi", ""), + "pmid": rec.get("pmid", ""), + "title": rec.get("title", ""), + "authors": rec.get("authors", []), + "journal": rec.get("journal", ""), + "journal_detail": rec.get("journal_detail", ""), + "year": rec.get("year", ""), + "license": rec.get("license", ""), + }, + + "patient_context": { + "age_years": case0.get("age"), + "sex": (case0.get("gender") or "unknown").lower(), + "immunocompromised": e.get("immunocompromised", "unknown"), + "chief_complaint": e.get("chief_complaint"), + "symptom_duration": e.get("symptom_duration"), + "comorbidities": e.get("comorbidities", []), + "medications_used": e.get("medications_used"), + }, + + "study": { + "modality": "CXR", + "view_position": normalise_view(first_cxr.get("radiology_view")), + "body_region": first_cxr.get("radiology_region") or "chest", + }, + + "images": [ + { + "image_id": str(uuid.uuid5(uuid.NAMESPACE_URL, img.get("file_id", img.get("file", "")))), + "local_image_path": img.get("local_image_path", ""), + "file_type": Path(img.get("file", "")).suffix.lstrip(".") or "webp", + "image_type": img.get("image_type", ""), + "image_subtype": img.get("image_subtype", ""), + "is_chest_xray": img.get("is_chest_xray", False), + "view_position": normalise_view(img.get("radiology_view")), + "radiology_region": img.get("radiology_region"), + "caption": img.get("caption", ""), + "text_references": img.get("text_references", []), + "ml_labels": img.get("ml_labels", ""), + "gt_labels": img.get("gt_labels", ""), + } + for img in images + ], + + "clinical_text": { + "raw_note": case0.get("case_text", ""), + "structured_note": { + "hpi": e.get("clinical_note_hpi"), + "pmh": e.get("clinical_note_pmh"), + "meds": e.get("clinical_note_meds"), + "allergies": e.get("clinical_note_allergies"), + }, + }, + + "radiology_labels": { + "primary_suspected": e.get("primary_suspected", []), + "differential": e.get("differential", []), + "urgency": e.get("urgency", "routine"), + "infectious_concern": e.get("infectious_concern", "unknown"), + "icu_candidate": e.get("icu_candidate", "unknown"), + }, + + "findings_structured": { + "lungs": { + "consolidation": { + "present": e.get("lungs_consolidation_present", "unknown"), + "location": e.get("lungs_consolidation_location", []), + "extent": e.get("lungs_consolidation_extent", "unknown"), + }, + "atelectasis": { + "present": e.get("lungs_atelectasis_present", "unknown"), + "location": e.get("lungs_atelectasis_location", []), + }, + "edema": { + "present": e.get("lungs_edema_present", "unknown"), + "pattern": e.get("lungs_edema_pattern", "unknown"), + }, + }, + "pleura": { + "effusion": { + "present": e.get("pleura_effusion_present", "unknown"), + "side": e.get("pleura_effusion_side", "unknown"), + "size": e.get("pleura_effusion_size", "unknown"), + }, + "pneumothorax": { + "present": e.get("pleura_pneumothorax_present", "unknown"), + "side": e.get("pleura_pneumothorax_side", "unknown"), + }, + }, + "cardiomediastinal": { + "cardiomegaly": e.get("cardiomegaly", "unknown"), + "mediastinal_widening": e.get("mediastinal_widening", "unknown"), + }, + "devices": { + "lines_tubes_present": e.get("lines_tubes_present", "unknown"), + "device_list": e.get("device_list", []), + }, + }, + + "text_outputs": { + "case_card": { + "summary_1_2_lines": e.get("summary_1_2_lines") or rec.get("abstract", ""), + "bullets": e.get("bullets", []), + "red_flags": e.get("red_flags", []), + "uncertainties": e.get("uncertainties", []), + }, + }, + + "outcomes": { + "label_type": "real", + "success": e.get("outcome_success", "unknown"), + "outcome_detail": e.get("outcome_detail"), + "ground_truth_diagnosis": e.get("ground_truth_diagnosis"), + "ground_truth_source": e.get("ground_truth_source", "case_text"), + }, + + "audit": { + "created_at": f"{rec.get('year', '2023')}-01-01T00:00:00Z", + "keywords": rec.get("keywords", []), + "mesh_terms": rec.get("mesh_terms", []), + }, + } + +# ── Main ─────────────────────────────────────────────────────────────────── +def main(): + with open(INPUT_JSON) as f: + source_data = json.load(f) + + print(f"Processing {len(source_data)} records with Gemini extraction...") + results = [] + errors = [] + + for rec in tqdm(source_data, unit="record"): + pmc_id = rec["pmc_id"] + cache_file = CACHE_DIR / f"{pmc_id}.json" + + # Use cache if already processed + if cache_file.exists(): + with open(cache_file) as f: + extracted = json.load(f) + else: + cases = rec.get("cases", []) + case0 = cases[0] if cases else {} + case_text = case0.get("case_text", "") + abstract = rec.get("abstract", "") + + extracted = call_gemini(abstract, case_text) + + # Save to cache + with open(cache_file, "w") as f: + json.dump(extracted, f) + + time.sleep(0.3) # gentle rate limiting + + try: + schema_record = build_record(rec, extracted) + results.append(schema_record) + except Exception as ex: + errors.append({"pmc_id": pmc_id, "error": str(ex)}) + + print(f"\nWriting {len(results)} records to {OUTPUT_JSON}...") + with open(OUTPUT_JSON, "w") as f: + json.dump(results, f, indent=2, default=str) + + if errors: + err_path = OUTPUT_DIR / "errors.json" + with open(err_path, "w") as f: + json.dump(errors, f, indent=2) + print(f"⚠ {len(errors)} errors written to {err_path}") + + print(f"\nDone. Output: {OUTPUT_JSON.resolve()}") + print(f" Records: {len(results)}") + +if __name__ == "__main__": + main() diff --git a/data_pipeline/cluster_cases.py b/data_pipeline/cluster_cases.py new file mode 100644 index 0000000..0bd8afb --- /dev/null +++ b/data_pipeline/cluster_cases.py @@ -0,0 +1,299 @@ +""" +Cluster the 750 cases in cxr_schema_dataset/dataset_per_case.json by clinical +similarity and write two output files: + + cxr_schema_dataset/clusters.json + Array of cluster objects: + { + "cluster_id": int, + "cluster_name": str, # auto-named from dominant features + "dominant_diagnosis": str, + "size": int, + "case_ids": [str, ...] + } + Outliers (HDBSCAN label -1) are collected into a single "Miscellaneous" cluster. + + cxr_schema_dataset/dataset_per_case.json (updated in-place) + cxr_schema_dataset/dataset_cxr_primary.json (updated in-place) + The `embeddings` field of every record is populated with the 384-dim + sentence-transformer vector for that case. +""" + +import json +import os +import warnings +from collections import Counter +from pathlib import Path + +import numpy as np + +warnings.filterwarnings("ignore") + +# Redirect HuggingFace cache to project dir (avoids home-dir space limits on HPC) +os.environ.setdefault( + "HF_HOME", + str(Path(__file__).parent / ".hf_cache"), +) + +# ── Paths ────────────────────────────────────────────────────────────────── +BASE = Path("cxr_schema_dataset") +PER_CASE_PATH = BASE / "dataset_per_case.json" +CXR_PRIMARY_PATH = BASE / "dataset_cxr_primary.json" +CLUSTERS_PATH = BASE / "clusters.json" + +# ── Embedding model ──────────────────────────────────────────────────────── +EMBED_MODEL = "all-MiniLM-L6-v2" # 384-dim, fast, good semantic similarity + +# ── UMAP / HDBSCAN tuning ────────────────────────────────────────────────── +UMAP_N_COMPONENTS = 20 +UMAP_N_NEIGHBORS = 15 +UMAP_MIN_DIST = 0.0 +HDBSCAN_MIN_CLUSTER = 5 # minimum cases per cluster +HDBSCAN_MIN_SAMPLES = 3 + + +# ── Text representation ──────────────────────────────────────────────────── +def case_to_text(r: dict) -> str: + """ + Build a single clinical text for a case by concatenating the most + semantically informative fields. + """ + parts = [] + + pat = r.get("patient") or {} + pres = r.get("presentation") or {} + asmt = r.get("assessment") or {} + find = r.get("findings") or {} + summ = r.get("summary") or {} + + # Demographics + chief complaint + age = pat.get("age_years") + sex = pat.get("sex", "") + if age: + parts.append(f"{int(age)}-year-old {sex}.") + + if pres.get("chief_complaint"): + parts.append(f"Chief complaint: {pres['chief_complaint']}.") + + # HPI (most information-dense field) + if pres.get("hpi"): + parts.append(pres["hpi"]) + + # PMH + if pres.get("pmh"): + parts.append(f"PMH: {pres['pmh']}.") + + # Comorbidities + comorbids = pat.get("comorbidities") or [] + if comorbids: + parts.append(f"Comorbidities: {', '.join(comorbids)}.") + + # Diagnosis + differentials + dx = asmt.get("diagnosis_primary") + if dx: + parts.append(f"Diagnosis: {dx}.") + suspected = asmt.get("suspected_primary") or [] + if suspected: + parts.append(f"Suspected: {', '.join(suspected)}.") + diff = asmt.get("differential") or [] + if diff: + parts.append(f"Differential: {', '.join(diff)}.") + + # Urgency / infection + if asmt.get("urgency"): + parts.append(f"Urgency: {asmt['urgency']}.") + if asmt.get("infectious_concern") == "yes": + parts.append("Infectious concern.") + if asmt.get("icu_candidate") == "yes": + parts.append("ICU candidate.") + + # Findings summary + lungs = find.get("lungs") or {} + pleura = find.get("pleura") or {} + cm = find.get("cardiomediastinal") or {} + + findings_parts = [] + if lungs.get("consolidation_present") == "yes": + locs = ", ".join(lungs.get("consolidation_locations") or []) + findings_parts.append(f"consolidation ({locs})" if locs else "consolidation") + if lungs.get("atelectasis_present") == "yes": + findings_parts.append("atelectasis") + if lungs.get("edema_present") == "yes": + findings_parts.append(f"edema ({lungs.get('edema_pattern', 'unknown')})") + if pleura.get("effusion_present") == "yes": + findings_parts.append(f"pleural effusion ({pleura.get('effusion_side', '')})") + if pleura.get("pneumothorax_present") == "yes": + findings_parts.append("pneumothorax") + if cm.get("cardiomegaly") == "yes": + findings_parts.append("cardiomegaly") + if findings_parts: + parts.append(f"Findings: {', '.join(findings_parts)}.") + + # Summary + if summ.get("one_liner"): + parts.append(summ["one_liner"]) + + key_pts = summ.get("key_points") or [] + if key_pts: + parts.append(" ".join(key_pts)) + + return " ".join(parts) + + +# ── Cluster naming ───────────────────────────────────────────────────────── +def name_cluster(cases) -> tuple: + """ + Return (cluster_name, dominant_diagnosis) for a group of case records. + Naming strategy: most common diagnosis + most common finding/urgency modifier. + """ + diagnoses = [ + r["assessment"]["diagnosis_primary"] + for r in cases + if r["assessment"].get("diagnosis_primary") + ] + findings_tokens = [] + for r in cases: + lungs = (r.get("findings") or {}).get("lungs") or {} + pleura = (r.get("findings") or {}).get("pleura") or {} + cm = (r.get("findings") or {}).get("cardiomediastinal") or {} + if lungs.get("consolidation_present") == "yes": + findings_tokens.append("consolidation") + if lungs.get("edema_present") == "yes": + findings_tokens.append("edema") + if pleura.get("effusion_present") == "yes": + findings_tokens.append("effusion") + if cm.get("cardiomegaly") == "yes": + findings_tokens.append("cardiomegaly") + + infectious = sum( + 1 for r in cases if r["assessment"].get("infectious_concern") == "yes" + ) + + if diagnoses: + top_dx = Counter(diagnoses).most_common(1)[0][0] + dominant = top_dx + else: + top_dx = "unknown" + dominant = "unknown" + + # Build a concise label + label_parts = [top_dx.title()] + if findings_tokens: + top_finding = Counter(findings_tokens).most_common(1)[0][0] + label_parts.append(f"({top_finding})") + if infectious > len(cases) // 2: + label_parts.append("[infectious]") + + return " ".join(label_parts), dominant + + +# ── Main ─────────────────────────────────────────────────────────────────── +def main(): + print("Loading per-case dataset ...") + with open(PER_CASE_PATH) as f: + cases = json.load(f) + + # Build text representations + print("Building text representations ...") + texts = [case_to_text(r) for r in cases] + case_ids = [r["case_id"] for r in cases] + + # Embed + print(f"Embedding {len(texts)} cases with '{EMBED_MODEL}' ...") + from sentence_transformers import SentenceTransformer + model = SentenceTransformer(EMBED_MODEL) + embeddings = model.encode(texts, show_progress_bar=True, batch_size=64) + print(f"Embedding shape: {embeddings.shape}") + + # UMAP dimensionality reduction + print(f"UMAP reduction to {UMAP_N_COMPONENTS} dims ...") + import umap + reducer = umap.UMAP( + n_components=UMAP_N_COMPONENTS, + n_neighbors=UMAP_N_NEIGHBORS, + min_dist=UMAP_MIN_DIST, + metric="cosine", + random_state=42, + ) + reduced = reducer.fit_transform(embeddings) + + # HDBSCAN clustering + print("HDBSCAN clustering ...") + import hdbscan + clusterer = hdbscan.HDBSCAN( + min_cluster_size=HDBSCAN_MIN_CLUSTER, + min_samples=HDBSCAN_MIN_SAMPLES, + metric="euclidean", + cluster_selection_method="eom", + ) + labels = clusterer.fit_predict(reduced) + + n_clusters = len(set(labels)) - (1 if -1 in labels else 0) + n_outliers = int(np.sum(labels == -1)) + print(f"Found {n_clusters} clusters, {n_outliers} outliers → 'Miscellaneous'") + + # Build case_id → embedding mapping (for backfill) + embed_map = {cid: emb.tolist() for cid, emb in zip(case_ids, embeddings)} + + # Build cluster → list of (case_id, case_record) pairs + cluster_map: dict[int, list] = {} + for idx, label in enumerate(labels): + cluster_map.setdefault(label, []).append(idx) + + # Assemble output clusters (sorted by size desc, outliers last) + cluster_objects = [] + cluster_id_counter = 0 + + for label in sorted(cluster_map.keys(), key=lambda l: (-len(cluster_map[l]), l)): + indices = cluster_map[label] + cluster_cases = [cases[i] for i in indices] + cluster_case_ids = [case_ids[i] for i in indices] + + if label == -1: + cname = "Miscellaneous" + dominant = "various" + else: + cname, dominant = name_cluster(cluster_cases) + + cluster_objects.append({ + "cluster_id": cluster_id_counter, + "cluster_name": cname, + "dominant_diagnosis": dominant, + "size": len(cluster_case_ids), + "case_ids": cluster_case_ids, + }) + cluster_id_counter += 1 + + print(f"\nCluster summary:") + for c in cluster_objects: + print(f" [{c['cluster_id']:2d}] {c['cluster_name']:<50s} n={c['size']}") + + # Write clusters.json + print(f"\nWriting {CLUSTERS_PATH} ...") + with open(CLUSTERS_PATH, "w") as f: + json.dump(cluster_objects, f, indent=2, ensure_ascii=False) + + # Backfill embeddings into per-case dataset + print(f"Backfilling embeddings into {PER_CASE_PATH} ...") + for r in cases: + r["embeddings"] = embed_map.get(r["case_id"], []) + with open(PER_CASE_PATH, "w") as f: + json.dump(cases, f, indent=2, ensure_ascii=False) + + # Backfill embeddings into cxr-primary dataset + print(f"Backfilling embeddings into {CXR_PRIMARY_PATH} ...") + with open(CXR_PRIMARY_PATH) as f: + cxr_primary = json.load(f) + for r in cxr_primary: + r["embeddings"] = embed_map.get(r["case_id"], []) + with open(CXR_PRIMARY_PATH, "w") as f: + json.dump(cxr_primary, f, indent=2, ensure_ascii=False) + + print("\nDone.") + print(f" clusters.json → {len(cluster_objects)} clusters") + print(f" dataset_per_case.json → embeddings backfilled") + print(f" dataset_cxr_primary.json → embeddings backfilled") + + +if __name__ == "__main__": + main() diff --git a/data_pipeline/diagnosis_graph.py b/data_pipeline/diagnosis_graph.py new file mode 100644 index 0000000..0fca2a2 --- /dev/null +++ b/data_pipeline/diagnosis_graph.py @@ -0,0 +1,1000 @@ +""" +MedGemma Differential Diagnosis System using LangGraph + MultiCaRe Dataset +""" + +import json +import base64 +import os +from pathlib import Path +from typing import TypedDict, Annotated, List, Optional +import operator + +from langgraph.graph import StateGraph, END +from langchain_core.messages import HumanMessage, AIMessage, SystemMessage + + +# ───────────────────────────────────────────── +# State Definition +# ───────────────────────────────────────────── + +class DiagnosisState(TypedDict): + case_id: str + case_data: dict + images_b64: List[dict] # [{image_id, b64, caption, subtype}] + initial_diagnosis: Optional[str] + differential_diagnoses: Optional[List[dict]] # [{condition, confidence, reasoning}] + bias_check_notes: Optional[str] + alternative_hypotheses: Optional[List[dict]] + final_report: Optional[str] + messages: Annotated[List, operator.add] + + +# ───────────────────────────────────────────── +# Data Loading +# ───────────────────────────────────────────── + +def load_case_node(state: DiagnosisState) -> DiagnosisState: + """Load case details from parsed MultiCaRe JSON.""" + case_id = state["case_id"] + + # Load from your parsed dataset — adjust path as needed + dataset_path = Path(os.getenv("MULTICARE_DATASET_PATH", "multicare_parsed.json")) + + if not dataset_path.exists(): + raise FileNotFoundError( + f"Dataset not found at {dataset_path}. " + "Set MULTICARE_DATASET_PATH env var to your parsed JSON file." + ) + + with open(dataset_path) as f: + dataset = json.load(f) + + # Support both list and dict-keyed formats + if isinstance(dataset, list): + cases = {c["case_id"]: c for c in dataset} + else: + cases = dataset + + if case_id not in cases: + raise ValueError(f"Case ID '{case_id}' not found in dataset.") + + case_data = cases[case_id] + print(f"[load_case] Loaded case: {case_data['provenance']['article_title']}") + + return { + **state, + "case_data": case_data, + "messages": [HumanMessage(content=f"Starting diagnosis for case {case_id}")] + } + + +def load_images_node(state: DiagnosisState) -> DiagnosisState: + """Load and base64-encode images for the case.""" + case_data = state["case_data"] + images_b64 = [] + + base_path = Path(os.getenv("MULTICARE_IMAGES_PATH", ".")) + + for img_meta in case_data.get("images", []): + img_path = base_path / img_meta["local_image_path"] + + if not img_path.exists(): + print(f"[load_images] Warning: Image not found: {img_path}") + continue + + with open(img_path, "rb") as f: + raw = f.read() + + b64 = base64.b64encode(raw).decode("utf-8") + + # Determine MIME type + ext = img_meta["file_type"].lower() + mime_map = { + "webp": "image/webp", + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "png": "image/png", + } + mime = mime_map.get(ext, "image/jpeg") + + images_b64.append({ + "image_id": img_meta["image_id"], + "b64": b64, + "mime": mime, + "caption": img_meta.get("caption", ""), + "subtype": img_meta.get("image_subtype", "unknown"), + "image_type": img_meta.get("image_type", "unknown"), + }) + print(f"[load_images] Loaded image: {img_meta['image_subtype']} ({img_meta['image_id'][:8]}...)") + + return {**state, "images_b64": images_b64} + + +# ───────────────────────────────────────────── +# MedGemma Helpers +# ───────────────────────────────────────────── +def _build_medgemma_prompt(case_data: dict, images_b64: list, task: str) -> list: + """Build a multimodal message list for MedGemma.""" + import json + + patient = case_data.get("patient", {}) + presentation = case_data.get("presentation", {}) + findings = case_data.get("findings", {}) + study = case_data.get("study", {}) + + # Clinical summary text + clinical_text = f""" +=== CLINICAL CASE === + +PATIENT: +- Age: {patient.get('age_years', 'unknown')} years, {patient.get('sex', 'unknown')} +- Immunocompromised: {patient.get('immunocompromised', 'unknown')} +- Comorbidities: {', '.join(patient.get('comorbidities', [])) or 'None documented'} +- Medications: {', '.join(patient.get('medications', [])) or 'None documented'} +- Allergies: {patient.get('allergies', 'None documented')} + +PRESENTATION: +- Chief Complaint: {presentation.get('chief_complaint', 'N/A')} +- Duration: {presentation.get('symptom_duration', 'N/A')} +- HPI: {presentation.get('hpi', 'N/A')} +- Past Medical History: {presentation.get('pmh', 'N/A')} + +IMAGING STUDY: +- Modality: {study.get('modality', 'N/A')} +- Region: {study.get('body_region', 'N/A')} +- View: {study.get('view_position', 'N/A')} + +RADIOLOGICAL FINDINGS: +""" + # Dynamically inject ALL available findings, regardless of body system + if findings: + for region, details in findings.items(): + clinical_text += f"- {region.upper()}: {json.dumps(details)}\n" + else: + clinical_text += "- No specific radiological findings documented.\n" + + # Image captions context + if images_b64: + clinical_text += "\nAVAILABLE IMAGES:\n" + for i, img in enumerate(images_b64, 1): + clinical_text += f" Image {i} ({img['subtype']}): {img['caption']}\n" + + content = [{"type": "text", "text": clinical_text + f"\n\n{task}"}] + + # Attach images + for img in images_b64: + content.append({ + "type": "image_url", + "image_url": {"url": f"data:{img['mime']};base64,{img['b64']}"} + }) + + return [HumanMessage(content=content)] + +def preload_model() -> None: + """ + Pre-warm the local MedGemma model so the first pipeline call isn't slow. + Call this once before graph.invoke() to load weights into GPU memory. + """ + use_local = os.getenv("USE_LOCAL_MEDGEMMA", "0") == "1" + if not use_local: + print("[preload] Using HF Inference API — no local model to preload.") + return + model_id = os.getenv("MEDGEMMA_MODEL", "google/medgemma-1.5-4b-it") + if hasattr(_call_medgemma_local, "_model") and _call_medgemma_local._model_id == model_id: + print(f"[preload] Model already loaded: {model_id}") + return + # Trigger the lazy-load by making a minimal text-only call + _call_medgemma_local( + [HumanMessage(content="Hello")], + system_prompt="You are a helpful assistant.", + max_new_tokens=8, + ) + print(f"[preload] Model ready: {model_id}") + + +def _call_medgemma(messages: list, system_prompt: str, max_new_tokens: int = 1024) -> str: + """ + Call MedGemma via HuggingFace Inference or local pipeline. + + Supports two backends: + 1. HuggingFace Inference API (set HF_TOKEN env var) + 2. Local transformers pipeline (set USE_LOCAL_MEDGEMMA=1) + """ + use_local = os.getenv("USE_LOCAL_MEDGEMMA", "0") == "1" + + if use_local: + return _call_medgemma_local(messages, system_prompt, max_new_tokens=max_new_tokens) + else: + return _call_medgemma_hf_api(messages, system_prompt, max_new_tokens=max_new_tokens) + + +def _call_medgemma_hf_api(messages: list, system_prompt: str, max_new_tokens: int = 1024) -> str: + """Call MedGemma via HuggingFace Inference API.""" + from huggingface_hub import InferenceClient + + token = os.getenv("HF_TOKEN") + if not token: + raise ValueError("HF_TOKEN environment variable not set.") + + model = os.getenv("MEDGEMMA_MODEL", "google/medgemma-1.5-4b-it") + client = InferenceClient(model=model, token=token) + + # Build HF-compatible message format + hf_messages = [{"role": "system", "content": system_prompt}] + + for msg in messages: + if isinstance(msg, HumanMessage): + if isinstance(msg.content, list): + # Multimodal + hf_content = [] + for part in msg.content: + if part["type"] == "text": + hf_content.append({"type": "text", "text": part["text"]}) + elif part["type"] == "image_url": + hf_content.append({ + "type": "image_url", + "image_url": part["image_url"] + }) + hf_messages.append({"role": "user", "content": hf_content}) + else: + hf_messages.append({"role": "user", "content": msg.content}) + + response = client.chat_completion( + messages=hf_messages, + max_tokens=max_new_tokens, + temperature=0.3, + ) + + return response.choices[0].message.content + + +def _call_medgemma_local(messages: list, system_prompt: str, max_new_tokens: int = 1024) -> str: + """ + Call MedGemma locally using AutoProcessor + AutoModelForImageTextToText. + Requires a GPU with sufficient VRAM (≥16 GB for 4b, ≥40 GB for 27b). + Set MEDGEMMA_MODEL env var to the HuggingFace model ID. + If the model is gated, set HF_TOKEN env var for authentication. + """ + from transformers import AutoProcessor, AutoModelForImageTextToText + from PIL import Image + import torch + import io + + model_id = os.getenv("MEDGEMMA_MODEL", "google/medgemma-1.5-4b-it") + hf_token = os.getenv("HF_TOKEN", None) + + # ── Lazy-load model + processor (cached on the function object) ────────── + if not hasattr(_call_medgemma_local, "_model") or _call_medgemma_local._model_id != model_id: + print(f"[medgemma] Loading local model: {model_id} (this may take a minute)") + _call_medgemma_local._processor = AutoProcessor.from_pretrained( + model_id, token=hf_token + ) + _call_medgemma_local._model = AutoModelForImageTextToText.from_pretrained( + model_id, + token=hf_token, + torch_dtype=torch.bfloat16, + device_map="auto", + ) + _call_medgemma_local._model_id = model_id + _call_medgemma_local._model.eval() + print("[medgemma] Model loaded.") + + processor = _call_medgemma_local._processor + model = _call_medgemma_local._model + + # ── Collect PIL images and build chat messages ──────────────────────────── + pil_images = [] + chat_messages = [{"role": "system", "content": system_prompt}] + + for msg in messages: + if isinstance(msg, HumanMessage): + if isinstance(msg.content, list): + content = [] + for part in msg.content: + if part["type"] == "text": + content.append({"type": "text", "text": part["text"]}) + elif part["type"] == "image_url": + data_url = part["image_url"]["url"] + b64_data = data_url.split(",", 1)[1] + img_bytes = base64.b64decode(b64_data) + pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB") + pil_images.append(pil_img) + content.append({"type": "image"}) + chat_messages.append({"role": "user", "content": content}) + else: + chat_messages.append( + {"role": "user", "content": [{"type": "text", "text": msg.content}]} + ) + + # ── Tokenise ───────────────────────────────────────────────────────────── + prompt_text = processor.apply_chat_template( + chat_messages, + tokenize=False, + add_generation_prompt=True, + ) + + inputs = processor( + text=prompt_text, + images=pil_images if pil_images else None, + return_tensors="pt", + ).to(model.device) + + # ── Generate ───────────────────────────────────────────────────────────── + with torch.inference_mode(): + output_ids = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + ) + + # Decode only the newly generated tokens + input_len = inputs["input_ids"].shape[-1] + generated = output_ids[0][input_len:] + return processor.decode(generated, skip_special_tokens=True) + + +# ───────────────────────────────────────────── +# Diagnosis Nodes +# ───────────────────────────────────────────── + +SYSTEM_PROMPT = """You are an expert AI medical diagnostic assistant. +You analyze clinical cases with the rigor of a senior physician, considering all +available information: patient history, symptoms, medications, imaging, and lab findings. +Be systematic, evidence-based, and transparent about your reasoning and uncertainty. +When asked for JSON output, respond with ONLY valid JSON — no markdown fences, no commentary before or after.""" + + +def _extract_text_diagnosis(text: str) -> dict | None: + """ + Fallback parser for when model outputs prose instead of JSON. + Attempts to extract diagnosis info from markdown/text format. + """ + import re + + result = {} + + # Extract primary/main condition + condition_patterns = [ + r'(?:PRIMARY DIAGNOSIS|Condition|Main Diagnosis)[:\s]*[-\*]?\s*([^\n]+)', + r'(?:The (?:primary|main) diagnosis is)[:\s]*([^\n\.]+)', + ] + for pat in condition_patterns: + m = re.search(pat, text, re.IGNORECASE) + if m: + condition = m.group(1).strip().strip('*').strip() + if condition and len(condition) > 2: + result['acute_complication'] = {'condition': condition, 'confidence': 85, 'reasoning': 'Extracted from text output'} + break + + # Extract confidence if present + conf_match = re.search(r'Confidence[:\s]*([\d]+)', text, re.IGNORECASE) + if conf_match and 'acute_complication' in result: + try: + result['acute_complication']['confidence'] = int(conf_match.group(1)) + except ValueError: + pass + + # Extract differentials from numbered or bulleted lists + diff_section = re.search(r'DIFFERENTIAL DIAGNOSES?[:\s]*(.+?)(?:CRITICAL|CLINICAL REASONING|$)', text, re.DOTALL | re.IGNORECASE) + if diff_section: + diff_text = diff_section.group(1) + # Find numbered items like "1. Condition" or "- Condition" + diff_items = re.findall(r'(?:\d+\.\s*\*\*|\d+\.\s*|[-\*]\s*\*\*)([^\*\n]+)', diff_text) + differentials = [] + for item in diff_items[:5]: # Limit to 5 + cond = item.strip().strip('*').strip() + if cond and len(cond) > 2 and not cond.lower().startswith(('confidence', 'supporting', 'against')): + differentials.append({ + 'condition': cond, + 'confidence': 50, + 'supporting_evidence': [], + 'against_evidence': [] + }) + if differentials: + result['differentials'] = differentials + + # Extract critical findings + crit_section = re.search(r'CRITICAL FINDINGS?[:\s]*(.+?)(?:CLINICAL REASONING|$)', text, re.DOTALL | re.IGNORECASE) + if crit_section: + findings = re.findall(r'[-\*]\s*\*\*([^\*]+)\*\*|[-\*]\s*([^\n]+)', crit_section.group(1)) + result['critical_findings'] = [f[0] or f[1] for f in findings if (f[0] or f[1]).strip()][:5] + + return result if result else None + + +def _extract_json_object(text: str) -> dict | None: + """Robustly extract a JSON object from model output.""" + import re + # 1. Try fenced ```json ... ``` block + fenced = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL) + if fenced: + try: + return json.loads(fenced.group(1)) + except json.JSONDecodeError: + pass + # 2. Find all '{' positions and try parsing from each (last to first) + candidates = [m.start() for m in re.finditer(r'\{', text)] + for start in reversed(candidates): + try: + return json.loads(text[start:]) + except json.JSONDecodeError: + # Try finding the matching closing brace + depth = 0 + for i, ch in enumerate(text[start:]): + if ch == '{': + depth += 1 + elif ch == '}': + depth -= 1 + if depth == 0: + try: + return json.loads(text[start:start + i + 1]) + except json.JSONDecodeError: + break + + # 3. Fallback: try to extract from prose/text format + text_result = _extract_text_diagnosis(text) + if text_result: + return text_result + + return None + + +def _extract_json_array(text: str) -> list | None: + """Robustly extract a JSON array from model output.""" + import re + # 1. Try fenced ```json ... ``` block + fenced = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL) + if fenced: + try: + return json.loads(fenced.group(1)) + except json.JSONDecodeError: + pass + # 2. Find all '[' positions and try parsing from each (last to first) + candidates = [m.start() for m in re.finditer(r'\[', text)] + for start in reversed(candidates): + try: + return json.loads(text[start:]) + except json.JSONDecodeError: + depth = 0 + for i, ch in enumerate(text[start:]): + if ch == '[': + depth += 1 + elif ch == ']': + depth -= 1 + if depth == 0: + try: + return json.loads(text[start:start + i + 1]) + except json.JSONDecodeError: + break + return None + +def initial_diagnosis_node(state: DiagnosisState) -> DiagnosisState: + """Generate initial diagnosis with differential and confidence scores.""" + print("[diagnosis] Generating initial diagnosis...") + + task = """Analyze the clinical case and images above. Respond with ONLY a JSON object (no other text). + +CRITICAL INSTRUCTIONS: +1. Separate the underlying etiology (root cause) from the acute complication or primary focus that requires immediate intervention. +2. MEDICAL REALITY CHECK: Never assign 100% confidence to any single condition. Always distribute probabilities to leave room for uncertainty and differentials. + +Required JSON structure: +{ + "underlying_etiology": {"condition": "", "confidence": <0-99>}, + "acute_complication": {"condition": "", "confidence": <0-99>, "requires_intervention": true, "reasoning": "<1-2 sentences>"}, + "differentials": [ + {"condition": "", "confidence": <0-99>, "supporting_evidence": [""], "against_evidence": [""]} + ], + "critical_findings": [""], + "clinical_reasoning": "" +} + +Provide 3-5 differentials. Be specific and evidence-based.""" + + # Retry logic: try up to 2 times if JSON parsing fails + max_retries = 2 + diagnosis_data = None + response = "" + + for attempt in range(max_retries): + messages = _build_medgemma_prompt(state["case_data"], state["images_b64"], task) + response = _call_medgemma(messages, SYSTEM_PROMPT, max_new_tokens=2048) + + # Debug: show raw response + print(f"[diagnosis] Attempt {attempt + 1} - Raw response ({len(response)} chars):") + print(response[:500]) + if len(response) > 500: + print(f"... ({len(response) - 500} more chars)") + + # Robust JSON extraction (includes text fallback) + diagnosis_data = _extract_json_object(response) + + if diagnosis_data is not None and diagnosis_data.get('acute_complication') or diagnosis_data.get('differentials'): + print(f"[diagnosis] JSON extraction successful on attempt {attempt + 1}") + break + elif attempt < max_retries - 1: + print(f"[diagnosis] JSON extraction failed, retrying with stricter prompt...") + # Make the task prompt stricter for retry + task = """IMPORTANT: You MUST respond with ONLY valid JSON. No explanatory text before or after. + +Analyze the clinical case. Return this exact JSON structure: +{"underlying_etiology": {"condition": "", "confidence": <0-99>}, "acute_complication": {"condition": "", "confidence": <0-99>, "requires_intervention": true, "reasoning": "<1-2 sentences>"}, "differentials": [{"condition": "", "confidence": <0-99>, "supporting_evidence": [""], "against_evidence": [""]}], "critical_findings": [""], "clinical_reasoning": ""} + +Provide 3-5 differentials. Start your response with { and end with }.""" + + if diagnosis_data is None: + print("[diagnosis] WARNING: Could not parse JSON from response after retries, storing raw text.") + diagnosis_data = {"raw_response": response} + + return { + **state, + "initial_diagnosis": json.dumps(diagnosis_data, indent=2), + "differential_diagnoses": diagnosis_data.get("differentials", []), + "messages": [AIMessage(content=f"Initial diagnosis generated: {diagnosis_data.get('acute_complication', {}).get('condition', 'unknown')}")] + } + +def bias_check_node(state: DiagnosisState) -> DiagnosisState: + """Check for diagnostic anchoring bias and cognitive shortcuts.""" + print("[bias_check] Performing cognitive bias check...") + + initial = state["initial_diagnosis"] + + task = f""" +You have generated an initial diagnosis. Now perform a COGNITIVE BIAS AUDIT: + +Initial diagnosis was: +{initial} + +Check for these specific biases: +1. ANCHORING BIAS: Are you too fixated on the first or most obvious finding? +2. AVAILABILITY BIAS: Are common diagnoses being over-weighted? +3. PREMATURE CLOSURE: Have you stopped considering alternatives too early? +4. FRAMING EFFECT: How is the case presentation framing your thinking? +5. REPRESENTATIVE HEURISTIC: Are you pattern-matching too quickly? +6. RARE DISEASE NEGLECT: What uncommon but serious conditions are being missed? +7. CONFIRMATION BIAS: Are you selectively weighing evidence? + +For each bias found, explain: +- What the bias is in this case +- How it might be skewing the diagnosis +- What to reconsider + +Also identify: What diagnoses might a physician MISS due to these biases? +""" + + messages = _build_medgemma_prompt(state["case_data"], [], task) # No images needed for bias check + response = _call_medgemma(messages, SYSTEM_PROMPT, max_new_tokens=1024) + + print(f"[bias_check] Response ({len(response)} chars)") + + return { + **state, + "bias_check_notes": response, + "messages": [AIMessage(content="Bias check completed.")] + } + + +def _extract_text_alternatives(text: str) -> list | None: + """ + Fallback parser for extracting alternative diagnoses from prose output. + """ + import re + + alternatives = [] + + # Look for numbered items like "1. Condition Name" or "**Condition Name**" + patterns = [ + r'\d+\.\s*\*\*([^\*]+)\*\*', # 1. **Condition** + r'\d+\.\s*([^:\n]+?)(?::|\n)', # 1. Condition: + r'[-\*]\s*\*\*([^\*]+)\*\*', # - **Condition** + ] + + for pat in patterns: + matches = re.findall(pat, text) + for match in matches: + cond = match.strip() + # Filter out common non-diagnosis phrases + if (cond and len(cond) > 3 and len(cond) < 100 and + not any(skip in cond.lower() for skip in + ['confidence', 'evidence', 'missed', 'why', 'risk', 'test', 'supporting'])): + # Try to extract confidence for this condition + conf_match = re.search(rf'{re.escape(cond)}.*?(?:confidence|Confidence)[:\s]*(\d+)', text, re.DOTALL) + confidence = int(conf_match.group(1)) if conf_match else 50 + + # Try to extract risk level + risk_match = re.search(rf'{re.escape(cond)}.*?(?:risk|Risk)[:\s]*([\w]+)', text, re.DOTALL) + risk = risk_match.group(1).lower() if risk_match else 'medium' + if risk not in ['low', 'medium', 'high', 'critical']: + risk = 'medium' + + alternatives.append({ + 'condition': cond, + 'confidence': confidence, + 'why_missed': 'Extracted from text output', + 'supporting_evidence': [], + 'confirmatory_tests': [], + 'risk_if_missed': risk + }) + if alternatives: + break + + return alternatives[:5] if alternatives else None + + +def alternative_hypotheses_node(state: DiagnosisState) -> DiagnosisState: + """Generate alternative diagnoses that might be missed.""" + print("[alternatives] Generating alternative hypotheses...") + + bias_notes = state.get("bias_check_notes", "") + initial = state["initial_diagnosis"] + + task = f"""Given this initial diagnosis and bias analysis, generate 3-5 alternative diagnoses that may have been missed. + +Initial diagnosis: +{initial} + +Bias analysis highlights: +{bias_notes[:800] if bias_notes else 'N/A'} + +Consider: atypical presentations, rare but serious conditions, dual pathology, mimickers, and systemic diseases. + +Respond with ONLY a JSON array (no other text). Each element must have: +{{"condition":"","why_missed":"","supporting_evidence":[""],"confirmatory_tests":[""],"risk_if_missed":"high|medium|low|critical","confidence":<0-100>}}""" + + messages = _build_medgemma_prompt(state["case_data"], state["images_b64"], task) + response = _call_medgemma(messages, SYSTEM_PROMPT, max_new_tokens=2048) + + # Debug: show raw response + print(f"[alternatives] Raw response ({len(response)} chars):") + print(response[:500]) + if len(response) > 500: + print(f"... ({len(response) - 500} more chars)") + + # Robust JSON extraction + alts = _extract_json_array(response) + if alts is None: + print("[alternatives] WARNING: Could not parse JSON array, trying object fallback.") + obj = _extract_json_object(response) + if obj and isinstance(obj, dict): + # Model may have wrapped array in an object + for v in obj.values(): + if isinstance(v, list): + alts = v + break + if alts is None: + # Try text extraction fallback + print("[alternatives] Trying text extraction fallback...") + alts = _extract_text_alternatives(response) + if alts is None: + alts = [{"raw_response": response, "condition": "See raw output", "confidence": 0, "risk_if_missed": "unknown"}] + + return { + **state, + "alternative_hypotheses": alts, + "messages": [AIMessage(content=f"Generated {len(alts)} alternative hypotheses.")] + } + + +def final_report_node(state: DiagnosisState) -> DiagnosisState: + """Compile everything into a final structured diagnostic report.""" + print("[report] Compiling final diagnostic report...") + + case = state["case_data"] + patient = case.get("patient", {}) + prov = case.get("provenance", {}) + + initial_data = {} + try: + initial_data = json.loads(state.get("initial_diagnosis", "{}")) + except: + pass + + alts = state.get("alternative_hypotheses", []) + + # Format differentials + def fmt_list(items, key): + return "\n".join(f" - {i.get(key, str(i))}" for i in items) if items else " None" + + report = f""" +╔══════════════════════════════════════════════════════════════════════╗ +║ MEDGEMMA DIFFERENTIAL DIAGNOSIS REPORT ║ +╚══════════════════════════════════════════════════════════════════════╝ + +CASE ID: {state['case_id']} +SOURCE: {prov.get('article_title', 'N/A')} ({prov.get('journal', 'N/A')}, {prov.get('year', 'N/A')}) +PMC ID: {prov.get('pmc_id', 'N/A')} + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +PATIENT SUMMARY +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +{case.get('summary', {}).get('one_liner', 'N/A')} + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +PRIMARY DIAGNOSIS +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +Condition: {initial_data.get('acute_complication', initial_data.get('underlying_etiology', {})).get('condition', 'See raw output')} +Confidence: {initial_data.get('acute_complication', initial_data.get('underlying_etiology', {})).get('confidence', '?')}% +Reasoning: {initial_data.get('acute_complication', initial_data.get('underlying_etiology', {})).get('reasoning', 'N/A')} + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +UNDERLYING ETIOLOGY (if different from primary) +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +Condition: {initial_data.get('underlying_etiology', {}).get('condition', 'Same as primary or N/A')} +Confidence: {initial_data.get('underlying_etiology', {}).get('confidence', 'N/A')}% + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +DIFFERENTIAL DIAGNOSES +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━""" + + for i, diff in enumerate(state.get("differential_diagnoses", []), 1): + report += f""" + [{i}] {diff.get('condition', 'Unknown')} — Confidence: {diff.get('confidence', '?')}% + Supporting: {'; '.join(diff.get('supporting_evidence', [])[:2])} + Against: {'; '.join(diff.get('against_evidence', [])[:2])}""" + + report += f""" + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +COGNITIVE BIAS ANALYSIS +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +{state.get('bias_check_notes', 'Not performed')[:800]}... + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +MISSED DIAGNOSES / ALTERNATIVE HYPOTHESES +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━""" + + for i, alt in enumerate(alts, 1): + if isinstance(alt, dict) and "condition" in alt: + report += f""" + [{i}] {alt.get('condition')} (Risk if missed: {alt.get('risk_if_missed', '?').upper()}) + Why missed: {alt.get('why_missed', 'N/A')} + Evidence: {'; '.join(alt.get('supporting_evidence', [])[:2])} + Confirm via: {'; '.join(alt.get('confirmatory_tests', [])[:3])}""" + + # Ground truth comparison + actual = case.get("assessment", {}) + report += f""" + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +GROUND TRUTH (from case record) +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +Actual Primary Diagnosis: {actual.get('diagnosis_primary', 'N/A')} +Suspected: {', '.join(actual.get('suspected_primary', []))} +Documented Differentials: {', '.join(actual.get('differential', []))} +Urgency: {actual.get('urgency', 'N/A')} +Outcome: {case.get('outcome', {}).get('detail', 'N/A')} + +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +⚠ DISCLAIMER: For research/educational use only. Not for clinical decisions. +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +""" + + return { + **state, + "final_report": report, + "messages": [AIMessage(content="Final report compiled.")] + } + + +# ───────────────────────────────────────────── +# Graph Construction +# ───────────────────────────────────────────── + +def build_diagnosis_graph() -> StateGraph: + graph = StateGraph(DiagnosisState) + + graph.add_node("load_case", load_case_node) + graph.add_node("load_images", load_images_node) + graph.add_node("initial_diagnosis", initial_diagnosis_node) + graph.add_node("bias_check", bias_check_node) + graph.add_node("alternative_hypotheses", alternative_hypotheses_node) + graph.add_node("final_report", final_report_node) + + graph.set_entry_point("load_case") + graph.add_edge("load_case", "load_images") + graph.add_edge("load_images", "initial_diagnosis") + graph.add_edge("initial_diagnosis", "bias_check") + graph.add_edge("bias_check", "alternative_hypotheses") + graph.add_edge("alternative_hypotheses", "final_report") + graph.add_edge("final_report", END) + + return graph.compile() + + +# ───────────────────────────────────────────── +# Entry Point +# ───────────────────────────────────────────── + +def run_diagnosis(case_id: str) -> dict: + """Run the full diagnosis pipeline for a given case ID.""" + graph = build_diagnosis_graph() + + initial_state: DiagnosisState = { + "case_id": case_id, + "case_data": {}, + "images_b64": [], + "initial_diagnosis": None, + "differential_diagnoses": None, + "bias_check_notes": None, + "alternative_hypotheses": None, + "final_report": None, + "messages": [], + } + + final_state = graph.invoke(initial_state) + + print("\n" + "="*70) + print(final_state["final_report"]) + print("="*70) + + return final_state + + +if __name__ == "__main__": + import sys + case_id = sys.argv[1] if len(sys.argv) > 1 else "e0a8f078-fb8f-5281-ae67-256e060d0ef0" + run_diagnosis(case_id) + + +# ───────────────────────────────────────────── +# Interactive Chat with Case Context +# ───────────────────────────────────────────── + +class CaseChat: + """ + Interactive chat session with patient case context. + Maintains conversation history for multi-turn dialogue. + """ + + def __init__(self, final_state: dict): + """ + Initialize chat with a completed diagnosis pipeline state. + + Args: + final_state: The dict returned by graph.invoke() containing + case_data, images_b64, initial_diagnosis, etc. + """ + self.state = final_state + self.case_data = final_state["case_data"] + self.images_b64 = final_state.get("images_b64", []) + self.diagnosis = final_state.get("initial_diagnosis", "{}") + self.differentials = final_state.get("differential_diagnoses", []) + self.alternatives = final_state.get("alternative_hypotheses", []) + self.history: List[dict] = [] # [{"role": "user/assistant", "content": "..."}] + + def _build_context_prompt(self) -> str: + """Build a compact case summary for chat context.""" + patient = self.case_data.get("patient", {}) + presentation = self.case_data.get("presentation", {}) + assessment = self.case_data.get("assessment", {}) + + return f"""You are a medical AI assistant discussing a specific patient case. Use the case context below to answer questions accurately. + +=== PATIENT CASE CONTEXT === +Patient: {patient.get('age_years', '?')} y/o {patient.get('sex', '?')} +Chief Complaint: {presentation.get('chief_complaint', 'N/A')} +HPI: {presentation.get('hpi', 'N/A')[:500]} +Comorbidities: {', '.join(patient.get('comorbidities', [])) or 'None'} +Medications: {', '.join(patient.get('medications', [])) or 'None'} + +=== AI DIAGNOSIS === +{self.diagnosis} + +=== GROUND TRUTH === +Actual Diagnosis: {assessment.get('diagnosis_primary', 'N/A')} +Differentials: {', '.join(assessment.get('differential', []))} + +Answer the user's questions about this case. Be specific, cite findings from the case, and explain your reasoning. If asked about treatment or prognosis, note that this is for educational discussion only.""" + + def ask(self, question: str, include_images: bool = False) -> str: + """ + Send a question about the case and get a response. + + Args: + question: User's question about the case + include_images: Whether to include case images in the context (slower but more accurate for imaging questions) + + Returns: + AI response string + """ + # Build messages with conversation history + system_prompt = self._build_context_prompt() + + # Construct content - text only or multimodal + if include_images and self.images_b64: + content = [{"type": "text", "text": question}] + for img in self.images_b64[:4]: # Limit to 4 images + content.append({ + "type": "image_url", + "image_url": {"url": f"data:{img['mime']};base64,{img['b64']}"} + }) + user_msg = HumanMessage(content=content) + else: + user_msg = HumanMessage(content=question) + + # Include recent history (last 6 turns to avoid context overflow) + messages = [] + for turn in self.history[-6:]: + if turn["role"] == "user": + messages.append(HumanMessage(content=turn["content"])) + else: + messages.append(AIMessage(content=turn["content"])) + messages.append(user_msg) + + # Call MedGemma + response = _call_medgemma(messages, system_prompt, max_new_tokens=1024) + + # Update history + self.history.append({"role": "user", "content": question}) + self.history.append({"role": "assistant", "content": response}) + + return response + + def clear_history(self): + """Reset conversation history.""" + self.history = [] + print("Chat history cleared.") + + def get_diagnosis_summary(self) -> dict: + """Return a structured summary of diagnoses for display.""" + try: + diag_data = json.loads(self.diagnosis) + except: + diag_data = {} + + # Support both schema variants: acute_complication (correct) or primary_diagnosis (legacy) + primary = diag_data.get("acute_complication", diag_data.get("primary_diagnosis", {})) + # Also check for underlying_etiology as a secondary primary source + if not primary.get("condition"): + primary = diag_data.get("underlying_etiology", {}) + diffs = diag_data.get("differentials", []) + + # Build ranked list + ranked = [] + if primary.get("condition"): + ranked.append({ + "rank": 1, + "condition": primary["condition"], + "confidence": primary.get("confidence", 0), + "reasoning": primary.get("reasoning", ""), + "type": "primary" + }) + + for i, d in enumerate(diffs, start=2): + ranked.append({ + "rank": i, + "condition": d.get("condition", "Unknown"), + "confidence": d.get("confidence", 0), + "reasoning": "; ".join(d.get("supporting_evidence", [])[:2]), + "type": "differential" + }) + + # Add alternatives + for alt in self.alternatives: + if isinstance(alt, dict) and alt.get("condition"): + ranked.append({ + "rank": len(ranked) + 1, + "condition": alt["condition"], + "confidence": alt.get("confidence", 0), + "reasoning": alt.get("why_missed", ""), + "type": "alternative", + "risk_if_missed": alt.get("risk_if_missed", "unknown") + }) + + return { + "primary": primary, + "ranked_list": ranked, + "ground_truth": self.case_data.get("assessment", {}).get("diagnosis_primary", "N/A") + } + + +def create_chat_session(final_state: dict) -> CaseChat: + """ + Create an interactive chat session from a completed diagnosis pipeline. + + Usage: + final_state = graph.invoke(initial_state) + chat = create_chat_session(final_state) + response = chat.ask("Why did you rule out tuberculosis?") + """ + return CaseChat(final_state) diff --git a/data_pipeline/schema.json b/data_pipeline/schema.json new file mode 100644 index 0000000..c500ad7 --- /dev/null +++ b/data_pipeline/schema.json @@ -0,0 +1,130 @@ +{ + "profile_id": "59ba3305-ec05-5873-9a45-4a766c6609fc:9a7b2e88-0b85-5e88-a5d7-77bd298ee714", + "case_id": "59ba3305-ec05-5873-9a45-4a766c6609fc", + "image_id": "9a7b2e88-0b85-5e88-a5d7-77bd298ee714", + + "patient": { + "age_years": 69, + "sex": "female", + "immunocompromised": "no", + "weight_kg": null, + "comorbidities": [ + "hypertension", + "type 2 diabetes", + "atrial fibrillation" + ], + "medications": [ + "Metoprolol 25 mg", + "Rivaroxiban" + ], + "allergies": null + }, + + "presentation": { + "chief_complaint": "scheduled for living donor liver transplantation (LDLT)", + "symptom_duration": null, + "hpi": "A 69-year-old female with a history of hypertension, type 2 diabetes, atrial fibrillation and hepatocellular carcinoma with liver cirrhosis was scheduled for living donor liver transplantation (LDLT). Pre-operative assessment revealed no cardio-respiratory symptoms, except an irregularly irregular pulse. Chest X-ray showed an incidental scimitar sign along with a smaller right lung shadow, and a marked displacement of the mediastinum to the right side.", + "pmh": "hypertension, type 2 diabetes, atrial fibrillation (AF), and interval progression of hepatocellular carcinoma with liver cirrhosis" + }, + + "study": { + "modality": "CXR", + "body_region": "thorax", + "view_position": "PA", + "radiology_region": "thorax", + "caption": "Chest X-ray showing Scimitar sign.", + "image_type": "radiology", + "image_subtype": "x_ray", + "image_url": "https://storage.googleapis.com/casetwin-xrays/chest-xrays/PMC10077807_SJA-17-101-g001_undivided_1_1.webp", + "storage_path": "gs://casetwin-xrays/chest-xrays/PMC10077807_SJA-17-101-g001_undivided_1_1.webp" + }, + + "assessment": { + "diagnosis_primary": "scimitar syndrome", + "suspected_primary": [ + "scimitar syndrome", + "hepatocellular carcinoma", + "liver cirrhosis" + ], + "differential": [], + "urgency": "routine", + "infectious_concern": "no", + "icu_candidate": "yes" + }, + + "findings": { + "lungs": { + "consolidation_present": "no", + "consolidation_locations": [], + "consolidation_extent": "unknown", + "atelectasis_present": "no", + "atelectasis_locations": [], + "edema_present": "no", + "edema_pattern": "unknown" + }, + "pleura": { + "effusion_present": "no", + "effusion_side": "unknown", + "effusion_size": "unknown", + "pneumothorax_present": "no", + "pneumothorax_side": "unknown" + }, + "cardiomediastinal": { + "cardiomegaly": "yes", + "mediastinal_widening": "no" + }, + "devices": { + "lines_tubes_present": "yes", + "device_list": [ + "ett", + "radial arterial line", + "femoral arterial line", + "pulmonary artery catheter" + ] + } + }, + + "summary": { + "one_liner": "69-year-old female with scimitar syndrome, hypertension, type 2 diabetes, atrial fibrillation and hepatocellular carcinoma with liver cirrhosis who underwent living donor liver transplantation.", + "key_points": [ + "Incidental scimitar sign on chest X-ray", + "Hypoplastic right lung with venous drainage to the IVC", + "Cardiomegaly with enlarged pulmonary vessels and bicaval dilatation" + ], + "red_flags": [ + "Potential air embolism in the presence of PFO", + "Hemodynamic instability during neohepatic phase" + ] + }, + + "outcome": { + "success": "yes", + "detail": "The patient tolerated the liver transplantation well and was discharged from ICU on the 4th post-op day." + }, + + "provenance": { + "dataset_name": "multicare", + "pmc_id": "PMC10077807", + "pmid": "37032661", + "doi": "10.4103/sja.sja_553_22", + "article_title": "Anesthesia management of living donor liver transplantation in a patient with scimitar syndrome", + "journal": "Saudi J Anaesth", + "year": 2023, + "authors": [ + "Muhammad Shabbir", + "Amer Majeed", + "Mudassir A Baig", + "Matloob A Shajar", + "Tahir Iqbal" + ], + "license": "CC BY-NC-SA", + "source_url": "https://pubmed.ncbi.nlm.nih.gov/37032661/" + }, + + "tags": { + "ml_labels": ["thorax", "radiology", "frontal", "x_ray"], + "gt_labels": ["radiology", "thorax", "x_ray"], + "keywords": ["aberrant pulmonary veins", "scimitar", "liver transplantation"], + "mesh_terms": ["Case Reports"] + } +} diff --git a/data_pipeline/transform_to_schema.py b/data_pipeline/transform_to_schema.py new file mode 100644 index 0000000..456f711 --- /dev/null +++ b/data_pipeline/transform_to_schema.py @@ -0,0 +1,322 @@ +""" +Produces two dataset variants from cxr_schema_dataset/dataset.json: + +1. dataset.json (per-case) + One record per case (750). `images` contains ALL images from the case + (CXR, CT, pathology, photos, etc.), each tagged with `is_chest_xray`. + +2. dataset_cxr_primary.json (per-CXR-image, CXR-primary) + One record per CXR image (1 002). The primary CXR fields mirror + schema.json. A `related_images` array holds the non-CXR images from + the same case for multi-modal context. + +Case-level clinical fields come from the Gemini-enriched flat records +already stored in cxr_schema_dataset/dataset.json (the per-case version +is written first and then used as input for the second variant). +""" + +import json +import ast +import re +import uuid +from pathlib import Path + +FLAT_PATH = Path("cxr_schema_dataset/dataset.json") # Gemini-enriched CXR flat records +SOURCE_PATH = Path("cxr_full_dataset/dataset.json") # original all-image source +OUTPUT_PATH = Path("cxr_schema_dataset/dataset.json") +OUTPUT_CXR_PRIMARY = Path("cxr_schema_dataset/dataset_cxr_primary.json") + + +def parse_label_list(raw) -> list: + """Convert a stringified Python list to a real list, or return as-is.""" + if raw is None: + return [] + if isinstance(raw, list): + return raw + if isinstance(raw, str): + try: + result = ast.literal_eval(raw) + return result if isinstance(result, list) else [] + except Exception: + return [] + return [] + + +def safe_int(val): + try: + return int(val) + except (TypeError, ValueError): + return val + + +def split_medications(meds_raw) -> list: + """Split a comma-separated medications string into a list.""" + if isinstance(meds_raw, list): + return meds_raw + if isinstance(meds_raw, str) and meds_raw: + return [m.strip() for m in re.split(r",(?![^()]*\))", meds_raw) if m.strip()] + return [] + + +VIEW_MAP = { + "frontal": "PA", + "sagittal": "LATERAL", + "axial": "UNKNOWN", + "oblique": "UNKNOWN", + None: "UNKNOWN", +} + +def normalise_view(v): + if v is None: + return "UNKNOWN" + return VIEW_MAP.get(v.lower(), v.upper()) + + +def build_image_entry(img: dict) -> dict: + """Build a single image object from a cxr_full_dataset image record.""" + file_key = img.get("file_id") or img.get("file") or "" + image_id = str(uuid.uuid5(uuid.NAMESPACE_URL, file_key)) + return { + "image_id": image_id, + "local_image_path": img.get("local_image_path", ""), + "file_type": Path(img.get("file", "")).suffix.lstrip(".") or "webp", + "image_type": img.get("image_type", ""), + "image_subtype": img.get("image_subtype", ""), + "is_chest_xray": img.get("is_chest_xray", False), + "view_position": normalise_view(img.get("radiology_view")), + "radiology_region": img.get("radiology_region"), + "caption": img.get("caption", ""), + "text_references": img.get("text_references", []), + "ml_labels": parse_label_list(img.get("ml_labels")), + "gt_labels": parse_label_list(img.get("gt_labels")), + } + + +def build_case_record(flat: dict, src_images: list) -> dict: + """ + Merge Gemini-enriched case fields (from flat CXR record) with the full + image list from the original source. + + flat – one representative flat record for this case (all case-level + Gemini fields are identical across CXR records of the same case) + src_images – all images for this case from cxr_full_dataset/dataset.json + """ + pat = flat.get("patient") or {} + pres = flat.get("presentation") or {} + study = flat.get("study") or {} + asmt = flat.get("assessment") or {} + find = flat.get("findings") or {} + summ = flat.get("summary") or {} + oc = flat.get("outcome") or {} + prov = flat.get("provenance") or {} + tags = flat.get("tags") or {} + + return { + "case_id": flat.get("case_id"), + + "patient": { + "age_years": pat.get("age_years"), + "sex": pat.get("sex"), + "immunocompromised": pat.get("immunocompromised"), + "weight_kg": pat.get("weight_kg"), + "comorbidities": pat.get("comorbidities") or [], + "medications": pat.get("medications") or [], + "allergies": pat.get("allergies"), + }, + + "presentation": { + "chief_complaint": pres.get("chief_complaint"), + "symptom_duration": pres.get("symptom_duration"), + "hpi": pres.get("hpi"), + "pmh": pres.get("pmh"), + }, + + "study": { + "modality": "CXR", + "body_region": study.get("body_region") or "thorax", + "view_position": study.get("view_position"), + }, + + "images": [build_image_entry(img) for img in src_images], + + "assessment": { + "diagnosis_primary": asmt.get("diagnosis_primary"), + "suspected_primary": asmt.get("suspected_primary") or [], + "differential": asmt.get("differential") or [], + "urgency": asmt.get("urgency"), + "infectious_concern": asmt.get("infectious_concern"), + "icu_candidate": asmt.get("icu_candidate"), + }, + + "findings": find, + + "summary": { + "one_liner": summ.get("one_liner"), + "key_points": summ.get("key_points") or [], + "red_flags": summ.get("red_flags") or [], + }, + + "outcome": { + "success": oc.get("success"), + "detail": oc.get("detail"), + }, + + "provenance": { + "dataset_name": prov.get("dataset_name"), + "pmc_id": prov.get("pmc_id"), + "pmid": prov.get("pmid"), + "doi": prov.get("doi"), + "article_title": prov.get("article_title"), + "journal": prov.get("journal"), + "year": prov.get("year"), + "authors": prov.get("authors") or [], + "license": prov.get("license"), + "source_url": prov.get("source_url"), + }, + + "tags": { + "ml_labels": tags.get("ml_labels") or [], + "gt_labels": tags.get("gt_labels") or [], + "keywords": tags.get("keywords") or [], + "mesh_terms": tags.get("mesh_terms") or [], + }, + + "embeddings": [], + } + + +def build_cxr_primary_record(case_record: dict, non_cxr_images: list) -> list: + """ + Given one per-case record (already in final schema) and the list of + non-CXR source images, return a list of flat per-CXR-image records. + Each has the CXR as the primary image (top-level `study` fields) and + all non-CXR images nested under `related_images`. + """ + related = [ + { + "image_id": img["image_id"], + "local_image_path": img["local_image_path"], + "file_type": img["file_type"], + "image_type": img["image_type"], + "image_subtype": img["image_subtype"], + "view_position": img["view_position"], + "radiology_region": img["radiology_region"], + "caption": img["caption"], + "ml_labels": img["ml_labels"], + "gt_labels": img["gt_labels"], + } + for img in non_cxr_images + ] + + records = [] + for img in case_record["images"]: + if not img["is_chest_xray"]: + continue + records.append({ + "profile_id": f"{case_record['case_id']}:{img['image_id']}", + "case_id": case_record["case_id"], + "image_id": img["image_id"], + + "patient": case_record["patient"], + "presentation": case_record["presentation"], + + "study": { + "modality": img["image_subtype"], + "body_region": img["radiology_region"] or case_record["study"].get("body_region"), + "view_position": img["view_position"], + "radiology_region": img["radiology_region"], + "caption": img["caption"], + "image_type": img["image_type"], + "image_subtype": img["image_subtype"], + "storage_path": img["local_image_path"], + }, + + "related_images": related, + + "assessment": case_record["assessment"], + "findings": case_record["findings"], + "summary": case_record["summary"], + "outcome": case_record["outcome"], + "provenance": case_record["provenance"], + + "tags": { + "ml_labels": img["ml_labels"], + "gt_labels": img["gt_labels"], + "keywords": case_record["tags"].get("keywords") or [], + "mesh_terms": case_record["tags"].get("mesh_terms") or [], + }, + + "embeddings": [], + }) + return records + + +def main(): + # ── Load Gemini-enriched flat records (CXR only) and group by PMC ID ── + print(f"Loading flat Gemini-enriched records from {FLAT_PATH} ...") + with open(FLAT_PATH, "r") as f: + flat_records = json.load(f) + + # Index: pmc_id → first flat record for that case (case-level fields are + # identical across all CXR records belonging to the same case) + flat_by_pmc: dict[str, dict] = {} + for rec in flat_records: + pmc = rec.get("provenance", {}).get("pmc_id") + if pmc and pmc not in flat_by_pmc: + flat_by_pmc[pmc] = rec + + # ── Load original source (all images) ── + print(f"Loading original full-image source from {SOURCE_PATH} ...") + with open(SOURCE_PATH, "r") as f: + source_cases = json.load(f) + + # ── Merge ── + output = [] + missing = [] + for src_case in source_cases: + pmc_id = src_case.get("pmc_id") + flat = flat_by_pmc.get(pmc_id) + if flat is None: + missing.append(pmc_id) + continue + record = build_case_record(flat, src_case.get("images", [])) + output.append(record) + + total_images = sum(len(r["images"]) for r in output) + cxr_images = sum(sum(1 for img in r["images"] if img["is_chest_xray"]) for r in output) + + print(f"Cases produced : {len(output)}") + print(f"Cases missing : {len(missing)}") + print(f"Total images : {total_images} ({cxr_images} CXR, {total_images - cxr_images} non-CXR)") + + print(f"Writing {OUTPUT_PATH} ...") + with open(OUTPUT_PATH, "w") as f: + json.dump(output, f, indent=2, ensure_ascii=False) + + # ── Build CXR-primary variant ── + print("\nBuilding CXR-primary variant ...") + + # Index source cases by pmc_id for quick lookup of non-CXR images + src_by_pmc = {c["pmc_id"]: c for c in source_cases} + + cxr_primary_records = [] + for case_record in output: + pmc_id = case_record["provenance"].get("pmc_id") + src_case = src_by_pmc.get(pmc_id, {}) + non_cxr = [ + build_image_entry(img) + for img in src_case.get("images", []) + if not img.get("is_chest_xray") + ] + cxr_primary_records.extend(build_cxr_primary_record(case_record, non_cxr)) + + print(f"CXR-primary records: {len(cxr_primary_records)}") + print(f"Writing {OUTPUT_CXR_PRIMARY} ...") + with open(OUTPUT_CXR_PRIMARY, "w") as f: + json.dump(cxr_primary_records, f, indent=2, ensure_ascii=False) + + print("Done.") + + +if __name__ == "__main__": + main()