From bf9f63869581f3f10f4609cd5133103b59dafd3d Mon Sep 17 00:00:00 2001 From: kuuhaku Date: Fri, 17 Jul 2026 11:34:57 +0800 Subject: [PATCH] =?UTF-8?q?feat:=E6=96=B0=E5=A2=9E=E5=8F=82=E6=95=B0?= =?UTF-8?q?=E7=9F=A9=E9=98=B5=E6=B5=8B=E8=AF=95=E7=A8=8B=E5=BA=8Fbenchmark?= =?UTF-8?q?.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 2 + README.md | 40 ++- benchmark.py | 584 ++++++++++++++++++++++++++++++++++++++++ benchmarks/README.md | 18 ++ tests/test_benchmark.py | 79 ++++++ 5 files changed, 722 insertions(+), 1 deletion(-) create mode 100644 benchmark.py create mode 100644 tests/test_benchmark.py diff --git a/.gitignore b/.gitignore index d06fa47..5b90b64 100644 --- a/.gitignore +++ b/.gitignore @@ -17,5 +17,7 @@ logs/* !logs/.gitkeep benchmarks/cpu/*.json benchmarks/gpu/*.json +benchmarks/parameter-runs/ +benchmarks/参数测试报告-*.md !benchmarks/cpu/.gitkeep !benchmarks/gpu/.gitkeep diff --git a/README.md b/README.md index 92b20aa..2ac8310 100644 --- a/README.md +++ b/README.md @@ -187,10 +187,48 @@ outputs/ JSON 中保留坐标，调用方可以按具体业务继续处理。 +## 参数速度测试 + +使用根目录 `benchmark.py` 对 `data/` 中的图片和 PDF 执行参数矩阵测试，并生成 Markdown 对比报告： + +```bash +# 默认测试 tiny/small/medium × fast/standard，PDF 强制 OCR +python benchmark.py + +# 显式指定模型、预处理配置、DPI 和 CPU 线程 +python benchmark.py \ + --models tiny small medium \ + --profiles fast standard robust \ + --dpis 120 144 200 \ + --threads auto 4 8 \ + --det-limit-side-lens 64 960 \ + --det-thresholds 0.2 0.3 \ + --det-box-thresholds 0.5 0.6 \ + --det-unclip-ratios 1.5 2.0 \ + --rec-score-thresholds 0.0 0.5 \ + --rec-batch-sizes 1 6 + +# 指定报告路径 +python benchmark.py --output benchmarks/参数测试报告.md +``` + +默认场景共 6 组： + +- 三种模型：`tiny`、`small`、`medium`； +- 两种预处理配置：`fast`、`standard`； +- PDF 使用 `ocr` 模式，确保不同模型真正参与 PDF 测试； +- 每个场景在独立进程中执行，避免模型缓存到内存造成场景间干扰。 + +可测试的参数维度包括模型规格、预处理配置、PDF DPI、CPU 线程、检测边长及限制方式、检测阈值、文本框阈值、扩张比例、识别阈值和识别 Batch Size。所有列表参数会组成笛卡尔积。 + +报告统计墙钟耗时、模型初始化、图片推理、PDF OCR、纯 OCR 耗时、吞吐、识别行数和平均置信度。原始输出、日志及汇总 JSON 保存在 `benchmarks/parameter-runs/`。 + +> `robust` 会启用文档去畸变，测试显著更慢；多组 DPI 和线程参数会产生笛卡尔积，请控制测试规模。 + ## 测试 ```bash uv run --project cpu pytest -q ``` -测试使用假模型覆盖结果适配、输出路由和 PDF 混合流程，不要求每次测试都下载 OCR 模型。 +测试使用假模型覆盖结果适配、输出路由、参数测试程序和 PDF 混合流程，不要求每次测试都下载 OCR 模型。 diff --git a/benchmark.py b/benchmark.py new file mode 100644 index 0000000..f73e67e --- /dev/null +++ b/benchmark.py @@ -0,0 +1,584 @@ +"""Run PP-OCRv6 parameter combinations against data/ and write a Markdown report.""" + +from __future__ import annotations + +import argparse +import itertools +import json +import shutil +import statistics +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from datetime import datetime +from pathlib import Path +from typing import Any + +ROOT = Path(__file__).resolve().parent +DEFAULT_DATA = ROOT / "data" +DEFAULT_BENCHMARK_ROOT = ROOT / "benchmarks" / "parameter-runs" +MODEL_SIZES = ("tiny", "small", "medium") +PROFILE_OPTIONS = { + "fast": { + "doc_orientation": False, + "doc_unwarping": False, + "textline_orientation": False, + "description": "关闭文档方向、去畸变和文本行方向，优先速度", + }, + "standard": { + "doc_orientation": True, + "doc_unwarping": False, + "textline_orientation": True, + "description": "开启文档方向和文本行方向，关闭去畸变", + }, + "robust": { + "doc_orientation": True, + "doc_unwarping": True, + "textline_orientation": True, + "description": "开启方向处理和文档去畸变，优先复杂输入适应性", + }, +} + + +@dataclass(frozen=True) +class Scenario: + model_size: str + profile: str + dpi: int + threads: int | None + det_limit_side_len: int = 64 + det_limit_type: str = "min" + det_thresh: float = 0.3 + det_box_thresh: float = 0.6 + det_unclip_ratio: float = 1.5 + rec_score_thresh: float = 0.0 + rec_batch_size: int = 6 + + @staticmethod + def _slug(value: float) -> str: + return f"{value:g}".replace("-", "m").replace(".", "p") + + @property + def name(self) -> str: + thread_name = "auto" if self.threads is None else str(self.threads) + return ( + f"{self.model_size}-{self.profile}-dpi{self.dpi}-th{thread_name}" + f"-s{self.det_limit_side_len}{self.det_limit_type}" + f"-d{self._slug(self.det_thresh)}-b{self._slug(self.det_box_thresh)}" + f"-u{self._slug(self.det_unclip_ratio)}-r{self._slug(self.rec_score_thresh)}" + f"-rb{self.rec_batch_size}" + ) + + +@dataclass +class ScenarioResult: + scenario: Scenario + status: str + exit_code: int + wall_seconds: float + model_init_seconds: float | None + pure_ocr_seconds: float + image_inference_seconds: float + pdf_ocr_seconds: float + processed_files: int + image_files: int + pdf_files: int + ocr_pages: int + recognized_lines: int + mean_confidence: float | None + units_per_second: float | None + output_dir: str + log_file: str + command: list[str] + error: str | None = None + + +def parse_threads(values: list[str]) -> list[int | None]: + parsed: list[int | None] = [] + for value in values: + normalized = value.strip().lower() + if normalized == "auto": + item = None + else: + item = int(normalized) + if item < 1: + raise ValueError("线程数必须大于等于 1") + if item not in parsed: + parsed.append(item) + return parsed + + +def build_scenarios( + model_sizes: list[str], + profiles: list[str], + dpis: list[int], + threads: list[int | None], + det_limit_side_lens: list[int] | None = None, + det_limit_types: list[str] | None = None, + det_thresholds: list[float] | None = None, + det_box_thresholds: list[float] | None = None, + det_unclip_ratios: list[float] | None = None, + rec_score_thresholds: list[float] | None = None, + rec_batch_sizes: list[int] | None = None, +) -> list[Scenario]: + dimensions = itertools.product( + model_sizes, + profiles, + dpis, + threads, + det_limit_side_lens or [64], + det_limit_types or ["min"], + det_thresholds or [0.3], + det_box_thresholds or [0.6], + det_unclip_ratios or [1.5], + rec_score_thresholds or [0.0], + rec_batch_sizes or [6], + ) + return [Scenario(*values) for values in dimensions] + + +def _bool_flag(name: str, enabled: bool) -> str: + return f"--{name}" if enabled else f"--no-{name}" + + +def build_command( + scenario: Scenario, + *, + data_dir: Path, + output_dir: Path, + device: str, + device_id: int, + warmup: int, + rounds: int, + pdf_mode: str, +) -> list[str]: + profile = PROFILE_OPTIONS[scenario.profile] + command = [ + sys.executable, + str(ROOT / "ocr.py"), + str(data_dir), + "--recursive", + "--device", + device, + "--device-id", + str(device_id), + "--model-size", + scenario.model_size, + "--pdf-mode", + pdf_mode, + "--dpi", + str(scenario.dpi), + "--warmup", + str(warmup), + "--rounds", + str(rounds), + "--text-det-limit-side-len", + str(scenario.det_limit_side_len), + "--text-det-limit-type", + scenario.det_limit_type, + "--text-det-thresh", + str(scenario.det_thresh), + "--text-det-box-thresh", + str(scenario.det_box_thresh), + "--text-det-unclip-ratio", + str(scenario.det_unclip_ratio), + "--text-rec-score-thresh", + str(scenario.rec_score_thresh), + "--text-recognition-batch-size", + str(scenario.rec_batch_size), + "--output", + str(output_dir), + "--overwrite", + "--no-result", + _bool_flag("doc-orientation-classify", profile["doc_orientation"]), + _bool_flag("doc-unwarping", profile["doc_unwarping"]), + _bool_flag("textline-orientation", profile["textline_orientation"]), + ] + if scenario.threads is not None: + command.extend(["--threads", str(scenario.threads)]) + return command + + +def _read_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _collect_confidences(payload: dict[str, Any]) -> list[float]: + scores: list[float] = [] + for line in payload.get("lines", []): + value = line.get("score") + if isinstance(value, (int, float)): + scores.append(float(value)) + return scores + + +def collect_scenario_metrics( + scenario: Scenario, + *, + output_dir: Path, + log_file: Path, + command: list[str], + exit_code: int, + wall_seconds: float, + error: str | None, +) -> ScenarioResult: + batch_manifests = sorted((output_dir / "batches").glob("*.json")) + batch = _read_json(batch_manifests[-1]) if batch_manifests else {} + image_benchmarks = list((output_dir / "images").rglob("benchmark.json")) + image_results = list((output_dir / "images").rglob("result.json")) + pdf_manifests = list((output_dir / "pdfs").rglob("manifest.json")) + pdf_page_results = list((output_dir / "pdfs").rglob("pages/page-*.json")) + + image_inference = 0.0 + model_init_values: list[float] = [] + recognized_lines = 0 + confidences: list[float] = [] + for path in image_benchmarks: + payload = _read_json(path) + image_inference += sum(float(value) for value in payload.get("inference_seconds", {}).get("all", [])) + value = payload.get("model_init_seconds") + if isinstance(value, (int, float)): + model_init_values.append(float(value)) + for path in image_results: + payload = _read_json(path) + recognized_lines += int(payload.get("summary", {}).get("non_empty_lines", 0)) + confidences.extend(_collect_confidences(payload)) + + pdf_ocr_seconds = 0.0 + ocr_pages = 0 + for path in pdf_manifests: + payload = _read_json(path) + summary = payload.get("summary", {}) + ocr_pages += int(summary.get("ocr_pages", 0)) + metadata = payload.get("run_metadata", {}) + value = metadata.get("model_init_seconds") + if isinstance(value, (int, float)): + model_init_values.append(float(value)) + for page in payload.get("pages", {}).values(): + pdf_ocr_seconds += float(page.get("ocr_seconds", 0.0) or 0.0) + if page.get("source_type") == "text": + recognized_lines += int(page.get("detected_lines", 0) or 0) + for path in pdf_page_results: + payload = _read_json(path) + if "lines" in payload: + recognized_lines += int(payload.get("summary", {}).get("non_empty_lines", 0)) + confidences.extend(_collect_confidences(payload)) + + pure_ocr_seconds = image_inference + pdf_ocr_seconds + image_files = len(image_results) + pdf_files = len(pdf_manifests) + units = image_files + ocr_pages + units_per_second = units / pure_ocr_seconds if pure_ocr_seconds > 0 else None + status = "completed" if exit_code == 0 else "failed" + if exit_code == 0 and not batch: + status = "incomplete" + error = error or "未找到批处理 manifest" + + return ScenarioResult( + scenario=scenario, + status=status, + exit_code=exit_code, + wall_seconds=wall_seconds, + model_init_seconds=max(model_init_values) if model_init_values else None, + pure_ocr_seconds=pure_ocr_seconds, + image_inference_seconds=image_inference, + pdf_ocr_seconds=pdf_ocr_seconds, + processed_files=int(batch.get("completed_files", image_files + pdf_files)), + image_files=image_files, + pdf_files=pdf_files, + ocr_pages=ocr_pages, + recognized_lines=recognized_lines, + mean_confidence=statistics.fmean(confidences) if confidences else None, + units_per_second=units_per_second, + output_dir=str(output_dir), + log_file=str(log_file), + command=command, + error=error, + ) + + +def _format_seconds(value: float | None) -> str: + return "-" if value is None else f"{value:.3f}" + + +def _format_float(value: float | None, digits: int = 4) -> str: + return "-" if value is None else f"{value:.{digits}f}" + + +def render_report( + results: list[ScenarioResult], + *, + data_dir: Path, + device: str, + pdf_mode: str, + warmup: int, + rounds: int, + started_at: datetime, + finished_at: datetime, +) -> str: + completed = [result for result in results if result.status == "completed"] + by_ocr = sorted(completed, key=lambda item: item.pure_ocr_seconds or float("inf")) + by_wall = sorted(completed, key=lambda item: item.wall_seconds) + lines = [ + "# PP-OCRv6 参数测试报告", + "", + "## 测试信息", + "", + f"- 数据目录：`{data_dir}`", + f"- 运行设备：`{device}`", + f"- PDF 模式：`{pdf_mode}`", + f"- 图片预热轮数：`{warmup}`", + f"- 图片推理轮数：`{rounds}`", + f"- 场景数量：`{len(results)}`", + f"- 开始时间：`{started_at.astimezone().isoformat()}`", + f"- 结束时间：`{finished_at.astimezone().isoformat()}`", + f"- 总耗时：`{(finished_at - started_at).total_seconds():.3f}s`", + "", + "> 纯 OCR 耗时为图片推理耗时与 PDF OCR 页面耗时之和，不含模型初始化、文件扫描、PDF 文本层提取和结果导出。首次下载模型会显著增加墙钟耗时，应优先参考缓存模型后的结果。", + "", + "## 汇总对比", + "", + "| 排名 | 场景 | 状态 | 墙钟耗时(s) | 模型初始化(s) | 纯OCR耗时(s) | 图片推理(s) | PDF OCR(s) | OCR单位/秒 | 识别行数 | 平均置信度 |", + "|---:|---|---|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + rank_map = {result.scenario.name: index + 1 for index, result in enumerate(by_ocr)} + for result in results: + lines.append( + "| {rank} | `{name}` | {status} | {wall} | {init} | {pure} | {image} | {pdf} | {rate} | {count} | {confidence} |".format( + rank=rank_map.get(result.scenario.name, "-"), + name=result.scenario.name, + status=result.status, + wall=_format_seconds(result.wall_seconds), + init=_format_seconds(result.model_init_seconds), + pure=_format_seconds(result.pure_ocr_seconds), + image=_format_seconds(result.image_inference_seconds), + pdf=_format_seconds(result.pdf_ocr_seconds), + rate=_format_float(result.units_per_second, 3), + count=result.recognized_lines, + confidence=_format_float(result.mean_confidence), + ) + ) + + lines.extend(["", "## 结论", ""]) + if by_ocr: + fastest = by_ocr[0] + lines.append( + f"- 纯 OCR 耗时最短：`{fastest.scenario.name}`，耗时 `{fastest.pure_ocr_seconds:.3f}s`，吞吐 `{_format_float(fastest.units_per_second, 3)}` OCR 单位/秒。" + ) + if by_wall: + fastest_wall = by_wall[0] + lines.append( + f"- 墙钟耗时最短：`{fastest_wall.scenario.name}`，耗时 `{fastest_wall.wall_seconds:.3f}s`。" + ) + if len(by_ocr) > 1 and by_ocr[-1].pure_ocr_seconds > 0: + speedup = by_ocr[-1].pure_ocr_seconds / by_ocr[0].pure_ocr_seconds + lines.append(f"- 最快与最慢完成场景的纯 OCR 速度差约为 `{speedup:.2f}x`。") + lines.append("- 平均置信度仅作为结果稳定性的辅助观察值，不等价于有标注数据集上的识别准确率。") + + lines.extend(["", "## 参数说明", ""]) + for name, profile in PROFILE_OPTIONS.items(): + lines.append(f"- `{name}`：{profile['description']}。") + lines.extend( + [ + "", + "## 场景明细", + "", + ] + ) + for result in results: + scenario = result.scenario + lines.extend( + [ + f"### {scenario.name}", + "", + f"- 状态：`{result.status}`；退出码：`{result.exit_code}`", + f"- 模型规格：`{scenario.model_size}`", + f"- 预处理配置：`{scenario.profile}`", + f"- PDF DPI：`{scenario.dpi}`", + f"- CPU 线程：`{'auto' if scenario.threads is None else scenario.threads}`", + f"- 检测边长：`{scenario.det_limit_side_len}`；限制方式：`{scenario.det_limit_type}`", + f"- 检测阈值：`{scenario.det_thresh}`；文本框阈值：`{scenario.det_box_thresh}`；扩张比例：`{scenario.det_unclip_ratio}`", + f"- 识别阈值：`{scenario.rec_score_thresh}`；识别 Batch Size：`{scenario.rec_batch_size}`", + f"- 完成文件：`{result.processed_files}`（图片 `{result.image_files}`，PDF `{result.pdf_files}`，OCR 页 `{result.ocr_pages}`）", + f"- 墙钟耗时：`{result.wall_seconds:.3f}s`", + f"- 纯 OCR 耗时：`{result.pure_ocr_seconds:.3f}s`", + f"- 输出目录：`{result.output_dir}`", + f"- 运行日志：`{result.log_file}`", + f"- 命令：`{' '.join(result.command)}`", + ] + ) + if result.error: + lines.append(f"- 错误：`{result.error}`") + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="对 data 目录执行 PP-OCRv6 参数矩阵测试并生成 Markdown 报告", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--data", type=Path, default=DEFAULT_DATA, help="测试数据目录") + parser.add_argument("--device", choices=("cpu", "gpu"), default="cpu", help="测试设备") + parser.add_argument("--device-id", type=int, default=0, help="GPU 编号") + parser.add_argument("--models", nargs="+", choices=MODEL_SIZES, default=list(MODEL_SIZES), help="模型规格列表") + parser.add_argument("--profiles", nargs="+", choices=tuple(PROFILE_OPTIONS), default=["fast", "standard"], help="预处理配置列表") + parser.add_argument("--dpis", nargs="+", type=int, default=[144], help="PDF 渲染 DPI 列表") + parser.add_argument("--threads", nargs="+", default=["auto"], help="CPU 线程列表，可使用 auto") + parser.add_argument("--det-limit-side-lens", nargs="+", type=int, default=[64], help="文本检测边长列表") + parser.add_argument("--det-limit-types", nargs="+", choices=("min", "max"), default=["min"], help="文本检测边长限制方式列表") + parser.add_argument("--det-thresholds", nargs="+", type=float, default=[0.3], help="文本检测像素阈值列表") + parser.add_argument("--det-box-thresholds", nargs="+", type=float, default=[0.6], help="文本框阈值列表") + parser.add_argument("--det-unclip-ratios", nargs="+", type=float, default=[1.5], help="文本框扩张比例列表") + parser.add_argument("--rec-score-thresholds", nargs="+", type=float, default=[0.0], help="识别置信度阈值列表") + parser.add_argument("--rec-batch-sizes", nargs="+", type=int, default=[6], help="文本识别 Batch Size 列表") + parser.add_argument("--warmup", type=int, default=1, help="每个场景首张图片预热轮数") + parser.add_argument("--rounds", type=int, default=1, help="每张图片推理轮数") + parser.add_argument("--pdf-mode", choices=("ocr", "hybrid", "text"), default="ocr", help="PDF 测试模式；速度对比推荐 ocr") + parser.add_argument("--output", type=Path, default=None, help="Markdown 报告路径") + parser.add_argument("--work-dir", type=Path, default=None, help="各场景原始输出目录") + parser.add_argument("--overwrite", action="store_true", help="允许删除已存在的 work-dir") + parser.add_argument("--fail-fast", action="store_true", help="任一场景失败后停止") + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + data_dir = args.data.expanduser().resolve() + if not data_dir.is_dir(): + print(f"测试数据目录不存在: {data_dir}", file=sys.stderr) + return 2 + if args.warmup < 0 or args.rounds < 1: + print("--warmup 必须 >= 0，--rounds 必须 >= 1", file=sys.stderr) + return 2 + if any(dpi < 72 or dpi > 600 for dpi in args.dpis): + print("--dpis 必须在 72 到 600 之间", file=sys.stderr) + return 2 + if any(value < 1 for value in args.det_limit_side_lens): + print("--det-limit-side-lens 必须大于等于 1", file=sys.stderr) + return 2 + if any(value < 1 for value in args.rec_batch_sizes): + print("--rec-batch-sizes 必须大于等于 1", file=sys.stderr) + return 2 + try: + thread_values = parse_threads(args.threads) + except ValueError as exc: + print(f"无效线程参数: {exc}", file=sys.stderr) + return 2 + + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + work_dir = (args.work_dir or DEFAULT_BENCHMARK_ROOT / timestamp).expanduser().resolve() + report_path = (args.output or ROOT / "benchmarks" / f"参数测试报告-{timestamp}.md").expanduser().resolve() + if work_dir.exists(): + if not args.overwrite: + print(f"测试输出目录已存在: {work_dir}；请使用 --overwrite", file=sys.stderr) + return 2 + shutil.rmtree(work_dir) + work_dir.mkdir(parents=True, exist_ok=True) + report_path.parent.mkdir(parents=True, exist_ok=True) + + scenarios = build_scenarios( + args.models, + args.profiles, + args.dpis, + thread_values, + args.det_limit_side_lens, + args.det_limit_types, + args.det_thresholds, + args.det_box_thresholds, + args.det_unclip_ratios, + args.rec_score_thresholds, + args.rec_batch_sizes, + ) + print(f"计划执行 {len(scenarios)} 个参数场景，报告将写入: {report_path}") + started_at = datetime.now().astimezone() + results: list[ScenarioResult] = [] + for index, scenario in enumerate(scenarios, 1): + scenario_dir = work_dir / scenario.name + output_dir = scenario_dir / "outputs" + log_file = scenario_dir / "run.log" + scenario_dir.mkdir(parents=True, exist_ok=True) + command = build_command( + scenario, + data_dir=data_dir, + output_dir=output_dir, + device=args.device, + device_id=args.device_id, + warmup=args.warmup, + rounds=args.rounds, + pdf_mode=args.pdf_mode, + ) + print(f"[{index}/{len(scenarios)}] {scenario.name}") + wall_started = time.perf_counter() + error = None + try: + completed = subprocess.run( + command, + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=False, + ) + raw_output = completed.stdout or b"" + log_file.write_bytes(raw_output) + exit_code = completed.returncode + if exit_code != 0: + text = raw_output.decode("utf-8", errors="replace") + error = next((line for line in reversed(text.splitlines()) if line.strip()), f"退出码 {exit_code}") + except Exception as exc: + exit_code = 1 + error = f"{type(exc).__name__}: {exc}" + log_file.write_text(error + "\n", encoding="utf-8") + wall_seconds = time.perf_counter() - wall_started + result = collect_scenario_metrics( + scenario, + output_dir=output_dir, + log_file=log_file, + command=command, + exit_code=exit_code, + wall_seconds=wall_seconds, + error=error, + ) + results.append(result) + print( + f" 状态={result.status} 墙钟={result.wall_seconds:.3f}s " + f"纯OCR={result.pure_ocr_seconds:.3f}s" + ) + if result.status != "completed" and args.fail_fast: + break + + finished_at = datetime.now().astimezone() + report = render_report( + results, + data_dir=data_dir, + device=args.device, + pdf_mode=args.pdf_mode, + warmup=args.warmup, + rounds=args.rounds, + started_at=started_at, + finished_at=finished_at, + ) + report_path.write_text(report, encoding="utf-8") + summary_path = work_dir / "summary.json" + summary_path.write_text( + json.dumps( + [ + { + **asdict(result), + "scenario": asdict(result.scenario), + } + for result in results + ], + ensure_ascii=False, + indent=2, + ), + encoding="utf-8", + ) + print(f"测试完成，Markdown 报告: {report_path}") + print(f"原始汇总 JSON: {summary_path}") + return 0 if all(result.status == "completed" for result in results) else 3 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/README.md b/benchmarks/README.md index f6a0f5d..d0f1d2d 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -17,3 +17,21 @@ python ocr.py data/images/手写01.png --warmup 1 --rounds 3 ```bash python ocr.py data/images/手写01.png --benchmark-json benchmarks/cpu/手写01.json ``` + +## 参数矩阵测试 + +```bash +python benchmark.py +``` + +默认比较 `tiny/small/medium` 与 `fast/standard` 预处理配置，报告生成到： + +```text +benchmarks/参数测试报告-<时间戳>.md +``` + +各场景原始结果和汇总 JSON 保存在： + +```text +benchmarks/parameter-runs/<时间戳>/ +``` diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py new file mode 100644 index 0000000..f03ea78 --- /dev/null +++ b/tests/test_benchmark.py @@ -0,0 +1,79 @@ +from pathlib import Path + +from benchmark import Scenario, build_command, build_scenarios, parse_threads, render_report, ScenarioResult + + +def test_build_default_matrix(): + scenarios = build_scenarios( + ["tiny", "small", "medium"], + ["fast", "standard"], + [144], + [None], + ) + assert len(scenarios) == 6 + assert scenarios[0].name == "tiny-fast-dpi144-thauto-s64min-d0p3-b0p6-u1p5-r0-rb6" + assert scenarios[-1].name == "medium-standard-dpi144-thauto-s64min-d0p3-b0p6-u1p5-r0-rb6" + + +def test_parse_threads(): + assert parse_threads(["auto", "4", "4"]) == [None, 4] + + +def test_command_contains_model_and_profile_flags(tmp_path): + scenario = Scenario("small", "fast", 120, 4) + command = build_command( + scenario, + data_dir=tmp_path / "data", + output_dir=tmp_path / "output", + device="cpu", + device_id=0, + warmup=1, + rounds=2, + pdf_mode="ocr", + ) + assert command[1].endswith("ocr.py") + assert command[command.index("--model-size") + 1] == "small" + assert "--no-doc-orientation-classify" in command + assert "--no-textline-orientation" in command + assert command[command.index("--threads") + 1] == "4" + assert command[command.index("--text-det-thresh") + 1] == "0.3" + assert command[command.index("--text-recognition-batch-size") + 1] == "6" + + +def test_report_contains_comparison_table(tmp_path): + scenario = Scenario("tiny", "fast", 144, None) + result = ScenarioResult( + scenario=scenario, + status="completed", + exit_code=0, + wall_seconds=3.0, + model_init_seconds=1.0, + pure_ocr_seconds=2.0, + image_inference_seconds=1.0, + pdf_ocr_seconds=1.0, + processed_files=2, + image_files=1, + pdf_files=1, + ocr_pages=1, + recognized_lines=10, + mean_confidence=0.95, + units_per_second=1.0, + output_dir=str(tmp_path / "output"), + log_file=str(tmp_path / "run.log"), + command=["python", "ocr.py"], + ) + from datetime import datetime + + report = render_report( + [result], + data_dir=Path("data"), + device="cpu", + pdf_mode="ocr", + warmup=1, + rounds=1, + started_at=datetime.now().astimezone(), + finished_at=datetime.now().astimezone(), + ) + assert "# PP-OCRv6 参数测试报告" in report + assert "tiny-fast-dpi144-thauto-s64min-d0p3-b0p6-u1p5-r0-rb6" in report + assert "纯 OCR 耗时最短" in report