317 lines
10 KiB
Python
317 lines
10 KiB
Python
"""
|
|
天目 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()
|