scripts/benchmark.py

"""Performance benchmark script for PageParse.

Usage: python scripts/benchmark.py <path> [--iterations N] [--output results.json]
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np


def benchmark_pipeline(path: str, iterations: int = 3) -> dict:
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

    from pageparse.ingest import load_image
    from pageparse.preprocess import preprocess, adaptive_preprocess
    from pageparse.ocr.handwriting import HandwritingOCR
    from pageparse.ocr.printed import PrintedOCR
    from pageparse.extract import Extractor

    print(f"Loading: {path}")
    img = load_image(path)
    h, w = img.shape[:2]
    print(f"  Image size: {w}x{h}")

    ocr = HandwritingOCR()
    printed_ocr = PrintedOCR()
    extractor = Extractor()

    results = {
        "image": str(path),
        "dimensions": f"{w}x{h}",
        "iterations": iterations,
        "stages": {},
    }

    for stage_name, stage_fn in [
        ("preprocess_basic", lambda: preprocess(img)),
        ("preprocess_adaptive", lambda: adaptive_preprocess(img)),
        ("handwriting_ocr", lambda: ocr.recognize(preprocess(img))),
        ("printed_ocr", lambda: printed_ocr.recognize(preprocess(img))),
        ("extract_todo", lambda: extractor.extract("Task 1: Buy milk\nTask 2: Review code", "benchmark.txt", "todo")),
    ]:
        durations = []
        errors = 0
        for i in range(iterations):
            t0 = time.perf_counter()
            try:
                stage_fn()
                durations.append(time.perf_counter() - t0)
            except Exception as e:
                errors += 1
                print(f"  {stage_name}: error: {e}")

        if durations:
            avg = float(np.mean(durations))
            std = float(np.std(durations))
            results["stages"][stage_name] = {
                "avg_seconds": round(avg, 4),
                "std_seconds": round(std, 4),
                "min_seconds": round(float(np.min(durations)), 4),
                "max_seconds": round(float(np.max(durations)), 4),
                "errors": errors,
            }
            print(f"  {stage_name}: avg={avg:.3f}s (std={std:.3f}s)")

    return results


def main():
    parser = argparse.ArgumentParser(description="PageParse benchmark")
    parser.add_argument("path", help="Path to image file")
    parser.add_argument("--iterations", "-n", type=int, default=3, help="Number of iterations")
    parser.add_argument("--output", "-o", type=str, help="Output JSON file")
    args = parser.parse_args()

    results = benchmark_pipeline(args.path, args.iterations)

    if args.output:
        with open(args.output, "w") as f:
            json.dump(results, f, indent=2)
        print(f"Results saved to: {args.output}")
    else:
        print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()