初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,316 @@
|
||||
"""
|
||||
天目 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()
|
||||
Reference in new issue
Block a user