249 lines
7.6 KiB
Python
249 lines
7.6 KiB
Python
"""
|
|
天目 OCR 布局质量集成测试。
|
|
|
|
生成模拟真实文档的图片（含标题、正文、列表、编号项、分隔线等），
|
|
验证：
|
|
1. PP-OCRv6 正确识别文字内容
|
|
2. 布局后处理生成结构化 Markdown（标题级别、段落分组）
|
|
3. 双栏排版识别
|
|
|
|
使用方法：
|
|
1. 先启动天目服务：uv run tianmu
|
|
2. 运行本脚本：python tests/test_integration_layout.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import urllib.request
|
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
API_BASE = os.environ.get("TIANMU_API_BASE", "http://localhost:8101")
|
|
|
|
|
|
# ── 字体加载 ──────────────────────────────────────────────
|
|
|
|
def _load_font(size: int) -> ImageFont.FreeTypeFont | None:
|
|
paths = [
|
|
"/System/Library/Fonts/PingFang.ttc",
|
|
"/System/Library/Fonts/STHeiti Light.ttc",
|
|
"/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
|
|
"C:/Windows/Fonts/msyh.ttc",
|
|
]
|
|
for fp in paths:
|
|
if os.path.exists(fp):
|
|
try:
|
|
return ImageFont.truetype(fp, size)
|
|
except Exception:
|
|
continue
|
|
return None
|
|
|
|
|
|
# ── 测试场景 1：标准文档布局 ────────────────────────────
|
|
|
|
def create_standard_doc() -> bytes:
|
|
"""生成一个含标题、正文、列表的标准文档图片。"""
|
|
font_title = _load_font(28)
|
|
font_body = _load_font(16)
|
|
|
|
img = Image.new("RGB", (600, 400), color="white")
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
y = 20
|
|
lines = []
|
|
title = "项目实施方案"
|
|
lines.append((title, font_title, (20, y)))
|
|
y += 45
|
|
|
|
lines.append(("", None, None)) # 空行
|
|
|
|
h2 = "一、项目背景"
|
|
lines.append((h2, font_body, (20, y)))
|
|
y += 30
|
|
|
|
body_texts = [
|
|
"本项目旨在构建一套基于视觉感知的智能文档处理系统。",
|
|
"系统通过本地OCR引擎实现文字提取，支持多种文档类型。",
|
|
]
|
|
for t in body_texts:
|
|
lines.append((t, font_body, (20, y)))
|
|
y += 25
|
|
|
|
lines.append(("", None, None))
|
|
y += 5
|
|
|
|
items = [
|
|
"• 核心技术栈：Python + FastAPI + ONNX Runtime",
|
|
"• 模型支持：PP-OCRv6 + PaddleOCR-VL",
|
|
"• 部署方式：Docker 容器化微服务",
|
|
]
|
|
for item in items:
|
|
lines.append((item, font_body, (20, y)))
|
|
y += 25
|
|
|
|
for text, font, pos in lines:
|
|
if text and font and pos:
|
|
draw.text(pos, text, fill="black", font=font)
|
|
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
# ── 测试场景 2：双栏布局 ────────────────────────────────
|
|
|
|
def create_two_column_doc() -> bytes:
|
|
"""生成一个模拟双栏排版的图片。"""
|
|
font_body = _load_font(14)
|
|
|
|
img = Image.new("RGB", (700, 300), color="white")
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
# 左栏 (x=20~300)
|
|
left_items = [
|
|
"左栏项目一：需求分析",
|
|
"左栏项目二：方案设计",
|
|
"左栏项目三：开发实施",
|
|
]
|
|
y = 20
|
|
for item in left_items:
|
|
draw.text((20, y), item, fill="black", font=font_body)
|
|
y += 30
|
|
|
|
# 右栏 (x=380~680)
|
|
right_items = [
|
|
"右栏概述：本栏说明验收标准",
|
|
"右栏条目1：功能完整性",
|
|
"右栏条目2：性能达标",
|
|
]
|
|
y = 20
|
|
for item in right_items:
|
|
draw.text((380, y), item, fill="black", font=font_body)
|
|
y += 30
|
|
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
# ── 测试场景 3：编号 + TOC 目录 ─────────────────────────
|
|
|
|
def create_toc_doc() -> bytes:
|
|
"""生成含编号和 TOC 目录的图片。"""
|
|
font = _load_font(16)
|
|
|
|
img = Image.new("RGB", (500, 250), color="white")
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
items = [
|
|
("1. 引言", 20),
|
|
("2. 相关工作", 50),
|
|
(" 2.1 传统方法", 75),
|
|
(" 2.2 深度学习方法", 100),
|
|
("3. 实验与分析", 130),
|
|
(" 3.1 数据集", 155),
|
|
(" 3.2 实验结果", 180),
|
|
]
|
|
for text, y in items:
|
|
draw.text((20, y), text, fill="black", font=font)
|
|
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
# ── 运行测试 ──────────────────────────────────────────────
|
|
|
|
def run_ocr(image_bytes: bytes, engine: str = "ppocrv6") -> dict | None:
|
|
b64 = base64.b64encode(image_bytes).decode("ascii")
|
|
body = json.dumps({"image": b64, "engine": engine}).encode("utf-8")
|
|
try:
|
|
req = urllib.request.Request(
|
|
f"{API_BASE}/v1/ocr",
|
|
data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=30) as resp:
|
|
return json.loads(resp.read())
|
|
except Exception as e:
|
|
print(f" ✗ API 请求失败: {e}")
|
|
return None
|
|
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print(" 天目 OCR 布局质量集成测试")
|
|
print(f" API: {API_BASE}")
|
|
print("=" * 60)
|
|
|
|
# 1. 标准文档布局
|
|
print("\n▶ 1/3 标准文档（标题 + 正文 + 列表）")
|
|
img1 = create_standard_doc()
|
|
result = run_ocr(img1)
|
|
if result and result.get("text"):
|
|
text = result["text"]
|
|
print(f" 识别文本 ({len(text)} 字符):")
|
|
for line in text.split("\n")[:10]:
|
|
if line.strip():
|
|
print(f" {line.strip()[:80]}")
|
|
if "项目实施方案" in text:
|
|
print(" ✓ 识别到标题")
|
|
if "•" in text or "- 核心" in text:
|
|
print(" ✓ 识别到列表项")
|
|
print(f" 耗时: {result.get('duration_ms', 0):.0f} ms")
|
|
else:
|
|
err = result.get("error") if result else "无响应"
|
|
print(f" ✗ 识别失败: {err}")
|
|
|
|
# 2. 双栏布局
|
|
print("\n▶ 2/3 双栏排版")
|
|
img2 = create_two_column_doc()
|
|
result = run_ocr(img2)
|
|
if result and result.get("text"):
|
|
text = result["text"]
|
|
print(f" 识别文本 ({len(text)} 字符):")
|
|
for line in text.split("\n")[:10]:
|
|
if line.strip():
|
|
print(f" {line.strip()[:80]}")
|
|
if "左栏项目" in text and "右栏概述" in text:
|
|
print(" ✓ 双栏内容均被识别")
|
|
print(f" 耗时: {result.get('duration_ms', 0):.0f} ms")
|
|
else:
|
|
err = result.get("error") if result else "无响应"
|
|
print(f" ✗ 识别失败: {err}")
|
|
|
|
# 3. 编号 + 目录
|
|
print("\n▶ 3/3 编号项 + TOC 目录")
|
|
img3 = create_toc_doc()
|
|
result = run_ocr(img3)
|
|
if result and result.get("text"):
|
|
text = result["text"]
|
|
print(f" 识别文本 ({len(text)} 字符):")
|
|
for line in text.split("\n")[:10]:
|
|
if line.strip():
|
|
print(f" {line.strip()[:80]}")
|
|
checks = [
|
|
("1. 引言" in text, "识别编号 1."),
|
|
("2. 相关" in text, "识别编号 2."),
|
|
("3. 实验" in text, "识别编号 3."),
|
|
]
|
|
for ok, desc in checks:
|
|
print(f" {'✓' if ok else '✗'} {desc}")
|
|
print(f" 耗时: {result.get('duration_ms', 0):.0f} ms")
|
|
else:
|
|
err = result.get("error") if result else "无响应"
|
|
print(f" ✗ 识别失败: {err}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print(" 测试完成")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|