data/build_test100_from_hf.py

"""
build_test100_from_hf.py
-------------------------
Dựng package TEST-100 (100 studies có bệnh + VQA, đã lọc trong
data/test100_manifest.json) bằng cách TRÍCH ảnh từ shard RESIZED trên HF
(hieu3636/cxr-vlm-data : MIMIC-CXR_resized/shards/cxr-*.tar) — KHÔNG cần PhysioNet.

Ảnh trong shard resized đã 518px (đúng như lúc train). Tên trong tar = "Study_N.jpg",
map study_id ↔ Study_N lấy từ MIMIC-CXR_resized/manifest_test.csv (nguồn chuẩn).

Output (giống mini MIMIC-CXR_resized, dùng lại được builder/evaluate):
    <OUT>/
      files/pXX/pSUBJ/sSTUDY/<dicom>.jpg     ← 100 ảnh resized
      files/pXX/pSUBJ/sSTUDY.txt             ← report (findings+impression GT)
      manifest_test.csv                       ← cột chuẩn resized
      vqa/vqa_test.json                       ← câu VQA
      test100.json                            ← GT gom theo study

Chạy:
    python data/build_test100_from_hf.py --out "D:/USTH/KLTN/MIMIC-CXR_test100"
Repo data public → không cần token. Nếu private: --hf_token <token> hoặc env HF_TOKEN.
"""
import argparse, csv, io, json, os, sys, tarfile
from collections import Counter
from pathlib import Path

try:
    sys.stdout.reconfigure(encoding="utf-8")   # Windows console hay là cp1252
except Exception:
    pass

LABELS = ["Atelectasis","Cardiomegaly","Consolidation","Edema","Enlarged Cardiomediastinum",
  "Fracture","Lung Lesion","Lung Opacity","No Finding","Pleural Effusion","Pleural Other",
  "Pneumonia","Pneumothorax","Support Devices"]

