#!/usr/bin/env python3 """ ============================================================================== Anatomy-Aware DRR Pipeline — Phase 1 & Phase 2 Batch Processing ============================================================================== Phase 1: Automated Table Removal (3D) - Run TotalSegmentator "body" task on raw CT - Replace voxels outside body mask with -1000 HU - Save CT_clean.nii.gz Phase 2: Full 3D Anatomy Segmentation - Run TotalSegmentator "total" task on CT_clean.nii.gz - Produce 118-class multilabel mask.nii.gz Follows the directory layout specified in ANATOMY_AWARE_DRR_PIPELINE_DESIGN.md. Supports resumption: already-processed cases are skipped automatically. Usage: python run_phase1_phase2.py [--device gpu] [--fast] [--dry-run] """ import argparse import csv import datetime import json import logging import os import sys import time import traceback from pathlib import Path from typing import Dict, List, Optional, Tuple import matplotlib matplotlib.use("Agg") # Non-interactive backend; must be set before pyplot import import matplotlib.pyplot as plt import nibabel as nib import numpy as np from tqdm import tqdm # ============================================================================ # Constants # ============================================================================ # Root paths (relative to this script) SCRIPT_DIR = Path(__file__).resolve().parent RAW_CT_DIR = SCRIPT_DIR / "data" / "lidc_nii_data_test" PIPELINE_DIR = SCRIPT_DIR / "data" / "lidc_TotalSeg_test" MANIFEST_DIR = PIPELINE_DIR / "manifests" # TotalSegmentator v2 "total" label dictionary (118 classes) TOTALSEG_V2_LABELS: Dict[int, str] = { 1: "spleen", 2: "kidney_right", 3: "kidney_left", 4: "gallbladder", 5: "liver", 6: "stomach", 7: "pancreas", 8: "adrenal_gland_right", 9: "adrenal_gland_left", 10: "lung_upper_lobe_left", 11: "lung_lower_lobe_left", 12: "lung_upper_lobe_right", 13: "lung_middle_lobe_right", 14: "lung_lower_lobe_right", 15: "esophagus", 16: "trachea", 17: "thyroid_gland", 18: "small_bowel", 19: "duodenum", 20: "colon", 21: "urinary_bladder", 22: "prostate", 23: "kidney_cyst_left", 24: "kidney_cyst_right", 25: "sacrum", 26: "vertebrae_S1", 27: "vertebrae_L5", 28: "vertebrae_L4", 29: "vertebrae_L3", 30: "vertebrae_L2", 31: "vertebrae_L1", 32: "vertebrae_T12", 33: "vertebrae_T11", 34: "vertebrae_T10", 35: "vertebrae_T9", 36: "vertebrae_T8", 37: "vertebrae_T7", 38: "vertebrae_T6", 39: "vertebrae_T5", 40: "vertebrae_T4", 41: "vertebrae_T3", 42: "vertebrae_T2", 43: "vertebrae_T1", 44: "vertebrae_C7", 45: "vertebrae_C6", 46: "vertebrae_C5", 47: "vertebrae_C4", 48: "vertebrae_C3", 49: "vertebrae_C2", 50: "vertebrae_C1", 51: "heart", 52: "aorta", 53: "pulmonary_vein", 54: "brachiocephalic_trunk", 55: "subclavian_artery_right", 56: "subclavian_artery_left", 57: "common_carotid_artery_right", 58: "common_carotid_artery_left", 59: "brachiocephalic_vein_left", 60: "brachiocephalic_vein_right", 61: "atrial_appendage_left", 62: "superior_vena_cava", 63: "inferior_vena_cava", 64: "portal_vein_and_splenic_vein", 65: "iliac_artery_left", 66: "iliac_artery_right", 67: "iliac_vena_left", 68: "iliac_vena_right", 69: "humerus_left", 70: "humerus_right", 71: "scapula_left", 72: "scapula_right", 73: "clavicula_left", 74: "clavicula_right", 75: "femur_left", 76: "femur_right", 77: "hip_left", 78: "hip_right", 79: "spinal_cord", 80: "gluteus_maximus_left", 81: "gluteus_maximus_right", 82: "gluteus_medius_left", 83: "gluteus_medius_right", 84: "gluteus_minimus_left", 85: "gluteus_minimus_right", 86: "autochthon_left", 87: "autochthon_right", 88: "iliopsoas_left", 89: "iliopsoas_right", 90: "brain", 91: "skull", 92: "rib_left_1", 93: "rib_left_2", 94: "rib_left_3", 95: "rib_left_4", 96: "rib_left_5", 97: "rib_left_6", 98: "rib_left_7", 99: "rib_left_8", 100: "rib_left_9", 101: "rib_left_10", 102: "rib_left_11", 103: "rib_left_12", 104: "rib_right_1", 105: "rib_right_2", 106: "rib_right_3", 107: "rib_right_4", 108: "rib_right_5", 109: "rib_right_6", 110: "rib_right_7", 111: "rib_right_8", 112: "rib_right_9", 113: "rib_right_10", 114: "rib_right_11", 115: "rib_right_12", 116: "sternum", 117: "carpal", 118: "costal_cartilages", } # Key thoracic structures that should be present in a chest CT EXPECTED_THORACIC_LABELS = { "lung_upper_lobe_left", "lung_lower_lobe_left", "lung_upper_lobe_right", "lung_middle_lobe_right", "lung_lower_lobe_right", "heart", "aorta", "trachea", } # Air HU value for table removal AIR_HU = -1000 # Logging setup logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[logging.StreamHandler(sys.stdout)], ) logger = logging.getLogger("pipeline") # ============================================================================ # Helper: Path builders for per-case directory layout # ============================================================================ def case_dir(case_id: str) -> Path: return PIPELINE_DIR / case_id def phase1_paths(case_id: str) -> dict: d = case_dir(case_id) / "01_body" return { "dir": d, "body_mask": d / "body_mask.nii.gz", "ct_clean": d / "CT_clean.nii.gz", "metadata": d / "phase1_metadata.json", } def phase2_paths(case_id: str) -> dict: d = case_dir(case_id) / "02_totalseg" return { "dir": d, "segments_dir": d / "segments", "mask": d / "mask.nii.gz", "metadata": d / "phase2_metadata.json", } def qc_dir(case_id: str) -> Path: return case_dir(case_id) / "qc_visuals" # ============================================================================ # Phase 1: Automated table removal # ============================================================================ def run_phase1(case_id: str, raw_ct_path: Path, device: str, fast: bool) -> dict: """ Phase 1: Body segmentation + CT table removal. Returns a stats dict for metadata logging. """ from totalsegmentator.python_api import totalsegmentator p = phase1_paths(case_id) p["dir"].mkdir(parents=True, exist_ok=True) # --- Step 1a: Load raw CT --- logger.info(f" [Phase 1] Loading raw CT: {raw_ct_path.name}") ct_img = nib.load(str(raw_ct_path)) ct_data = ct_img.get_fdata(dtype=np.float32) # --- Step 1b: Run TotalSegmentator body task --- logger.info(f" [Phase 1] Running TotalSegmentator body task...") body_seg_img = totalsegmentator( input=str(raw_ct_path), output=str(p["dir"] / "body_tmp"), task="body", ml=True, device=device, fast=fast, quiet=True, ) body_data = np.asarray(body_seg_img.dataobj) # body task labels: 1=body_trunc, 2=body_extremities # Combine into binary body mask (anything > 0) body_mask = (body_data > 0).astype(np.uint8) # --- QC: body mask sanity --- body_voxel_count = int(body_mask.sum()) total_voxel_count = int(np.prod(body_mask.shape)) body_fraction = body_voxel_count / total_voxel_count if body_voxel_count == 0: raise RuntimeError(f"Body mask is empty for {case_id}!") if body_fraction < 0.01: logger.warning(f" [Phase 1] Body mask occupies only {body_fraction:.4f} of volume — suspiciously small.") # Save body mask body_mask_img = nib.Nifti1Image(body_mask, ct_img.affine, ct_img.header) nib.save(body_mask_img, str(p["body_mask"])) logger.info(f" [Phase 1] Body mask saved. Fraction: {body_fraction:.4f}") # --- Step 1c: Apply table removal --- ct_clean = ct_data.copy() ct_clean[body_mask == 0] = AIR_HU # Verify shapes match assert ct_clean.shape == ct_data.shape, "Shape mismatch after cleaning!" # Save as int16 for consistency and space efficiency ct_clean_int16 = np.clip(ct_clean, -1024, 3071).astype(np.int16) ct_clean_img = nib.Nifti1Image(ct_clean_int16, ct_img.affine, ct_img.header) ct_clean_img.header.set_data_dtype(np.int16) nib.save(ct_clean_img, str(p["ct_clean"])) logger.info(f" [Phase 1] CT_clean.nii.gz saved.") # Clean up temp directory body_tmp = p["dir"] / "body_tmp" if body_tmp.exists(): import shutil shutil.rmtree(body_tmp, ignore_errors=True) # --- Build metadata --- stats = { "case_id": case_id, "phase": 1, "timestamp": datetime.datetime.now().isoformat(), "raw_ct_path": str(raw_ct_path), "ct_shape": list(ct_data.shape), "voxel_spacing": [float(v) for v in ct_img.header.get_zooms()[:3]], "body_voxel_count": body_voxel_count, "body_fraction": round(body_fraction, 6), "output_dtype": "int16", "air_hu_value": AIR_HU, } with open(p["metadata"], "w") as f: json.dump(stats, f, indent=2) return stats # ============================================================================ # Phase 2: Full 3D anatomy segmentation (118-class) # ============================================================================ def run_phase2(case_id: str, device: str, fast: bool) -> dict: """ Phase 2: Run full TotalSegmentator on CT_clean to produce 118-class mask. Returns a stats dict for metadata logging. """ from totalsegmentator.python_api import totalsegmentator p1 = phase1_paths(case_id) p2 = phase2_paths(case_id) p2["dir"].mkdir(parents=True, exist_ok=True) ct_clean_path = p1["ct_clean"] if not ct_clean_path.exists(): raise FileNotFoundError(f"CT_clean.nii.gz not found for {case_id}. Run Phase 1 first.") # --- Step 2a: Run TotalSegmentator full task with multilabel output --- logger.info(f" [Phase 2] Running TotalSegmentator full (total) task...") mask_img = totalsegmentator( input=str(ct_clean_path), output=str(p2["mask"]), task="total", ml=True, device=device, fast=fast, quiet=True, ) # --- Step 2b: Validate the mask --- mask_data = np.asarray(mask_img.dataobj) ct_clean_img = nib.load(str(ct_clean_path)) # Shape check if mask_data.shape != ct_clean_img.shape: raise RuntimeError( f"Shape mismatch: mask {mask_data.shape} vs CT_clean {ct_clean_img.shape}" ) # Label validity check # unique_labels = set(np.unique(mask_data).astype(int)) unique_labels = set(int(x) for x in np.unique(mask_data)) unique_labels.discard(0) # background valid_labels = set(TOTALSEG_V2_LABELS.keys()) invalid = unique_labels - valid_labels if invalid: logger.warning(f" [Phase 2] Found unexpected label IDs: {invalid}") # Thoracic structures check present_names = {TOTALSEG_V2_LABELS.get(lbl, "?") for lbl in unique_labels} missing_thoracic = EXPECTED_THORACIC_LABELS - present_names if missing_thoracic: logger.warning(f" [Phase 2] Missing expected thoracic structures: {missing_thoracic}") num_classes_found = len(unique_labels) logger.info(f" [Phase 2] mask.nii.gz saved. {num_classes_found} classes found.") # --- Build metadata --- stats = { "case_id": case_id, "phase": 2, "timestamp": datetime.datetime.now().isoformat(), "mask_shape": list(mask_data.shape), "num_classes_found": num_classes_found, "labels_found": sorted(list(unique_labels)), "label_names_found": sorted(list(present_names)), "missing_thoracic": sorted(list(missing_thoracic)), } with open(p2["metadata"], "w") as f: json.dump(stats, f, indent=2) return stats # ============================================================================ # QC Visualization: Middle-slice extraction for Raw CT, Clean CT, Mask overlay # ============================================================================ def generate_qc_visuals(case_id: str, raw_ct_path: Path) -> None: """ Extract axial, coronal, and sagittal middle slices from: 1. Raw CT 2. Clean CT 3. Clean CT + 3D mask overlay Save as PNG images in the case's qc_visuals/ folder. """ p1 = phase1_paths(case_id) p2 = phase2_paths(case_id) out_dir = qc_dir(case_id) out_dir.mkdir(parents=True, exist_ok=True) # Load volumes raw_img = nib.load(str(raw_ct_path)) raw_data = raw_img.get_fdata(dtype=np.float32) dx, dy, dz = raw_img.header.get_zooms()[:3] clean_img = nib.load(str(p1["ct_clean"])) clean_data = clean_img.get_fdata(dtype=np.float32) mask_img = nib.load(str(p2["mask"])) mask_data = np.asarray(mask_img.dataobj).astype(np.int16) shape = raw_data.shape mid = (shape[0] // 2, shape[1] // 2, shape[2] // 2) # Window for chest CT display (W=1500, L=-500 => range [-1250, 250]) ct_vmin, ct_vmax = -1250, 250 # slice_specs = [ # ("axial", lambda vol: vol[:, :, mid[2]]), # ("coronal", lambda vol: vol[:, mid[1], :]), # ("sagittal", lambda vol: vol[mid[0], :, :]), # ] slice_specs = [ ("axial", lambda vol: vol[:, :, mid[2]], dy / dx), ("coronal", lambda vol: vol[:, mid[1], :], dz / dx), ("sagittal", lambda vol: vol[mid[0], :, :], dz / dy), ] # 【修改 for 循环解包】 for plane_name, slicer, aspect_ratio in slice_specs: raw_slice = np.rot90(slicer(raw_data)) clean_slice = np.rot90(slicer(clean_data)) mask_slice = np.rot90(slicer(mask_data)) fig, axes = plt.subplots(1, 3, figsize=(18, 6), facecolor="black") # Panel 1: Raw CT # 【新增 aspect=aspect_ratio】 axes[0].imshow(raw_slice, cmap="gray", vmin=ct_vmin, vmax=ct_vmax, aspect=aspect_ratio) axes[0].set_title("Raw CT", color="white", fontsize=14) axes[0].axis("off") # Panel 2: Clean CT # 【新增 aspect=aspect_ratio】 axes[1].imshow(clean_slice, cmap="gray", vmin=ct_vmin, vmax=ct_vmax, aspect=aspect_ratio) axes[1].set_title("Clean CT (table removed)", color="white", fontsize=14) axes[1].axis("off") # Panel 3: Clean CT + Mask overlay # 【新增 aspect=aspect_ratio】 axes[2].imshow(clean_slice, cmap="gray", vmin=ct_vmin, vmax=ct_vmax, aspect=aspect_ratio) mask_overlay = np.ma.masked_where(mask_slice == 0, mask_slice) # 【蒙版层也要加 aspect=aspect_ratio】 axes[2].imshow(mask_overlay, cmap="nipy_spectral", alpha=0.45, vmin=0, vmax=118, interpolation="nearest", aspect=aspect_ratio) axes[2].set_title("Clean CT + Mask Overlay", color="white", fontsize=14) axes[2].axis("off") # for plane_name, slicer in slice_specs: # raw_slice = np.rot90(slicer(raw_data)) # clean_slice = np.rot90(slicer(clean_data)) # mask_slice = np.rot90(slicer(mask_data)) # fig, axes = plt.subplots(1, 3, figsize=(18, 6), facecolor="black") # # Panel 1: Raw CT # axes[0].imshow(raw_slice, cmap="gray", vmin=ct_vmin, vmax=ct_vmax) # axes[0].set_title("Raw CT", color="white", fontsize=14) # axes[0].axis("off") # # Panel 2: Clean CT # axes[1].imshow(clean_slice, cmap="gray", vmin=ct_vmin, vmax=ct_vmax) # axes[1].set_title("Clean CT (table removed)", color="white", fontsize=14) # axes[1].axis("off") # # Panel 3: Clean CT + Mask overlay # axes[2].imshow(clean_slice, cmap="gray", vmin=ct_vmin, vmax=ct_vmax) # # Create a colored overlay where mask > 0 (use nipy_spectral for distinct colors) # mask_overlay = np.ma.masked_where(mask_slice == 0, mask_slice) # axes[2].imshow(mask_overlay, cmap="nipy_spectral", alpha=0.45, # vmin=0, vmax=118, interpolation="nearest") # axes[2].set_title("Clean CT + Mask Overlay", color="white", fontsize=14) # axes[2].axis("off") fig.suptitle(f"{case_id} — {plane_name} (middle slice)", color="white", fontsize=16, fontweight="bold") plt.tight_layout(rect=[0, 0, 1, 0.95]) out_path = out_dir / f"{plane_name}.png" fig.savefig(str(out_path), dpi=120, bbox_inches="tight", facecolor="black", edgecolor="none") plt.close(fig) logger.info(f" [QC] Saved 3 visualization PNGs to {out_dir}") # ============================================================================ # Manifest / CSV management # ============================================================================ def write_csv_row(csv_path: Path, row: dict, fieldnames: List[str]) -> None: """Append one row to a CSV. Create with header if file does not exist.""" file_exists = csv_path.exists() with open(csv_path, "a", newline="") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) if not file_exists: writer.writeheader() writer.writerow(row) SUCCESS_FIELDS = ["case_id", "timestamp", "body_fraction", "num_classes_found", "missing_thoracic"] FAILED_FIELDS = ["case_id", "timestamp", "phase_failed", "error"] def save_label_dictionary() -> None: """Write the canonical label dictionary JSON to manifests/.""" MANIFEST_DIR.mkdir(parents=True, exist_ok=True) label_path = MANIFEST_DIR / "labels_totalseg_118.json" if not label_path.exists(): payload = { "description": "TotalSegmentator v2 'total' task — 118 anatomical classes", "background_label": 0, "labels": {str(k): v for k, v in TOTALSEG_V2_LABELS.items()}, } with open(label_path, "w") as f: json.dump(payload, f, indent=2) logger.info(f"Label dictionary saved: {label_path}") # ============================================================================ # Resumability: detect already-processed cases # ============================================================================ def is_case_complete(case_id: str) -> bool: """Return True if both Phase 1 and Phase 2 outputs exist for this case.""" p1 = phase1_paths(case_id) p2 = phase2_paths(case_id) return ( p1["ct_clean"].exists() and p1["body_mask"].exists() and p1["metadata"].exists() and p2["mask"].exists() and p2["metadata"].exists() ) # ============================================================================ # Main batch loop # ============================================================================ def discover_cases() -> List[Tuple[str, Path]]: """Discover all .nii.gz files in RAW_CT_DIR and return (case_id, path) pairs.""" if not RAW_CT_DIR.exists(): raise FileNotFoundError(f"Raw CT directory not found: {RAW_CT_DIR}") cases = [] for f in sorted(RAW_CT_DIR.glob("*.nii.gz")): case_id = f.name.replace(".nii.gz", "") cases.append((case_id, f)) if not cases: raise FileNotFoundError(f"No .nii.gz files found in {RAW_CT_DIR}") return cases def run_pipeline(device: str, fast: bool, dry_run: bool = False) -> None: """Run full Phase 1 + Phase 2 batch with resumability.""" cases = discover_cases() logger.info(f"Discovered {len(cases)} CT volumes in {RAW_CT_DIR}") # Prepare manifests MANIFEST_DIR.mkdir(parents=True, exist_ok=True) save_label_dictionary() # Write cases_all.csv cases_all_path = MANIFEST_DIR / "cases_all.csv" if not cases_all_path.exists(): with open(cases_all_path, "w", newline="") as f: writer = csv.writer(f) writer.writerow(["case_id", "raw_ct_path"]) for cid, cpath in cases: writer.writerow([cid, str(cpath)]) logger.info(f" Written {cases_all_path}") success_csv = MANIFEST_DIR / "cases_success.csv" failed_csv = MANIFEST_DIR / "cases_failed.csv" # Determine which cases to process skip_count = 0 to_process = [] for case_id, raw_path in cases: if is_case_complete(case_id): skip_count += 1 else: to_process.append((case_id, raw_path)) logger.info(f"Already complete: {skip_count} | To process: {len(to_process)}") if dry_run: logger.info("=== DRY RUN MODE — No processing will be performed ===") logger.info(f"Would process {len(to_process)} cases.") for i, (cid, rp) in enumerate(to_process[:10]): p1 = phase1_paths(cid) p2 = phase2_paths(cid) logger.info(f" [{i+1}] {cid}") logger.info(f" Input: {rp}") logger.info(f" CT_clean: {p1['ct_clean']}") logger.info(f" body_mask: {p1['body_mask']}") logger.info(f" mask: {p2['mask']}") logger.info(f" QC dir: {qc_dir(cid)}") if len(to_process) > 10: logger.info(f" ... and {len(to_process) - 10} more cases.") # Show directory structure preview logger.info("") logger.info("=== OUTPUT DIRECTORY STRUCTURE PREVIEW ===") logger.info(f" {PIPELINE_DIR}/") logger.info(f" ├── manifests/") logger.info(f" │ ├── cases_all.csv") logger.info(f" │ ├── cases_success.csv") logger.info(f" │ ├── cases_failed.csv") logger.info(f" │ └── labels_totalseg_118.json") if to_process: cid0 = to_process[0][0] logger.info(f" ├── {cid0}/") logger.info(f" │ ├── 01_body/") logger.info(f" │ │ ├── body_mask.nii.gz") logger.info(f" │ │ ├── CT_clean.nii.gz") logger.info(f" │ │ └── phase1_metadata.json") logger.info(f" │ ├── 02_totalseg/") logger.info(f" │ │ ├── mask.nii.gz") logger.info(f" │ │ └── phase2_metadata.json") logger.info(f" │ └── qc_visuals/") logger.info(f" │ ├── axial.png") logger.info(f" │ ├── coronal.png") logger.info(f" │ └── sagittal.png") logger.info(f" └── ... ({len(to_process)} cases total)") logger.info("") logger.info("=== DRY RUN COMPLETE — Pass without --dry-run to execute ===") return # Process each case n_success = 0 n_failed = 0 for case_id, raw_path in tqdm(to_process, desc="Phase 1+2", unit="case"): logger.info(f"\n{'='*60}") logger.info(f"Processing: {case_id}") logger.info(f"{'='*60}") try: t0 = time.time() # Phase 1 p1_stats = run_phase1(case_id, raw_path, device=device, fast=fast) # Phase 2 p2_stats = run_phase2(case_id, device=device, fast=fast) # QC Visuals logger.info(f" [QC] Generating middle-slice visualizations...") generate_qc_visuals(case_id, raw_path) elapsed = time.time() - t0 logger.info(f" DONE: {case_id} ({elapsed:.1f}s)") # Record success write_csv_row(success_csv, { "case_id": case_id, "timestamp": datetime.datetime.now().isoformat(), "body_fraction": p1_stats["body_fraction"], "num_classes_found": p2_stats["num_classes_found"], "missing_thoracic": "|".join(p2_stats["missing_thoracic"]), }, SUCCESS_FIELDS) n_success += 1 except Exception as e: logger.error(f" FAILED: {case_id} — {e}") logger.error(traceback.format_exc()) # Determine which phase failed p1 = phase1_paths(case_id) phase_failed = 1 if not p1["ct_clean"].exists() else 2 write_csv_row(failed_csv, { "case_id": case_id, "timestamp": datetime.datetime.now().isoformat(), "phase_failed": phase_failed, "error": str(e), }, FAILED_FIELDS) n_failed += 1 # Final summary logger.info(f"\n{'='*60}") logger.info(f"BATCH COMPLETE") logger.info(f" Processed: {n_success + n_failed}") logger.info(f" Success: {n_success}") logger.info(f" Failed: {n_failed}") logger.info(f" Skipped: {skip_count} (already complete)") logger.info(f" Manifests: {MANIFEST_DIR}") logger.info(f"{'='*60}") # ============================================================================ # CLI entry point # ============================================================================ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Anatomy-Aware DRR Pipeline: Phase 1 (table removal) + Phase 2 (3D segmentation)" ) parser.add_argument( "--device", type=str, default="gpu", help="Device for TotalSegmentator: 'gpu', 'cpu', 'gpu:0', 'gpu:1', etc. (default: gpu)", ) parser.add_argument( "--fast", action="store_true", help="Use TotalSegmentator fast mode (3mm resolution). Faster but lower quality.", ) parser.add_argument( "--dry-run", action="store_true", help="Show what would be processed without actually running.", ) return parser.parse_args() if __name__ == "__main__": args = parse_args() run_pipeline(device=args.device, fast=args.fast, dry_run=args.dry_run)