初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,247 @@
|
||||
# 天目 · 本地 OCR 引擎
|
||||
|
||||
天目（Tiānmù）是魂镜的**本地 OCR 推理引擎**，基于 RapidOCR（PP-OCRv6）和 PaddleOCR-VL-1.6，纯 Python 实现，直接加载 ONNX / PaddlePaddle 模型进行本地文字识别和文档解析。
|
||||
|
||||
> 「天目」为观天象、察文字之眼，映照万物纹理；在魂镜中承载本地 OCR 感知能力，为人物的视觉感知提供文字提取与文档理解。
|
||||
|
||||
---
|
||||
|
||||
## 架构概览
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
lingtai[灵台 · Backend]
|
||||
tianmu[天目 · OCR Local Engine]
|
||||
models[(models/ocr/ 目录)]
|
||||
|
||||
lingtai -->|POST /v1/ocr| tianmu
|
||||
lingtai -->|POST /v1/ocr/vl| tianmu
|
||||
lingtai -->|GET /healthz| tianmu
|
||||
tianmu -->|RapidOCR ONNX Runtime| models
|
||||
tianmu -->|PaddleOCR-VL 直接加载| models
|
||||
```
|
||||
|
||||
## 提供的 API
|
||||
|
||||
| 端点 | 能力 | 引擎 |
|
||||
| --- | --- | --- |
|
||||
| `POST /v1/ocr` | 通用文字识别 | RapidOCR（PP-OCRv6） |
|
||||
| `POST /v1/ocr/vl` | 复杂文档解析（Phase 3） | PaddleOCR-VL-1.6 |
|
||||
| `GET /healthz` | 健康检查 | — |
|
||||
|
||||
### POST /v1/ocr
|
||||
|
||||
请求：
|
||||
|
||||
```json
|
||||
{
|
||||
"image": "<base64 或 HTTP URL>",
|
||||
"engine": "auto",
|
||||
"lang": ["ch", "en"]
|
||||
}
|
||||
```
|
||||
|
||||
响应：
|
||||
|
||||
```json
|
||||
{
|
||||
"engine": "ppocrv6",
|
||||
"text": "识别后的文本...",
|
||||
"blocks": [
|
||||
{
|
||||
"text": "字",
|
||||
"confidence": 0.98,
|
||||
"bbox": { "points": [[x1,y1],[x2,y1],[x2,y2],[x1,y2]] }
|
||||
}
|
||||
],
|
||||
"duration_ms": 423,
|
||||
"error": null
|
||||
}
|
||||
```
|
||||
|
||||
## 模型文件结构
|
||||
|
||||
所有模型统一放在项目根目录 `models/ocr/` 下，已预下载完毕：
|
||||
|
||||
```text
|
||||
models/ocr/
|
||||
├── ppocrv6/ # PP-OCRv6 ONNX 模型（168 MB）
|
||||
│ ├── det/
|
||||
│ │ ├── PP-OCRv6_det_tiny.onnx # 轻量检测（1.8 MB）
|
||||
│ │ ├── PP-OCRv6_det_small.onnx # 默认检测（9.5 MB）
|
||||
│ │ └── PP-OCRv6_det_medium.onnx # 高精度检测（59 MB）
|
||||
│ ├── rec/
|
||||
│ │ ├── PP-OCRv6_rec_tiny.onnx # 轻量识别（4.3 MB）
|
||||
│ │ ├── PP-OCRv6_rec_small.onnx # 默认识别（20 MB）
|
||||
│ │ ├── PP-OCRv6_rec_medium.onnx # 高精度识别（73 MB）
|
||||
│ │ ├── ppocrv6_tiny_dict.txt # tiny 词表（27 KB）
|
||||
│ │ └── ppocrv6_dict.txt # small/medium 词表（75 KB）
|
||||
│ └── cls/
|
||||
│ └── ch_ppocr_mobile_v2.0_cls_mobile.onnx # 方向分类（0.57 MB）
|
||||
└── paddleocr-vl/ # PaddleOCR-VL-1.6（1.8 GB）
|
||||
├── model.safetensors
|
||||
├── config.json
|
||||
├── tokenizer.json / tokenizer.model
|
||||
└── modeling_paddleocr_vl.py
|
||||
```
|
||||
|
||||
各模型来源：
|
||||
|
||||
| 模型 | 来源 | 下载方式 |
|
||||
| --- | --- | --------- |
|
||||
| PP-OCRv6（det/rec/cls） | [RapidAI/RapidOCR - ModelScope](https://www.modelscope.cn/models/RapidAI/RapidOCR) | 预下载或自动下载 |
|
||||
| PaddleOCR-VL-1.6 | [PaddlePaddle/PaddleOCR-VL-1.6 - HuggingFace](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.6) | 已预下载 |
|
||||
|
||||
## 快速开始
|
||||
|
||||
```bash
|
||||
# 1. 进入天目目录
|
||||
cd tianmu
|
||||
|
||||
# 2. 创建虚拟环境并安装依赖
|
||||
uv venv -p 3.12
|
||||
source .venv/bin/activate
|
||||
uv sync --extra rapidocr
|
||||
|
||||
# 3. 启动天目（推荐）
|
||||
uv run tianmu
|
||||
|
||||
# 4. 验证健康检查
|
||||
curl http://localhost:8101/healthz
|
||||
|
||||
# 5. 测试 OCR（base64 图片）
|
||||
python3 tests/test_ocr.py
|
||||
```
|
||||
|
||||
## 启动方式对比
|
||||
|
||||
| 命令 | 适用场景 | reload | 泄漏警告 |
|
||||
|------|---------|--------|---------|
|
||||
| `uv run tianmu` | 开发（推荐） | ❌ | ❌ |
|
||||
| `uv run uvicorn src.main:app --port 8101 --reload` | 开发热重载 | ✅ | ⚠️ reloader 进程可能打印 resource_tracker 警告 |
|
||||
| `python3 -m src.main` | 快速测试 | ❌ | ❌ |
|
||||
|
||||
> `--reload` 下的 `resource_tracker` 警告是 uvicorn WatchFiles reloader 进程本身的多进程行为，不影响天目功能，纯属视觉问题。生产环境（Docker）无此问题。
|
||||
|
||||
## 启动和测试
|
||||
|
||||
天目提供两种测试方式：
|
||||
|
||||
### 方式一：远程图片测试
|
||||
|
||||
```bash
|
||||
curl -s -X POST http://localhost:8101/v1/ocr \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"image": "https://example.com/test-image.jpg"}' \
|
||||
| python3 -m json.tool
|
||||
```
|
||||
|
||||
### 方式二：base64 图片测试
|
||||
|
||||
```bash
|
||||
# 使用内置测试脚本生成测试图片并调用
|
||||
cd tianmu && source .venv/bin/activate
|
||||
python3 tests/test_ocr.py
|
||||
```
|
||||
|
||||
### 方式三：本地文件测试
|
||||
|
||||
```bash
|
||||
curl -s -X POST http://localhost:8101/v1/ocr \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "{\"image\": \"$(base64 -i /path/to/image.jpg | tr -d '\n')\"}" \
|
||||
| python3 -m json.tool
|
||||
```
|
||||
|
||||
## 配置
|
||||
|
||||
通过 `TIANMU_*` 环境变量配置（详见 `.env.example`）：
|
||||
|
||||
| 变量 | 默认值 | 说明 |
|
||||
| --- | --- | --- |
|
||||
| `TIANMU_HOST` | `0.0.0.0` | 监听地址 |
|
||||
| `TIANMU_PORT` | `8101` | 监听端口 |
|
||||
| `TIANMU_MODELS_DIR` | `../models` | 模型文件根目录 |
|
||||
| `TIANMU_DEVICE` | `auto` | 推理设备：auto / cuda / mps / cpu |
|
||||
| `TIANMU_LOG_LEVEL` | `info` | 日志级别 |
|
||||
|
||||
## 项目结构
|
||||
|
||||
```text
|
||||
tianmu/ # 源代码
|
||||
├── pyproject.toml # 依赖（rapidocr / paddlepaddle）
|
||||
├── README.md # 本文件
|
||||
├── src/
|
||||
│ ├── main.py # FastAPI 入口
|
||||
│ ├── config.py # 配置加载（环境变量 → models/ 路径解析）
|
||||
│ ├── endpoints/
|
||||
│ │ ├── ocr.py # OCR 端点（/v1/ocr + /v1/ocr/vl）
|
||||
│ │ └── health.py # 健康检查端点（/healthz）
|
||||
│ ├── engine/
|
||||
│ │ ├── rapidocr.py # RapidOCR（PP-OCRv6）推理引擎
|
||||
│ │ └── paddleocr_vl.py # PaddleOCR-VL-1.6 引擎（预留）
|
||||
│ ├── schemas/
|
||||
│ │ └── ocr.py # Pydantic 请求/响应模型
|
||||
│ └── utils/
|
||||
│ └── logging.py # loguru 日志配置
|
||||
└── tests/
|
||||
├── test_ocr.py # OCR 集成测试（生成图片 + 调用 API）
|
||||
└── __init__.py
|
||||
|
||||
docker/tianmu/ # Docker 构建文件
|
||||
├── Dockerfile
|
||||
└── .dockerignore
|
||||
```
|
||||
|
||||
## 引擎选择
|
||||
|
||||
天目提供两个引擎，按需切换：
|
||||
|
||||
| 引擎 | 场景 | 速度 | 精度 | 依赖 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PP-OCRv6（via RapidOCR） | 通用文字提取 | ~400ms（CPU） | ⭐⭐⭐⭐ | `rapidocr` |
|
||||
| PaddleOCR-VL-1.6 | 复杂文档/表格/古籍 | ~3-12s（GPU） | ⭐⭐⭐⭐⭐ | `paddlepaddle` |
|
||||
|
||||
调用时通过 `engine` 参数指定：
|
||||
|
||||
```json
|
||||
// 强制使用 PP-OCRv6
|
||||
{"image": "...", "engine": "ppocrv6"}
|
||||
|
||||
// 强制使用 PaddleOCR-VL
|
||||
{"image": "...", "engine": "paddleocr-vl"}
|
||||
```
|
||||
|
||||
`engine` 不传或传 `auto` 时默认走 PP-OCRv6。
|
||||
|
||||
## Docker 部署
|
||||
|
||||
详见 `docker/README.md` 及 `docker-compose.app.yml` 中 `hunjing-tianmu` 服务定义。
|
||||
|
||||
```bash
|
||||
# Docker 启动
|
||||
docker compose -f docker/docker-compose.app.yml --profile tianmu up -d --build hunjing-tianmu
|
||||
```
|
||||
|
||||
## 依赖分组
|
||||
|
||||
```bash
|
||||
# 仅 PP-OCRv6（默认推荐）
|
||||
uv sync --extra rapidocr
|
||||
|
||||
# PP-OCRv6 + PaddleOCR-VL（需要复杂文档解析）
|
||||
uv sync --extra rapidocr --extra vl
|
||||
|
||||
# 所有可选
|
||||
uv sync --extra all
|
||||
```
|
||||
|
||||
## 测试
|
||||
|
||||
```bash
|
||||
# 基础测试（无需启动服务）
|
||||
pytest tests/ -v
|
||||
|
||||
# 集成测试（需先启动服务）
|
||||
python3 tests/test_ocr.py
|
||||
```
|
||||
@@ -0,0 +1,109 @@
|
||||
[project]
|
||||
name = "tianmu"
|
||||
version = "0.2.0"
|
||||
description = "天目 · 魂镜本地 OCR 引擎 — PP-OCRv6（via RapidOCR）+ PaddleOCR-VL-1.6"
|
||||
authors = [{ name = "chenxin", email = "2685475410@qq.com" }]
|
||||
license = "MIT"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"fastapi>=0.139.0",
|
||||
"httpx[http2]>=0.28.1",
|
||||
"loguru>=0.7.3",
|
||||
"pillow>=10.0.0",
|
||||
"pydantic>=2.13.4",
|
||||
"pydantic-settings>=2.14.2",
|
||||
"python-multipart>=0.0.32",
|
||||
"uvicorn[standard]>=0.51.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# PP-OCRv6 引擎（RapidOCR + ONNX Runtime）
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# CPU 推理（所有平台通用，包括 M4 的 CoreML 实验性支持）
|
||||
rapidocr = [
|
||||
"rapidocr>=3.9.0",
|
||||
"onnxruntime>=1.20.0",
|
||||
]
|
||||
|
||||
# GPU 推理（NVIDIA CUDA，仅 Linux/Windows）
|
||||
rapidocr-gpu = [
|
||||
"onnxruntime-gpu>=1.20.0",
|
||||
]
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# PaddleOCR-VL 引擎（HuggingFace Transformers）
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# CPU/MPS（macOS 自动获取 CPU/MPS 版本，Linux/Windows 自动走 CPU 版本）
|
||||
vl = [
|
||||
"torch>=2.0.0",
|
||||
"torchvision>=0.15.0",
|
||||
"transformers>=5.0.0",
|
||||
]
|
||||
|
||||
# GPU（通过 [tool.uv.sources] 自动选择 CUDA 版本，无需重复声明）
|
||||
# 见下方 pytorch-cu128 源映射：Linux/Windows 自动使用 CUDA 12.8 版本
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
all = ["tianmu[rapidocr,rapidocr-gpu,vl]"]
|
||||
dev = [
|
||||
"pre-commit>=4.6.0",
|
||||
"pytest>=9.1.1",
|
||||
"pytest-asyncio>=1.4.0",
|
||||
"ruff>=0.15.21",
|
||||
]
|
||||
lint = ["black>=26.5.1", "mypy>=2.3.0", "ruff>=0.15.21"]
|
||||
test = ["pytest>=9.1.1", "pytest-asyncio>=1.4.0", "pytest-cov>=7.1.0"]
|
||||
dev-all = ["tianmu[dev,lint,test]"]
|
||||
|
||||
[project.scripts]
|
||||
tianmu = "src.main:main_cli"
|
||||
|
||||
# =============================================
|
||||
# uv 配置 - 跨平台 PyTorch 自动适配
|
||||
# =============================================
|
||||
|
||||
# 默认使用中科大镜像（所有平台通用）
|
||||
[[tool.uv.index]]
|
||||
name = "ustc"
|
||||
url = "https://pypi.mirrors.ustc.edu.cn/simple/"
|
||||
default = true
|
||||
|
||||
# PyTorch CUDA 12.8 专用源（仅 Linux/Windows 使用）
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cu128"
|
||||
url = "https://mirrors.nju.edu.cn/pytorch/whl/cu128/"
|
||||
explicit = true
|
||||
|
||||
# === 包来源映射：根据操作系统自动选择 ===
|
||||
[tool.uv.sources]
|
||||
# Linux/Windows：使用 CUDA 12.8 版本
|
||||
torch = { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }
|
||||
torchvision = { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }
|
||||
torchaudio = { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }
|
||||
tianmu = { workspace = true }
|
||||
|
||||
# macOS：从默认源（中科大）自动获取 CPU/MPS 版本
|
||||
# 不需要额外配置，会自动回退到 default = true 的源
|
||||
|
||||
# 验证命令：
|
||||
# macOS: python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'MPS available: {torch.backends.mps.is_available()}')"
|
||||
# Linux: python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}')"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling>=1.25.0"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src"]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
target-version = "py312"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "N", "W"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
@@ -0,0 +1,2 @@
|
||||
"""天目"""
|
||||
__version__ = "0.2.0"
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
天目配置管理。
|
||||
|
||||
从环境变量加载配置，遵循魂镜统一配置规范。
|
||||
模型文件统一存放在项目根目录 models/ 下，各引擎从对应子目录加载。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
_DEFAULT_MODELS_DIR = _PROJECT_ROOT / "models"
|
||||
|
||||
|
||||
def detect_device(preferred: str = "auto") -> str:
|
||||
"""检测当前可用的最佳推理设备。
|
||||
|
||||
使用 PyTorch 检测（torch 在 vl 组中，未安装时回退 CPU）。
|
||||
检测顺序：CUDA → MPS（Apple Silicon）→ CPU
|
||||
|
||||
Args:
|
||||
preferred: 用户偏好（auto / cuda / mps / cpu）。
|
||||
|
||||
Returns:
|
||||
实际可用的设备名称。
|
||||
"""
|
||||
if preferred == "cpu":
|
||||
return "cpu"
|
||||
|
||||
# 使用 PyTorch 检测（vl 组可选，未安装时回退 CPU）
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return "cpu"
|
||||
|
||||
if preferred in ("auto", "cuda") and torch.cuda.is_available():
|
||||
cuda_ver = torch.version.cuda or "unknown"
|
||||
return f"cuda ({cuda_ver})"
|
||||
|
||||
if preferred in ("auto", "mps") and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
|
||||
return "cpu"
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""天目应用配置。"""
|
||||
|
||||
# 服务
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8101
|
||||
|
||||
# ── 模型文件路径 ────────────────────────────────────
|
||||
models_dir: str = str(_DEFAULT_MODELS_DIR)
|
||||
|
||||
# PP-OCRv6 模型：det / rec / cls 子目录或具体 ONNX 路径
|
||||
# 为空时自动扫描 $models_dir/ocr/ppocrv6/{det,rec,cls}/
|
||||
ppocrv6_det_model: str = ""
|
||||
ppocrv6_rec_model: str = ""
|
||||
ppocrv6_cls_model: str = ""
|
||||
|
||||
# 识别词表路径（默认 auto 在 rec 目录下查找 ppocrv6_dict.txt）
|
||||
ppocrv6_dict_path: str = ""
|
||||
|
||||
# PaddleOCR-VL 模型目录（默认 $models_dir/ocr/paddleocr-vl/）
|
||||
paddleocr_vl_model_dir: str = ""
|
||||
|
||||
# ── 输出路径 ────────────────────────────────────────
|
||||
# OCR 结果输出目录（Markdown 文件等），默认项目根目录 output/
|
||||
output_dir: str = ""
|
||||
|
||||
# 推理设备：auto / cuda / mps / cpu
|
||||
device: str = "auto"
|
||||
|
||||
# ── 网络与超时 ──────────────────────────────────────
|
||||
# 下载远端图片的 HTTP 超时（秒）
|
||||
http_timeout: int = 30
|
||||
|
||||
# ── PaddleOCR-VL 推理参数 ───────────────────────────
|
||||
# 最大生成 tokens（复杂文档可能需要更长）
|
||||
vl_max_new_tokens: int = 1024
|
||||
# 图像最大像素数（影响内存占用，越大越吃显存）
|
||||
# 默认 1280*28*28 ≈ 1,003,520，即长边约 1280px
|
||||
vl_max_pixels: int = 1280 * 28 * 28
|
||||
|
||||
# 日志级别
|
||||
log_level: str = "info"
|
||||
|
||||
# ── 运行时属性（非配置项）──
|
||||
_detected_device: str | None = None
|
||||
|
||||
@property
|
||||
def detected_device(self) -> str:
|
||||
"""实际检测到的设备（懒加载）。"""
|
||||
if self._detected_device is None:
|
||||
self._detected_device = detect_device(self.device)
|
||||
return self._detected_device
|
||||
|
||||
@property
|
||||
def models_dir_resolved(self) -> Path:
|
||||
return Path(self.models_dir).resolve()
|
||||
|
||||
@property
|
||||
def ppocrv6_dir(self) -> Path:
|
||||
"""PP-OCRv6 模型根目录。"""
|
||||
return self.models_dir_resolved / "ocr" / "ppocrv6"
|
||||
|
||||
@property
|
||||
def paddleocr_vl_dir(self) -> Path:
|
||||
"""PaddleOCR-VL 模型目录。"""
|
||||
if self.paddleocr_vl_model_dir:
|
||||
return Path(self.paddleocr_vl_model_dir).resolve()
|
||||
return self.models_dir_resolved / "ocr" / "paddleocr-vl"
|
||||
|
||||
@property
|
||||
def output_dir_resolved(self) -> Path:
|
||||
"""OCR 输出目录（Markdown 文件等）。
|
||||
|
||||
优先使用 ``TIANMU_OUTPUT_DIR`` 环境变量，
|
||||
否则默认为项目根目录的 ``output/``。
|
||||
"""
|
||||
if self.output_dir:
|
||||
return Path(self.output_dir).resolve()
|
||||
# 项目根目录 = models_dir 的父目录
|
||||
return self.models_dir_resolved.parent / "output"
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="tianmu_",
|
||||
env_file=None,
|
||||
)
|
||||
|
||||
|
||||
settings = Settings()
|
||||
@@ -0,0 +1 @@
|
||||
"""天目 API 端点。"""
|
||||
@@ -0,0 +1,55 @@
|
||||
"""健康检查端点。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from src import __version__
|
||||
from src.engine.rapidocr import RapidOCREngine
|
||||
from src.engine.paddleocr_vl import PaddleOCRVLEngine
|
||||
from src.schemas.ocr import EngineStatus, HealthResponse
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# 全局引擎实例（由 main.py startup 注入）
|
||||
rapidocr: RapidOCREngine | None = None
|
||||
paddleocr_vl: PaddleOCRVLEngine | None = None
|
||||
|
||||
|
||||
def init(
|
||||
rapidocr_engine: RapidOCREngine,
|
||||
paddleocr_vl_engine: PaddleOCRVLEngine,
|
||||
) -> None:
|
||||
"""注入引擎实例。"""
|
||||
global rapidocr, paddleocr_vl
|
||||
rapidocr = rapidocr_engine
|
||||
paddleocr_vl = paddleocr_vl_engine
|
||||
|
||||
|
||||
@router.get("/healthz", response_model=HealthResponse)
|
||||
async def health() -> HealthResponse:
|
||||
engines_status: dict = {}
|
||||
|
||||
for name, engine in [("ppocrv6", rapidocr), ("paddleocr-vl", paddleocr_vl)]:
|
||||
if engine is None:
|
||||
engines_status[name] = EngineStatus(status="unavailable")
|
||||
else:
|
||||
h = await engine.health()
|
||||
engines_status[name] = EngineStatus(status=h.get("status", "error"))
|
||||
|
||||
# 整体状态
|
||||
ok_count = sum(1 for e in engines_status.values() if e.status == "ok")
|
||||
unavailable_count = sum(1 for e in engines_status.values() if e.status == "unavailable")
|
||||
total = len(engines_status)
|
||||
if ok_count == total:
|
||||
overall = "ok"
|
||||
elif unavailable_count == total:
|
||||
overall = "error"
|
||||
else:
|
||||
overall = "degraded"
|
||||
|
||||
return HealthResponse(
|
||||
status=overall,
|
||||
engines=engines_status,
|
||||
version=__version__,
|
||||
)
|
||||
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
OCR 推理端点。
|
||||
|
||||
职责：
|
||||
1. 接收 OCR / VL 请求
|
||||
2. 解析图片输入（base64 / URL / 文件路径）
|
||||
3. 按需懒加载引擎
|
||||
4. 执行推理并返回结构化结果
|
||||
|
||||
注意：
|
||||
- PP-OCRv6 引擎在 startup 时加载（有模型目录则加载，失败不阻塞）
|
||||
- PaddleOCR-VL 引擎在首次请求时懒加载（模型大，启动慢）
|
||||
- 两条入口（/v1/ocr?engine=paddleocr-vl 和 /v1/ocr/vl）走同一懒加载逻辑
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
from src.config import settings
|
||||
from src.engine.rapidocr import RapidOCREngine
|
||||
from src.engine.paddleocr_vl import PaddleOCRVLEngine
|
||||
from src.schemas.ocr import OCRRequest, OCRResponse, OCRVLRequest, OCRVLResponse
|
||||
from src.utils.logging import logger
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# 全局引擎实例（由 main.py startup 注入）
|
||||
rapidocr: RapidOCREngine | None = None
|
||||
paddleocr_vl: PaddleOCRVLEngine | None = None
|
||||
|
||||
|
||||
def init(
|
||||
rapidocr_engine: RapidOCREngine,
|
||||
paddleocr_vl_engine: PaddleOCRVLEngine,
|
||||
) -> None:
|
||||
"""注入引擎实例。"""
|
||||
global rapidocr, paddleocr_vl
|
||||
rapidocr = rapidocr_engine
|
||||
paddleocr_vl = paddleocr_vl_engine
|
||||
|
||||
|
||||
async def _ensure_paddleocr_vl() -> bool:
|
||||
"""确保 PaddleOCR-VL 引擎已加载（懒加载）。
|
||||
|
||||
该引擎模型约 1.8 GB，启动慢，所以在首次请求时加载而非 startup。
|
||||
|
||||
Returns:
|
||||
True 表示引擎就绪，False 表示加载失败或不可用。
|
||||
"""
|
||||
global paddleocr_vl
|
||||
if paddleocr_vl is None:
|
||||
logger.warning("PaddleOCR-VL 引擎实例未初始化")
|
||||
return False
|
||||
if paddleocr_vl.loaded:
|
||||
return True
|
||||
try:
|
||||
await paddleocr_vl.load()
|
||||
logger.info("PaddleOCR-VL 引擎懒加载完成")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.exception("PaddleOCR-VL 引擎加载失败")
|
||||
return False
|
||||
|
||||
|
||||
@router.post("/v1/ocr", response_model=OCRResponse)
|
||||
async def ocr(req: OCRRequest) -> OCRResponse:
|
||||
"""通用文字识别。
|
||||
|
||||
支持 base64 编码图片和 HTTP(s) URL。
|
||||
默认使用 PP-OCRv6 引擎，速度快、精度高。
|
||||
如有复杂文档需求，可使用 /v1/ocr/vl 端点。
|
||||
"""
|
||||
engine_name = req.engine
|
||||
if engine_name == "auto":
|
||||
engine_name = "ppocrv6" # 默认走 RapidOCR
|
||||
|
||||
# 解析图片
|
||||
try:
|
||||
image_data = await _resolve_image(req.image)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
if req.resize:
|
||||
image_data = _resize_image(image_data, req.resize)
|
||||
|
||||
# 选择引擎（统一走懒加载）
|
||||
if engine_name == "ppocrv6":
|
||||
result = await _run_rapidocr(image_data, req)
|
||||
elif engine_name == "paddleocr-vl":
|
||||
result = await _run_paddleocr_vl_with_lazy_load(image_data, req)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail=f"未知引擎: {engine_name}")
|
||||
|
||||
# 保存 Markdown 文件
|
||||
if req.save_md and result.text:
|
||||
md_path = _save_md(result.text, "ppocrv6" if engine_name == "ppocrv6" else "paddleocr-vl")
|
||||
result.saved_md_path = str(md_path)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@router.post("/v1/ocr/vl", response_model=OCRVLResponse)
|
||||
async def ocr_vl(req: OCRVLRequest) -> OCRVLResponse:
|
||||
"""复杂文档解析（PaddleOCR-VL）。
|
||||
|
||||
适用于表格、公式、古籍、印章等复杂场景。
|
||||
输出 Markdown / JSON 结构化文本。
|
||||
"""
|
||||
try:
|
||||
image_data = await _resolve_image(req.image)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
|
||||
# 懒加载 PaddleOCR-VL
|
||||
ready = await _ensure_paddleocr_vl()
|
||||
if not ready:
|
||||
return OCRVLResponse(
|
||||
engine="paddleocr-vl",
|
||||
text="",
|
||||
format=req.output_format,
|
||||
duration_ms=0,
|
||||
error="PaddleOCR-VL 引擎不可用或加载失败",
|
||||
)
|
||||
|
||||
if paddleocr_vl is None:
|
||||
return OCRVLResponse(
|
||||
engine="paddleocr-vl",
|
||||
text="",
|
||||
format=req.output_format,
|
||||
duration_ms=0,
|
||||
error="引擎不可用",
|
||||
)
|
||||
|
||||
result = await paddleocr_vl.parse_document(image_data, req.output_format)
|
||||
|
||||
# 保存 Markdown 文件
|
||||
md_path = None
|
||||
if req.save_md and result.text:
|
||||
md_path = str(_save_md(result.text, "paddleocr-vl"))
|
||||
logger.info("VL 结果已保存 | path={} chars={}", md_path, len(result.text))
|
||||
|
||||
# 重建响应以包含 saved_md_path（防止 Pydantic 序列化问题）
|
||||
return OCRVLResponse(
|
||||
engine=result.engine,
|
||||
text=result.text,
|
||||
format=result.format,
|
||||
duration_ms=result.duration_ms,
|
||||
error=result.error,
|
||||
saved_md_path=md_path,
|
||||
)
|
||||
|
||||
|
||||
# ── 辅助函数 ──────────────────────────────────────────────
|
||||
|
||||
|
||||
async def _resolve_image(image: str) -> bytes:
|
||||
"""解析图片输入：支持 base64 或 HTTP URL。
|
||||
|
||||
Args:
|
||||
image: base64 编码的图片数据或 HTTP(s) URL。
|
||||
|
||||
Returns:
|
||||
图片的二进制数据。
|
||||
"""
|
||||
# URL 下载
|
||||
if image.startswith(("http://", "https://")):
|
||||
async with httpx.AsyncClient(timeout=settings.http_timeout) as client:
|
||||
resp = await client.get(image)
|
||||
resp.raise_for_status()
|
||||
return resp.content
|
||||
|
||||
# 本地文件路径
|
||||
if os.path.isfile(image):
|
||||
with open(image, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
# Base64 解码（支持 data:image/...;base64,... 格式）
|
||||
if "," in image:
|
||||
image = image.split(",", 1)[1]
|
||||
try:
|
||||
return base64.b64decode(image)
|
||||
except Exception:
|
||||
raise ValueError("图片输入格式无效：请提供 base64、HTTP URL 或文件路径")
|
||||
|
||||
|
||||
def _resize_image(data: bytes, max_side: int) -> bytes:
|
||||
"""等比缩放图片，确保最长边不超过 max_side（px）。
|
||||
|
||||
Args:
|
||||
data: 原图二进制数据。
|
||||
max_side: 最大边长（像素）。
|
||||
|
||||
Returns:
|
||||
缩放后的图片二进制数据。
|
||||
"""
|
||||
try:
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
logger.warning("PIL 不可用，跳过图片缩放")
|
||||
return data
|
||||
|
||||
import io
|
||||
|
||||
try:
|
||||
img = Image.open(io.BytesIO(data))
|
||||
w, h = img.size
|
||||
if max(w, h) <= max_side:
|
||||
return data
|
||||
|
||||
if w >= h:
|
||||
new_w = max_side
|
||||
new_h = int(h * max_side / w)
|
||||
else:
|
||||
new_h = max_side
|
||||
new_w = int(w * max_side / h)
|
||||
|
||||
img = img.resize((new_w, new_h), Image.LANCZOS)
|
||||
buf = io.BytesIO()
|
||||
fmt = img.format or "JPEG"
|
||||
img.save(buf, format=fmt)
|
||||
return buf.getvalue()
|
||||
except Exception as e:
|
||||
logger.warning("图片缩放失败: {}", e)
|
||||
return data
|
||||
|
||||
|
||||
async def _run_rapidocr(image_data: bytes, req: OCRRequest) -> OCRResponse:
|
||||
"""使用 PP-OCRv6 引擎执行 OCR。"""
|
||||
if rapidocr is None or not rapidocr.loaded:
|
||||
return OCRResponse(engine="ppocrv6", error="PP-OCRv6 引擎未加载")
|
||||
|
||||
result = await rapidocr.recognize(image_data)
|
||||
return result
|
||||
|
||||
|
||||
async def _run_paddleocr_vl_with_lazy_load(image_data: bytes, req: OCRRequest) -> OCRResponse:
|
||||
"""使用 PaddleOCR-VL 引擎（含懒加载）。"""
|
||||
ready = await _ensure_paddleocr_vl()
|
||||
if not ready:
|
||||
return OCRResponse(engine="paddleocr-vl", error="PaddleOCR-VL 引擎不可用或加载失败")
|
||||
|
||||
if paddleocr_vl is None:
|
||||
return OCRResponse(engine="paddleocr-vl", error="引擎不可用")
|
||||
|
||||
vl_result = await paddleocr_vl.parse_document(image_data, "text")
|
||||
return OCRResponse(
|
||||
engine="paddleocr-vl",
|
||||
text=vl_result.text,
|
||||
blocks=[],
|
||||
duration_ms=vl_result.duration_ms,
|
||||
error=vl_result.error,
|
||||
)
|
||||
|
||||
|
||||
def _save_md(text: str, engine_prefix: str = "ocr") -> Path:
|
||||
"""将 OCR 结果保存为 Markdown 文件。
|
||||
|
||||
文件保存在 ``{TIANMU_OUTPUT_DIR}/md/`` 下（默认项目根目录 output/md/），
|
||||
文件名格式为 ``{engine_prefix}_{datetime}_{uuid}.md``。
|
||||
使用 uuid4 保证并发不冲突，自动创建目录（如不存在）。
|
||||
|
||||
Args:
|
||||
text: OCR 识别文本（已按布局排序）。
|
||||
engine_prefix: 引擎前缀标识（如 ocr / ppocrv6 / paddleocr-vl），用于文件名。
|
||||
|
||||
Returns:
|
||||
保存的文件路径。
|
||||
"""
|
||||
output_dir = settings.output_dir_resolved / "md"
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
file_id = uuid.uuid4().hex[:12]
|
||||
safe_prefix = engine_prefix.replace("/", "-").replace("_", "-")
|
||||
file_path = output_dir / f"{safe_prefix}_{timestamp}_{file_id}.md"
|
||||
|
||||
file_path.write_text(text, encoding="utf-8")
|
||||
logger.info("OCR 结果已保存到 {}", file_path)
|
||||
return file_path
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
天目引擎协议（Protocol）。
|
||||
|
||||
定义统一的引擎接口，两个引擎实现（RapidOCR / PaddleOCR-VL）均遵循此协议。
|
||||
调用方只需依赖协议，不依赖具体实现。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class OCREngine(Protocol):
|
||||
"""OCR 引擎统一协议。"""
|
||||
|
||||
loaded: bool
|
||||
|
||||
async def load(self) -> None:
|
||||
"""初始化并加载模型。"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放模型资源。"""
|
||||
...
|
||||
|
||||
async def health(self) -> dict:
|
||||
"""返回引擎健康状态。"""
|
||||
...
|
||||
|
||||
|
||||
class RapidOCRProtocol(OCREngine, Protocol):
|
||||
"""RapidOCR 引擎协议（含 recognize 方法）。"""
|
||||
|
||||
async def recognize(self, image_data: bytes | str) -> object:
|
||||
"""执行 OCR 识别，返回结构化响应。"""
|
||||
...
|
||||
|
||||
|
||||
class PaddleOCRVlProtocol(OCREngine, Protocol):
|
||||
"""PaddleOCR-VL 引擎协议（含 parse_document 方法）。"""
|
||||
|
||||
async def parse_document(
|
||||
self,
|
||||
image_data: bytes | str,
|
||||
output_format: str = "markdown",
|
||||
) -> object:
|
||||
"""解析文档图片，返回结构化文本。"""
|
||||
...
|
||||
@@ -0,0 +1,489 @@
|
||||
"""
|
||||
OCR 结果布局后处理（纯逻辑，无 IO）。
|
||||
|
||||
对 RapidOCR 检测到的文本框做阅读顺序修复和段落分组，
|
||||
将扁平 blocks 转换为结构化文本（Markdown 兼容）。
|
||||
|
||||
能力范围（Phase 2，简单排版）：
|
||||
1. 阅读顺序排序：Y→X 排序
|
||||
2. 行归组：Y 坐标容差内归为同一行
|
||||
3. 段落分组：行间距 > 阈值时分段
|
||||
4. 基础格式推断：全数字行、短标题行标记
|
||||
5. 多栏检测：X 范围不重叠的 blocks 按栏分组排序
|
||||
6. 标题层级：基于 block 高度相对中位高度的倍数推断
|
||||
|
||||
依赖：
|
||||
- 无 IO，无第三方库（仅标准库）
|
||||
- 可独立测试，不依赖引擎实例
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import statistics
|
||||
from typing import Any
|
||||
|
||||
from src.schemas.ocr import BBox, TextBlock
|
||||
|
||||
|
||||
def _detect_columns(blocks: list[TextBlock], page_width: float = 800) -> list[list[TextBlock]]:
|
||||
"""多栏检测：将 blocks 按水平范围分成左右栏。
|
||||
|
||||
检测策略：
|
||||
1. 计算所有 block 的 X 范围（左边界、右边界）
|
||||
2. 如果存在明显的水平间隙（gap > page_width * 0.1），则按间隙分栏
|
||||
3. 每栏内部按 Y 排序
|
||||
4. 按左 → 右顺序返回
|
||||
|
||||
Args:
|
||||
blocks: 原始 OCR blocks。
|
||||
page_width: 页面参考宽度。
|
||||
|
||||
Returns:
|
||||
分栏后的 blocks 列表（左栏 → 右栏，每栏内部 Y 排序）。
|
||||
如果未检测到多栏，返回 [blocks]（单栏）。
|
||||
"""
|
||||
if len(blocks) < 4:
|
||||
return [blocks]
|
||||
|
||||
# 收集所有 block 的 X 范围
|
||||
x_ranges: list[tuple[float, float, float]] = [] # (left, right, center)
|
||||
for b in blocks:
|
||||
pts = b.bbox.points
|
||||
if len(pts) >= 2:
|
||||
xs = [p[0] for p in pts]
|
||||
x_ranges.append((min(xs), max(xs), statistics.mean(xs)))
|
||||
else:
|
||||
x_ranges.append((0, page_width, page_width / 2))
|
||||
|
||||
# 计算中位宽度，用于检测间隙
|
||||
widths = [r[1] - r[0] for r in x_ranges]
|
||||
median_width = statistics.median(widths) if widths else page_width
|
||||
|
||||
# 对所有右边界排序，找间隙
|
||||
right_edges = sorted(set(r[1] for r in x_ranges))
|
||||
left_edges = sorted(set(r[0] for r in x_ranges))
|
||||
|
||||
# 找大间隙：如果某列右边界与下一列左边界的间隙 > median_width
|
||||
split_x: list[float] = []
|
||||
for re in right_edges:
|
||||
candidates = [le for le in left_edges if le > re + median_width * 1.5]
|
||||
if candidates:
|
||||
gap = candidates[0] - re
|
||||
if gap > page_width * 0.08: # 间隙 > 页面 8% 宽度
|
||||
split_x.append(re + gap / 2)
|
||||
|
||||
if not split_x:
|
||||
return [blocks]
|
||||
|
||||
# 按间隙分栏（取最大间隙位置）
|
||||
split_pos = max(split_x)
|
||||
|
||||
col_left = [b for b, r in zip(blocks, x_ranges) if r[0] < split_pos]
|
||||
col_right = [b for b, r in zip(blocks, x_ranges) if r[0] >= split_pos]
|
||||
|
||||
# 每栏内按 Y 排序
|
||||
col_left.sort(key=_block_center_y)
|
||||
col_right.sort(key=_block_center_y)
|
||||
|
||||
result = []
|
||||
if col_left:
|
||||
result.append(col_left)
|
||||
if col_right:
|
||||
result.append(col_right)
|
||||
return result
|
||||
|
||||
|
||||
def sort_and_group(blocks: list[TextBlock], page_width: float = 800) -> list[list[TextBlock]]:
|
||||
"""对 OCR blocks 做阅读顺序排序和行归组。
|
||||
|
||||
支持多栏检测：如果 blocks 可以按水平范围分成左右栏，
|
||||
则先按栏分组，每栏内部 Y→X 排序，再左→右合并。
|
||||
|
||||
Args:
|
||||
blocks: 原始 OCR 结果块。
|
||||
page_width: 页面参考宽度，用于多栏检测和居中判断。
|
||||
|
||||
Returns:
|
||||
按阅读顺序排列的行列表，每行为该行内已按 X 排序的 blocks。
|
||||
"""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
# 多栏检测（有间隙时分栏，每栏独立 Y 排序）
|
||||
columns = _detect_columns(blocks, page_width)
|
||||
if len(columns) > 1:
|
||||
# 多栏：逐栏处理行归组，再左→右合并行
|
||||
all_lines: list[list[TextBlock]] = []
|
||||
for col in columns:
|
||||
col_lines = _sort_and_group_single(col)
|
||||
all_lines.extend(col_lines)
|
||||
return all_lines
|
||||
|
||||
return _sort_and_group_single(blocks)
|
||||
|
||||
|
||||
def _sort_and_group_single(blocks: list[TextBlock]) -> list[list[TextBlock]]:
|
||||
"""单栏场景的 Y→X 排序和行归组。"""
|
||||
if not blocks:
|
||||
return []
|
||||
|
||||
# 1. 计算 Y 容差（基于所有 bbox 的中位高度）
|
||||
heights = [_block_height(b) for b in blocks]
|
||||
if not heights:
|
||||
return [[b] for b in blocks]
|
||||
|
||||
median_height = statistics.median(heights) if len(heights) > 1 else heights[0]
|
||||
y_tolerance = median_height * 0.5 # 半行高以内视为同一行
|
||||
|
||||
# 2. 按 Y 中心点排序（全局从上到下）
|
||||
sorted_blocks = sorted(blocks, key=_block_center_y)
|
||||
|
||||
# 3. 行归组
|
||||
lines: list[list[TextBlock]] = []
|
||||
current_line: list[TextBlock] = [sorted_blocks[0]]
|
||||
current_y = _block_center_y(sorted_blocks[0])
|
||||
|
||||
for block in sorted_blocks[1:]:
|
||||
y_center = _block_center_y(block)
|
||||
if abs(y_center - current_y) <= y_tolerance:
|
||||
current_line.append(block)
|
||||
else:
|
||||
lines.append(current_line)
|
||||
current_line = [block]
|
||||
current_y = y_center
|
||||
|
||||
if current_line:
|
||||
lines.append(current_line)
|
||||
|
||||
# 4. 每行内按 X 从左到右排序
|
||||
for line in lines:
|
||||
line.sort(key=_block_center_x)
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def group_paragraphs(lines: list[list[TextBlock]], page_width: float = 800) -> list[dict[str, Any]]:
|
||||
"""将行归组为段落。
|
||||
|
||||
基于行间距判断段落边界。行间距明显大于中位行距时视为新段落。
|
||||
|
||||
Args:
|
||||
lines: 已按阅读顺序排列的行列表。
|
||||
page_width: 页面参考宽度（px），用于居中检测，默认 800。
|
||||
|
||||
Returns:
|
||||
段落列表，每段包含文本和类型。
|
||||
"""
|
||||
if not lines:
|
||||
return []
|
||||
|
||||
# 计算中位行间距
|
||||
if len(lines) > 1:
|
||||
spacings: list[float] = []
|
||||
for i in range(len(lines) - 1):
|
||||
bottom = max(_block_bottom(b) for b in lines[i])
|
||||
top = min(_block_top(b) for b in lines[i + 1])
|
||||
spacings.append(top - bottom)
|
||||
median_spacing = statistics.median(spacings) if spacings else 8.0
|
||||
paragraph_threshold = median_spacing * 1.8 # 1.8 倍以上视为段落分隔
|
||||
else:
|
||||
paragraph_threshold = 20.0 # 默认值
|
||||
|
||||
paragraphs: list[dict[str, Any]] = []
|
||||
current_para: list[str] = []
|
||||
current_type = "text"
|
||||
current_level: int = 3 # 当前标题级别（用于连续标题的场景）
|
||||
|
||||
# 所有行 block 的高度（用于计算中位高度比例，辅助标题检测）
|
||||
all_heights = [max(_block_height(b) for b in line) for line in lines if line]
|
||||
median_line_height = statistics.median(all_heights) if len(all_heights) > 1 else (all_heights[0] if all_heights else 16.0)
|
||||
|
||||
for i, line_blocks in enumerate(lines):
|
||||
line_text = "".join(b.text for b in line_blocks)
|
||||
|
||||
# 计算行信息（居中 + 相对高度）
|
||||
avg_center_x = statistics.mean(
|
||||
_block_center_x(b) for b in line_blocks
|
||||
) if line_blocks else 0
|
||||
center_x_ratio = avg_center_x / page_width if page_width > 0 and avg_center_x > 0 else None
|
||||
line_height = max(_block_height(b) for b in line_blocks) if line_blocks else 0
|
||||
rel_height = line_height / median_line_height if median_line_height > 0 and line_height > 0 else None
|
||||
is_centered = center_x_ratio is not None and 0.25 <= center_x_ratio <= 0.75
|
||||
|
||||
line_type = _detect_line_type(
|
||||
line_text,
|
||||
center_x_ratio=center_x_ratio,
|
||||
rel_height=rel_height,
|
||||
)
|
||||
|
||||
# 标题级别推断
|
||||
level = _detect_heading_level(
|
||||
rel_height=rel_height,
|
||||
is_centered=is_centered,
|
||||
text=line_text,
|
||||
) if line_type == "heading" else 3
|
||||
|
||||
# 判断是否分段
|
||||
is_new_paragraph = False
|
||||
|
||||
# 空行强制分段
|
||||
if not line_text.strip():
|
||||
is_new_paragraph = True
|
||||
else:
|
||||
# 行间距大于阈值（仅正间距，防止重叠/同 Y 误分段）
|
||||
if i > 0:
|
||||
prev_bottom = max(_block_bottom(b) for b in lines[i - 1])
|
||||
curr_top = min(_block_top(b) for b in line_blocks)
|
||||
gap = curr_top - prev_bottom
|
||||
if gap > 0 and gap > paragraph_threshold:
|
||||
is_new_paragraph = True
|
||||
|
||||
# 行类型变化（如文字→标题，独立于间距判断）
|
||||
if current_para and line_type != current_type:
|
||||
is_new_paragraph = True
|
||||
|
||||
if is_new_paragraph and current_para:
|
||||
paragraphs.append({
|
||||
"type": current_type,
|
||||
"text": "\n".join(current_para),
|
||||
"level": current_level,
|
||||
})
|
||||
current_para = []
|
||||
current_type = "text"
|
||||
current_level = 3
|
||||
|
||||
if line_text.strip():
|
||||
current_para.append(line_text)
|
||||
current_type = line_type
|
||||
current_level = level
|
||||
|
||||
if current_para:
|
||||
paragraphs.append({
|
||||
"type": current_type,
|
||||
"text": "\n".join(current_para),
|
||||
"level": current_level,
|
||||
})
|
||||
|
||||
return paragraphs
|
||||
|
||||
|
||||
def build_markdown_text(
|
||||
blocks: list[TextBlock],
|
||||
page_width: float = 800,
|
||||
) -> tuple[str, list[list[TextBlock]], list[dict[str, Any]]]:
|
||||
"""将 OCR blocks 转换为 markdown 格式文本。
|
||||
|
||||
支持段落类型：
|
||||
- heading：标题（## 前缀）
|
||||
- list_item：无序列表（- 前缀）
|
||||
- numbered：有序列表
|
||||
- code：代码块（``` 包裹）
|
||||
- separator：分隔线（---）
|
||||
- text：普通段落
|
||||
|
||||
Args:
|
||||
blocks: 原始 OCR 识别结果块。
|
||||
page_width: 页面参考宽度（px），用于居中检测。
|
||||
|
||||
Returns:
|
||||
(markdown 文本, 行列表, 段落列表)
|
||||
"""
|
||||
lines = sort_and_group(blocks)
|
||||
paragraphs = group_paragraphs(lines, page_width=page_width)
|
||||
|
||||
# 构建 markdown
|
||||
md_parts: list[str] = []
|
||||
code_block_open = False
|
||||
|
||||
for para in paragraphs:
|
||||
para_type = para["type"]
|
||||
para_text = para["text"]
|
||||
|
||||
if para_type == "heading":
|
||||
level = para.get("level", 2)
|
||||
level = max(1, min(level, 3)) # 只输出 1~3 级
|
||||
prefix = "#" * level
|
||||
md_parts.append(f"\n{prefix} {para_text.strip()}\n")
|
||||
elif para_type == "list_item":
|
||||
# 每行加 - 前缀
|
||||
items = para_text.strip().split("\n")
|
||||
md_parts.append("\n".join(f"- {line.strip()}" for line in items))
|
||||
elif para_type == "numbered":
|
||||
md_parts.append(f"\n{para_text.strip()}\n")
|
||||
elif para_type == "code":
|
||||
md_parts.append(f"\n```\n{para_text.strip()}\n```\n")
|
||||
elif para_type == "separator":
|
||||
md_parts.append("\n---\n")
|
||||
else:
|
||||
md_parts.append(para_text)
|
||||
|
||||
markdown = "\n\n".join(p for p in md_parts if p.strip())
|
||||
|
||||
return markdown, lines, paragraphs
|
||||
|
||||
|
||||
# ── 辅助函数 ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def _block_height(block: TextBlock) -> float:
|
||||
"""计算文本框高度。"""
|
||||
try:
|
||||
pts = block.bbox.points
|
||||
return abs(pts[2][1] - pts[0][1]) if len(pts) >= 4 else 0.0
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _block_center_y(block: TextBlock) -> float:
|
||||
"""文本框 Y 中心点。"""
|
||||
try:
|
||||
pts = block.bbox.points
|
||||
ys = [p[1] for p in pts]
|
||||
return statistics.mean(ys)
|
||||
except (IndexError, TypeError, AttributeError, statistics.StatisticsError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _block_center_x(block: TextBlock) -> float:
|
||||
"""文本框 X 中心点。"""
|
||||
try:
|
||||
pts = block.bbox.points
|
||||
xs = [p[0] for p in pts]
|
||||
return statistics.mean(xs)
|
||||
except (IndexError, TypeError, AttributeError, statistics.StatisticsError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _block_bottom(block: TextBlock) -> float:
|
||||
"""文本框底部 Y 坐标。"""
|
||||
try:
|
||||
pts = block.bbox.points
|
||||
return max(p[1] for p in pts)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _block_top(block: TextBlock) -> float:
|
||||
"""文本框顶部 Y 坐标。"""
|
||||
try:
|
||||
pts = block.bbox.points
|
||||
return min(p[1] for p in pts)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _detect_line_type(
|
||||
text: str,
|
||||
center_x_ratio: float | None = None,
|
||||
rel_height: float | None = None,
|
||||
) -> str:
|
||||
"""检测行类型。
|
||||
|
||||
Args:
|
||||
text: 行文本。
|
||||
center_x_ratio: 文本框中心 X 相对页面宽度的比例（0~1），
|
||||
用于检测居中对齐（标题候选）。
|
||||
rel_height: 文本框高度相对中位高度的比例，
|
||||
用于检测标题级别（更大的字 → 更高级别标题）。
|
||||
|
||||
Returns:
|
||||
类型：heading / list_item / numbered / code / separator / text
|
||||
"""
|
||||
stripped = text.strip()
|
||||
|
||||
if not stripped:
|
||||
return "text"
|
||||
|
||||
# 分隔线（纯符号）
|
||||
if re.match(r'^[-=*_]{3,}$', stripped):
|
||||
return "separator"
|
||||
|
||||
# 代码行（检查原文以保留前导空格，stripped 会去除它们）
|
||||
if text.startswith((" ", "\t")):
|
||||
return "code"
|
||||
|
||||
# 列表项（以 - * • ▪ ▸ → · 开头）
|
||||
if stripped.startswith(("- ", "* ", "• ", "▪ ", "▸ ", "→ ", "· ")):
|
||||
return "list_item"
|
||||
|
||||
# 独立 bullet 字符（单个符号作为独立行）
|
||||
if stripped in ("•", "▪", "▸", "→", "·", "-", "*"):
|
||||
return "list_item"
|
||||
|
||||
# 编号项（支持 1.内容 / 1)内容 / (1) 等，点后可不带空格）
|
||||
if re.match(r'^\d+[\.\)]\s*', stripped):
|
||||
return "numbered"
|
||||
if re.match(r'^\(\d+\)\s*', stripped):
|
||||
return "numbered"
|
||||
if re.match(r'^[①②③④⑤⑥⑦⑧⑨⑩]', stripped):
|
||||
return "numbered"
|
||||
|
||||
# TOC 目录行："第一章 .............. 1" — 含连续点号的行
|
||||
# 特征：短行(10-50 char) + 含连续 3+ 点号 + 非标题
|
||||
if 10 <= len(stripped) < 60 and "..." in stripped and "。" not in stripped:
|
||||
return "numbered"
|
||||
|
||||
# 标题候选：短行 + 居中或短文本特征
|
||||
# 标题通常满足：较短 + 不以句号/逗号/分号/冒号结尾
|
||||
is_short = len(stripped) < 40
|
||||
no_punct_end = not stripped.endswith(("。", "，", "；", "：", "、", "！", "？", "…"))
|
||||
# 居中检测：如果 center_x_ratio 在 0.25~0.75 之间，视为居中
|
||||
is_centered = center_x_ratio is not None and 0.25 <= center_x_ratio <= 0.75
|
||||
# 字体较大：rel_height > 1.3 说明字号明显大于中位文本
|
||||
is_larger_font = rel_height is not None and rel_height > 1.3
|
||||
|
||||
if is_short and no_punct_end:
|
||||
# 全大写英文短行 → 标题
|
||||
if stripped.isupper() and stripped.replace(" ", "").isalpha():
|
||||
return "heading"
|
||||
# 居中短行 → 标题
|
||||
if is_centered:
|
||||
return "heading"
|
||||
# 大字短行 → 标题
|
||||
if is_larger_font:
|
||||
return "heading"
|
||||
# 极短的中文行（≤ 8 字符）→ 标题（如"第一章"、"摘要"）
|
||||
cn_chars = sum(1 for c in stripped if '一' <= c <= '鿿')
|
||||
cn_ratio = cn_chars / len(stripped) if len(stripped) > 0 else 0
|
||||
if cn_ratio > 0.6 and len(stripped) <= 8:
|
||||
return "heading"
|
||||
# 纯数字短行（页码、编号）
|
||||
if stripped.isdigit():
|
||||
return "numbered"
|
||||
|
||||
return "text"
|
||||
|
||||
|
||||
def _detect_heading_level(
|
||||
rel_height: float | None = None,
|
||||
is_centered: bool = False,
|
||||
text: str = "",
|
||||
) -> int:
|
||||
"""推断标题级别。
|
||||
|
||||
级别规则：
|
||||
- h1（1）：字体很大(rel_height ≥ 1.8) 或 居中 + 大字
|
||||
- h2（2）：字体中等(rel_height ≥ 1.3) 或 居中
|
||||
- h3（3）：其他标题
|
||||
|
||||
Args:
|
||||
rel_height: 相对中位字体高度的比例。
|
||||
is_centered: 是否居中。
|
||||
text: 标题文本（用于额外判断）。
|
||||
|
||||
Returns:
|
||||
1 / 2 / 3——对应 markdown 的 # / ## / ###。
|
||||
"""
|
||||
# 居中的大字 → h1
|
||||
if is_centered and rel_height is not None and rel_height >= 1.5:
|
||||
return 1
|
||||
# 大字 → h2
|
||||
if rel_height is not None and rel_height >= 1.8:
|
||||
return 1
|
||||
if rel_height is not None and rel_height >= 1.3:
|
||||
return 2
|
||||
# 居中 → h2
|
||||
if is_centered:
|
||||
return 2
|
||||
return 3
|
||||
@@ -0,0 +1,252 @@
|
||||
"""
|
||||
PaddleOCR-VL 引擎封装。
|
||||
|
||||
基于 PaddleOCR-VL-1.6（via HuggingFace Transformers）的端到端文档解析引擎。
|
||||
支持表格/公式/古籍/印章等复杂文档的 Markdown 结构化输出。
|
||||
|
||||
依赖：
|
||||
- torch >= 2.0（推理后端）
|
||||
- transformers >= 4.40（模型加载，内置 PaddleOCR-VL 支持）
|
||||
- Pillow（图片处理）
|
||||
|
||||
设备：MPS (Apple Silicon) / CUDA / CPU。
|
||||
注意：CPU 推理极慢（~数分钟/张），强烈建议在有 GPU 的环境使用。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from src.config import settings
|
||||
from src.schemas.ocr import OCRVLResponse
|
||||
from src.utils.logging import logger
|
||||
|
||||
|
||||
class PaddleOCRVLEngine:
|
||||
"""PaddleOCR-VL-1.6 文档解析引擎。
|
||||
|
||||
通过 HuggingFace Transformers 加载，使用 chat template 消息格式
|
||||
进行结构化文档解析，输出 Markdown / JSON / 纯文本。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._processor: Any | None = None
|
||||
self._model: Any | None = None
|
||||
self._device: str = ""
|
||||
self.loaded = False
|
||||
|
||||
async def load(self) -> None:
|
||||
"""加载 PaddleOCR-VL 模型和处理器。
|
||||
|
||||
从本地路径加载，不使用 trust_remote_code（transformers 5.13 已内置）。
|
||||
"""
|
||||
model_path = str(settings.paddleocr_vl_dir)
|
||||
logger.info("加载 PaddleOCR-VL 模型 | dir={}", model_path)
|
||||
|
||||
if not Path(model_path).exists():
|
||||
raise FileNotFoundError(f"PaddleOCR-VL 模型目录不存在: {model_path}")
|
||||
|
||||
# 检测设备（尊重 TIANMU_DEVICE 显式设置）
|
||||
import torch
|
||||
|
||||
user_pref = settings.device # "auto" / "cuda" / "mps" / "cpu"
|
||||
|
||||
if user_pref == "cpu":
|
||||
self._device = "cpu"
|
||||
elif user_pref in ("auto", "cuda") and torch.cuda.is_available():
|
||||
self._device = "cuda"
|
||||
elif user_pref in ("auto", "mps") and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
self._device = "mps"
|
||||
else:
|
||||
self._device = "cpu"
|
||||
|
||||
logger.info(
|
||||
" device={} | TIANMU_DEVICE={} | CUDA={} MPS={}",
|
||||
self._device, user_pref,
|
||||
torch.cuda.is_available(),
|
||||
torch.backends.mps.is_available() if hasattr(torch.backends, "mps") else False,
|
||||
)
|
||||
|
||||
from transformers import AutoModelForImageTextToText, AutoProcessor
|
||||
|
||||
# 加载 processor（需 trust_remote_code，因为处理有自定义逻辑）
|
||||
self._processor = AutoProcessor.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
# 加载模型（使用内置 transformers 实现，不 trust_remote_code）
|
||||
self._model = AutoModelForImageTextToText.from_pretrained(
|
||||
model_path,
|
||||
dtype=torch.bfloat16,
|
||||
).to(self._device).eval()
|
||||
|
||||
self.loaded = True
|
||||
logger.info("PaddleOCR-VL 引擎就绪 | device={}", self._device)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放模型和处理器。"""
|
||||
self._model = None
|
||||
self._processor = None
|
||||
self.loaded = False
|
||||
logger.info("PaddleOCR-VL 引擎已关闭")
|
||||
|
||||
import gc, torch
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
async def health(self) -> dict:
|
||||
"""返回引擎健康状态。"""
|
||||
return {
|
||||
"status": "ok" if self.loaded else "unavailable",
|
||||
"model": "paddleocr-vl-1.6",
|
||||
"device": self._device if self.loaded else "N/A",
|
||||
}
|
||||
|
||||
async def parse_document(
|
||||
self,
|
||||
image_data: bytes | str | Path,
|
||||
output_format: str = "markdown",
|
||||
) -> OCRVLResponse:
|
||||
"""解析文档图片，返回结构化文本。
|
||||
|
||||
Args:
|
||||
image_data: 图片二进制数据或文件路径。
|
||||
output_format: 输出格式（markdown / json / text）。
|
||||
|
||||
Returns:
|
||||
OCRVLResponse: 结构化解析结果。
|
||||
"""
|
||||
if not self.loaded or self._model is None or self._processor is None:
|
||||
return OCRVLResponse(
|
||||
engine="paddleocr-vl",
|
||||
text="",
|
||||
format=output_format,
|
||||
error="引擎未加载",
|
||||
)
|
||||
|
||||
# 解析图片
|
||||
try:
|
||||
img = self._resolve_image(image_data)
|
||||
except Exception as e:
|
||||
logger.exception("图片解析失败")
|
||||
return OCRVLResponse(
|
||||
engine="paddleocr-vl", text="", format=output_format, error=f"图片解析失败: {e}"
|
||||
)
|
||||
|
||||
# 构建消息（官方格式）
|
||||
task = self._task_from_format(output_format)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "image": img},
|
||||
{"type": "text", "text": task},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
# 使用官方的 apply_chat_template
|
||||
max_pixels = getattr(
|
||||
self._processor.image_processor, "max_pixels", settings.vl_max_pixels
|
||||
)
|
||||
|
||||
inputs = self._processor.apply_chat_template(
|
||||
messages,
|
||||
add_generation_prompt=True,
|
||||
tokenize=True,
|
||||
return_dict=True,
|
||||
return_tensors="pt",
|
||||
processor_kwargs={
|
||||
"images_kwargs": {
|
||||
"size": {
|
||||
"shortest_edge": self._processor.image_processor.min_pixels,
|
||||
"longest_edge": max_pixels,
|
||||
},
|
||||
},
|
||||
},
|
||||
).to(self._device)
|
||||
|
||||
# 推理
|
||||
import torch
|
||||
logger.info("PaddleOCR-VL 推理开始 | img={} outputs={}", img.size, output_format)
|
||||
start = time.perf_counter()
|
||||
|
||||
try:
|
||||
with torch.no_grad():
|
||||
outputs = self._model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=settings.vl_max_new_tokens,
|
||||
do_sample=False,
|
||||
temperature=None,
|
||||
top_p=None,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("PaddleOCR-VL 推理失败")
|
||||
return OCRVLResponse(
|
||||
engine="paddleocr-vl", text="", format=output_format, error=f"推理失败: {e}"
|
||||
)
|
||||
|
||||
elapsed = (time.perf_counter() - start) * 1000
|
||||
result_text = self._processor.decode(
|
||||
outputs[0][inputs["input_ids"].shape[-1]:-1]
|
||||
).strip()
|
||||
|
||||
logger.info("PaddleOCR-VL 推理完成 | {}ms | chars={}", round(elapsed), len(result_text))
|
||||
|
||||
return OCRVLResponse(
|
||||
engine="paddleocr-vl",
|
||||
text=result_text,
|
||||
format=output_format,
|
||||
duration_ms=round(elapsed, 1),
|
||||
error=None,
|
||||
)
|
||||
|
||||
# ── 辅助方法 ──────────────────────────────────────────
|
||||
|
||||
# 官方任务提示词
|
||||
_TASK_PROMPTS: dict[str, str] = {
|
||||
"markdown": "OCR:",
|
||||
"json": "OCR:",
|
||||
"text": "OCR:",
|
||||
"table": "Table Recognition:",
|
||||
"formula": "Formula Recognition:",
|
||||
"chart": "Chart Recognition:",
|
||||
"spotting": "Spotting:",
|
||||
"seal": "Seal Recognition:",
|
||||
}
|
||||
|
||||
def _task_from_format(self, output_format: str) -> str:
|
||||
"""将 tianmu 的 output_format 映射为官方任务提示词。
|
||||
|
||||
Args:
|
||||
output_format: markdown / json / text / table / formula / chart / spotting / seal。
|
||||
|
||||
Returns:
|
||||
官方任务提示词字符串。
|
||||
"""
|
||||
return self._TASK_PROMPTS.get(output_format, "OCR:")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_image(data: bytes | str | Path) -> Image.Image:
|
||||
"""将各种输入格式解析为 PIL Image。
|
||||
|
||||
Args:
|
||||
data: 图片数据（bytes）、文件路径（str/Path）。
|
||||
|
||||
Returns:
|
||||
RGB 模式的 PIL Image。
|
||||
"""
|
||||
if isinstance(data, (str, Path)):
|
||||
path = Path(data)
|
||||
if path.is_file():
|
||||
return Image.open(path).convert("RGB")
|
||||
|
||||
# bytes → PIL
|
||||
return Image.open(io.BytesIO(data if isinstance(data, bytes) else data.encode())).convert("RGB")
|
||||
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
RapidOCR 引擎封装。
|
||||
|
||||
基于 RapidOCR（PP-OCRv6）的本地 ONNX 推理引擎，用于快速文字提取。
|
||||
默认使用 small 模型组合，可通过环境变量切换模型尺寸。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from src.config import settings
|
||||
from src.engine.layout import build_markdown_text
|
||||
from src.schemas.ocr import BBox, OCRResponse, TextBlock
|
||||
from src.utils.logging import logger
|
||||
|
||||
|
||||
class RapidOCREngine:
|
||||
"""RapidOCR（PP-OCRv6）推理引擎。
|
||||
|
||||
使用 ONNX Runtime 进行 CPU/GPU 推理，支持自定义模型路径。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._engine: object | None = None
|
||||
self.loaded = False
|
||||
self._model_name = ""
|
||||
self._use_coreml = False
|
||||
|
||||
async def load(self) -> None:
|
||||
"""初始化 RapidOCR 引擎，加载 ONNX 模型。"""
|
||||
import numpy as np # noqa: F401 — rapidocr 依赖 numpy
|
||||
from rapidocr import RapidOCR
|
||||
|
||||
models_dir = settings.ppocrv6_dir
|
||||
logger.info("加载 PP-OCRv6 模型 | dir={}", models_dir)
|
||||
|
||||
# RapidOCR 的 params 使用 "section.key" 平铺键格式
|
||||
det_path = settings.ppocrv6_det_model or str(models_dir / "det" / "PP-OCRv6_det_small.onnx")
|
||||
rec_path = settings.ppocrv6_rec_model or str(models_dir / "rec" / "PP-OCRv6_rec_small.onnx")
|
||||
cls_path = settings.ppocrv6_cls_model or str(models_dir / "cls" / "ch_ppocr_mobile_v2.0_cls_mobile.onnx")
|
||||
dict_path = settings.ppocrv6_dict_path or str(models_dir / "rec" / "ppocrv6_dict.txt")
|
||||
|
||||
params: dict = {
|
||||
"Global.model_root_dir": str(models_dir),
|
||||
"Global.log_level": "warn",
|
||||
"Det.model_dir": det_path,
|
||||
"Rec.model_dir": rec_path,
|
||||
"Rec.dict_path": dict_path,
|
||||
"Cls.model_dir": cls_path,
|
||||
}
|
||||
|
||||
# Apple Silicon CoreML 加速（实验性）
|
||||
# M4/M 系列芯片可用 Neural Engine，但 RapidOCR 官方标注实测性能不如 CPU
|
||||
# 默认由 TIANMU_DEVICE 控制：auto 时尝试 CoreML，cpu 时强制 CPU
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
if settings.device in ("auto", "mps") and "CoreMLExecutionProvider" in ort.get_available_providers():
|
||||
params["EngineConfig.onnxruntime.use_coreml"] = True
|
||||
self._use_coreml = True
|
||||
logger.info(" ✅ Apple CoreML 加速已启用（实验性）")
|
||||
else:
|
||||
logger.info(" CPU 推理（CoreML 未启用）")
|
||||
except ImportError:
|
||||
logger.info(" CPU 推理（onnxruntime 未安装）")
|
||||
|
||||
# 验证模型文件存在
|
||||
for name, fp in [("det", det_path), ("rec", rec_path), ("cls", cls_path), ("dict", dict_path)]:
|
||||
logger.info(" {}: {} | exists={}", name, fp, Path(fp).exists())
|
||||
|
||||
self._engine = RapidOCR(params=params)
|
||||
self.loaded = True
|
||||
self._model_name = "rapidocr"
|
||||
logger.info("PP-OCRv6 引擎就绪")
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放引擎资源。"""
|
||||
self._engine = None
|
||||
self.loaded = False
|
||||
logger.info("PP-OCRv6 引擎已关闭")
|
||||
|
||||
async def health(self) -> dict:
|
||||
"""返回引擎健康状态。"""
|
||||
device = "coreml" if self._use_coreml else "cpu"
|
||||
return {
|
||||
"status": "ok" if self.loaded else "unavailable",
|
||||
"model": self._model_name,
|
||||
"device": device,
|
||||
}
|
||||
|
||||
async def recognize(self, image_data: bytes | str) -> OCRResponse:
|
||||
"""执行 OCR 识别。
|
||||
|
||||
Args:
|
||||
image_data: 图片二进制数据或文件路径字符串。
|
||||
|
||||
Returns:
|
||||
OCRResponse: 包含识别结果的结构化响应。
|
||||
"""
|
||||
if not self.loaded or self._engine is None:
|
||||
return OCRResponse(engine="ppocrv6", error="引擎未加载")
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
try:
|
||||
result = self._engine(image_data)
|
||||
except Exception as e:
|
||||
logger.exception("PP-OCRv6 推理失败")
|
||||
return OCRResponse(engine="ppocrv6", error=str(e))
|
||||
|
||||
elapsed = (time.perf_counter() - start) * 1000
|
||||
|
||||
# 解析返回结果
|
||||
if result is None or len(result) == 0:
|
||||
logger.info("PP-OCRv6 未识别到文字")
|
||||
return OCRResponse(
|
||||
engine="ppocrv6",
|
||||
text="",
|
||||
blocks=[],
|
||||
duration_ms=round(elapsed, 1),
|
||||
)
|
||||
|
||||
# RapidOCROutput → 结构化响应
|
||||
return self._build_response(result, elapsed)
|
||||
|
||||
def _build_response(self, result: object, elapsed: float) -> OCRResponse:
|
||||
"""将 RapidOCROutput 转换为统一响应模型。
|
||||
|
||||
Args:
|
||||
result: RapidOCR 返回的 RapidOCROutput 实例。
|
||||
elapsed: 推理耗时（毫秒）。
|
||||
"""
|
||||
# 尝试获取属性 — RapidOCROutput 的结构
|
||||
boxes = getattr(result, "boxes", None)
|
||||
txts = getattr(result, "txts", None)
|
||||
scores = getattr(result, "scores", None)
|
||||
|
||||
if txts is None or boxes is None:
|
||||
return OCRResponse(
|
||||
engine="ppocrv6",
|
||||
text="",
|
||||
blocks=[],
|
||||
duration_ms=round(elapsed, 1),
|
||||
)
|
||||
|
||||
blocks: list[TextBlock] = []
|
||||
|
||||
for i, text in enumerate(txts):
|
||||
score = float(scores[i]) if scores is not None else 1.0
|
||||
box = boxes[i] if len(boxes) > i else []
|
||||
|
||||
# 标准化 bbox 为 4 顶点格式
|
||||
points = self._normalize_bbox(box)
|
||||
|
||||
blocks.append(
|
||||
TextBlock(text=text, confidence=round(score, 4), bbox=BBox(points=points))
|
||||
)
|
||||
|
||||
# 布局后处理：阅读顺序排序 → 行归组 → 段落分组 → Markdown
|
||||
markdown, lines, paragraphs = build_markdown_text(blocks)
|
||||
|
||||
return OCRResponse(
|
||||
engine="ppocrv6",
|
||||
text=markdown,
|
||||
blocks=blocks,
|
||||
duration_ms=round(elapsed, 1),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_bbox(box: object) -> list[list[float]]:
|
||||
"""将 RapidOCR 输出的 bbox 统一为 4 顶点格式。"""
|
||||
import numpy as np
|
||||
|
||||
if isinstance(box, (list, tuple)):
|
||||
if not box:
|
||||
return [[0, 0], [0, 0], [0, 0], [0, 0]]
|
||||
if len(box) == 4 and all(isinstance(p, (list, tuple)) and len(p) == 2 for p in box):
|
||||
return [[float(x), float(y)] for x, y in box]
|
||||
if len(box) == 8:
|
||||
return [[float(box[i]), float(box[i + 1])] for i in range(0, 8, 2)]
|
||||
return [[float(box[0]), float(box[1])]] * 4
|
||||
if isinstance(box, np.ndarray):
|
||||
return box.reshape(-1, 2).tolist() if box.size >= 8 else [[0, 0]] * 4
|
||||
return [[0, 0]] * 4
|
||||
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
天目 · 魂镜本地 OCR 引擎 — FastAPI 应用入口。
|
||||
|
||||
提供通用文字识别（PP-OCRv6，via RapidOCR）和复杂文档解析（PaddleOCR-VL）API。
|
||||
所有模型由本地引擎直接加载推理，不依赖外部服务。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# 抑制进程退出时 multiprocessing 信号量泄漏警告（强制覆盖，子进程继承）
|
||||
os.environ["PYTHONWARNINGS"] = "ignore::UserWarning"
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
import gc
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore", message="resource_tracker:")
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from src import __version__
|
||||
from src.config import settings
|
||||
from src.engine.rapidocr import RapidOCREngine
|
||||
from src.engine.paddleocr_vl import PaddleOCRVLEngine
|
||||
from src.endpoints import health, ocr
|
||||
from src.utils.logging import logger, setup_logging
|
||||
|
||||
setup_logging(settings.log_level)
|
||||
|
||||
app = FastAPI(
|
||||
title="天目 · 魂镜本地 OCR 引擎",
|
||||
version=__version__,
|
||||
description="本地 OCR 推理引擎 — PP-OCRv6（快速文字提取）+ PaddleOCR-VL（复杂文档解析）",
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# 全局引擎实例
|
||||
rapidocr_engine = RapidOCREngine()
|
||||
paddleocr_vl_engine = PaddleOCRVLEngine()
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup() -> None:
|
||||
logger.info(
|
||||
"天目启动 | host={} port={} device={}(detected={}) models_dir={}",
|
||||
settings.host,
|
||||
settings.port,
|
||||
settings.device,
|
||||
settings.detected_device,
|
||||
settings.models_dir_resolved,
|
||||
)
|
||||
# 注入引擎实例到端点模块
|
||||
health.init(rapidocr_engine, paddleocr_vl_engine)
|
||||
ocr.init(rapidocr_engine, paddleocr_vl_engine)
|
||||
|
||||
# 懒加载 RapidOCR 引擎（模型存在则加载，失败不阻塞启动）
|
||||
if settings.ppocrv6_dir.exists():
|
||||
try:
|
||||
await rapidocr_engine.load()
|
||||
except Exception as e:
|
||||
logger.warning("PP-OCRv6 引擎加载失败（可在运行时重试）: {}", e)
|
||||
else:
|
||||
logger.warning("PP-OCRv6 模型目录不存在: {}", settings.ppocrv6_dir)
|
||||
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown() -> None:
|
||||
"""优雅关闭：释放引擎资源，超时兜底。"""
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(
|
||||
rapidocr_engine.close(),
|
||||
paddleocr_vl_engine.close(),
|
||||
),
|
||||
timeout=5.0,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("引擎关闭超时，强制退出")
|
||||
except Exception as e:
|
||||
logger.warning("引擎关闭异常: {}", e)
|
||||
|
||||
gc.collect()
|
||||
logger.info("天目停止")
|
||||
|
||||
|
||||
# 注册路由
|
||||
app.include_router(health.router, tags=["health"])
|
||||
app.include_router(ocr.router, tags=["ocr"])
|
||||
|
||||
|
||||
def _run() -> None:
|
||||
"""启动 uvicorn 服务。"""
|
||||
import uvicorn
|
||||
uvicorn.run(
|
||||
"src.main:app",
|
||||
host=settings.host,
|
||||
port=settings.port,
|
||||
log_level=settings.log_level.lower(),
|
||||
)
|
||||
|
||||
|
||||
def main_cli() -> None:
|
||||
"""uv 命令行入口（uv run tianmu），确保 PYTHONWARNINGS 在 reloader 之前设置。"""
|
||||
import sys
|
||||
os.environ["PYTHONWARNINGS"] = "ignore::UserWarning"
|
||||
sys.argv = [a for a in sys.argv if a != "tianmu"]
|
||||
_run()
|
||||
|
||||
|
||||
# 直接运行入口（python -m src.main）
|
||||
if __name__ == "__main__":
|
||||
_run()
|
||||
@@ -0,0 +1 @@
|
||||
"""天目数据模型。"""
|
||||
@@ -0,0 +1,104 @@
|
||||
"""
|
||||
OCR 请求 / 响应 Pydantic 模型。
|
||||
|
||||
定义天目微服务内部的 API 契约，供 endpoints 和 engine 使用。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class OCRRequest(BaseModel):
|
||||
"""通用 OCR 请求。"""
|
||||
|
||||
image: str = Field(description="图片数据：base64 编码字符串或 HTTP(s) URL")
|
||||
engine: str = Field(
|
||||
default="auto",
|
||||
description="引擎选择：auto（自动选择）/ ppocrv6 / paddleocr-vl",
|
||||
)
|
||||
lang: list[str] = Field(
|
||||
default_factory=lambda: ["ch", "en"],
|
||||
description="识别语言列表，默认中英文。仅对 ppocrv6 引擎生效",
|
||||
)
|
||||
resize: int | None = Field(
|
||||
default=None,
|
||||
description="可选：图片最大边长（px），超过则等比缩放。减少传输 / 推理耗时",
|
||||
)
|
||||
save_md: bool = Field(
|
||||
default=False,
|
||||
description="是否将识别结果保存为 Markdown 文件到 output/md/ 目录",
|
||||
)
|
||||
|
||||
|
||||
class BBox(BaseModel):
|
||||
"""文本框坐标。"""
|
||||
|
||||
points: list[list[float]] = Field(description="四个顶点坐标 [[x1,y1],[x2,y1],[x3,y2],[x4,y2]]")
|
||||
|
||||
|
||||
class TextBlock(BaseModel):
|
||||
"""单行/单块识别结果。"""
|
||||
|
||||
text: str = Field(description="识别文字")
|
||||
confidence: float = Field(description="置信度 0~1", ge=0, le=1)
|
||||
bbox: BBox = Field(description="文本框坐标")
|
||||
|
||||
|
||||
class OCRResponse(BaseModel):
|
||||
"""通用 OCR 响应。"""
|
||||
|
||||
engine: str = Field(description="实际使用的引擎，如 ppocrv6 / paddleocr-vl")
|
||||
text: str = Field(default="", description="拼接后的全文")
|
||||
blocks: list[TextBlock] = Field(default_factory=list, description="逐块识别结果")
|
||||
duration_ms: float = Field(default=0, description="推理耗时（毫秒）")
|
||||
error: str | None = Field(default=None, description="错误信息，成功时为 null")
|
||||
saved_md_path: str | None = Field(
|
||||
default=None, description="保存的 Markdown 文件路径（仅 save_md=true 时有值）"
|
||||
)
|
||||
|
||||
|
||||
class OCRVLRequest(BaseModel):
|
||||
"""PaddleOCR-VL 复杂文档解析请求。"""
|
||||
|
||||
image: str = Field(description="图片数据：base64 编码字符串或 HTTP(s) URL")
|
||||
output_format: str = Field(
|
||||
default="markdown",
|
||||
description="输出格式：markdown / json / text",
|
||||
)
|
||||
save_md: bool = Field(
|
||||
default=False,
|
||||
description="是否将识别结果保存为 Markdown 文件到 output/md/ 目录",
|
||||
)
|
||||
|
||||
|
||||
class OCRVLResponse(BaseModel):
|
||||
"""PaddleOCR-VL 响应。"""
|
||||
|
||||
engine: str = Field(default="paddleocr-vl")
|
||||
text: str = Field(default="", description="结构化输出文本（Markdown / JSON / 纯文本）")
|
||||
format: str = Field(default="markdown", description="实际输出格式")
|
||||
duration_ms: float = Field(default=0, description="推理耗时（毫秒）")
|
||||
error: str | None = Field(default=None, description="错误信息，成功时为 null")
|
||||
saved_md_path: str | None = Field(
|
||||
default=None, description="保存的 Markdown 文件路径（仅 save_md=true 时有值）"
|
||||
)
|
||||
|
||||
|
||||
# ── 健康检查 ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class EngineStatus(BaseModel):
|
||||
"""引擎状态。"""
|
||||
|
||||
status: str = Field(description="ok / degraded / error / unavailable")
|
||||
|
||||
|
||||
class HealthResponse(BaseModel):
|
||||
"""健康检查响应。"""
|
||||
|
||||
status: str = Field(description="整体状态：ok / degraded / error")
|
||||
engines: dict[str, Any] = Field(default_factory=dict, description="各引擎状态")
|
||||
version: str = Field(description="天目版本")
|
||||
@@ -0,0 +1 @@
|
||||
"""天目工具模块。"""
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
统一日志配置。
|
||||
|
||||
所有模块统一从这里导入 ``logger``，避免 loguru 配置散落在各处。
|
||||
在应用启动时调用 ``setup_logging()`` 完成初始化。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import loguru
|
||||
|
||||
logger = loguru.logger
|
||||
|
||||
|
||||
def setup_logging(level: str = "INFO") -> None:
|
||||
"""配置 loguru：清空默认 handler，写入 stderr（结构化 + 彩色输出）。
|
||||
|
||||
Args:
|
||||
level: 日志级别（DEBUG / INFO / WARNING / ERROR）。
|
||||
"""
|
||||
logger.remove()
|
||||
logger.add(
|
||||
sink=sys.stderr,
|
||||
format=(
|
||||
"<green>{time:YYYY-MM-DD HH:mm:ss}</green> "
|
||||
"| <level>{level: <7}</level> "
|
||||
"| <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> "
|
||||
"- <level>{message}</level>"
|
||||
),
|
||||
level=level.upper(),
|
||||
colorize=True,
|
||||
enqueue=True,
|
||||
)
|
||||
Whitespace-only changes.
@@ -0,0 +1,84 @@
|
||||
"""
|
||||
天目 pytest 配置。
|
||||
|
||||
提供服务启停夹具，方便运行集成测试。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
API_BASE = os.environ.get("TIANMU_API_BASE", "http://localhost:8101")
|
||||
TIANMU_DIR = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
def _wait_for_service(url: str, timeout: float = 15.0, interval: float = 0.5) -> bool:
|
||||
"""轮询等待服务就绪。
|
||||
|
||||
Args:
|
||||
url: 健康检查 URL。
|
||||
timeout: 最长等待秒数。
|
||||
interval: 轮询间隔秒数。
|
||||
|
||||
Returns:
|
||||
True 表示服务就绪，False 表示超时。
|
||||
"""
|
||||
start = time.monotonic()
|
||||
while time.monotonic() - start < timeout:
|
||||
try:
|
||||
req = urllib.request.Request(url, method="GET")
|
||||
with urllib.request.urlopen(req, timeout=2) as resp:
|
||||
if resp.status == 200:
|
||||
return True
|
||||
except (urllib.error.URLError, ConnectionError, OSError):
|
||||
pass
|
||||
time.sleep(interval)
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def tianmu_server():
|
||||
"""启动天目服务（仅一次，整个 session 共用）。
|
||||
|
||||
如果环境变量 TIANMU_API_BASE 已设置，跳过启动（连接外部服务）。
|
||||
启动后轮询 /healthz，就绪后才继续运行测试。
|
||||
"""
|
||||
if "TIANMU_API_BASE" in os.environ:
|
||||
yield
|
||||
return
|
||||
|
||||
proc = subprocess.Popen(
|
||||
["python3", "-m", "src.main"],
|
||||
cwd=str(TIANMU_DIR),
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
ready = _wait_for_service(f"{API_BASE}/healthz")
|
||||
if not ready:
|
||||
proc.terminate()
|
||||
proc.wait()
|
||||
raise RuntimeError(
|
||||
f"天目服务启动超时（{API_BASE}/healthz 15s 内未就绪）"
|
||||
)
|
||||
|
||||
yield
|
||||
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def api_base() -> str:
|
||||
return API_BASE
|
||||
@@ -0,0 +1,78 @@
|
||||
"""
|
||||
天目健康检查测试（pytest）。
|
||||
|
||||
需先启动天目服务：
|
||||
cd tianmu && source .venv/bin/activate && python3 -m src.main
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
|
||||
API_BASE = os.environ.get("TIANMU_API_BASE", "http://localhost:8101")
|
||||
|
||||
|
||||
def test_healthz():
|
||||
"""健康检查应返回 200 并包含引擎状态。"""
|
||||
req = urllib.request.Request(f"{API_BASE}/healthz")
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
assert resp.status == 200
|
||||
data = json.loads(resp.read())
|
||||
assert "status" in data
|
||||
assert "engines" in data
|
||||
assert "version" in data
|
||||
assert data["version"] == "0.2.0"
|
||||
assert "ppocrv6" in data["engines"]
|
||||
|
||||
|
||||
def test_ocr_empty_image():
|
||||
"""发送空数据应得到有效的错误响应（而非崩溃）。"""
|
||||
body = json.dumps({"image": ""}).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
f"{API_BASE}/v1/ocr",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
assert resp.status == 200
|
||||
data = json.loads(resp.read())
|
||||
# 空图片不会触发 422，而是返回 engine 处理结果或错误
|
||||
assert "engine" in data
|
||||
assert "error" in data # 可能是 None 或错误消息，但字段必须存在
|
||||
|
||||
|
||||
def test_ocr_png_image():
|
||||
"""PP-OCRv6 应正确识别简单图片中的文字。"""
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
# 生成白底黑字图片
|
||||
img = Image.new("RGB", (400, 60), color="white")
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw.text((10, 20), "Hello World 123", fill="black")
|
||||
|
||||
import io
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
b64 = base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
body = json.dumps({"image": b64, "engine": "ppocrv6"}).encode("utf-8")
|
||||
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:
|
||||
data = json.loads(resp.read())
|
||||
assert data["engine"] == "ppocrv6"
|
||||
assert data["error"] is None
|
||||
assert len(data["text"]) > 0 # 识别到文字
|
||||
assert len(data["blocks"]) > 0
|
||||
assert data["blocks"][0]["confidence"] > 0.5
|
||||
assert data["duration_ms"] > 0
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
天目 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()
|
||||
@@ -0,0 +1,531 @@
|
||||
"""
|
||||
布局后处理单元测试（纯逻辑，不依赖引擎和服务）。
|
||||
|
||||
测试 `src.engine.layout` 中的纯函数：
|
||||
- sort_and_group：阅读顺序排序 + 行归组
|
||||
- group_paragraphs：段落分组
|
||||
- build_markdown_text：Markdown 格式化
|
||||
- 辅助函数
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from src.engine.layout import (
|
||||
_block_center_x,
|
||||
_block_center_y,
|
||||
_block_height,
|
||||
_detect_columns,
|
||||
_detect_heading_level,
|
||||
_detect_line_type,
|
||||
_sort_and_group_single,
|
||||
build_markdown_text,
|
||||
group_paragraphs,
|
||||
sort_and_group,
|
||||
)
|
||||
from src.schemas.ocr import BBox, TextBlock
|
||||
|
||||
|
||||
# ── 工厂函数 ──────────────────────────────────────────────
|
||||
|
||||
def _make_block(
|
||||
text: str,
|
||||
x1: float = 0,
|
||||
y1: float = 0,
|
||||
x2: float = 100,
|
||||
y2: float = 30,
|
||||
confidence: float = 0.95,
|
||||
) -> TextBlock:
|
||||
"""创建一个测试用的 TextBlock。"""
|
||||
return TextBlock(
|
||||
text=text,
|
||||
confidence=confidence,
|
||||
bbox=BBox(points=[[x1, y1], [x2, y1], [x2, y2], [x1, y2]]),
|
||||
)
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# 辅助函数
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestHelpers:
|
||||
"""辅助函数测试。"""
|
||||
|
||||
def test_block_height(self) -> None:
|
||||
block = _make_block("测试", y1=10, y2=50)
|
||||
assert _block_height(block) == 40.0
|
||||
|
||||
def test_block_height_empty(self) -> None:
|
||||
block = _make_block("测试", x1=0, y1=0, x2=0, y2=0)
|
||||
assert _block_height(block) == 0.0
|
||||
|
||||
def test_block_center_y(self) -> None:
|
||||
block = _make_block("测试", y1=10, y2=50)
|
||||
assert _block_center_y(block) == 30.0
|
||||
|
||||
def test_block_center_x(self) -> None:
|
||||
block = _make_block("测试", x1=20, x2=100)
|
||||
assert _block_center_x(block) == 60.0
|
||||
|
||||
def test_detect_heading_chinese(self) -> None:
|
||||
"""中文短行 → heading。"""
|
||||
assert _detect_line_type("第一章 绪论") == "heading"
|
||||
|
||||
def test_detect_heading_short_no_punct(self) -> None:
|
||||
"""短行不以句号结尾 → heading。"""
|
||||
assert _detect_line_type("摘要") == "heading"
|
||||
|
||||
def test_detect_not_heading_with_punct(self) -> None:
|
||||
"""以句号结尾的短行 → text。"""
|
||||
assert _detect_line_type("这是一个句子。") == "text"
|
||||
|
||||
def test_detect_long_text(self) -> None:
|
||||
"""长行（>= 40 字符）→ text。"""
|
||||
long = "一二三四五六七八九十" * 4
|
||||
assert _detect_line_type(long) == "text"
|
||||
|
||||
def test_detect_list_item_dash(self) -> None:
|
||||
assert _detect_line_type("- 列表项") == "list_item"
|
||||
|
||||
def test_detect_list_item_bullet(self) -> None:
|
||||
assert _detect_line_type("• 项目") == "list_item"
|
||||
|
||||
def test_detect_numbered(self) -> None:
|
||||
assert _detect_line_type("1. 首先") == "numbered"
|
||||
assert _detect_line_type("2) 其次") == "numbered"
|
||||
|
||||
def test_detect_separator(self) -> None:
|
||||
assert _detect_line_type("---") == "separator"
|
||||
assert _detect_line_type("===") == "separator"
|
||||
assert _detect_line_type("***") == "separator"
|
||||
|
||||
def test_detect_code(self) -> None:
|
||||
assert _detect_line_type(" def foo():") == "code"
|
||||
assert _detect_line_type("\t\tindented") == "code"
|
||||
|
||||
def test_detect_blockquote_as_text(self) -> None:
|
||||
"""> 前缀在 OCR 语境中通常是排版伪影，不作为代码块。"""
|
||||
assert _detect_line_type("> 引用") == "text"
|
||||
|
||||
def test_detect_empty_text(self) -> None:
|
||||
assert _detect_line_type("") == "text"
|
||||
assert _detect_line_type(" ") == "text"
|
||||
|
||||
def test_detect_centered_heading(self) -> None:
|
||||
"""居中短行 → 即使非中文也应视为 heading。"""
|
||||
assert _detect_line_type("Hello World", center_x_ratio=0.5) == "heading"
|
||||
|
||||
def test_detect_not_centered(self) -> None:
|
||||
"""不居中短行 + 非中文锁 → 不回退。"""
|
||||
# "hello world" 非中文、非大写、不居中 → text
|
||||
assert _detect_line_type("hello world", center_x_ratio=0.1) == "text"
|
||||
|
||||
def test_detect_uppercase_heading(self) -> None:
|
||||
"""全大写短行 → heading。"""
|
||||
assert _detect_line_type("INTRODUCTION") == "heading"
|
||||
|
||||
# ── 新增：编号变体 ────────────────────────────────
|
||||
|
||||
def test_detect_numbered_without_space(self) -> None:
|
||||
"""编号后无空格（"1.内容"）→ numbered。"""
|
||||
assert _detect_line_type("1.内容") == "numbered"
|
||||
assert _detect_line_type("1)内容") == "numbered"
|
||||
|
||||
def test_detect_numbered_parenthesized(self) -> None:
|
||||
"""括号编号 → numbered。"""
|
||||
assert _detect_line_type("(1) 内容") == "numbered"
|
||||
assert _detect_line_type("(2)内容") == "numbered"
|
||||
|
||||
def test_detect_numbered_chinese_numeral(self) -> None:
|
||||
"""中文数字序号 → numbered。"""
|
||||
assert _detect_line_type("① 第一项") == "numbered"
|
||||
assert _detect_line_type("②第二项") == "numbered"
|
||||
|
||||
# ── 新增：TOC 目录行 ──────────────────────────────
|
||||
|
||||
def test_detect_toc_line(self) -> None:
|
||||
"""TOC 目录行（含连续点号）→ numbered。"""
|
||||
assert _detect_line_type("第一章 .............. 1") == "numbered"
|
||||
assert _detect_line_type("1.1 概述 ............ 10") == "numbered"
|
||||
|
||||
# ── 新增：更多 bullet 变体 ────────────────────────
|
||||
|
||||
def test_detect_bullet_variants(self) -> None:
|
||||
"""更多 bullet 变体 → list_item。"""
|
||||
assert _detect_line_type("▪ 项目") == "list_item"
|
||||
assert _detect_line_type("▸ 子项") == "list_item"
|
||||
assert _detect_line_type("→ 链接") == "list_item"
|
||||
assert _detect_line_type("· 条目") == "list_item"
|
||||
|
||||
def test_detect_standalone_bullet(self) -> None:
|
||||
"""独立 bullet 字符 → list_item。"""
|
||||
assert _detect_line_type("•") == "list_item"
|
||||
assert _detect_line_type("▪") == "list_item"
|
||||
assert _detect_line_type("→") == "list_item"
|
||||
|
||||
# ── 新增：字体大小辅助标题检测 ────────────────────
|
||||
|
||||
def test_detect_larger_font_heading(self) -> None:
|
||||
"""大字（rel_height > 1.3）→ heading。"""
|
||||
assert _detect_line_type("这是一段长文本但字号大", rel_height=1.5) == "heading"
|
||||
# 中号字（rel_height < 1.3）且不居中 → text
|
||||
assert _detect_line_type("这是一段长文本字号普通", rel_height=1.1) == "text"
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# 标题级别检测
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestHeadingLevel:
|
||||
"""标题级别检测。"""
|
||||
|
||||
def test_h1_centered_large(self) -> None:
|
||||
assert _detect_heading_level(rel_height=1.5, is_centered=True) == 1
|
||||
|
||||
def test_h1_large_alone(self) -> None:
|
||||
assert _detect_heading_level(rel_height=1.8) == 1
|
||||
|
||||
def test_h2_large(self) -> None:
|
||||
assert _detect_heading_level(rel_height=1.4) == 2
|
||||
|
||||
def test_h2_centered(self) -> None:
|
||||
assert _detect_heading_level(is_centered=True) == 2
|
||||
|
||||
def test_h3_default(self) -> None:
|
||||
assert _detect_heading_level() == 3
|
||||
assert _detect_heading_level(rel_height=1.1) == 3
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# sort_and_group
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestSortAndGroup:
|
||||
"""阅读顺序排序 + 行归组。"""
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert sort_and_group([]) == []
|
||||
|
||||
def test_single_block(self) -> None:
|
||||
blocks = [_make_block("你好")]
|
||||
lines = sort_and_group(blocks)
|
||||
assert len(lines) == 1
|
||||
assert len(lines[0]) == 1
|
||||
assert lines[0][0].text == "你好"
|
||||
|
||||
def test_top_to_bottom_sort(self) -> None:
|
||||
"""按 Y 坐标从上到下排序。"""
|
||||
blocks = [
|
||||
_make_block("第二行", y1=100, y2=130),
|
||||
_make_block("第一行", y1=10, y2=40),
|
||||
]
|
||||
lines = sort_and_group(blocks)
|
||||
assert len(lines) == 2
|
||||
assert lines[0][0].text == "第一行"
|
||||
assert lines[1][0].text == "第二行"
|
||||
|
||||
def test_same_line_grouping(self) -> None:
|
||||
"""Y 坐标接近的 blocks 归为同一行，按 X 排序。"""
|
||||
blocks = [
|
||||
_make_block("右", x1=100, x2=130, y1=10, y2=30),
|
||||
_make_block("左", x1=0, x2=30, y1=10, y2=30),
|
||||
]
|
||||
lines = sort_and_group(blocks)
|
||||
assert len(lines) == 1
|
||||
assert len(lines[0]) == 2
|
||||
assert lines[0][0].text == "左"
|
||||
assert lines[0][1].text == "右"
|
||||
|
||||
def test_multi_line_grouping(self) -> None:
|
||||
"""多行正确归组。"""
|
||||
blocks = [
|
||||
_make_block("B1", x1=0, x2=30, y1=100, y2=130),
|
||||
_make_block("A1", x1=0, x2=30, y1=10, y2=40),
|
||||
_make_block("A2", x1=40, x2=70, y1=10, y2=40),
|
||||
]
|
||||
lines = sort_and_group(blocks)
|
||||
assert len(lines) == 2
|
||||
# 第一行：A1、A2
|
||||
assert len(lines[0]) == 2
|
||||
assert lines[0][0].text == "A1"
|
||||
assert lines[0][1].text == "A2"
|
||||
# 第二行：B1
|
||||
assert len(lines[1]) == 1
|
||||
assert lines[1][0].text == "B1"
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# 多栏检测
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestMultiColumn:
|
||||
"""多栏检测。"""
|
||||
|
||||
def assert_one_column(self, blocks: list, page_width: float = 800) -> None:
|
||||
cols = _detect_columns(blocks, page_width)
|
||||
assert len(cols) == 1, f"应检测为单栏，实际 {len(cols)} 栏"
|
||||
|
||||
def assert_two_columns(self, blocks: list, page_width: float = 800) -> None:
|
||||
cols = _detect_columns(blocks, page_width)
|
||||
assert len(cols) == 2, f"应检测为双栏，实际 {len(cols)} 栏"
|
||||
assert len(cols[0]) > 0
|
||||
assert len(cols[1]) > 0
|
||||
|
||||
def test_single_column_few_blocks(self) -> None:
|
||||
"""少于 4 个 block 不触发多栏检测。"""
|
||||
self.assert_one_column([_make_block("A")])
|
||||
|
||||
def test_single_column_no_gap(self) -> None:
|
||||
"""无间隙 → 单栏。"""
|
||||
blocks = [
|
||||
_make_block("左1", x1=0, x2=100, y1=10, y2=30),
|
||||
_make_block("右1", x1=100, x2=200, y1=10, y2=30),
|
||||
_make_block("左2", x1=0, x2=100, y1=40, y2=60),
|
||||
_make_block("右2", x1=100, x2=200, y1=40, y2=60),
|
||||
]
|
||||
self.assert_one_column(blocks)
|
||||
|
||||
def test_two_columns_clear_gap(self) -> None:
|
||||
"""有明显水平间隙 → 双栏。"""
|
||||
blocks = [
|
||||
_make_block("左栏 A", x1=0, x2=150, y1=10, y2=30),
|
||||
_make_block("左栏 B", x1=0, x2=150, y1=40, y2=60),
|
||||
_make_block("右栏 A", x1=500, x2=650, y1=15, y2=35),
|
||||
_make_block("右栏 B", x1=500, x2=650, y1=50, y2=70),
|
||||
]
|
||||
self.assert_two_columns(blocks, page_width=800)
|
||||
|
||||
def test_sort_and_group_two_columns(self) -> None:
|
||||
"""双栏排版：左栏两行，右栏两行。"""
|
||||
blocks = [
|
||||
_make_block("右A", x1=500, x2=650, y1=50, y2=70),
|
||||
_make_block("左B", x1=0, x2=150, y1=40, y2=60),
|
||||
_make_block("左A", x1=0, x2=150, y1=10, y2=30),
|
||||
_make_block("右B", x1=500, x2=650, y1=80, y2=100),
|
||||
]
|
||||
lines = sort_and_group(blocks, page_width=800)
|
||||
# 左栏两行 → 右栏两行
|
||||
assert len(lines) == 4
|
||||
# 左栏第一行
|
||||
assert lines[0][0].text == "左A"
|
||||
# 左栏第二行
|
||||
assert lines[1][0].text == "左B"
|
||||
# 右栏第一行
|
||||
assert lines[2][0].text == "右A"
|
||||
# 右栏第二行
|
||||
assert lines[3][0].text == "右B"
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# group_paragraphs
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestGroupParagraphs:
|
||||
"""段落分组。"""
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert group_paragraphs([]) == []
|
||||
|
||||
def test_single_para(self) -> None:
|
||||
"""多行同段落（用长文本避免触发标题检测）。"""
|
||||
lines = [
|
||||
[_make_block("第一行正文内容这里是长文本")],
|
||||
[_make_block("第二行正文内容再接一些字")],
|
||||
]
|
||||
paras = group_paragraphs(lines)
|
||||
assert len(paras) == 1
|
||||
assert paras[0]["type"] == "text"
|
||||
assert paras[0]["type"] == "text"
|
||||
|
||||
def test_paragraph_break_by_gap(self) -> None:
|
||||
"""行间距大 → 分段（需要 3+ 行，单行间距无法超过自身中位值的 1.8 倍）。"""
|
||||
lines = [
|
||||
[_make_block("第一段正文在这里第一行啊", y1=10, y2=30)],
|
||||
[_make_block("第一段正文在这里第二行", y1=40, y2=60)],
|
||||
[_make_block("第二段正文这里间距很大", y1=200, y2=230)],
|
||||
]
|
||||
paras = group_paragraphs(lines, page_width=800)
|
||||
assert len(paras) == 2
|
||||
# 第一段：两行
|
||||
assert paras[0]["type"] == "text"
|
||||
assert "第一段正文在这里" in paras[0]["text"]
|
||||
# 第二段：一行
|
||||
assert "第二段正文" in paras[1]["text"]
|
||||
|
||||
def test_heading_paragraph(self) -> None:
|
||||
"""短行 → 标题段。"""
|
||||
lines = [[_make_block("第一章")], [_make_block("正文内容")]]
|
||||
paras = group_paragraphs(lines)
|
||||
assert len(paras) >= 1
|
||||
# 标题段
|
||||
assert paras[0]["type"] == "heading"
|
||||
|
||||
def test_list_paragraph(self) -> None:
|
||||
"""列表项段落。"""
|
||||
lines = [
|
||||
[_make_block("- 项目一", y1=10, y2=30)],
|
||||
[_make_block("- 项目二", y1=40, y2=60)],
|
||||
]
|
||||
paras = group_paragraphs(lines)
|
||||
assert len(paras) == 1
|
||||
assert paras[0]["type"] == "list_item"
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# build_markdown_text
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestBuildMarkdown:
|
||||
"""完整的 Markdown 转换。"""
|
||||
|
||||
def test_empty(self) -> None:
|
||||
md, lines, paras = build_markdown_text([])
|
||||
assert md == ""
|
||||
assert lines == []
|
||||
assert paras == []
|
||||
|
||||
def test_simple_text(self) -> None:
|
||||
"""简单段落转换。"""
|
||||
blocks = [
|
||||
_make_block("你好世界"),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "你好世界" in md
|
||||
|
||||
def test_heading_conversion(self) -> None:
|
||||
"""标题应加 # 前缀（级别基于字体大小）。"""
|
||||
# 标题 block 更大 → 推断为 h2
|
||||
blocks = [
|
||||
_make_block("第一章", x1=0, x2=100, y1=10, y2=50), # 40px 高
|
||||
_make_block("正文内容长一些避免被检测为标题", x1=0, x2=200, y1=60, y2=80), # 20px 高
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "### 第一章" in md or "## 第一章" in md
|
||||
assert "正文内容长一些避免被检测为标题" in md
|
||||
|
||||
def test_list_conversion(self) -> None:
|
||||
"""列表项应加 - 前缀。"""
|
||||
blocks = [
|
||||
_make_block("- 苹果"),
|
||||
_make_block("- 香蕉"),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "- 苹果" in md
|
||||
assert "- 香蕉" in md
|
||||
|
||||
def test_centered_heading(self) -> None:
|
||||
"""居中的短行 → heading。"""
|
||||
# 中心 X 在页面宽度中间
|
||||
blocks = [
|
||||
_make_block("欢迎光临", x1=300, x2=500, y1=10, y2=40),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks, page_width=800)
|
||||
assert "欢迎光临" in md
|
||||
|
||||
def test_heading_level_h1(self) -> None:
|
||||
"""居中大字 → h1（# 前缀）。"""
|
||||
blocks = [
|
||||
_make_block("年度报告", x1=300, x2=500, y1=10, y2=50),
|
||||
_make_block("正文内容较小", x1=0, x2=100, y1=60, y2=75),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks, page_width=800)
|
||||
assert "# 年度报告" in md or "## 年度报告" in md
|
||||
|
||||
def test_heading_level_h3(self) -> None:
|
||||
"""等大字体短行 → h3（### 前缀）。"""
|
||||
blocks = [
|
||||
_make_block("摘要", x1=0, x2=100, y1=10, y2=30),
|
||||
_make_block("第一章内容文字", x1=0, x2=100, y1=40, y2=70),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks, page_width=800)
|
||||
# "摘要" ≤ 8 字无标点 → heading, 但字体与正文同大 → h3
|
||||
assert "### 摘要" in md
|
||||
|
||||
def test_separator_conversion(self) -> None:
|
||||
"""分隔线 → ---。"""
|
||||
blocks = [
|
||||
_make_block("---"),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "---" in md
|
||||
|
||||
def test_multi_line_same_para(self) -> None:
|
||||
"""多行近距离内容合并在一个段落（长文本避免标题检测）。"""
|
||||
blocks = [
|
||||
_make_block("第一行长文本避免被检测为标题行", y1=10, y2=30),
|
||||
_make_block("第二行避免被检测为标题", y1=35, y2=55),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "第一行长文本避免被检测为标题行" in md
|
||||
assert "第二行避免被检测为标题" in md
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
# 边缘情况
|
||||
# ═══════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""边缘情况。"""
|
||||
|
||||
def test_single_block_only(self) -> None:
|
||||
blocks = [_make_block("Hello")]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert md == "Hello"
|
||||
|
||||
def test_blocks_with_same_y(self) -> None:
|
||||
"""完全同一 Y 位置的多个 blocks（罕见但需处理）。"""
|
||||
blocks = [
|
||||
_make_block("A", x1=0, x2=30, y1=10, y2=30),
|
||||
_make_block("B", x1=40, x2=70, y1=10, y2=30),
|
||||
]
|
||||
lines = sort_and_group(blocks)
|
||||
assert len(lines) == 1
|
||||
assert len(lines[0]) == 2
|
||||
|
||||
def test_blocks_with_negligible_height(self) -> None:
|
||||
"""高度为 0 的 block 不应崩。"""
|
||||
block = TextBlock(
|
||||
text="边缘",
|
||||
confidence=0.5,
|
||||
bbox=BBox(points=[[0, 0], [0, 0], [0, 0], [0, 0]]),
|
||||
)
|
||||
assert _block_height(block) == 0.0
|
||||
lines = sort_and_group([block])
|
||||
assert len(lines) == 1
|
||||
|
||||
def test_confidence_range(self) -> None:
|
||||
"""置信度在 0~1 范围内。"""
|
||||
block = _make_block("test", confidence=1.0)
|
||||
assert 0 <= block.confidence <= 1
|
||||
|
||||
def test_long_line_content(self) -> None:
|
||||
"""长行（> 40 字符）→ text 而非 heading。"""
|
||||
long_text = "A" * 50
|
||||
blocks = [_make_block(long_text)]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert long_text in md
|
||||
assert "##" not in md
|
||||
|
||||
def test_empty_text_blocks(self) -> None:
|
||||
"""空文本 block 不应造成影响。"""
|
||||
blocks = [
|
||||
_make_block(""),
|
||||
_make_block("实际内容"),
|
||||
]
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "实际内容" in md
|
||||
|
||||
def test_numeric_only(self) -> None:
|
||||
"""纯数字行 → numbered。"""
|
||||
blocks = [_make_block("42")]
|
||||
# 短行数字 → numbered（先于 heading 触发）
|
||||
md, lines, paras = build_markdown_text(blocks)
|
||||
assert "42" in md
|
||||
@@ -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()
|
||||
Generated
+2444
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user