total_Seg_test.py

#!/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)