src/veil_pgd/eval/harness.py

"""Run the full protect pipeline over a labeled image set and report metrics.

Also evaluates robustness by re-querying the black-box targets on the protected
image after the `hard_transforms` ceiling set.
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

from PIL import Image

from veil_pgd.config import STEALTH_PRESETS, get_settings, stealth_preset
from veil_pgd.eval.metrics import summarize
from veil_pgd.pipeline import protect_image
from veil_pgd.targets.registry import Registry


def _log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True, file=sys.stderr)


def _read_rows(manifest_csv: str) -> list[tuple[str, str]]:
    rows = []
    for line in Path(manifest_csv).read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        path, truth = line.split(",", 1)
        rows.append((path.strip(), truth.strip()))
    return rows


def run_eval(manifest_csv: str, out_dir: str = "artifacts/eval") -> dict:
    """Single-level eval at the configured stealth gate. CSV: `image_path,truth`."""
    s = get_settings()
    rows = _read_rows(manifest_csv)
    distances: list[float] = []
    stealth_passes: list[bool] = []
    per_image = []
    reg = Registry(s)
    try:
        for path, truth in rows:
            res = protect_image(path, truth, settings=s, reg=reg)
            best_dist = max(res.manifest["per_model_distance"].values(), default=0.0)
            distances.append(best_dist)
            st = res.manifest.get("stealth") or {}
            stealth_passes.append(bool(st.get("passed")))
            per_image.append({"image": path, "truth": truth, "distance": best_dist,
                              "spec": res.manifest["spec"]})
    finally:
        reg.close()

    summary = summarize(distances, stealth_passes)
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    report = {"summary": summary.__dict__, "per_image": per_image}
    (out / "eval_report.json").write_text(json.dumps(report, indent=2))
    return report


def run_sweep(
    manifest_csv: str,
    levels: list[str] | None = None,
    out_dir: str = "artifacts/eval",
    skip_blackbox: bool = False,
) -> dict:
    """Stealth sweep: run every image at each stealth level to map the
    attack-strength vs human-visibility frontier.

    Returns {level: {summary, per_image}} plus a compact frontier table.
    """
    s = get_settings()
    levels = levels or list(STEALTH_PRESETS.keys())
    rows = _read_rows(manifest_csv)

    results: dict[str, dict] = {}
    frontier = []
    reg = Registry(s)
    try:
        for level in levels:
            th = stealth_preset(level)
            distances: list[float] = []
            stealth_passes: list[bool] = []
            per_image = []
            _log(f"level={level}: {len(rows)} images")
            for i, (path, truth) in enumerate(rows, 1):
                t0 = time.time()
                try:
                    res = protect_image(
                        path, truth, settings=s, stealth_thresholds=th, reg=reg,
                        skip_blackbox=skip_blackbox,
                    )
                    best_dist = max(res.manifest["per_model_distance"].values(), default=0.0)
                    st = res.manifest.get("stealth") or {}
                    passed = bool(st.get("passed"))
                    spec = res.manifest["spec"]
                    _log(f"  [{i}/{len(rows)}] {Path(path).name} truth={truth!r} "
                         f"dist={best_dist:.3f} stealth={'pass' if passed else 'FAIL'} "
                         f"text={spec.get('text')!r} ({time.time()-t0:.0f}s)")
                except Exception as e:  # noqa: BLE001 - record failure, keep sweeping
                    best_dist, passed, spec = 0.0, False, {"error": str(e)}
                    _log(f"  [{i}/{len(rows)}] {Path(path).name} FAILED: {e} "
                         f"({time.time()-t0:.0f}s)")
                distances.append(best_dist)
                stealth_passes.append(passed)
                per_image.append({"image": path, "truth": truth,
                                  "distance": best_dist, "spec": spec})
            summary = summarize(distances, stealth_passes)
            results[level] = {"summary": summary.__dict__, "per_image": per_image}
            frontier.append({
                "level": level,
                "mean_distance": summary.mean_distance,
                "success@0.5": summary.poison_success_rate.get(0.5, 0.0),
                "stealth_pass_rate": summary.stealth_pass_rate,
            })
    finally:
        reg.close()

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    report = {"levels": results, "frontier": frontier}
    (out / "sweep_report.json").write_text(json.dumps(report, indent=2))
    return report