Upload scripts/run_pipeline.py with huggingface_hub

scripts/run_pipeline.py

@@ -1,22 +1,9 @@
 #!/usr/bin/env python3
-"""End-to-end coronary segmentation pipeline.
-
-CT image (DICOM/NIfTI/NRRD) → binary coronary mask �� segmental coronary labels.
-
-Chains two 5-fold ensemble inferences:
-  1. Binary coronary segmentation (CT → vessel mask)
-  2. Segmental coronary segmentation (vessel mask → 21-class labels)
-
-Works both locally and inside Docker. Detects bundle directories automatically:
-  - HF clone:  <root>/ct_binary_coronary_segmentation/
-  - Git repo:  <root>/bundle/ct_binary_coronary_segmentation/
+"""End-to-end coronary segmentation: CT image → binary mask → segmental labels.
 
 Usage:
-    # Local
     python scripts/run_pipeline.py --input /path/to/dicoms --output /path/to/results
-
-    # Docker
-    docker run --gpus all -v /data/patient:/input -v /data/output:/output ct-heart-seg
+    docker run --gpus all -v /data/in:/input -v /data/out:/output ct-heart-seg
 """
 
 import argparse
@@ -29,140 +16,82 @@ import tempfile
 import time
 from pathlib import Path
 
-
-def setup_logging(output_dir: Path):
-    output_dir.mkdir(parents=True, exist_ok=True)
-    handlers = [
-        logging.StreamHandler(),
-        logging.FileHandler(output_dir / "pipeline.log"),
-    ]
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s [%(levelname)s] %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-        handlers=handlers,
-        force=True,
-    )
-
-
 logger = logging.getLogger(__name__)
 
 BINARY_BUNDLE = "ct_binary_coronary_segmentation"
 SEGMENTAL_BUNDLE = "ct_segmental_coronary_segmentation"
 
 
-def find_bundle_dir(root: Path, bundle_name: str) -> Path:
-    """Find bundle directory — supports both HF clone and git repo layouts."""
-    # HF clone: bundles at root level
-    if (root / bundle_name / "configs").is_dir():
-        return root / bundle_name
-    # Git repo: bundles under bundle/
-    if (root / "bundle" / bundle_name / "configs").is_dir():
-        return root / "bundle" / bundle_name
-    raise FileNotFoundError(
-        f"Cannot find {bundle_name}/ in {root} or {root / 'bundle'}"
-    )
+def find_bundle_dir(root: Path, name: str) -> Path:
+    for prefix in ["", "bundle"]:
+        d = root / prefix / name if prefix else root / name
+        if (d / "configs").is_dir():
+            return d
+    raise FileNotFoundError(f"Cannot find {name}/ in {root} or {root / 'bundle'}")
 
 
-def run_inference(bundle_dir: Path, config: str, extra_args: list[str], gpu: int, label: str):
-    """Run MONAI bundle inference as subprocess."""
+def run_inference(bundle_dir: Path, extra_args: list[str], gpu: int, label: str):
     env = os.environ.copy()
     env["CUDA_VISIBLE_DEVICES"] = str(gpu)
-
     cmd = [
         sys.executable, "-m", "monai.bundle", "run", "inference",
-        "--config_file", config,
+        "--config_file", "configs/ensemble_inference.yaml",
         *extra_args,
     ]
-
-    logger.info("[%s] Starting inference", label)
-    logger.info("[%s] cwd: %s", label, bundle_dir)
-    logger.info("[%s] cmd: %s", label, " ".join(cmd))
-
+    logger.info("[%s] %s (cwd=%s)", label, " ".join(cmd), bundle_dir)
     t0 = time.time()
-    result = subprocess.run(
-        cmd,
-        cwd=str(bundle_dir),
-        env=env,
-        capture_output=True,
-        text=True,
-    )
-
+    result = subprocess.run(cmd, cwd=str(bundle_dir), env=env, capture_output=True, text=True)
     elapsed = time.time() - t0
+
     if result.stdout:
         for line in result.stdout.strip().split("\n"):
             logger.info("[%s] %s", label, line)
     if result.stderr:
         for line in result.stderr.strip().split("\n"):
             logger.info("[%s] %s", label, line)
-
     if result.returncode != 0:
-        logger.error("[%s] FAILED (exit code %d) after %.1fs", label, result.returncode, elapsed)
+        logger.error("[%s] FAILED (exit %d) after %.1fs", label, result.returncode, elapsed)
         sys.exit(result.returncode)
-
     logger.info("[%s] Done in %.1fs", label, elapsed)
 
 
 def main():
-    parser = argparse.ArgumentParser(
-        description="End-to-end coronary segmentation: CT → binary mask → segmental labels"
-    )
-    parser.add_argument("--input", required=True, help="Input directory (DICOM/NIfTI/NRRD files)")
-    parser.add_argument("--output", required=True, help="Output directory")
-    parser.add_argument("--gpu", type=int, default=0, help="GPU index (default: 0)")
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--input", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--gpu", type=int, default=0)
     args = parser.parse_args()
 
     output_dir = Path(args.output)
-    setup_logging(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    logging.basicConfig(
+        level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        handlers=[logging.StreamHandler(), logging.FileHandler(output_dir / "pipeline.log")],
+        force=True,
+    )
 
-    # Find repo root (parent of scripts/)
     repo_root = Path(__file__).resolve().parent.parent
     binary_dir = find_bundle_dir(repo_root, BINARY_BUNDLE)
     segmental_dir = find_bundle_dir(repo_root, SEGMENTAL_BUNDLE)
-
-    logger.info("Input: %s", args.input)
-    logger.info("Output: %s", output_dir)
-    logger.info("Binary bundle: %s", binary_dir)
-    logger.info("Segmental bundle: %s", segmental_dir)
-    logger.info("GPU: %d", args.gpu)
+    logger.info("Input: %s | Output: %s | GPU: %d", args.input, output_dir, args.gpu)
 
     t_start = time.time()
-
-    # Stage 1: Binary coronary segmentation
     with tempfile.TemporaryDirectory(prefix="binary_output_") as tmp_binary:
-        run_inference(
-            bundle_dir=binary_dir,
-            config="configs/ensemble_inference.yaml",
-            extra_args=["--dataset_dir", str(Path(args.input).resolve()),
-                        "--output_dir", tmp_binary],
-            gpu=args.gpu,
-            label="binary",
-        )
-
-        # Copy binary output to final location
+        run_inference(binary_dir, ["--dataset_dir", str(Path(args.input).resolve()), "--output_dir", tmp_binary], args.gpu, "binary")
+
         binary_out = output_dir / "binary"
         binary_out.mkdir(parents=True, exist_ok=True)
         for f in Path(tmp_binary).glob("*"):
             shutil.copy2(f, binary_out / f.name)
-        n_binary = len(list(binary_out.glob("*.nii.gz")))
-        logger.info("Binary output: %d files in %s", n_binary, binary_out)
+        logger.info("Binary: %d files → %s", len(list(binary_out.glob("*.nii.gz"))), binary_out)
 
-        # Stage 2: Segmental coronary segmentation (uses binary output as input)
         segmental_out = output_dir / "segmental"
         segmental_out.mkdir(parents=True, exist_ok=True)
-        run_inference(
-            bundle_dir=segmental_dir,
-            config="configs/ensemble_inference.yaml",
-            extra_args=["--binary_label_dir", tmp_binary,
-                        "--output_dir", str(segmental_out)],
-            gpu=args.gpu,
-            label="segmental",
-        )
-        n_segmental = len(list(segmental_out.glob("*.nii.gz")))
-        logger.info("Segmental output: %d files in %s", n_segmental, segmental_out)
-
-    total = time.time() - t_start
-    logger.info("Pipeline complete in %.1fs. Output: %s", total, output_dir)
+        run_inference(segmental_dir, ["--binary_label_dir", tmp_binary, "--output_dir", str(segmental_out)], args.gpu, "segmental")
+        logger.info("Segmental: %d files → %s", len(list(segmental_out.glob("*.nii.gz"))), segmental_out)
+
+    logger.info("Pipeline complete in %.1fs", time.time() - t_start)
 
 
 if __name__ == "__main__":