HF_REPO   = "hieu3636/cxr-vlm-data"
RESIZED   = "MIMIC-CXR_resized"
SHARDS    = [f"{RESIZED}/shards/cxr-0000.tar", f"{RESIZED}/shards/cxr-0001.tar"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", default=str(Path(__file__).parent / "test100_manifest.json"))
    ap.add_argument("--out", default=r"D:\USTH\KLTN\MIMIC-CXR_test100")
    ap.add_argument("--hf_token", default=os.environ.get("HF_TOKEN"))
    args = ap.parse_args()

    from huggingface_hub import hf_hub_download

    studies = json.load(open(args.manifest, encoding="utf-8"))
    OUT = Path(args.out); OUT.mkdir(parents=True, exist_ok=True)
    print(f"studies: {len(studies)} | out: {OUT}")

    # ── 1) map study_id → "Study_N.jpg" từ manifest_test.csv của bản resized ──
    mpath = hf_hub_download(HF_REPO, repo_type="dataset",
                            filename=f"{RESIZED}/manifest_test.csv", token=args.hf_token)
    sid2member = {}
    for r in csv.DictReader(io.open(mpath, encoding="utf-8")):
        sid2member[str(r["study_id"])] = r["image_relpath"]   # "Study_N.jpg"
    want = {}   # tar member name -> our study dict
    miss_map = []
    for s in studies:
        m = sid2member.get(str(s["study_id"]))
        if m is None:
            miss_map.append(s["study_name"]); continue
        want[m] = s
    print(f"map study_id→Study_N.jpg: matched {len(want)}/{len(studies)}"
          + (f" | KHÔNG map được: {miss_map[:10]}" if miss_map else ""))
    if miss_map:
        print("  (study_id không có trong resized manifest_test — bỏ qua)")

    # ── 2) tải 2 shard resized (HF cache lại, lần sau khỏi tải) ───────────────
    tar_paths = []
    for sh in SHARDS:
        print(f"  tải {sh} (cache nếu đã có)…", flush=True)
        tar_paths.append(hf_hub_download(HF_REPO, repo_type="dataset",
                                         filename=sh, token=args.hf_token))

    # ── 3) trích ảnh cần thiết, ghi vào OUT theo image_relpath gốc của mình ───
    found = {}
    for tp in tar_paths:
        with tarfile.open(tp, "r") as tf:
            for mem in tf:
                if not mem.isfile():
                    continue
                name = os.path.basename(mem.name)            # "Study_N.jpg"
                s = want.get(name)
                if s is None:
                    continue
                dst = OUT / s["image_relpath"]               # files/pXX/.../<dicom>.jpg
                dst.parent.mkdir(parents=True, exist_ok=True)
                with tf.extractfile(mem) as src, open(dst, "wb") as out:
                    out.write(src.read())
                found[name] = s
        print(f"  {Path(tp).name}: đã trích {len(found)}/{len(want)} ảnh", flush=True)
    missing = [s["study_name"] for m, s in want.items() if m not in found]
    if missing:
        print(f"!! còn thiếu {len(missing)} ảnh trong shard: {missing[:10]}")

    # ── 4) manifest_test.csv (cột chuẩn resized) ─────────────────────────────
    man_cols = (["study_name","split","subject_id","study_id","subset","dicom_id",
                 "image_filename","view","image_relpath","report_relpath","jpg_url","has_vqa"]
                + [f"chex_{l}" for l in LABELS])
    with open(OUT / "manifest_test.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=man_cols); w.writeheader()
        for s in studies:
            row = {"study_name":s["study_name"],"split":"test","subject_id":s["subject_id"],
                   "study_id":s["study_id"],"subset":s["subset"],"dicom_id":s["dicom_id"],
                   "image_filename":f"{s['dicom_id']}.jpg","view":s["view"],
                   "image_relpath":s["image_relpath"],"report_relpath":s["report_relpath"],
                   "jpg_url":s["jpg_url"],"has_vqa":True}
            row.update({f"chex_{l}": s["chex"].get(l, "") for l in LABELS})
            w.writerow(row)

    # ── 5) vqa/vqa_test.json ─────────────────────────────────────────────────
    (OUT / "vqa").mkdir(exist_ok=True)
    vqa_rows = []
    for s in studies:
        for q in s["vqa"]:
            vqa_rows.append({"study_name":s["study_name"],"image_path":s["image_relpath"],
                "question":q["question"],"answer":[a.strip() for a in q["answer"].split(",")],
                "semantic_type":q.get("semantic_type"),"content_type":q.get("content_type"),
                "subject_id":s["subject_id"],"study_id":s["study_id"],"image_id":q.get("image_id")})
    json.dump(vqa_rows, open(OUT / "vqa" / "vqa_test.json", "w", encoding="utf-8"), indent=1)

    # ── 6) test100.json (GT gom theo study) ──────────────────────────────────
    preview = [{"study_name":s["study_name"],"study_id":s["study_id"],
                "image_path":s["image_relpath"],"positive_labels":s["positive_labels"],
                "structured_findings":s["structured_findings"],
                "findings":s["findings"],"impression":s["impression"],
                "vqa":[{"question":q["question"],"answer":q["answer"]} for q in s["vqa"]]}
               for s in studies]
    json.dump(preview, open(OUT / "test100.json", "w", encoding="utf-8"),
              ensure_ascii=False, indent=1)

    # ── 7) reports (findings+impression đã parse) ────────────────────────────
    for s in studies:
        rp = OUT / s["report_relpath"]
        rp.parent.mkdir(parents=True, exist_ok=True)
        rp.write_text(f"FINDINGS: {s['findings']}\n\nIMPRESSION: {s['impression']}\n",
                      encoding="utf-8")

    n_img = sum(1 for s in studies if (OUT / s["image_relpath"]).exists())
    print("\n=== PACKAGE DONE ===")
    print(f"  ảnh        : {n_img}/{len(studies)}")
    print(f"  manifest   : {len(studies)} rows")
    print(f"  vqa        : {len(vqa_rows)} câu")
    print(f"  test100    : {len(preview)} studies")
    print(f"  package    : {OUT}")


if __name__ == "__main__":
    main()