feat: parallelized cli batch processing

tests/test_cli.py

@@ -6,16 +6,62 @@
 import pytest
 from PIL import Image
 
-from vesskel.cli import (
-    _discover_input_paths,
+from vesskel._batch import (
     _load_image,
     _sanitize_for_csv,
     _save_radius,
     _save_skeleton,
     _write_csv,
 )
+from vesskel.cli import _discover_input_paths
 from vesskel.config import CONFIG_SCHEMA_VERSION, PipelineConfig
 
+HEAVY_MODULES = frozenset(
+    {"numpy", "PIL", "PIL.Image", "vesskel.pipeline", "vesskel._batch"}
+)
+
+
+class TestCompletionSpeed:
+    """Shell completions must exit before importing heavy modules (numpy/PIL)."""
+
+    def test_completions_skip_heavy_imports(self):
+        import subprocess
+        import sys
+
+        heavy_modules = sorted(HEAVY_MODULES)
+        probe = f"""
+import os, sys
+
+def _fake_exit(code=0):
+    raise SystemExit(code)
+os._exit = _fake_exit
+
+os.environ["_ARGCOMPLETE"] = "1"
+sys.argv = ["vesskel", "complete", "zsh"]
+
+try:
+    import vesskel.cli
+except SystemExit:
+    pass
+
+heavy = [m for m in {heavy_modules!r} if m in sys.modules]
+if heavy:
+    sys.stdout.write("HEAVY:" + ",".join(heavy))
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", probe],
+            capture_output=True,
+            text=True,
+            timeout=15,
+        )
+        assert result.returncode == 0, f"Subprocess failed (stderr): {result.stderr}"
+        assert (
+            "HEAVY:" not in result.stdout
+        ), f"Heavy modules loaded during completions: {result.stdout}"
+        assert (
+            "HEAVY:" not in result.stderr
+        ), f"Heavy modules leaked to stderr: {result.stderr}"
+
 
 class TestDiscoverInputPaths:
     def test_single_png_file(self, tmp_path):
@@ -285,6 +331,7 @@ def test_run_batch_processes_single_image(self, tmp_path):
             config=str(config_path),
             out=str(out_dir),
             recursive=False,
+            jobs=1,
         )
 
         exit_code = _run_batch(args)
@@ -311,11 +358,112 @@ def test_run_batch_no_input_files_raises(self, tmp_path):
             config=str(config_path),
             out=str(tmp_path / "out"),
             recursive=False,
+            jobs=1,
         )
 
         with pytest.raises(ValueError, match="No input files found"):
             _run_batch(args)
 
+    def test_run_batch_parallel_multiple_images(self, tmp_path):
+        from vesskel.cli import _run_batch
+
+        img_path = tmp_path / "input"
+        img_path.mkdir()
+        for name in ("a", "b"):
+            image = np.zeros((16, 16), dtype=np.uint8)
+            image[8, 4:12] = 1
+            image[4:12, 8] = 1
+            Image.fromarray(image * 255).save(img_path / f"{name}.png")
+
+        config_path = tmp_path / "pipeline.json"
+        config = PipelineConfig.from_dict(
+            {"schema_version": CONFIG_SCHEMA_VERSION, "extraction": {}, "output": {}}
+        )
+        config_path.write_text(json.dumps(config.to_dict(), indent=2))
+
+        out_dir = tmp_path / "output"
+
+        args = _make_args(
+            input=[str(img_path)],
+            config=str(config_path),
+            out=str(out_dir),
+            recursive=False,
+            jobs=2,
+        )
+
+        exit_code = _run_batch(args)
+        assert exit_code == 0
+        assert (out_dir / "a").is_dir()
+        assert (out_dir / "b").is_dir()
+        assert (out_dir / "a" / "a_skeleton.npy").exists()
+        assert (out_dir / "b" / "b_skeleton.npy").exists()
+        assert (out_dir / "summary.csv").exists()
+
+        with open(out_dir / "summary.csv") as f:
+            reader = csv.DictReader(f)
+            data = list(reader)
+        assert len(data) == 2
+
+    def test_run_batch_parallel_failure_raises(self, tmp_path):
+        from vesskel.cli import _run_batch
+
+        img_path = tmp_path / "input"
+        img_path.mkdir()
+        image = np.zeros((16, 16), dtype=np.uint8)
+        image[8, 4:12] = 1
+        image[4:12, 8] = 1
+        Image.fromarray(image * 255).save(img_path / "good.png")
+        np.save(img_path / "bad.npy", np.zeros((2, 2, 2, 2), dtype=np.uint8))
+
+        config_path = tmp_path / "pipeline.json"
+        config = PipelineConfig.from_dict(
+            {"schema_version": CONFIG_SCHEMA_VERSION, "extraction": {}, "output": {}}
+        )
+        config_path.write_text(json.dumps(config.to_dict(), indent=2))
+
+        out_dir = tmp_path / "output"
+
+        args = _make_args(
+            input=[str(img_path)],
+            config=str(config_path),
+            out=str(out_dir),
+            recursive=False,
+            jobs=2,
+        )
+
+        with pytest.raises(RuntimeError, match=r"image\(s\) failed"):
+            _run_batch(args)
+
+    def test_run_batch_jobs_zero_auto_detect(self, tmp_path):
+        from vesskel.cli import _run_batch
+
+        img_path = tmp_path / "input"
+        img_path.mkdir()
+        image = np.zeros((16, 16), dtype=np.uint8)
+        image[8, 4:12] = 1
+        image[4:12, 8] = 1
+        Image.fromarray(image * 255).save(img_path / "cross.png")
+
+        config_path = tmp_path / "pipeline.json"
+        config = PipelineConfig.from_dict(
+            {"schema_version": CONFIG_SCHEMA_VERSION, "extraction": {}, "output": {}}
+        )
+        config_path.write_text(json.dumps(config.to_dict(), indent=2))
+
+        out_dir = tmp_path / "output"
+
+        args = _make_args(
+            input=[str(img_path)],
+            config=str(config_path),
+            out=str(out_dir),
+            recursive=False,
+            jobs=0,
+        )
+
+        exit_code = _run_batch(args)
+        assert exit_code == 0
+        assert (out_dir / "cross" / "cross_skeleton.npy").exists()
+
     def test_main_dispatch_config_init(self, tmp_path, monkeypatch):
         import sys
 

vesskel/_batch.py

@@ -0,0 +1,101 @@
+"""Batch I/O helpers and worker function for multiprocessing."""
+
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+from PIL import Image
+
+from vesskel.config import PipelineConfig
+from vesskel.pipeline import analyze_binary_image
+
+
+def _load_image(path: Path) -> np.ndarray:
+    if path.suffix.lower() == ".npy":
+        arr = np.load(path)
+    else:
+        with Image.open(path) as im:
+            arr = np.asarray(im)
+
+    if arr.ndim == 0:
+        raise ValueError("Scalar input is not supported")
+
+    if arr.ndim == 3 and arr.shape[-1] in (3, 4):
+        arr = np.max(arr[..., :3], axis=-1)
+
+    if arr.ndim not in (2, 3):
+        raise ValueError(f"Expected 2D or 3D image, got shape={arr.shape}")
+
+    return arr
+
+
+def _sanitize_for_csv(value: object) -> object:
+    if isinstance(value, (np.generic,)):
+        return value.item()
+    return value
+
+
+def _write_csv(path: Path, rows: Iterable[dict[str, object]]) -> None:
+    rows = list(rows)
+    if not rows:
+        return
+
+    fieldnames = sorted({key for row in rows for key in row.keys()})
+    with path.open("w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for row in rows:
+            writer.writerow({k: _sanitize_for_csv(v) for k, v in row.items()})
+
+
+def _save_skeleton(
+    path: Path,
+    skeleton: np.ndarray,
+    *,
+    npy: bool = True,
+    png: bool = False,
+) -> None:
+    if npy:
+        np.save(path.with_suffix(".npy"), skeleton.astype(np.uint8))
+    if png:
+        if skeleton.ndim != 2:
+            raise ValueError("PNG skeleton output is only supported for 2D images")
+        img = Image.fromarray((skeleton > 0).astype(np.uint8) * 255)
+        img.save(path.with_suffix(".png"))
+
+
+def _save_radius(path: Path, radius_matrix: np.ndarray) -> None:
+    np.save(path.with_suffix(".npy"), radius_matrix.astype(np.float64))
+
+
+def process_one(
+    in_path: Path,
+    safe_name: str,
+    out_dir: Path,
+    config: PipelineConfig,
+) -> dict[str, object]:
+    """Load, analyse, save one image. Returns summary row for agg CSV."""
+    image = _load_image(in_path)
+    result = analyze_binary_image(image=image, base_name=in_path.stem, config=config)
+
+    image_out_dir = out_dir / safe_name
+    image_out_dir.mkdir(parents=True, exist_ok=True)
+
+    if config.output.write_skeleton_npy or config.output.write_skeleton_png:
+        _save_skeleton(
+            image_out_dir / f"{safe_name}_skeleton",
+            result.skeleton,
+            npy=config.output.write_skeleton_npy,
+            png=config.output.write_skeleton_png,
+        )
+
+    if config.output.write_radius and result.radius_matrix is not None:
+        _save_radius(image_out_dir / f"{safe_name}_radius", result.radius_matrix)
+
+    if config.output.write_branch_csv and result.branch_records:
+        _write_csv(image_out_dir / f"{safe_name}_branches.csv", result.branch_records)
+
+    return {"image": in_path.name, **result.summary_features}