From 1cf2c729f35899e0723a8fd0cce4523c56367e32 Mon Sep 17 00:00:00 2001 From: kuuhaku Date: Fri, 17 Jul 2026 11:35:47 +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 | 13 +- README.md | 126 +++- benchmark.py | 711 +++++++++++++++++++++++ benchmarks/parameter-matrix.example.json | 58 ++ ocr.py | 5 + ocr_app/cli.py | 8 +- ocr_app/commands.py | 18 +- tests/test_benchmark.py | 127 ++++ 8 files changed, 1055 insertions(+), 11 deletions(-) create mode 100644 benchmark.py create mode 100644 benchmarks/parameter-matrix.example.json create mode 100644 tests/test_benchmark.py diff --git a/.gitignore b/.gitignore index 00cd6bf..dd4af31 100644 --- a/.gitignore +++ b/.gitignore @@ -15,15 +15,16 @@ gpu/.gpu-ready # Generated benchmark results benchmarks/cpu/*.json benchmarks/gpu/*.json +benchmarks/runs/ !benchmarks/cpu/.gitkeep !benchmarks/gpu/.gitkeep # OCR outputs -# outputs/ +outputs/ # Generated structured logs (legacy logs directly under logs/ remain tracked) -# logs/input/ -# logs/verify/ -# logs/image/ -# logs/batch/ -# logs/pdf/ +logs/input/ +logs/verify/ +logs/image/ +logs/batch/ +logs/pdf/ diff --git a/README.md b/README.md index 1a8491b..d231da6 100644 --- a/README.md +++ b/README.md @@ -11,12 +11,14 @@ PDF 默认使用 **文本提取 + OCR 混合模式**：优先提取 PDF 原始 ```text ocr-VL1.6/ ├── ocr.py # 唯一用户入口，自动选择 CPU/GPU 子环境 +├── benchmark.py # data 参数矩阵测试与 Markdown 报告 ├── ocr_app/ # CPU/GPU 共享业务代码 │ ├── cli.py # image/batch/pdf/verify 命令 │ ├── commands.py # 命令实现 │ ├── runtime.py # 设备验证、模型延迟加载 │ ├── pdf.py # 混合 PDF、断点续传、导出 │ ├── pdf_text.py # 文本层提取与质量评估 +│ ├── output.py # 图片/PDF/批次统一输出路径和保存 │ └── logging_utils.py # UTF-8 结构化日志 ├── cpu/ │ ├── pyproject.toml # CPU 独立依赖 @@ -388,7 +390,127 @@ logs/legacy/ 日志使用 UTF-8。Windows 控制台即使显示乱码，日志文件中的中文仍正常。 -## 测试 +## 参数矩阵性能测试 + +新增 `benchmark.py`，会对 `data/` 中的图片和 PDF 按多组参数运行，统计速度并增量生成 Markdown 报告。 + +也可通过统一入口调用： + +```bash +python ocr.py benchmark --profile smoke --device cpu +``` + +默认报告： + +```text +benchmarks/parameter-comparison.md +``` + +每个用例的日志、输出和批次 manifest 保存在： + +```text +benchmarks/runs/<运行时间>// +``` + +### 预设矩阵 + +```bash +# 较快：图片 token=64/128、PDF hybrid/text +python ocr.py benchmark --profile smoke --device cpu + +# 常用对比：增加 CPU 线程、默认 token、PDF 阈值和 OCR DPI +python ocr.py benchmark --profile standard --device cpu + +# 完整矩阵：更多线程、token 和 DPI，耗时最长 +python ocr.py benchmark --profile full --device cpu +``` + +真实 CPU OCR 很慢，`standard/full` 可能运行数小时。建议先运行 `smoke`，或通过 `--case` 单独运行： + +```bash +python ocr.py benchmark \ + --profile standard \ + --case image-threads-8 \ + --case image-threads-safe \ + --device cpu +``` + +查看用例列表而不执行： + +```bash +python ocr.py benchmark --profile standard --list-cases +``` + +仅生成计划和空报告： + +```bash +python ocr.py benchmark --profile standard --dry-run +``` + +### 自定义参数矩阵 + +示例配置： + +```text +benchmarks/parameter-matrix.example.json +``` + +运行： + +```bash +python ocr.py benchmark \ + --config benchmarks/parameter-matrix.example.json \ + --device cpu \ + --report benchmarks/custom-comparison.md +``` + +每个 case 支持统一入口中的参数，例如： + +```json +{ + "id": "image-threads-8", + "title": "图片 8 线程", + "group": "image-threads", + "kind": "image", + "params": { + "threads": 8, + "max_new_tokens": 64, + "rounds": 1 + } +} +``` + +### Markdown 报告统计 + +报告按参数组比较： + +- 系统、CPU、逻辑核心和内存信息 +- 测试文件清单、类型和大小 +- Wall time +- 模型初始化耗时 +- 图片核心推理 / PDF OCR 总耗时 +- 平均每文件耗时 +- 每分钟处理文件数 +- 相对基准加速比 +- 输出文本字符数 +- 文本 SHA-256 是否一致 +- PDF Text/OCR 页数 +- PDF 路由原因统计 +- 完整命令、日志和输出路径 + +报告每完成一个 case 就原子更新；按 `Ctrl+C` 中断时，已经完成的结果不会丢失。 + +其他参数： + +```text +--data PATH 测试数据目录 +--report PATH Markdown 报告位置 +--run-root PATH 详细运行数据目录 +--timeout SEC 单个用例超时 +--case ID 只执行指定用例，可重复使用 +``` + +## 自动化测试 ```bash uv run --project cpu pytest -q @@ -405,7 +527,7 @@ uv run --project cpu pytest -q 当前结果： ```text -17 passed +21 passed ``` ## 已验证状态 diff --git a/benchmark.py b/benchmark.py new file mode 100644 index 0000000..07b7f05 --- /dev/null +++ b/benchmark.py @@ -0,0 +1,711 @@ +"""Run OCR parameter matrices against data/ and generate an incremental Markdown report.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import platform +import re +import signal +import subprocess +import sys +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any + +ROOT = Path(__file__).resolve().parent +DEFAULT_DATA = ROOT / "data" +DEFAULT_REPORT = ROOT / "benchmarks" / "parameter-comparison.md" + + +@dataclass +class BenchmarkCase: + id: str + title: str + group: str + kind: str + params: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class CaseResult: + case: BenchmarkCase + status: str + exit_code: int | None + command: list[str] + started_at: str + wall_seconds: float + output_dir: str + log_file: str + batch_manifest: str | None = None + discovered_files: int = 0 + completed_files: int = 0 + failed_files: int = 0 + model_init_seconds: float = 0.0 + program_total_seconds: float = 0.0 + total_file_seconds: float = 0.0 + core_processing_seconds: float = 0.0 + average_file_seconds: float = 0.0 + throughput_files_per_minute: float = 0.0 + recognized_chars: int = 0 + result_hash: str | None = None + text_pages: int = 0 + ocr_pages: int = 0 + routing_reasons: dict[str, int] = field(default_factory=dict) + error: str | None = None + + +def safe_name(value: str) -> str: + cleaned = re.sub(r"[^\w.-]+", "_", value, flags=re.UNICODE).strip("._") + return cleaned or "case" + + +def profile_cases(profile: str, device: str) -> list[BenchmarkCase]: + safe_threads = max(1, (os.cpu_count() or 4) - 2) + if device == "gpu": + smoke = [ + BenchmarkCase("image-baseline", "图片基准", "image-rounds", "image", {"rounds": 1}), + BenchmarkCase("image-token-64", "图片 max_new_tokens=64", "image-tokens", "image", {"max_new_tokens": 64}), + BenchmarkCase("image-token-128", "图片 max_new_tokens=128", "image-tokens", "image", {"max_new_tokens": 128}), + BenchmarkCase("pdf-hybrid", "PDF 混合模式", "pdf-mode", "pdf", {"pdf_mode": "hybrid"}), + BenchmarkCase("pdf-text", "PDF 文本模式", "pdf-mode", "pdf", {"pdf_mode": "text"}), + ] + else: + smoke = [ + BenchmarkCase("image-token-64", "图片 max_new_tokens=64", "image-tokens", "image", {"threads": safe_threads, "max_new_tokens": 64}), + BenchmarkCase("image-token-128", "图片 max_new_tokens=128", "image-tokens", "image", {"threads": safe_threads, "max_new_tokens": 128}), + BenchmarkCase("pdf-hybrid", "PDF 混合模式", "pdf-mode", "pdf", {"pdf_mode": "hybrid"}), + BenchmarkCase("pdf-text", "PDF 文本模式", "pdf-mode", "pdf", {"pdf_mode": "text"}), + ] + + if profile == "smoke": + return smoke + + standard = list(smoke) + if device == "cpu": + standard.extend( + [ + BenchmarkCase("image-threads-1", "图片 1 线程（token=64）", "image-threads", "image", {"threads": 1, "max_new_tokens": 64}), + BenchmarkCase("image-threads-8", "图片 8 线程（token=64）", "image-threads", "image", {"threads": 8, "max_new_tokens": 64}), + BenchmarkCase("image-threads-safe", f"图片 {safe_threads} 线程（token=64）", "image-threads", "image", {"threads": safe_threads, "max_new_tokens": 64}), + BenchmarkCase("image-token-default", "图片默认 token", "image-tokens", "image", {"threads": safe_threads}), + ] + ) + standard.extend( + [ + BenchmarkCase("pdf-hybrid-threshold-default", "PDF 混合默认阈值", "pdf-threshold", "pdf", {"pdf_mode": "hybrid"}), + BenchmarkCase("pdf-hybrid-minchars-20", "PDF 混合 min_chars=20", "pdf-threshold", "pdf", {"pdf_mode": "hybrid", "text_min_chars": 20}), + BenchmarkCase("pdf-hybrid-minchars-100", "PDF 混合 min_chars=100", "pdf-threshold", "pdf", {"pdf_mode": "hybrid", "text_min_chars": 100}), + BenchmarkCase("pdf-ocr-dpi-72", "PDF 强制 OCR DPI 72（每份第 1 页）", "pdf-ocr-dpi", "pdf", {"pdf_mode": "ocr", "pages": "1", "dpi": 72, "max_new_tokens": 64}), + BenchmarkCase("pdf-ocr-dpi-144", "PDF 强制 OCR DPI 144（每份第 1 页）", "pdf-ocr-dpi", "pdf", {"pdf_mode": "ocr", "pages": "1", "dpi": 144, "max_new_tokens": 64}), + ] + ) + if profile == "standard": + return _deduplicate_cases(standard) + + full = list(standard) + if device == "cpu": + for threads in (2, 4, 12, 16): + full.append( + BenchmarkCase( + f"image-threads-{threads}", + f"图片 {threads} 线程", + "image-threads", + "image", + {"threads": threads, "max_new_tokens": 64}, + ) + ) + for dpi in (120, 200): + full.append( + BenchmarkCase( + f"pdf-ocr-dpi-{dpi}", + f"PDF 强制 OCR DPI {dpi}（每份第 1 页）", + "pdf-ocr-dpi", + "pdf", + {"pdf_mode": "ocr", "pages": "1", "dpi": dpi, "max_new_tokens": 64}, + ) + ) + for max_tokens in (32, 256): + full.append( + BenchmarkCase( + f"image-token-{max_tokens}", + f"图片 max_new_tokens={max_tokens}", + "image-tokens", + "image", + {"max_new_tokens": max_tokens, **({"threads": safe_threads} if device == "cpu" else {})}, + ) + ) + return _deduplicate_cases(full) + + +def _deduplicate_cases(cases: list[BenchmarkCase]) -> list[BenchmarkCase]: + seen: set[str] = set() + result = [] + for case in cases: + if case.id not in seen: + seen.add(case.id) + result.append(case) + return result + + +def load_config(path: Path) -> list[BenchmarkCase]: + data = json.loads(path.read_text(encoding="utf-8")) + raw_cases = data.get("cases", data) if isinstance(data, dict) else data + if not isinstance(raw_cases, list): + raise ValueError("配置必须是 case 数组，或包含 cases 数组的对象") + cases = [] + for item in raw_cases: + case = BenchmarkCase( + id=str(item["id"]), + title=str(item.get("title", item["id"])), + group=str(item.get("group", "custom")), + kind=str(item["kind"]), + params=dict(item.get("params", {})), + ) + if case.kind not in {"image", "pdf"}: + raise ValueError(f"case {case.id} 的 kind 必须是 image/pdf") + cases.append(case) + ids = [case.id for case in cases] + if len(ids) != len(set(ids)): + raise ValueError("case id 必须唯一") + return cases + + +def param_to_cli(key: str, value: Any) -> list[str]: + option = "--" + key.replace("_", "-") + if isinstance(value, bool): + return [option] if value else [] + if value is None: + return [] + return [option, str(value)] + + +def build_command( + case: BenchmarkCase, + *, + data_dir: Path, + device: str, + output_dir: Path, + log_file: Path, +) -> list[str]: + command = [ + sys.executable, + str(ROOT / "ocr.py"), + str(data_dir), + "--recursive", + "--only", + case.kind, + "--device", + device, + "--output", + str(output_dir), + "--log-file", + str(log_file), + "--no-result", + "--overwrite", + ] + for key, value in case.params.items(): + command.extend(param_to_cli(key, value)) + return command + + +def file_text_and_hash(paths: list[Path]) -> tuple[int, str | None]: + digest = hashlib.sha256() + character_count = 0 + has_content = False + for path in sorted(paths, key=lambda value: str(value).casefold()): + if not path.is_file(): + continue + text = path.read_text(encoding="utf-8", errors="replace") + normalized = "\n".join(line.rstrip() for line in text.replace("\r\n", "\n").split("\n")).strip() + character_count += len("".join(normalized.split())) + digest.update(str(path.name).encode("utf-8")) + digest.update(normalized.encode("utf-8")) + has_content = True + return character_count, digest.hexdigest() if has_content else None + + +def latest_batch_manifest(output_dir: Path) -> Path | None: + files = sorted((output_dir / "batches").glob("*.json"), key=lambda path: path.stat().st_mtime) + return files[-1] if files else None + + +def parse_case_result( + case: BenchmarkCase, + *, + command: list[str], + exit_code: int, + started_at: str, + wall_seconds: float, + output_dir: Path, + log_file: Path, + stderr_tail: str, +) -> CaseResult: + manifest_path = latest_batch_manifest(output_dir) + result = CaseResult( + case=case, + status="completed" if exit_code == 0 else "failed", + exit_code=exit_code, + command=command, + started_at=started_at, + wall_seconds=wall_seconds, + output_dir=str(output_dir), + log_file=str(log_file), + batch_manifest=str(manifest_path) if manifest_path else None, + error=stderr_tail or None if exit_code else None, + ) + if manifest_path is None: + return result + + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + result.discovered_files = int(manifest.get("discovered_files", 0)) + result.completed_files = int(manifest.get("completed_files", 0)) + result.failed_files = int(manifest.get("failed_files", 0)) + result.model_init_seconds = float(manifest.get("model_init_seconds", 0.0)) + result.program_total_seconds = float(manifest.get("program_total_seconds", 0.0)) + result.total_file_seconds = float(manifest.get("total_file_seconds", 0.0)) + result.average_file_seconds = ( + result.total_file_seconds / result.completed_files if result.completed_files else 0.0 + ) + result.throughput_files_per_minute = ( + result.completed_files * 60.0 / wall_seconds if wall_seconds else 0.0 + ) + + text_paths: list[Path] = [] + core_seconds = 0.0 + for item in manifest.get("results", []): + outputs = item.get("outputs", {}) + if item.get("kind") == "image": + benchmark_path = Path(outputs.get("benchmark", "")) + text_path = Path(outputs.get("text", "")) + if benchmark_path.is_file(): + benchmark = json.loads(benchmark_path.read_text(encoding="utf-8")) + core_seconds += sum(benchmark.get("inference_seconds", {}).get("all", [])) + if text_path.is_file(): + text_paths.append(text_path) + elif item.get("kind") == "pdf": + pdf_manifest_path = Path(outputs.get("manifest_path", "")) + document_dir = Path(outputs.get("document_dir", "")) + if pdf_manifest_path.is_file(): + pdf_manifest = json.loads(pdf_manifest_path.read_text(encoding="utf-8")) + summary = pdf_manifest.get("summary", {}) + result.text_pages += int(summary.get("text_pages", 0)) + result.ocr_pages += int(summary.get("ocr_pages", 0)) + timing = summary.get("timing", {}) + core_seconds += float(timing.get("ocr_total_seconds", 0.0)) + for page in pdf_manifest.get("pages", {}).values(): + reason = page.get("routing_reason", "unknown") + result.routing_reasons[reason] = result.routing_reasons.get(reason, 0) + 1 + document_md = document_dir / "document.md" + if document_md.is_file(): + text_paths.append(document_md) + + result.core_processing_seconds = core_seconds + result.recognized_chars, result.result_hash = file_text_and_hash(text_paths) + return result + + +def format_seconds(value: float) -> str: + if value < 60: + return f"{value:.2f}s" + return f"{value / 60:.2f}min" + + +def params_text(params: dict[str, Any]) -> str: + return ", ".join(f"{key}={value}" for key, value in params.items()) or "默认" + + +def md_escape(value: Any) -> str: + return str(value).replace("|", "\\|").replace("\n", "
") + + +def data_inventory(data_dir: Path) -> list[dict[str, Any]]: + supported = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp", ".pdf"} + inventory = [] + for path in sorted(data_dir.rglob("*"), key=lambda value: str(value).casefold()): + if path.is_file() and path.suffix.lower() in supported: + inventory.append( + { + "path": path.relative_to(data_dir).as_posix(), + "type": "pdf" if path.suffix.lower() == ".pdf" else "image", + "size_bytes": path.stat().st_size, + } + ) + return inventory + + +def memory_info() -> tuple[str, str]: + try: + import psutil + + memory = psutil.virtual_memory() + return f"{memory.total / (1024**3):.1f} GB", f"{memory.available / (1024**3):.1f} GB" + except Exception: + pass + + if sys.platform == "win32": + try: + import ctypes + + class MemoryStatus(ctypes.Structure): + _fields_ = [ + ("dwLength", ctypes.c_ulong), + ("dwMemoryLoad", ctypes.c_ulong), + ("ullTotalPhys", ctypes.c_ulonglong), + ("ullAvailPhys", ctypes.c_ulonglong), + ("ullTotalPageFile", ctypes.c_ulonglong), + ("ullAvailPageFile", ctypes.c_ulonglong), + ("ullTotalVirtual", ctypes.c_ulonglong), + ("ullAvailVirtual", ctypes.c_ulonglong), + ("ullAvailExtendedVirtual", ctypes.c_ulonglong), + ] + + status = MemoryStatus() + status.dwLength = ctypes.sizeof(MemoryStatus) + if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)): + return ( + f"{status.ullTotalPhys / (1024**3):.1f} GB", + f"{status.ullAvailPhys / (1024**3):.1f} GB", + ) + except Exception: + pass + return "unknown", "unknown" + + +def render_report( + *, + report_path: Path, + run_id: str, + profile: str, + device: str, + data_dir: Path, + cases: list[BenchmarkCase], + results: list[CaseResult], + run_root: Path, + started_at: str, +) -> None: + by_id = {result.case.id: result for result in results} + inventory = data_inventory(data_dir) + total_memory, available_memory = memory_info() + groups: dict[str, list[BenchmarkCase]] = {} + for case in cases: + groups.setdefault(case.group, []).append(case) + + lines = [ + "# PaddleOCR-VL 参数性能对比", + "", + "## 运行信息", + "", + f"- 运行 ID：`{run_id}`", + f"- 开始时间：`{started_at}`", + f"- 更新时间：`{datetime.now().astimezone().isoformat()}`", + f"- Profile：`{profile}`", + f"- Device：`{device}`", + f"- Data：`{data_dir}`", + f"- 系统：`{platform.platform()}`", + f"- Python：`{platform.python_version()}`", + f"- CPU：`{platform.processor() or 'unknown'}`", + f"- CPU 逻辑核心：`{os.cpu_count()}`", + f"- 内存：总计 `{total_memory}`，当前可用 `{available_memory}`", + f"- 测试文件：`{len(inventory)}`（图片 `{sum(item['type'] == 'image' for item in inventory)}`，PDF `{sum(item['type'] == 'pdf' for item in inventory)}`）", + f"- 运行目录：`{run_root}`", + f"- 完成进度：`{len(results)}/{len(cases)}`", + "", + "### 测试数据", + "", + "| 文件 | 类型 | 大小 |", + "|---|---:|---:|", + *[ + f"| {md_escape(item['path'])} | {item['type']} | {item['size_bytes'] / 1024:.1f} KiB |" + for item in inventory + ], + "", + "> 加速比以同组第一个成功用例的 wall time 为基准。结果一致性使用输出文本 SHA-256 比较；仅表示文本完全一致，不代表 OCR 准确率。", + "", + ] + + for group, group_cases in groups.items(): + lines.extend([f"## 参数组：{group}", ""]) + baseline = next( + (by_id[case.id] for case in group_cases if case.id in by_id and by_id[case.id].status == "completed"), + None, + ) + lines.extend( + [ + "| 用例 | 参数 | 状态 | 文件 | Wall time | 模型初始化 | 核心推理/OCR | 平均每文件 | 吞吐(files/min) | 相对加速 | 文本字符 | 结果一致 | Text/OCR 页 |", + "|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + ) + for case in group_cases: + result = by_id.get(case.id) + if result is None: + lines.append( + f"| {md_escape(case.title)} | {md_escape(params_text(case.params))} | 待运行 | - | - | - | - | - | - | - | - | - | - |" + ) + continue + speedup = baseline.wall_seconds / result.wall_seconds if baseline and result.wall_seconds else 0.0 + consistent = ( + "基准" + if baseline is result + else "是" + if baseline and baseline.result_hash and baseline.result_hash == result.result_hash + else "否" + if baseline and baseline.result_hash and result.result_hash + else "未知" + ) + lines.append( + "| {title} | {params} | {status} | {files} | {wall} | {init} | {core} | {avg} | {throughput:.2f} | {speedup:.2f}x | {chars} | {consistent} | {text}/{ocr} |".format( + title=md_escape(case.title), + params=md_escape(params_text(case.params)), + status=result.status, + files=f"{result.completed_files}/{result.discovered_files}", + wall=format_seconds(result.wall_seconds), + init=format_seconds(result.model_init_seconds), + core=format_seconds(result.core_processing_seconds), + avg=format_seconds(result.average_file_seconds), + throughput=result.throughput_files_per_minute, + speedup=speedup, + chars=result.recognized_chars, + consistent=consistent, + text=result.text_pages, + ocr=result.ocr_pages, + ) + ) + lines.append("") + + lines.extend(["## 用例详情", ""]) + for case in cases: + result = by_id.get(case.id) + lines.extend([f"### {case.id} — {case.title}", "", f"- 类型：`{case.kind}`", f"- 参数：`{params_text(case.params)}`"]) + if result is None: + lines.extend(["- 状态：待运行", ""]) + continue + lines.extend( + [ + f"- 状态：`{result.status}`", + f"- Exit code：`{result.exit_code}`", + f"- Wall time：`{format_seconds(result.wall_seconds)}`", + f"- 输出目录：`{result.output_dir}`", + f"- 日志：`{result.log_file}`", + f"- 批次 manifest：`{result.batch_manifest}`", + f"- 结果 SHA-256：`{result.result_hash}`", + "- 命令：", + "", + "```text", + subprocess.list2cmdline(result.command), + "```", + ] + ) + if result.routing_reasons: + lines.append(f"- PDF 路由原因：`{json.dumps(result.routing_reasons, ensure_ascii=False)}`") + if result.error: + lines.extend(["- 错误摘要：", "", "```text", result.error[-3000:], "```"]) + lines.append("") + + report_path.parent.mkdir(parents=True, exist_ok=True) + temporary = report_path.with_name(f".{report_path.name}.tmp") + temporary.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8") + temporary.replace(report_path) + + +def run_case( + case: BenchmarkCase, + *, + data_dir: Path, + device: str, + run_root: Path, + timeout: float | None, +) -> CaseResult: + case_root = run_root / safe_name(case.id) + output_dir = case_root / "outputs" + log_file = case_root / "run.log" + console_file = case_root / "console.log" + case_root.mkdir(parents=True, exist_ok=True) + command = build_command( + case, + data_dir=data_dir, + device=device, + output_dir=output_dir, + log_file=log_file, + ) + started_at = datetime.now().astimezone().isoformat() + started = time.perf_counter() + print(f"\n[{case.id}] {case.title}") + print(" " + subprocess.list2cmdline(command)) + environment = os.environ.copy() + environment["PYTHONUTF8"] = "1" + environment["PYTHONIOENCODING"] = "utf-8" + stderr_tail = "" + exit_code = -1 + try: + creation_flags = ( + subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0 + ) + with console_file.open("w", encoding="utf-8") as console: + process = subprocess.Popen( + command, + cwd=ROOT, + stdout=console, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + env=environment, + creationflags=creation_flags, + start_new_session=sys.platform != "win32", + ) + try: + exit_code = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + if sys.platform == "win32": + subprocess.run( + ["taskkill", "/PID", str(process.pid), "/T", "/F"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + else: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + exit_code = 124 + stderr_tail = f"case timeout after {timeout}s" + console_text = console_file.read_text(encoding="utf-8", errors="replace") + console_lines = console_text.splitlines() + for line in console_lines[-20:]: + print(" " + line) + if exit_code != 124: + stderr_tail = "\n".join(console_lines[-100:]) + except KeyboardInterrupt: + raise + except Exception as exc: + exit_code = 1 + stderr_tail = f"{type(exc).__name__}: {exc}" + wall_seconds = time.perf_counter() - started + return parse_case_result( + case, + command=command, + exit_code=exit_code, + started_at=started_at, + wall_seconds=wall_seconds, + output_dir=output_dir, + log_file=log_file, + stderr_tail=stderr_tail, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="对 data 下图片/PDF执行参数矩阵测试并生成 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("--profile", choices=("smoke", "standard", "full"), default="smoke", help="预设参数矩阵") + parser.add_argument("--config", type=Path, default=None, help="自定义 JSON 参数矩阵") + parser.add_argument("--report", type=Path, default=DEFAULT_REPORT, help="Markdown 报告路径") + parser.add_argument("--run-root", type=Path, default=ROOT / "benchmarks" / "runs", help="详细运行数据根目录") + parser.add_argument("--case", action="append", default=[], help="只运行指定 case id，可重复使用") + parser.add_argument("--timeout", type=float, default=None, help="每个用例超时秒数") + parser.add_argument("--dry-run", action="store_true", help="只生成计划和 Markdown，不执行") + parser.add_argument("--list-cases", action="store_true", help="列出用例后退出") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + data_dir = args.data.expanduser().resolve() + if not data_dir.is_dir(): + print(f"数据目录不存在: {data_dir}", file=sys.stderr) + return 1 + cases = load_config(args.config.expanduser().resolve()) if args.config else profile_cases(args.profile, args.device) + if args.case: + requested = set(args.case) + cases = [case for case in cases if case.id in requested] + missing = requested - {case.id for case in cases} + if missing: + print(f"未知 case: {sorted(missing)}", file=sys.stderr) + return 2 + if args.list_cases: + for case in cases: + print(f"{case.id:30} {case.kind:5} {case.group:18} {params_text(case.params)}") + return 0 + if not cases: + print("没有可运行用例", file=sys.stderr) + return 2 + + run_id = datetime.now().strftime("%Y%m%d-%H%M%S") + run_root = args.run_root.expanduser().resolve() / run_id + report_path = args.report.expanduser().resolve() + started_at = datetime.now().astimezone().isoformat() + results: list[CaseResult] = [] + render_report( + report_path=report_path, + run_id=run_id, + profile=args.profile if not args.config else f"config:{args.config}", + device=args.device, + data_dir=data_dir, + cases=cases, + results=results, + run_root=run_root, + started_at=started_at, + ) + print(f"报告: {report_path}") + print(f"运行目录: {run_root}") + print(f"用例数: {len(cases)}") + if args.dry_run: + for case in cases: + command = build_command( + case, + data_dir=data_dir, + device=args.device, + output_dir=run_root / safe_name(case.id) / "outputs", + log_file=run_root / safe_name(case.id) / "run.log", + ) + print(f"[{case.id}] {subprocess.list2cmdline(command)}") + return 0 + + interrupted = False + try: + for index, case in enumerate(cases, 1): + print(f"\n=== Case {index}/{len(cases)} ===") + result = run_case( + case, + data_dir=data_dir, + device=args.device, + run_root=run_root, + timeout=args.timeout, + ) + results.append(result) + render_report( + report_path=report_path, + run_id=run_id, + profile=args.profile if not args.config else f"config:{args.config}", + device=args.device, + data_dir=data_dir, + cases=cases, + results=results, + run_root=run_root, + started_at=started_at, + ) + print( + f"完成: status={result.status}, wall={format_seconds(result.wall_seconds)}, " + f"files={result.completed_files}/{result.discovered_files}" + ) + except KeyboardInterrupt: + interrupted = True + print("\n测试被中断，已完成结果已写入 Markdown。", file=sys.stderr) + + failed = [result for result in results if result.status != "completed"] + print(f"\n报告已更新: {report_path}") + print(f"完成 {len(results)}/{len(cases)}，失败 {len(failed)}") + if interrupted: + return 130 + return 0 if not failed else 3 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/parameter-matrix.example.json b/benchmarks/parameter-matrix.example.json new file mode 100644 index 0000000..37627cd --- /dev/null +++ b/benchmarks/parameter-matrix.example.json @@ -0,0 +1,58 @@ +{ + "cases": [ + { + "id": "image-cpu-threads-8", + "title": "图片 CPU 8 线程", + "group": "image-threads", + "kind": "image", + "params": { + "threads": 8, + "max_new_tokens": 64, + "rounds": 1 + } + }, + { + "id": "image-cpu-threads-18", + "title": "图片 CPU 18 线程", + "group": "image-threads", + "kind": "image", + "params": { + "threads": 18, + "max_new_tokens": 64, + "rounds": 1 + } + }, + { + "id": "pdf-hybrid-default", + "title": "PDF 混合默认阈值", + "group": "pdf-routing", + "kind": "pdf", + "params": { + "pdf_mode": "hybrid" + } + }, + { + "id": "pdf-hybrid-strict", + "title": "PDF 混合严格文本阈值", + "group": "pdf-routing", + "kind": "pdf", + "params": { + "pdf_mode": "hybrid", + "text_min_chars": 100, + "text_min_content_ratio": 0.7 + } + }, + { + "id": "pdf-ocr-dpi-144", + "title": "PDF OCR DPI 144", + "group": "pdf-ocr-dpi", + "kind": "pdf", + "params": { + "pdf_mode": "ocr", + "pages": "1", + "dpi": 144, + "max_new_tokens": 64 + } + } + ] +} diff --git a/ocr.py b/ocr.py index 8aeb598..df49325 100644 --- a/ocr.py +++ b/ocr.py @@ -4,6 +4,7 @@ Examples: python ocr.py data/images/手写01.png --device cpu python ocr.py data/documents/sample.pdf --device cpu --pdf-mode hybrid python ocr.py data/ --recursive --device cpu + python ocr.py benchmark --profile smoke --device cpu python ocr.py verify --device gpu The launcher deliberately executes the selected isolated uv project, so CPU @@ -32,6 +33,10 @@ def _requested_device(argv: list[str]) -> str: def main() -> int: + if len(sys.argv) > 1 and sys.argv[1] == "benchmark": + command = [sys.executable, str(ROOT / "benchmark.py"), *sys.argv[2:]] + return subprocess.run(command, cwd=ROOT).returncode + try: device = _requested_device(sys.argv[1:]) except SystemExit as exc: diff --git a/ocr_app/cli.py b/ocr_app/cli.py index e13a7ef..8e5aaa1 100644 --- a/ocr_app/cli.py +++ b/ocr_app/cli.py @@ -30,7 +30,13 @@ def build_input_parser(device_override: str | None = None) -> argparse.ArgumentP ) parser.add_argument("input", type=Path, help="图片、PDF 或目录") parser.add_argument("--recursive", action="store_true", help="目录模式递归扫描子目录") - parser.add_argument("--output", type=Path, default=PROJECT_ROOT / "outputs", help="PDF 输出根目录") + parser.add_argument( + "--only", + choices=("all", "image", "pdf"), + default="all", + help="目录模式只处理指定类型", + ) + parser.add_argument("--output", type=Path, default=PROJECT_ROOT / "outputs", help="统一输出根目录") parser.add_argument("--fail-fast", action="store_true", help="单文件失败后立即停止目录任务") # Image options diff --git a/ocr_app/commands.py b/ocr_app/commands.py index 6a74f86..1ba511f 100644 --- a/ocr_app/commands.py +++ b/ocr_app/commands.py @@ -47,13 +47,24 @@ def detect_input_kind(path: Path) -> str: return "unsupported" -def discover_supported_files(directory: Path, *, recursive: bool, output_root: Path) -> list[Path]: +def discover_supported_files( + directory: Path, + *, + recursive: bool, + output_root: Path, + only: str = "all", +) -> list[Path]: iterator = directory.rglob("*") if recursive else directory.glob("*") output_root = output_root.expanduser().resolve() files: list[Path] = [] for path in iterator: if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS: continue + suffix = path.suffix.lower() + if only == "image" and suffix not in IMAGE_EXTENSIONS: + continue + if only == "pdf" and suffix not in PDF_EXTENSIONS: + continue resolved = path.resolve() try: resolved.relative_to(output_root) @@ -375,14 +386,16 @@ def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_ input_path, recursive=args.recursive, output_root=args.output, + only=args.only, ) if not files: logger.error("NO_SUPPORTED_FILES directory=%s recursive=%s", input_path, args.recursive) return 1 logger.info( - "DIRECTORY_PLAN directory=%s recursive=%s files=%d image_files=%d pdf_files=%d", + "DIRECTORY_PLAN directory=%s recursive=%s only=%s files=%d image_files=%d pdf_files=%d", input_path, args.recursive, + args.only, len(files), sum(path.suffix.lower() in IMAGE_EXTENSIONS for path in files), sum(path.suffix.lower() in PDF_EXTENSIONS for path in files), @@ -430,6 +443,7 @@ def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_ batch_manifest = { "input_directory": str(input_path), "recursive": args.recursive, + "only": args.only, "device": provider.resolved_device, "discovered_files": len(files), "completed_files": len(results), diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py new file mode 100644 index 0000000..3884479 --- /dev/null +++ b/tests/test_benchmark.py @@ -0,0 +1,127 @@ +import json +from pathlib import Path + +from benchmark import ( + BenchmarkCase, + CaseResult, + load_config, + profile_cases, + render_report, + run_case, +) + + +def test_smoke_profile_has_image_and_pdf_cases(): + cases = profile_cases("smoke", "cpu") + assert {case.kind for case in cases} == {"image", "pdf"} + assert len({case.id for case in cases}) == len(cases) + + +def test_load_custom_config(tmp_path): + config = tmp_path / "matrix.json" + config.write_text( + json.dumps( + { + "cases": [ + { + "id": "custom-image", + "title": "Custom", + "group": "custom", + "kind": "image", + "params": {"threads": 4}, + } + ] + } + ), + encoding="utf-8", + ) + cases = load_config(config) + assert cases == [ + BenchmarkCase( + id="custom-image", + title="Custom", + group="custom", + kind="image", + params={"threads": 4}, + ) + ] + + +def test_render_report_contains_speed_and_hash_comparison(tmp_path): + cases = [ + BenchmarkCase("base", "Base", "image", "image", {"threads": 1}), + BenchmarkCase("fast", "Fast", "image", "image", {"threads": 8}), + ] + results = [ + CaseResult( + case=cases[0], + status="completed", + exit_code=0, + command=["python", "ocr.py"], + started_at="2026-01-01T00:00:00+08:00", + wall_seconds=20.0, + output_dir="out/base", + log_file="base.log", + completed_files=2, + discovered_files=2, + average_file_seconds=8.0, + throughput_files_per_minute=6.0, + result_hash="same", + recognized_chars=100, + ), + CaseResult( + case=cases[1], + status="completed", + exit_code=0, + command=["python", "ocr.py"], + started_at="2026-01-01T00:00:00+08:00", + wall_seconds=10.0, + output_dir="out/fast", + log_file="fast.log", + completed_files=2, + discovered_files=2, + average_file_seconds=4.0, + throughput_files_per_minute=12.0, + result_hash="same", + recognized_chars=100, + ), + ] + report = tmp_path / "report.md" + render_report( + report_path=report, + run_id="test", + profile="test", + device="cpu", + data_dir=Path("data"), + cases=cases, + results=results, + run_root=Path("runs"), + started_at="2026-01-01T00:00:00+08:00", + ) + text = report.read_text(encoding="utf-8") + assert "2.00x" in text + assert "结果一致" in text + assert "| 是 |" in text + + +def test_run_case_timeout_returns_124(tmp_path, monkeypatch): + script = tmp_path / "slow.py" + script.write_text("import time; time.sleep(10)", encoding="utf-8") + case = BenchmarkCase("slow", "Slow", "timeout", "image", {}) + + def fake_build_command(*args, **kwargs): + import sys + + return [sys.executable, str(script)] + + monkeypatch.setattr("benchmark.build_command", fake_build_command) + result = run_case( + case, + data_dir=tmp_path, + device="cpu", + run_root=tmp_path / "runs", + timeout=0.1, + ) + assert result.exit_code == 124 + assert result.status == "failed" + assert "timeout" in (result.error or "")