Enhance device selection and logging for inference; add end-to-end tests

README.md

@@ -33,6 +33,7 @@ This is not the canonical upstream repository. The upstream project remains `Eye
 - The legacy `setup.py` and installed `vascx` console script were removed
 - Supported entrypoints are `./run.sh` and `python -m vascx_models`
 - Overlay generation can now be configured from the root `config.yaml`
+- Inference device selection is automatic by default and can be overridden explicitly
 - Local helper scripts and docs were updated to point at this fork instead of the upstream Hub repo
 - Generated outputs, caches, and other non-repository artifacts are excluded from version control
 
@@ -71,6 +72,7 @@ Run the full pipeline:
 
 ```bash
 INPUT_PATH=/path/to/images OUTPUT_PATH=/path/to/output N_JOBS=4 ./run.sh
+DEVICE=cpu INPUT_PATH=/path/to/images OUTPUT_PATH=/path/to/output ./run.sh
 ./run.sh --sample-run
 ```
 
@@ -94,10 +96,25 @@ Typical examples:
 python -m vascx_models run /path/to/images /path/to/output
 python -m vascx_models run /path/to/image_list.csv /path/to/output
 python -m vascx_models run /path/to/preprocessed/images /path/to/output --no-preprocess
+python -m vascx_models run /path/to/images /path/to/output --device auto
+python -m vascx_models run /path/to/images /path/to/output --device cpu
 python -m vascx_models run /path/to/images /path/to/output --no-disc --no-quality --no-fovea --no-overlay
 python -m vascx_models run /path/to/images /path/to/output --no-vessels
 ```
 
+## Device Selection
+
+Inference device selection is automatic by default.
+
+- `--device auto` is the default for `python -m vascx_models run`
+- `DEVICE=auto` is the default for `./run.sh`
+- Auto-selection priority is `cuda` first, then Apple Metal `mps`, then `cpu`
+- The CLI logs detected availability as `cuda=...`, `mps=...`, `cpu=True`
+- The CLI also logs the selected device for each run
+- You can force a backend with `--device cuda`, `--device mps`, or `--device cpu`
+- `./run.sh` forwards the `DEVICE` environment variable to the Python CLI
+- If you request `cuda` or `mps` explicitly and that backend is unavailable, the run exits with a clear error instead of silently falling back
+
 ## Configuration
 
 This fork adds a root-level `config.yaml` for overlay behavior, disc-circle generation, and vessel-width sampling.
@@ -197,10 +214,30 @@ Current measurement behavior is intentionally conservative:
 - `vascx_models/`: package source and CLI
 - `artery_vein/`, `disc/`, `fovea/`, `vessels/`, `quality/`, `odfd/`, `discedge/`: model artifacts
 - `config.yaml`: fork-specific overlay configuration
+- `pytest.ini`: pytest marker definitions for slow and end-to-end tests
 - `run.sh`: primary local runner
 - `tests/`: pytest suite
 - `notebooks/`: preprocessing and inference examples
 
+## Testing
+
+The test suite includes unit tests, CLI tests, and an opt-in real-model single-image end-to-end smoke test in `tests/test_e2e.py`.
+
+Useful commands:
+
+```bash
+conda run -n vascx-fork pytest
+KMP_DUPLICATE_LIB_OK=TRUE conda run -n vascx-fork pytest tests/test_e2e.py -q
+KMP_DUPLICATE_LIB_OK=TRUE VASCX_RUN_E2E=1 conda run -n vascx-fork pytest tests/test_e2e.py -q -k cpu
+```
+
+Explicitly tested in this fork as of April 21, 2026:
+
+- README and CLI/config behavior updates are covered by the regular pytest suite
+- device resolution priority and explicit unavailable-device failures are covered by unit tests
+- the real single-image end-to-end pipeline was run successfully on CPU with preprocessing enabled
+- the end-to-end test is parameterized for `cpu`, `cuda`, and `mps`, but actual `cuda` and `mps` execution were not exercised in this workspace because those backends were unavailable
+
 ## Upstream Reference
 
 Original upstream project:

