""" 天目 OCR 集成测试。 覆盖场景： - PP-OCRv6 文字识别（英文、中文、中英混合） - PaddleOCR-VL 文档解析（含引擎未加载时的优雅降级） - URL 图片下载识别 - resize 缩放功能 - save_md 保存 Markdown - 引擎选择（auto / 显式指定） - 无效输入（空数据、非法 base64） 使用方法： 1. 先启动天目服务： cd tianmu && source .venv/bin/activate && python3 -m src.main 2. 运行此脚本（另一终端）： python3 tests/run_ocr_test.py 3. 或指定图片： python3 tests/run_ocr_test.py /path/to/image.jpg """ from __future__ import annotations import base64 import io import json import os import sys import urllib.request import urllib.error from PIL import Image, ImageDraw, ImageFont API_BASE = os.environ.get("TIANMU_API_BASE", "http://localhost:8101") # ── 测试用例 ────────────────────────────────────────────── TEST_CASES = [ { "name": "英文数字", "text": "Hello OCR 12345", "size": (400, 60), "font_size": 20, }, { "name": "中文短句", "text": "你好世界 OCR 测试", "size": (400, 60), "font_size": 20, }, { "name": "中英混合", "text": "RapidOCR 天目 PP-OCRv6 识别", "size": (500, 60), "font_size": 20, }, { "name": "多行文本", "text": "第一行文字\n第二行文字\n第三行文字", "size": (400, 120), "font_size": 18, }, ] # ── PaddleOCR-VL 解析场景（仅当引擎就绪时测试） ── VL_TEST_PROMPTS = [ "markdown", "text", ] # ── 辅助函数 ────────────────────────────────────────────── def create_test_image(text: str, size: tuple[int, int] = (400, 60), font_size: int = 20) -> bytes: """生成用于测试的图片（白底黑字）。""" img = Image.new("RGB", size, color="white") draw = ImageDraw.Draw(img) # 尝试加载中文字体，回退默认 font = _load_font(font_size) # 多行支持 y_offset = 10 for line in text.split("\n"): draw.text((10, y_offset), line, fill="black", font=font) y_offset += font_size + 5 buf = io.BytesIO() img.save(buf, format="PNG") return buf.getvalue() def _load_font(font_size: int) -> ImageFont.FreeTypeFont | None: """加载系统字体（多平台回退）。""" font_paths = [ "/System/Library/Fonts/PingFang.ttc", "/System/Library/Fonts/STHeiti Light.ttc", "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc", "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", "C:/Windows/Fonts/msyh.ttc", ] for fp in font_paths: if os.path.exists(fp): try: return ImageFont.truetype(fp, font_size) except Exception: continue return None # ── 测试函数 ────────────────────────────────────────────── def check_health() -> dict | None: """测试健康检查端点。""" try: req = urllib.request.Request(f"{API_BASE}/healthz") with urllib.request.urlopen(req, timeout=5) as resp: data = json.loads(resp.read()) engine_status = data.get("engines", {}) ppocr = engine_status.get("ppocrv6", {}).get("status", "?") pvl = engine_status.get("paddleocr-vl", {}).get("status", "?") print(f" ✓ 健康检查 | status={data.get('status')} " f"| ppocrv6={ppocr} | paddleocr-vl={pvl}") return data except Exception as e: print(f" ✗ 健康检查失败: {e}") return None def run_ocr_test(image_bytes: bytes, engine: str = "ppocrv6", **kwargs) -> dict | None: """测试 OCR 端点。""" b64 = base64.b64encode(image_bytes).decode("ascii") body = {"image": b64, "engine": engine} body.update(kwargs) try: req = urllib.request.Request( f"{API_BASE}/v1/ocr", data=json.dumps(body).encode("utf-8"), headers={"Content-Type": "application/json"}, ) with urllib.request.urlopen(req, timeout=30) as resp: return json.loads(resp.read()) except urllib.error.HTTPError as e: print(f" ✗ HTTP {e.code}: {e.read().decode()[:200]}") return None except Exception as e: print(f" ✗ 请求失败: {e}") return None def run_ocr_vl_test(image_bytes: bytes, output_format: str = "markdown", **kwargs) -> dict | None: """测试 PaddleOCR-VL 端点（无需引擎就绪）。""" b64 = base64.b64encode(image_bytes).decode("ascii") body = {"image": b64, "output_format": output_format} body.update(kwargs) try: req = urllib.request.Request( f"{API_BASE}/v1/ocr/vl", data=json.dumps(body).encode("utf-8"), headers={"Content-Type": "application/json"}, ) with urllib.request.urlopen(req, timeout=10) as resp: return json.loads(resp.read()) except urllib.error.HTTPError as e: print(f" ✗ HTTP {e.code}: {e.read().decode()[:200]}") return None except Exception as e: print(f" ✗ 请求失败: {e}") return None def run_ocr_by_url(url: str) -> dict | None: """测试通过 URL 识别图片。""" body = json.dumps({"image": url, "engine": "ppocrv6"}) try: req = urllib.request.Request( f"{API_BASE}/v1/ocr", data=body.encode("utf-8"), headers={"Content-Type": "application/json"}, ) with urllib.request.urlopen(req, timeout=15) as resp: return json.loads(resp.read()) except Exception as e: print(f" ✗ URL 请求失败: {e}") return None # ── 结果格式化 ──────────────────────────────────────────── def print_result(case_name: str, result: dict | None) -> None: """格式化打印测试结果。""" if result is None: print(f" ✗ {case_name}: 无响应") return error = result.get("error") if error: print(f" ✗ {case_name}: {error}") return text = result.get("text", "") blocks = result.get("blocks", []) duration = result.get("duration_ms", 0) # 截断显示，防止控制台溢出 display_text = repr(text[:80]) + ("..." if len(text) > 80 else "") print(f" ✓ {case_name}") print(f" text: {display_text}") print(f" blocks: {len(blocks)} 块" if blocks else "") print(f" duration: {duration:.0f} ms") if blocks: print(f" confidence: {blocks[0]['confidence']:.3f}") # ── 主入口 ──────────────────────────────────────────────── def main(): print("=" * 60) print(" 天目 OCR 集成测试套件") print(f" API: {API_BASE}") print("=" * 60) print() # 1. 健康检查 print("▶ 1/6 健康检查") health = check_health() if health is None: print("\n⚠ 天目服务未启动，请先运行: cd tianmu && python3 -m src.main") sys.exit(1) print() # 2. PP-OCRv6 基础识别测试 print("▶ 2/6 PP-OCRv6 文字识别") for case in TEST_CASES: img_bytes = create_test_image(case["text"], case["size"], case["font_size"]) result = run_ocr_test(img_bytes, engine="ppocrv6") print_result(case["name"], result) print() # 3. PaddleOCR-VL 端点测试（引擎可能未加载，验证降级） print("▶ 3/6 PaddleOCR-VL 端点（含降级验证）") img_bytes = create_test_image("文档测试", (400, 60)) for fmt in VL_TEST_PROMPTS: result = run_ocr_test_vl(img_bytes, output_format=fmt) if result: err = result.get("error") if err: print(f" ~ {fmt}: 引擎未就绪 | {err[:80]}") else: print(f" ✓ {fmt}: {repr(result.get('text', '')[:60])}") else: print(f" ✗ {fmt}: 无响应") print() # 4. resize 缩放测试 print("▶ 4/6 resize 缩放测试") img_bytes = create_test_image("Resize Test", (800, 60), font_size=30) result_resized = run_ocr_test(img_bytes, engine="ppocrv6", resize=400) result_original = run_ocr_test(img_bytes, engine="ppocrv6") if result_resized: print(f" ✓ resize=400 | duration={result_resized.get('duration_ms', 0):.0f} ms") if result_original: print(f" ✓ 原始大小 | duration={result_original.get('duration_ms', 0):.0f} ms (对比)") print() # 5. save_md 保存功能测试 print("▶ 5/6 save_md 保存测试") img_bytes = create_test_image("保存测试 Markdown 文件", (500, 60)) result = run_ocr_test(img_bytes, engine="ppocrv6", save_md=True) if result: md_path = result.get("saved_md_path") if md_path: print(f" ✓ 已保存到: {md_path}") else: print(f" ~ 未返回路径 | text={repr(result.get('text', '')[:40])}") print() # 6. 自定义图片 if len(sys.argv) > 1: print("▶ 6/6 自定义图片") for img_path in sys.argv[1:]: if os.path.isfile(img_path): print(f" 📷 {img_path}") with open(img_path, "rb") as f: image_bytes = f.read() result = run_ocr_test(image_bytes, engine="auto") print_result(os.path.basename(img_path), result) else: print(f" ✗ 文件不存在: {img_path}") print() else: # 或者测试 auto 引擎选择 print("▶ 6/6 auto 引擎选择测试") img_bytes = create_test_image("引擎自动选择", (400, 60)) result = run_ocr_test(img_bytes, engine="auto") print_result("auto → ppocrv6", result) print() # 汇总 print("=" * 60) print(" 测试完成") print("=" * 60) if __name__ == "__main__": main()