""" export_split_by_pathology.py ---------------------------- Tải 1 split (mặc định: test) của MIMIC-CXR_resized từ HF, giải nén tar shards, rồi đổ ảnh vào thư mục theo 14 nhãn bệnh lý CheXpert. Vì sao phải làm thế này: - Trên HF, MIMIC-CXR_resized lưu ảnh dưới dạng tar shards (cxr-0000.tar, ...), KHÔNG tách theo split. Ảnh của 1 split nằm rải khắp các shard. - manifest_{train,val,test}.csv mới là nơi biết ảnh nào thuộc split nào, và chứa 14 cột `chex_` (giá trị U-MultiClass: 1/0/-1/blank). - Mỗi study là MULTI-LABEL: 1 ảnh có thể dương tính nhiều bệnh -> được copy vào NHIỀU thư mục (mỗi nhãn positive 1 thư mục). Quy ước nhãn (giống mimic_cxr_resized_builder._row_to_pnu): "1" / "1.0" -> positive -> bỏ vào thư mục / "-1" / "-1.0" -> uncertain -> tuỳ --uncertain "0"/"0.0"/blank -> negative -> bỏ qua Cấu trúc output (mặc định --uncertain separate): OUT/ Cardiomegaly/.jpg Pleural_Effusion/.jpg No_Finding/.jpg _uncertain/Atelectasis/.jpg (nếu --uncertain separate) _summary.csv (đếm ảnh mỗi nhãn) Chạy LOCAL (đã có data) hoặc trên Colab/Kaggle (tự tải từ HF): # tải từ HF rồi export test python data/export_split_by_pathology.py --out ./test_by_pathology # đã giải nén sẵn cây files/ ở đâu đó -> bỏ qua tải/giải nén python data/export_split_by_pathology.py --out ./test_by_pathology \ --extracted_root /content/data/MIMIC-CXR_resized \ --manifest /content/data/MIMIC-CXR_resized/manifest_test.csv # export cả split val, dùng hardlink cho đỡ tốn ổ python data/export_split_by_pathology.py --split val --link hardlink """ from __future__ import annotations import argparse import csv import os import sys import tarfile from collections import defaultdict from pathlib import Path # 14 nhãn CheXpert — single source of truth. try: from model.chexpert_classifier import PATHOLOGIES except Exception: # fallback nếu chạy ngoài project (giữ đúng thứ tự). PATHOLOGIES = [ "No Finding", "Enlarged Cardiomediastinum", "Cardiomegaly", "Lung Opacity", "Lung Lesion", "Edema", "Consolidation", "Pneumonia", "Atelectasis", "Pneumothorax", "Pleural Effusion", "Pleural Other", "Fracture", "Support Devices", ] _POS = {"1", "1.0"} _UNC = {"-1", "-1.0"} def _safe(name: str) -> str: """Tên thư mục an toàn: 'Pleural Effusion' -> 'Pleural_Effusion'.""" return name.replace(" ", "_") def _norm(p: str) -> str: """Chuẩn hoá path để so khớp tar member name <-> manifest image_relpath.""" return p.replace("\\", "/").lstrip("/") # ── Phase 1: tải từ HF (manifest + shards) ────────────────────────────────── def download_from_hf(repo_id: str, split: str, work: Path) -> tuple[Path, list[Path]]: """Tải manifest_.csv + toàn bộ tar shards về `work`. Trả về (manifest_path, [shard_paths]).""" from huggingface_hub import snapshot_download manifest_name = {"train": "manifest_train.csv", "val": "manifest_val.csv", "validate": "manifest_val.csv", "test": "manifest_test.csv"}[split] print(f"[download] snapshot_download {repo_id}:MIMIC-CXR_resized " f"(manifest + shards) -> {work}") snapshot_download( repo_id=repo_id, repo_type="dataset", local_dir=str(work), allow_patterns=[ f"MIMIC-CXR_resized/{manifest_name}", "MIMIC-CXR_resized/shards/*.tar", ], ) mr = work / "MIMIC-CXR_resized" manifest = mr / manifest_name shards = sorted((mr / "shards").glob("*.tar")) if not manifest.is_file(): sys.exit(f"ERROR: không thấy manifest sau khi tải: {manifest}") if not shards: sys.exit(f"ERROR: không thấy tar shard nào dưới {mr/'shards'}") print(f"[download] manifest={manifest.name} shards={len(shards)}") return manifest, shards # ── Phase 2: đọc manifest -> map ảnh test -> nhãn ─────────────────────────── def load_label_map(manifest: Path): """Trả về dict: image_relpath(norm) -> {'pos': set, 'unc': set, 'report': str|None}.""" label_map: dict[str, dict] = {} missing_cols = None with open(manifest, encoding="utf-8", newline="") as f: reader = csv.DictReader(f) cols = reader.fieldnames or [] chex_cols = {p: f"chex_{p}" for p in PATHOLOGIES if f"chex_{p}" in cols} missing_cols = [p for p in PATHOLOGIES if f"chex_{p}" not in cols] rel_col = "image_relpath" if "image_relpath" in cols else None if rel_col is None: sys.exit(f"ERROR: manifest thiếu cột 'image_relpath'. Có: {cols}") has_report = "report_relpath" in cols for row in reader: rel = _norm(str(row[rel_col]).strip()) pos, unc = set(), set() for path, col in chex_cols.items(): v = str(row.get(col, "")).strip() if v in _POS: pos.add(path) elif v in _UNC: unc.add(path) rep = _norm(str(row["report_relpath"]).strip()) if has_report else None label_map[rel] = {"pos": pos, "unc": unc, "report": rep or None} if missing_cols: print(f"[labels] CẢNH BÁO: manifest thiếu cột cho: {missing_cols}") if not has_report: print("[labels] CẢNH BÁO: manifest không có cột 'report_relpath' → bỏ qua report") print(f"[labels] {len(label_map):,} ảnh trong manifest") return label_map def gather_reports(shards: list[Path], report_set: set) -> dict: """Pass phụ: rút text của các report cần dùng từ tar (report nằm rải, gom 1 lượt). Report là file .txt nhỏ nên giữ trong RAM thoải mái.""" reports: dict[str, bytes] = {} if not report_set: return reports for shard in shards: with tarfile.open(shard, "r") as tf: for m in tf: if not m.isfile(): continue name = _norm(m.name) if name in report_set and name not in reports: reports[name] = tf.extractfile(m).read() print(f"[reports] rút được {len(reports):,} / {len(report_set):,} report") return reports # ── Phase 3: rút ảnh từ tar -> thư mục theo nhãn ──────────────────────────── def _place(data: bytes, dicom_name: str, paths: set, base: Path, counts: defaultdict, link_mode: str, report: bytes | None = None): """Ghi 1 ảnh (và report cùng tên .txt nếu có) vào nhiều thư mục nhãn.""" txt_name = Path(dicom_name).stem + ".txt" first_written: Path | None = None for lab in paths: d = base / _safe(lab) d.mkdir(parents=True, exist_ok=True) dst = d / dicom_name counts[lab] += 1 # report .txt đặt cạnh ảnh, cùng tên if report is not None: (d / txt_name).write_bytes(report) if dst.exists(): continue if link_mode == "copy" or first_written is None: dst.write_bytes(data) first_written = dst else: try: if link_mode == "hardlink": os.link(first_written, dst) else: # symlink os.symlink(os.path.abspath(first_written), dst) except OSError: dst.write_bytes(data) # fallback nếu FS không hỗ trợ link def export(shards: list[Path], label_map: dict, out: Path, uncertain: str, link_mode: str, with_report: bool = True): out.mkdir(parents=True, exist_ok=True) unc_base = out / "_uncertain" test_set = set(label_map.keys()) # Gom report cần dùng (1 pass phụ qua tar) trước khi rút ảnh. reports: dict = {} if with_report: report_set = {label_map[k]["report"] for k in test_set if label_map[k].get("report")} reports = gather_reports(shards, report_set) counts_pos: defaultdict = defaultdict(int) counts_unc: defaultdict = defaultdict(int) n_imgs = 0 n_no_report = 0 seen: set[str] = set() for si, shard in enumerate(shards, 1): print(f"[extract] [{si}/{len(shards)}] {shard.name}") with tarfile.open(shard, "r") as tf: for m in tf: if not m.isfile(): continue name = _norm(m.name) if name not in test_set: continue seen.add(name) ent = label_map[name] pos, unc = ent["pos"], ent["unc"] if not pos and not (uncertain != "skip" and unc): # không có nhãn positive (toàn negative) -> bỏ qua if not pos: continue data = tf.extractfile(m).read() dicom_name = Path(name).name rep = reports.get(ent.get("report")) if with_report else None if with_report and rep is None: n_no_report += 1 n_imgs += 1 if pos: _place(data, dicom_name, pos, out, counts_pos, link_mode, rep) if unc and uncertain != "skip": if uncertain == "merge": _place(data, dicom_name, unc, out, counts_unc, link_mode, rep) else: # separate _place(data, dicom_name, unc, unc_base, counts_unc, link_mode, rep) if with_report and n_no_report: print(f"[reports] CẢNH BÁO: {n_no_report:,} ảnh không tìm thấy report → chỉ có .jpg") missing = test_set - seen print(f"\n[done] ảnh test rút được: {n_imgs:,} / {len(test_set):,} trong manifest") if missing: print(f"[done] CẢNH BÁO: {len(missing):,} ảnh trong manifest không thấy trong shard " f"(ví dụ: {list(missing)[:3]})") print(f"[done] ảnh toàn-negative (không có positive): bỏ qua") # _summary.csv summ = out / "_summary.csv" with open(summ, "w", encoding="utf-8", newline="") as f: w = csv.writer(f) w.writerow(["pathology", "positive_images", "uncertain_images"]) for p in PATHOLOGIES: w.writerow([p, counts_pos.get(p, 0), counts_unc.get(p, 0)]) print(f"[done] thống kê -> {summ}") print("\n Nhãn positive uncertain") for p in PATHOLOGIES: print(f" {p:28s} {counts_pos.get(p,0):8d} {counts_unc.get(p,0):8d}") # ── CLI ───────────────────────────────────────────────────────────────────── def main(): ap = argparse.ArgumentParser( description="Export 1 split của MIMIC-CXR_resized thành thư mục theo 14 nhãn bệnh lý.") ap.add_argument("--out", required=True, help="Thư mục output.") ap.add_argument("--split", default="test", choices=["train", "val", "validate", "test"]) ap.add_argument("--repo_id", default="hieu3636/cxr-vlm-data") ap.add_argument("--work", default="./_hf_resized_dl", help="Thư mục cache tải từ HF (khi không dùng --extracted_root).") ap.add_argument("--extracted_root", default=None, help="Nếu đã có shards giải nén/tar sẵn ở local: trỏ tới thư mục " "MIMIC-CXR_resized (chứa shards/*.tar). Bỏ qua bước tải HF.") ap.add_argument("--manifest", default=None, help="Đường dẫn manifest_.csv (mặc định lấy trong dữ liệu đã tải).") ap.add_argument("--uncertain", default="separate", choices=["separate", "merge", "skip"], help="Xử lý nhãn uncertain: separate=thư mục _uncertain/