pytest.ini

@@ -0,0 +1,6 @@
+[pytest]
+markers =
+    e2e: real-model end-to-end integration tests
+    slow: long-running tests
+filterwarnings =
+    ignore:`torch\.jit\.interface` is deprecated\. Please use `torch\.compile` instead\.:DeprecationWarning

run.sh

@@ -10,6 +10,7 @@ TIMESTAMP="$(date +"%Y%m%d_%H%M%S")"
 DEFAULT_OUTPUT_PATH="$REPO_ROOT/output_$TIMESTAMP"
 OUTPUT_PATH="${OUTPUT_PATH:-$DEFAULT_OUTPUT_PATH}"
 N_JOBS="${N_JOBS:-1}"
+DEVICE="${DEVICE:-auto}"
 
 while [[ $# -gt 0 ]]; do
   case "$1" in
@@ -37,6 +38,7 @@ echo "  conda env:   $CONDA_ENV"
 echo "  input path:  $INPUT_PATH"
 echo "  output path: $OUTPUT_PATH"
 echo "  n_jobs:      $N_JOBS"
+echo "  device:      $DEVICE"
 
 CONDA_BASE="$(conda info --base)"
 # shellcheck disable=SC1091
@@ -44,4 +46,4 @@ source "$CONDA_BASE/etc/profile.d/conda.sh"
 conda activate "$CONDA_ENV"
 
 cd "$REPO_ROOT"
-exec python -m vascx_models run "$INPUT_PATH" "$OUTPUT_PATH" --n_jobs "$N_JOBS"
+exec python -m vascx_models run "$INPUT_PATH" "$OUTPUT_PATH" --n_jobs "$N_JOBS" --device "$DEVICE"

tests/test_cli.py

@@ -1,8 +1,9 @@
 from pathlib import Path
 
+import logging
 import pandas as pd
 from click.testing import CliRunner
-import logging
+import torch
 from PIL import Image
 
 from vascx_models.cli import cli
@@ -20,7 +21,15 @@ def test_cli_run_passes_measurement_config_and_data_to_overlays(
 
     calls: dict[str, object] = {}
 
-    monkeypatch.setattr("vascx_models.cli.preferred_device", lambda: "cpu")
+    def fake_available_device_types():
+        return {"cuda": False, "mps": False, "cpu": True}
+
+    def fake_resolve_device(device_name):
+        calls["device_name"] = device_name
+        return torch.device("cpu")
+
+    monkeypatch.setattr("vascx_models.cli.available_device_types", fake_available_device_types)
+    monkeypatch.setattr("vascx_models.cli.resolve_device", fake_resolve_device)
 
     def fake_run_segmentation_vessels_and_av(**kwargs):
         calls["run_segmentation_vessels_and_av"] = kwargs
@@ -128,9 +137,11 @@ def test_cli_run_passes_measurement_config_and_data_to_overlays(
     )
 
     assert result.exit_code == 0, result.output
+    assert calls["device_name"] == "auto"
     assert calls["run_segmentation_vessels_and_av"]["artery_color"] == (170, 0, 0)
     assert calls["run_segmentation_vessels_and_av"]["vein_color"] == (0, 0, 187)
     assert calls["run_segmentation_vessels_and_av"]["vessel_color"] == (0, 204, 0)
+    assert calls["run_segmentation_vessels_and_av"]["device"] == torch.device("cpu")
     assert calls["run_segmentation_disc"]["disc_color"] == (221, 221, 221)
     assert calls["measure_vessel_widths"]["inner_circle"].name == "2r"
     assert calls["measure_vessel_widths"]["outer_circle"].name == "3r"
@@ -167,3 +178,50 @@ def test_cli_run_reports_missing_path_column_in_csv(tmp_path: Path, caplog) -> N
 
     assert result.exit_code == 0
     assert "CSV must contain a 'path' column" in caplog.text
+
+
+def test_cli_run_accepts_explicit_device_and_logs_selection(
+    tmp_path: Path, monkeypatch, caplog
+) -> None:
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    input_dir.mkdir()
+    Image.new("RGB", (32, 32), color=(0, 0, 0)).save(input_dir / "sample.png")
+
+    calls: dict[str, object] = {}
+
+    monkeypatch.setattr(
+        "vascx_models.cli.available_device_types",
+        lambda: {"cuda": False, "mps": False, "cpu": True},
+    )
+
+    def fake_resolve_device(device_name):
+        calls["device_name"] = device_name
+        return torch.device("cpu")
+
+    monkeypatch.setattr("vascx_models.cli.resolve_device", fake_resolve_device)
+    monkeypatch.setattr("vascx_models.cli.run_quality_estimation", lambda **kwargs: pd.DataFrame())
+    monkeypatch.setattr("vascx_models.cli.run_fovea_detection", lambda **kwargs: pd.DataFrame())
+
+    with caplog.at_level(logging.INFO):
+        result = CliRunner().invoke(
+            cli,
+            [
+                "run",
+                str(input_dir),
+                str(output_dir),
+                "--no-preprocess",
+                "--no-vessels",
+                "--no-disc",
+                "--no-quality",
+                "--no-fovea",
+                "--no-overlay",
+                "--device",
+                "cpu",
+            ],
+        )
+
+    assert result.exit_code == 0, result.output
+    assert calls["device_name"] == "cpu"
+    assert "Device availability: cuda=False, mps=False, cpu=True" in caplog.text
+    assert "Using requested device 'cpu': cpu" in caplog.text

tests/test_e2e.py

@@ -0,0 +1,219 @@
+import os
+import shutil
+from pathlib import Path
+
+os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+from PIL import Image
+from rtnls_fundusprep.cli import _run_preprocessing
+
+from vascx_models.config import AppConfig
+from vascx_models.disc_circles import generate_disc_circles
+from vascx_models.inference import (
+    run_fovea_detection,
+    run_quality_estimation,
+    run_segmentation_disc,
+    run_segmentation_vessels_and_av,
+)
+from vascx_models.runtime import configure_runtime_environment
+from vascx_models.utils import batch_create_overlays
+from vascx_models.vessel_widths import (
+    measure_vessel_widths_between_disc_circle_pair,
+    resolve_vessel_width_circle_pair,
+)
+
+pytestmark = [pytest.mark.e2e, pytest.mark.slow]
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SAMPLE_IMAGE = REPO_ROOT / "samples" / "fundus" / "original" / "DRIVE_22.png"
+EXPECTED_VESSEL_WIDTH_COLUMNS = [
+    "image_id",
+    "inner_circle",
+    "outer_circle",
+    "inner_circle_radius_px",
+    "outer_circle_radius_px",
+    "connection_index",
+    "sample_index",
+    "x",
+    "y",
+    "width_px",
+    "x_start",
+    "y_start",
+    "x_end",
+    "y_end",
+    "vessel_type",
+]
+
+
+def _require_e2e_opt_in() -> None:
+    if os.environ.get("VASCX_RUN_E2E") != "1":
+        pytest.skip("Set VASCX_RUN_E2E=1 to run real-model end-to-end tests")
+
+
+def _device_or_skip(device_name: str) -> torch.device:
+    if device_name == "cpu":
+        return torch.device("cpu")
+    if device_name == "cuda":
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA is not available in this environment")
+        return torch.device("cuda:0")
+    if device_name == "mps":
+        if not torch.backends.mps.is_available():
+            pytest.skip("MPS is not available in this environment")
+        return torch.device("mps")
+    raise AssertionError(f"Unsupported device name: {device_name}")
+
+
+def _prepare_single_image_input(tmp_path: Path) -> tuple[str, Path, Path]:
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+
+    image_path = input_dir / SAMPLE_IMAGE.name
+    shutil.copy2(SAMPLE_IMAGE, image_path)
+    return SAMPLE_IMAGE.stem, image_path, input_dir
+
+
+def _assert_nonempty_mask(path: Path) -> None:
+    assert path.exists()
+    assert np.any(np.array(Image.open(path)) > 0)
+
+
+@pytest.mark.parametrize("device_name", ["cpu", "cuda", "mps"])
+def test_single_image_pipeline_smoke(tmp_path: Path, device_name: str) -> None:
+    _require_e2e_opt_in()
+    configure_runtime_environment()
+    device = _device_or_skip(device_name)
+    app_config = AppConfig()
+
+    image_id, image_path, input_dir = _prepare_single_image_input(tmp_path)
+
+    output_dir = tmp_path / "output"
+    output_dir.mkdir()
+    preprocessed_rgb_dir = output_dir / "preprocessed_rgb"
+    av_dir = output_dir / "artery_vein"
+    vessels_dir = output_dir / "vessels"
+    disc_dir = output_dir / "disc"
+    disc_circles_dir = output_dir / "disc_circles"
+    overlay_dir = output_dir / "overlays"
+    preprocessed_rgb_dir.mkdir()
+    av_dir.mkdir()
+    vessels_dir.mkdir()
+    disc_dir.mkdir()
+    overlay_dir.mkdir()
+
+    bounds_path = output_dir / "bounds.csv"
+    quality_path = output_dir / "quality.csv"
+    fovea_path = output_dir / "fovea.csv"
+    disc_geometry_path = output_dir / "disc_geometry.csv"
+    vessel_widths_path = output_dir / "vessel_widths.csv"
+
+    _run_preprocessing(
+        files=[image_path],
+        ids=[image_id],
+        rgb_path=preprocessed_rgb_dir,
+        bounds_path=bounds_path,
+        n_jobs=1,
+    )
+    preprocessed_image_path = preprocessed_rgb_dir / f"{image_id}.png"
+
+    df_quality = run_quality_estimation([preprocessed_image_path], ids=[image_id], device=device)
+    df_quality.to_csv(quality_path)
+
+    run_segmentation_vessels_and_av(
+        rgb_paths=[preprocessed_image_path],
+        ids=[image_id],
+        av_path=av_dir,
+        vessels_path=vessels_dir,
+        artery_color=app_config.overlay.colors.artery,
+        vein_color=app_config.overlay.colors.vein,
+        vessel_color=app_config.overlay.colors.vessel,
+        device=device,
+    )
+    run_segmentation_disc(
+        rgb_paths=[preprocessed_image_path],
+        ids=[image_id],
+        output_path=disc_dir,
+        disc_color=app_config.overlay.colors.disc,
+        device=device,
+    )
+
+    df_disc_geometry = generate_disc_circles(
+        disc_dir=disc_dir,
+        circle_output_dir=disc_circles_dir,
+        circles=app_config.overlay.circles,
+        measurements_path=disc_geometry_path,
+    )
+    inner_circle, outer_circle = resolve_vessel_width_circle_pair(
+        app_config.overlay.circles,
+        inner_circle_name=app_config.vessel_widths.inner_circle,
+        outer_circle_name=app_config.vessel_widths.outer_circle,
+    )
+    df_vessel_widths = measure_vessel_widths_between_disc_circle_pair(
+        vessels_dir=vessels_dir,
+        av_dir=av_dir,
+        disc_geometry_path=disc_geometry_path,
+        inner_circle=inner_circle,
+        outer_circle=outer_circle,
+        output_path=vessel_widths_path,
+        samples_per_connection=app_config.vessel_widths.samples_per_connection,
+    )
+
+    df_fovea = run_fovea_detection([preprocessed_image_path], ids=[image_id], device=device)
+    df_fovea.to_csv(fovea_path)
+
+    batch_create_overlays(
+        rgb_dir=preprocessed_rgb_dir,
+        output_dir=overlay_dir,
+        av_dir=av_dir,
+        disc_dir=disc_dir,
+        vessels_dir=vessels_dir,
+        circle_dirs={
+            circle.name: disc_circles_dir / circle.name
+            for circle in app_config.overlay.circles
+        },
+        vessel_width_data=df_vessel_widths,
+        fovea_data={
+            index: (row["x_fovea"], row["y_fovea"])
+            for index, row in df_fovea.iterrows()
+        },
+        overlay_config=app_config.overlay,
+    )
+
+    assert df_quality.index.tolist() == [image_id]
+    assert df_quality.columns.tolist() == ["q1", "q2", "q3"]
+    assert np.isfinite(df_quality.to_numpy()).all()
+    assert quality_path.exists()
+    assert bounds_path.exists()
+    assert preprocessed_image_path.exists()
+
+    _assert_nonempty_mask(av_dir / f"{image_id}.png")
+    _assert_nonempty_mask(vessels_dir / f"{image_id}.png")
+    _assert_nonempty_mask(disc_dir / f"{image_id}.png")
+
+    assert df_disc_geometry.index.tolist() == [image_id]
+    assert float(df_disc_geometry.loc[image_id, "disc_radius_px"]) > 0.0
+    assert disc_geometry_path.exists()
+    for circle in app_config.overlay.circles:
+        _assert_nonempty_mask(disc_circles_dir / circle.name / f"{image_id}.png")
+
+    assert vessel_widths_path.exists()
+    df_vessel_widths_disk = pd.read_csv(vessel_widths_path)
+    assert df_vessel_widths_disk.columns.tolist() == EXPECTED_VESSEL_WIDTH_COLUMNS
+    assert df_vessel_widths.columns.tolist() == EXPECTED_VESSEL_WIDTH_COLUMNS
+    if not df_vessel_widths.empty:
+        assert df_vessel_widths["image_id"].eq(image_id).all()
+        assert df_vessel_widths["vessel_type"].isin(["artery", "vein"]).all()
+        assert (df_vessel_widths["width_px"] > 0).all()
+
+    assert df_fovea.index.tolist() == [image_id]
+    assert df_fovea.columns.tolist() == ["x_fovea", "y_fovea"]
+    assert np.isfinite(df_fovea.to_numpy()).all()
+    assert fovea_path.exists()
+
+    overlay_path = overlay_dir / f"{image_id}.png"
+    assert overlay_path.exists()
+    assert Image.open(overlay_path).size == Image.open(preprocessed_image_path).size

tests/test_inference.py

@@ -0,0 +1,29 @@
+import pytest
+import torch
+
+from vascx_models import inference
+
+
+def test_resolve_device_auto_prefers_cuda_then_mps_then_cpu(monkeypatch) -> None:
+    monkeypatch.setattr(inference.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(inference.torch.backends.mps, "is_available", lambda: True)
+    assert inference.resolve_device("auto") == torch.device("cuda:0")
+
+    monkeypatch.setattr(inference.torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(inference.torch.backends.mps, "is_available", lambda: True)
+    assert inference.resolve_device("auto") == torch.device("mps")
+
+    monkeypatch.setattr(inference.torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(inference.torch.backends.mps, "is_available", lambda: False)
+    assert inference.resolve_device("auto") == torch.device("cpu")
+
+
+def test_resolve_device_rejects_unavailable_requested_accelerator(monkeypatch) -> None:
+    monkeypatch.setattr(inference.torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(inference.torch.backends.mps, "is_available", lambda: False)
+
+    with pytest.raises(RuntimeError, match="Requested device 'cuda' is not available"):
+        inference.resolve_device("cuda")
+
+    with pytest.raises(RuntimeError, match="Requested device 'mps' is not available"):
+        inference.resolve_device("mps")

vascx_models/cli.py

@@ -14,7 +14,8 @@ from .runtime import configure_runtime_environment
 configure_runtime_environment()
 
 from .inference import (
-    preferred_device,
+    available_device_types,
+    resolve_device,
     run_fovea_detection,
     run_quality_estimation,
     run_segmentation_disc,
@@ -75,6 +76,14 @@ def cli():
     default=None,
     help="Create visualization overlays. Defaults to the config value when set.",
 )
+@click.option(
+    "--device",
+    "device_name",
+    type=click.Choice(["auto", "cuda", "mps", "cpu"], case_sensitive=False),
+    default="auto",
+    show_default=True,
+    help="Inference device. 'auto' prefers CUDA first, then Apple Metal (MPS), then CPU.",
+)
 @click.option("--n_jobs", type=int, default=4, help="Number of preprocessing workers")
 def run(
     data_path,
@@ -86,6 +95,7 @@ def run(
     quality,
     fovea,
     overlay,
+    device_name,
     n_jobs,
 ):
     """Run the complete inference pipeline on fundus images.
@@ -185,9 +195,21 @@ def run(
     ids = [f.stem for f in preprocessed_files]
     logger.info("Prepared %d images for inference", len(preprocessed_files))
 
-    # Prefer hardware acceleration when the active torch build supports it.
-    device = preferred_device()
-    logger.info("Using device: %s", device)
+    available_devices = available_device_types()
+    logger.info(
+        "Device availability: cuda=%s, mps=%s, cpu=%s",
+        available_devices["cuda"],
+        available_devices["mps"],
+        available_devices["cpu"],
+    )
+    try:
+        device = resolve_device(device_name)
+    except (RuntimeError, ValueError) as exc:
+        raise click.ClickException(str(exc)) from exc
+    if device_name == "auto":
+        logger.info("Auto-selected device: %s", device)
+    else:
+        logger.info("Using requested device '%s': %s", device_name, device)
 
     # Step 2: Run quality estimation if requested
     if quality:

vascx_models/disc_circles.py

@@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
 
 
 def _save_visual_circle_mask(mask: np.ndarray, path: Path, color: tuple[int, int, int]) -> None:
-    image = Image.fromarray(mask.astype(np.uint8), mode="P")
+    image = Image.fromarray(mask.astype(np.uint8))
     palette = [0] * (256 * 3)
     palette[255 * 3 : 255 * 3 + 3] = list(color)
     image.putpalette(palette)

vascx_models/inference.py

@@ -23,7 +23,7 @@ logger = logging.getLogger(__name__)
 def _save_visual_mask(mask: np.ndarray, path: str, color_by_value: dict[int, tuple[int, int, int]]) -> None:
     """Save a label mask with a palette while preserving label values."""
     mask_uint8 = mask.squeeze().astype(np.uint8)
-    image = Image.fromarray(mask_uint8, mode="P")
+    image = Image.fromarray(mask_uint8)
     palette = [0] * (256 * 3)
     for value, color in color_by_value.items():
         start = int(value) * 3
@@ -32,12 +32,37 @@ def _save_visual_mask(mask: np.ndarray, path: str, color_by_value: dict[int, tup
     image.save(path)
 
 
+def available_device_types() -> dict[str, bool]:
+    return {
+        "cuda": torch.cuda.is_available(),
+        "mps": torch.backends.mps.is_available(),
+        "cpu": True,
+    }
+
+
 def preferred_device() -> torch.device:
-    if torch.cuda.is_available():
+    return resolve_device("auto")
+
+
+def resolve_device(device_name: str = "auto") -> torch.device:
+    available = available_device_types()
+    if device_name == "auto":
+        if available["cuda"]:
+            return torch.device("cuda:0")
+        if available["mps"]:
+            return torch.device("mps")
+        return torch.device("cpu")
+    if device_name == "cuda":
+        if not available["cuda"]:
+            raise RuntimeError("Requested device 'cuda' is not available")
         return torch.device("cuda:0")
-    if torch.backends.mps.is_available():
+    if device_name == "mps":
+        if not available["mps"]:
+            raise RuntimeError("Requested device 'mps' is not available")
         return torch.device("mps")
-    return torch.device("cpu")
+    if device_name == "cpu":
+        return torch.device("cpu")
+    raise ValueError(f"Unsupported device '{device_name}'")
 
 
 def _inference_num_workers(device: torch.device) -> int: