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,6 +358,7 @@ 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"):

vesskel/_batch.py

@@ -0,0 +1,105 @@
+"""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 AnalysisResult, analyze_binary_image
+
+_SUPPORTED_EXTENSIONS = frozenset(
+    {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".npy"}
+)
+
+
+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}