; " "merge=gộp chung

; skip=bỏ. Mặc định separate.") ap.add_argument("--link", default="copy", choices=["copy", "hardlink", "symlink"], help="copy (an toàn nhất) | hardlink/symlink (tiết kiệm ổ khi 1 ảnh nhiều nhãn).") ap.add_argument("--no_report", action="store_true", help="Không ghi report .txt cạnh ảnh (mặc định CÓ ghi).") a = ap.parse_args() out = Path(a.out) # 1) Lấy manifest + shards if a.extracted_root: mr = Path(a.extracted_root) shards = sorted((mr / "shards").glob("*.tar")) or sorted(mr.glob("*.tar")) if not shards: sys.exit(f"ERROR: không thấy *.tar dưới {mr} hoặc {mr/'shards'}") if a.manifest: manifest = Path(a.manifest) else: mname = {"train": "manifest_train.csv", "val": "manifest_val.csv", "validate": "manifest_val.csv", "test": "manifest_test.csv"}[a.split] manifest = mr / mname if not manifest.is_file(): sys.exit(f"ERROR: không thấy manifest: {manifest}") print(f"[local] manifest={manifest} shards={len(shards)}") else: manifest, shards = download_from_hf(a.repo_id, a.split, Path(a.work)) if a.manifest: manifest = Path(a.manifest) # 2) Đọc nhãn label_map = load_label_map(manifest) # 3) Rút ảnh (+ report) -> thư mục nhãn export(shards, label_map, out, a.uncertain, a.link, with_report=not a.no_report) print(f"\nXong. Output: {out.resolve()}") if __name__ == "__main__": main()