add benchmark script

This commit is contained in:
2026-07-02 12:26:20 +00:00
parent d24c6ed608
commit 7a87a00ca1
+68
View File
@@ -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()