diff --git a/scripts/bench_pipeline.py b/scripts/bench_pipeline.py new file mode 100644 index 0000000..8c6bfa4 --- /dev/null +++ b/scripts/bench_pipeline.py @@ -0,0 +1,68 @@ +"""Benchmark the default matting pipeline (excludes model-load time). + +Runs the default config (base ViTMatte + BiRefNet + directional trimap) with a +fixed screen_color over TestImage3/4, timing 10 in-memory runs after a warmup +that loads the models. Saves only the last output to Outputs/. + + python scripts/bench_pipeline.py +""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO)) + +from bgfilter.config import load_settings, override_settings +from bgfilter.io import load_rgb, save_rgba +from bgfilter.pipeline import MattingPipeline + +SCREEN_COLOR = "#CFEFFF" +IMAGES = ["TestImage3.png", "TestImage4.png"] +N = 10 + + +def main() -> None: + settings = override_settings( + load_settings(str(REPO / "configs" / "default.yaml")), + screen_color=SCREEN_COLOR, + ) + pipe = MattingPipeline(settings) + imgs = {name: load_rgb(REPO / "samples" / name) for name in IMAGES} + + # Warmup: loads ViTMatte + BiRefNet and triggers one-time JIT. NOT timed. + t0 = time.perf_counter() + pipe.run_rgb(imgs[IMAGES[0]]) + print(f"warmup (model load + first run): {time.perf_counter() - t0:.1f}s", flush=True) + + runs: list[tuple[str, float]] = [] + last = None + for i in range(N): + name = IMAGES[i % len(IMAGES)] + t0 = time.perf_counter() + res = pipe.run_rgb(imgs[name]) + dt = time.perf_counter() - t0 + runs.append((name, dt)) + last = (name, res) + print(f" run {i + 1:2d} {name:16s} {dt:6.2f}s", flush=True) + + ts = [d for _, d in runs] + total = sum(ts) + print(f"\n{N} runs: total {total:.2f}s avg {total / N:.2f}s " + f"min {min(ts):.2f}s max {max(ts):.2f}s") + for name in IMAGES: + d = [x for nm, x in runs if nm == name] + if d: + print(f" {name}: avg {sum(d) / len(d):.2f}s over {len(d)} runs") + + name, res = last + out = REPO / "Outputs" / (Path(name).stem + "_rgba.png") + save_rgba(out, res.rgb, res.alpha) + print(f"saved last output -> {out}", flush=True) + + +if __name__ == "__main__": + main()