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