|
|
@@ -1,29 +1,63 @@
|
|
|
"""
|
|
|
-混合全文提取实现 (HybridFullTextExtractor) - 飞浆版面分析版
|
|
|
+混合全文提取实现 (HybridFullTextExtractor) - GLM-OCR 版
|
|
|
|
|
|
-基于飞浆 RapidLayout 版面分析,检测 table 区域判断扫描件:
|
|
|
-1. 第一阶段:使用飞浆 RapidLayout 对所有页面进行版面分析
|
|
|
-2. 第二阶段:含有 table 区域的页面走 MinerU OCR,其余走本地提取
|
|
|
+【修改日期】2025-03-27
|
|
|
+【修改说明】OCR 引擎从 MinerU 替换为 GLM-OCR 本地 API
|
|
|
+- 版面分析阶段:保持不变(飞浆 RapidLayout)
|
|
|
+- OCR 阶段:改为 GLM-OCR 单页请求
|
|
|
+- 删除所有 MinerU 相关代码
|
|
|
+
|
|
|
+【请求格式】参考 glm_ocr_api_extractor.py 最终实现版本
|
|
|
+【API 地址】http://183.220.37.46:25429/v1/chat/completions
|
|
|
"""
|
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
+import base64
|
|
|
import io
|
|
|
+import time
|
|
|
+from typing import Any, Dict, List, Optional, Set
|
|
|
+
|
|
|
import fitz # PyMuPDF
|
|
|
-import os
|
|
|
-import tempfile
|
|
|
import numpy as np
|
|
|
-from typing import Any, Dict, List, Optional, Set
|
|
|
+import requests
|
|
|
|
|
|
from foundation.observability.logger.loggering import review_logger as logger
|
|
|
|
|
|
from ..config.provider import default_config_provider
|
|
|
from ..interfaces import DocumentSource, FullTextExtractor
|
|
|
from .fulltext_extractor import PdfFullTextExtractor
|
|
|
-from .mineru_extractor import LocalMinerUFullTextExtractor
|
|
|
-from foundation.observability.cachefiles.cache_manager import cache, CacheBaseDir
|
|
|
|
|
|
-# 尝试导入 RapidLayout,如果未安装则给出友好提示
|
|
|
+
|
|
|
+def _read_ini_config(section: str, key: str, default: Any = None) -> Any:
|
|
|
+ """从项目根目录的 config.ini 读取配置"""
|
|
|
+ try:
|
|
|
+ import configparser
|
|
|
+ from pathlib import Path
|
|
|
+
|
|
|
+ # 查找项目根目录的 config.ini
|
|
|
+ config_path = Path(__file__).parent.parent.parent.parent.parent.parent / "config" / "config.ini"
|
|
|
+ if not config_path.exists():
|
|
|
+ return default
|
|
|
+
|
|
|
+ config = configparser.ConfigParser()
|
|
|
+ config.read(config_path, encoding="utf-8")
|
|
|
+
|
|
|
+ if section in config and key in config[section]:
|
|
|
+ return config[section][key]
|
|
|
+ return default
|
|
|
+ except Exception:
|
|
|
+ return default
|
|
|
+
|
|
|
+# 尝试导入 PIL 用于图片压缩
|
|
|
+try:
|
|
|
+ from PIL import Image
|
|
|
+ PIL_AVAILABLE = True
|
|
|
+except ImportError:
|
|
|
+ PIL_AVAILABLE = False
|
|
|
+ logger.warning("PIL 未安装,GLM-OCR 图片压缩功能将不可用")
|
|
|
+
|
|
|
+# 尝试导入 RapidLayout
|
|
|
try:
|
|
|
from rapid_layout import RapidLayout
|
|
|
RAPID_LAYOUT_AVAILABLE = True
|
|
|
@@ -34,32 +68,81 @@ except ImportError:
|
|
|
|
|
|
class HybridFullTextExtractor(FullTextExtractor):
|
|
|
"""
|
|
|
- 混合提取器:基于飞浆版面分析检测 table 区域,智能路由扫描页到 MinerU OCR。
|
|
|
+ 混合提取器:基于飞浆版面分析检测 table 区域,智能路由扫描页到 GLM-OCR。
|
|
|
+
|
|
|
+ 【变更记录】
|
|
|
+ - 2025-03-27: OCR 引擎从 MinerU 切换为 GLM-OCR 本地 API
|
|
|
"""
|
|
|
|
|
|
+ # GLM-OCR 图片尺寸限制
|
|
|
+ MAX_SHORT_EDGE = 1024 # 短边最大 1024px
|
|
|
+ JPEG_QUALITY = 90 # 提高质量到 90,平衡识别效果和传输大小
|
|
|
+
|
|
|
def __init__(
|
|
|
self,
|
|
|
- layout_dpi: int = 180,
|
|
|
- ocr_dpi: int = 220,
|
|
|
- jpg_quality: int = 90
|
|
|
+ layout_dpi: int = 200, # 【优化】统一 DPI 为 200,兼顾版面分析和 OCR 质量
|
|
|
+ ocr_dpi: int = 200, # 【优化】与 layout_dpi 保持一致,避免重复渲染
|
|
|
+ jpg_quality: int = 90,
|
|
|
+ api_url: Optional[str] = None,
|
|
|
+ timeout: int = 600
|
|
|
) -> None:
|
|
|
self._cfg = default_config_provider
|
|
|
- # 复用已有的提取器
|
|
|
self.local_extractor = PdfFullTextExtractor()
|
|
|
- self.mineru_extractor = LocalMinerUFullTextExtractor() # 使用本地 MinerU
|
|
|
-
|
|
|
- # 飞浆版面分析配置(保守版优化参数)
|
|
|
- self.layout_dpi = layout_dpi # 版面分析 DPI:180(平衡检测精度和速度)
|
|
|
- self.ocr_dpi = ocr_dpi # OCR阶段 DPI:220(表格识别甜点值)
|
|
|
- self.jpg_quality = jpg_quality # JPEG质量:90(几乎无损,文件可控)
|
|
|
- self._layout_engine: Optional[Any] = None # 延迟初始化
|
|
|
-
|
|
|
- # 外部注入的进度状态字典(由 DocumentWorkflow 设置,心跳协程读取)
|
|
|
- # 格式:{'current': int(0-100), 'message': str}
|
|
|
- # 阶段一(版面分析):current 0→50,阶段二(OCR提取):current 50→100
|
|
|
+
|
|
|
+ # 【新增】OCR 引擎选择配置
|
|
|
+ # 优先级:config.ini [ocr] ENGINE > 默认 glm_ocr
|
|
|
+ # 同时支持 "glm_ocr"/"glm-ocr" 和 "mineru"/"mineru-ocr" 等多种写法
|
|
|
+ raw_engine = _read_ini_config("ocr", "engine", "glm_ocr")
|
|
|
+ self.ocr_engine = raw_engine.lower().strip() if raw_engine else "glm_ocr"
|
|
|
+
|
|
|
+ # 规范化引擎名称(统一转换为标准格式)
|
|
|
+ if self.ocr_engine in ("glm_ocr", "glm-ocr", "glmocr"):
|
|
|
+ self.ocr_engine_normalized = "glm_ocr"
|
|
|
+ elif self.ocr_engine in ("mineru", "mineru-ocr", "mineru_ocr"):
|
|
|
+ self.ocr_engine_normalized = "mineru"
|
|
|
+ else:
|
|
|
+ logger.warning(f"[HybridExtractor] 未知的 OCR 引擎 '{self.ocr_engine}',使用默认 glm_ocr")
|
|
|
+ self.ocr_engine_normalized = "glm_ocr"
|
|
|
+
|
|
|
+ logger.info(f"[HybridExtractor] OCR 引擎配置: '{self.ocr_engine}' -> 使用: '{self.ocr_engine_normalized}'")
|
|
|
+
|
|
|
+ # GLM-OCR 配置(从 config.ini 读取,兼容原有逻辑)
|
|
|
+ self.glm_api_url = api_url or _read_ini_config(
|
|
|
+ "ocr", "glm_ocr_api_url",
|
|
|
+ "http://183.220.37.46:25429/v1/chat/completions"
|
|
|
+ )
|
|
|
+ self.glm_timeout = int(_read_ini_config("ocr", "glm_ocr_timeout", "600"))
|
|
|
+
|
|
|
+ # 【新增】读取 GLM-OCR API Key(用于鉴权)
|
|
|
+ self.glm_api_key = _read_ini_config("ocr", "glm_ocr_api_key", "")
|
|
|
+
|
|
|
+ # 构建请求头,如果配置了 API Key 则添加 Authorization
|
|
|
+ self.glm_headers = {"Content-Type": "application/json"}
|
|
|
+ if self.glm_api_key:
|
|
|
+ self.glm_headers["Authorization"] = f"Bearer {self.glm_api_key}"
|
|
|
+ logger.debug(f"[HybridExtractor] GLM-OCR 已配置 API Key 鉴权")
|
|
|
+
|
|
|
+ # 【新增】MinerU 配置
|
|
|
+ self.mineru_api_url = _read_ini_config(
|
|
|
+ "ocr", "mineru_api_url",
|
|
|
+ "http://183.220.37.46:25428/file_parse"
|
|
|
+ )
|
|
|
+ self.mineru_timeout = int(_read_ini_config("ocr", "mineru_timeout", "300"))
|
|
|
+
|
|
|
+ # 【优化】飞浆版面分析配置 - DPI 统一为 200
|
|
|
+ # 原理:版面分析和 OCR 使用相同 DPI,第一阶段渲染的图片可直接复用
|
|
|
+ self.layout_dpi = layout_dpi
|
|
|
+ self.ocr_dpi = ocr_dpi
|
|
|
+ self.jpg_quality = jpg_quality
|
|
|
+ self._layout_engine: Optional[Any] = None
|
|
|
+
|
|
|
+ # 【优化】图片缓存:版面分析阶段缓存 table 页图片,供 OCR 阶段复用
|
|
|
+ # 格式: {page_num: (width, height, jpeg_bytes)}
|
|
|
+ self._image_cache: Dict[int, tuple] = {}
|
|
|
+
|
|
|
+ # 外部注入的进度状态字典
|
|
|
self._progress_state: Optional[dict] = None
|
|
|
|
|
|
- # 检查 RapidLayout 是否可用
|
|
|
if not RAPID_LAYOUT_AVAILABLE:
|
|
|
raise ImportError(
|
|
|
"RapidLayout 未安装。请在 doc_worker_venv 虚拟环境中运行:\n"
|
|
|
@@ -73,74 +156,90 @@ class HybridFullTextExtractor(FullTextExtractor):
|
|
|
self._layout_engine = RapidLayout()
|
|
|
return self._layout_engine
|
|
|
|
|
|
- def _detect_table_pages(self, doc: fitz.Document, dpi: int = 150) -> Set[int]:
|
|
|
+ def _detect_table_pages(self, doc: fitz.Document, dpi: int = 200) -> Set[int]:
|
|
|
"""
|
|
|
使用飞浆 RapidLayout 检测所有页面,返回包含 table 区域的页码集合。
|
|
|
|
|
|
- Args:
|
|
|
- doc: PyMuPDF 文档对象
|
|
|
- dpi: PDF 转图片的分辨率
|
|
|
-
|
|
|
- Returns:
|
|
|
- 包含 table 区域的页码集合 (1-based)
|
|
|
+ 【优化】检测到 table 的页面,将 JPEG 图片缓存到 self._image_cache
|
|
|
+ 供后续 OCR 阶段直接使用,避免重复渲染 PDF。
|
|
|
"""
|
|
|
table_pages: Set[int] = set()
|
|
|
layout_engine = self._get_layout_engine()
|
|
|
total_pages = len(doc)
|
|
|
+
|
|
|
+ # 清空图片缓存
|
|
|
+ self._image_cache.clear()
|
|
|
|
|
|
- logger.debug(f" [飞浆分析] 开始版面分析,共 {total_pages} 页...")
|
|
|
+ logger.info(f" [飞浆分析] 开始版面分析,共 {total_pages} 页,DPI={dpi}(图片缓存已启用)")
|
|
|
|
|
|
for page_num in range(1, total_pages + 1):
|
|
|
- page = doc[page_num - 1] # PyMuPDF 使用 0-based 索引
|
|
|
+ page = doc[page_num - 1]
|
|
|
|
|
|
- # 1. 将页面转换为图片
|
|
|
+ # 将页面转换为图片
|
|
|
pix = page.get_pixmap(dpi=dpi)
|
|
|
img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
|
|
|
|
|
|
- # 2. 飞浆版面分析
|
|
|
+ # 飞浆版面分析
|
|
|
try:
|
|
|
layout_output = layout_engine(img)
|
|
|
|
|
|
- # 3. 解析版面结果,检查是否有 table 区域
|
|
|
+ # 解析版面结果,检查是否有 table 区域
|
|
|
labels = []
|
|
|
if hasattr(layout_output, 'class_names'):
|
|
|
labels = list(layout_output.class_names)
|
|
|
elif hasattr(layout_output, 'boxes'):
|
|
|
- # 兼容不同版本的输出格式
|
|
|
labels = [
|
|
|
label for _, label, _
|
|
|
in zip(layout_output.boxes, layout_output.class_names, layout_output.scores)
|
|
|
]
|
|
|
|
|
|
- # 4. 判断是否包含 table
|
|
|
+ # 判断是否包含 table
|
|
|
if "table" in labels:
|
|
|
table_pages.add(page_num)
|
|
|
- logger.debug(f" 第 {page_num} 页: 检测到 table 区域 -> 将走 MinerU OCR")
|
|
|
+
|
|
|
+ # 【优化】缓存 table 页图片为 JPEG,供 OCR 阶段复用
|
|
|
+ try:
|
|
|
+ # 直接保存 Pixmap 的 JPEG 数据,无需 PIL 转换
|
|
|
+ jpeg_bytes = pix.tobytes("jpeg")
|
|
|
+ self._image_cache[page_num] = (pix.width, pix.height, jpeg_bytes)
|
|
|
+ logger.debug(f" 第 {page_num} 页: 检测到 table -> 缓存图片 "
|
|
|
+ f"({pix.width}x{pix.height}, {len(jpeg_bytes)/1024:.1f} KB)")
|
|
|
+ except Exception as cache_err:
|
|
|
+ logger.warning(f" 第 {page_num} 页: 图片缓存失败 ({cache_err})")
|
|
|
+
|
|
|
else:
|
|
|
region_types = ", ".join(set(labels)) if labels else "无"
|
|
|
logger.debug(f" 第 {page_num} 页: {region_types}")
|
|
|
|
|
|
except Exception as e:
|
|
|
logger.error(f" 第 {page_num} 页: 版面分析失败 ({e}),默认不走 OCR")
|
|
|
- # 分析失败时,保守起见不走 OCR
|
|
|
pass
|
|
|
|
|
|
- # 阶段一进度:已分析页 / 总页数 → 0% ~ 50%
|
|
|
+ # 阶段一进度
|
|
|
if self._progress_state is not None:
|
|
|
self._progress_state['current'] = int(page_num / total_pages * 50)
|
|
|
self._progress_state['message'] = f"版面分析中:已分析 {page_num}/{total_pages} 页"
|
|
|
|
|
|
- logger.debug(f" [飞浆分析] 完成,共 {len(table_pages)} 页包含 table 区域: {sorted(table_pages)}")
|
|
|
+ cache_size_mb = sum(len(data[2]) for data in self._image_cache.values()) / 1024 / 1024
|
|
|
+ logger.info(f" [飞浆分析] 完成: {len(table_pages)} 页 table,"
|
|
|
+ f"缓存 {len(self._image_cache)} 页图片 ({cache_size_mb:.1f} MB)")
|
|
|
return table_pages
|
|
|
|
|
|
def extract_full_text(self, source: DocumentSource) -> List[Dict[str, Any]]:
|
|
|
"""
|
|
|
执行混合提取流程:
|
|
|
1. 首先用飞浆 RapidLayout 检测所有页面的 table 区域
|
|
|
- 2. 含有 table 的页面走 MinerU OCR
|
|
|
+ 2. 含有 table 的页面走 GLM-OCR
|
|
|
3. 其他页面走本地 PyMuPDF 提取
|
|
|
+
|
|
|
+ 【统计信息】本方法会统计并输出总提取时间、OCR页数等信息
|
|
|
"""
|
|
|
- # 1. 打开文档
|
|
|
+ # 记录总开始时间
|
|
|
+ total_start_time = time.time()
|
|
|
+ layout_analysis_time = 0.0
|
|
|
+ ocr_total_time = 0.0
|
|
|
+
|
|
|
+ # 打开文档
|
|
|
if source.content is not None:
|
|
|
doc = fitz.open(stream=io.BytesIO(source.content))
|
|
|
source_file = "bytes_stream"
|
|
|
@@ -155,14 +254,31 @@ class HybridFullTextExtractor(FullTextExtractor):
|
|
|
|
|
|
try:
|
|
|
total_pages = len(doc)
|
|
|
- logger.debug(f"开始混合提取(飞浆版面分析 + 本地 MinerU),共 {total_pages} 页...")
|
|
|
+ ocr_page_count = 0 # 统计需要OCR的页数
|
|
|
+
|
|
|
+ # INFO级别:开始文档提取(方便查看主要流程)
|
|
|
+ current_engine = "GLM-OCR" if self.ocr_engine_normalized == "glm_ocr" else "MinerU"
|
|
|
+ logger.info(f"[文档提取] 开始处理,共 {total_pages} 页,OCR引擎: {current_engine}")
|
|
|
+ logger.debug(f"开始混合提取(飞浆版面分析 + {current_engine}),共 {total_pages} 页...")
|
|
|
|
|
|
if self._progress_state is not None:
|
|
|
self._progress_state['current'] = 0
|
|
|
self._progress_state['message'] = f"版面分析中:已分析 0/{total_pages} 页"
|
|
|
|
|
|
- # ========== 第一阶段:飞浆版面分析,检测 table 页 ==========
|
|
|
+ # ========== 第一阶段:飞浆版面分析 ==========
|
|
|
+ layout_start_time = time.time()
|
|
|
table_pages = self._detect_table_pages(doc, dpi=self.layout_dpi)
|
|
|
+ layout_analysis_time = time.time() - layout_start_time
|
|
|
+ ocr_page_count = len(table_pages)
|
|
|
+
|
|
|
+ # INFO级别:版面分析完成,显示OCR页数
|
|
|
+ if ocr_page_count > 0:
|
|
|
+ logger.info(f"[文档提取] 版面分析完成,共 {ocr_page_count} 页需要OCR识别,"
|
|
|
+ f"{total_pages - ocr_page_count} 页直接提取,"
|
|
|
+ f"版面分析耗时: {layout_analysis_time:.2f}s")
|
|
|
+ else:
|
|
|
+ logger.info(f"[文档提取] 版面分析完成,无扫描页,全部直接提取,"
|
|
|
+ f"版面分析耗时: {layout_analysis_time:.2f}s")
|
|
|
|
|
|
# ========== 第二阶段:分流处理 ==========
|
|
|
logger.debug(f"\n开始分流处理...")
|
|
|
@@ -170,25 +286,32 @@ class HybridFullTextExtractor(FullTextExtractor):
|
|
|
for i, page in enumerate(doc):
|
|
|
page_num = i + 1
|
|
|
|
|
|
- # 判断是否为 table 页(即扫描件)
|
|
|
if page_num in table_pages:
|
|
|
- logger.debug(f" [第 {page_num} 页] 检测到 table -> 走本地 MinerU OCR")
|
|
|
+ # 【修改】根据配置选择 OCR 引擎
|
|
|
+ # 使用规范化后的引擎名称(支持 glm_ocr/glm-ocr 和 mineru/mineru-ocr)
|
|
|
+ is_glm_ocr = self.ocr_engine_normalized == "glm_ocr"
|
|
|
+ ocr_name = "GLM-OCR" if is_glm_ocr else "MinerU"
|
|
|
+ logger.debug(f" [第 {page_num} 页] 检测到 table -> 走 {ocr_name}")
|
|
|
|
|
|
- # --- 扫描件处理 (MinerU OCR) ---
|
|
|
try:
|
|
|
- page_text = self._ocr_page(page, page_num, source_file)
|
|
|
+ # 根据配置调用不同的 OCR 引擎,并统计 OCR 时间
|
|
|
+ ocr_start_time = time.time()
|
|
|
+ if is_glm_ocr:
|
|
|
+ page_text = self._ocr_page_with_glm(page, page_num, source_file)
|
|
|
+ else:
|
|
|
+ page_text = self._ocr_page_with_mineru(doc, page_num, source_file)
|
|
|
+ ocr_total_time += time.time() - ocr_start_time
|
|
|
except Exception as e:
|
|
|
- logger.error(f" MinerU OCR 失败,回退到本地提取: {e}")
|
|
|
+ logger.error(f" {ocr_name} 失败,回退到本地提取: {e}")
|
|
|
raw_text = page.get_text()
|
|
|
page_text = self.local_extractor._filter_header_footer(raw_text)
|
|
|
else:
|
|
|
logger.debug(f" [第 {page_num} 页] 无 table -> 走本地 PyMuPDF 提取")
|
|
|
|
|
|
- # --- 电子版处理 (本地 PyMuPDF) ---
|
|
|
text_with_tables = self.local_extractor._extract_text_with_table_placeholders(page)
|
|
|
page_text = self.local_extractor._filter_header_footer(text_with_tables)
|
|
|
|
|
|
- # --- 组装结果 ---
|
|
|
+ # 组装结果
|
|
|
pages.append({
|
|
|
"page_num": page_num,
|
|
|
"text": page_text,
|
|
|
@@ -198,7 +321,7 @@ class HybridFullTextExtractor(FullTextExtractor):
|
|
|
})
|
|
|
current_pos += len(page_text)
|
|
|
|
|
|
- # 阶段二进度:已处理页 / 总页数 → 50% ~ 100%
|
|
|
+ # 阶段二进度
|
|
|
if self._progress_state is not None:
|
|
|
self._progress_state['current'] = 50 + int(page_num / total_pages * 50)
|
|
|
ocr_flag = "(OCR)" if page_num in table_pages else ""
|
|
|
@@ -206,8 +329,36 @@ class HybridFullTextExtractor(FullTextExtractor):
|
|
|
|
|
|
finally:
|
|
|
doc.close()
|
|
|
+ # 【优化】清理图片缓存,释放内存
|
|
|
+ if hasattr(self, '_image_cache'):
|
|
|
+ cache_size = len(self._image_cache)
|
|
|
+ self._image_cache.clear()
|
|
|
+ if cache_size > 0:
|
|
|
+ logger.debug(f" [缓存清理] 已清理 {cache_size} 页图片缓存")
|
|
|
+
|
|
|
+ # ========== 统计信息输出 ==========
|
|
|
+ # INFO级别:文档提取完成,输出详细统计
|
|
|
+ total_time = time.time() - total_start_time
|
|
|
+ total_chars = sum(len(page['text']) for page in pages)
|
|
|
+
|
|
|
+ # 计算各类时间占比
|
|
|
+ ocr_avg_time = ocr_total_time / ocr_page_count if ocr_page_count > 0 else 0
|
|
|
+ local_pages = total_pages - ocr_page_count
|
|
|
+
|
|
|
+ logger.info(
|
|
|
+ f"[文档提取] 完成统计 | "
|
|
|
+ f"总页数: {total_pages} | "
|
|
|
+ f"OCR页数: {ocr_page_count} | "
|
|
|
+ f"本地提取: {local_pages} | "
|
|
|
+ f"总耗时: {total_time:.2f}s | "
|
|
|
+ f"版面分析: {layout_analysis_time:.2f}s | "
|
|
|
+ f"OCR耗时: {ocr_total_time:.2f}s | "
|
|
|
+ f"OCR平均: {ocr_avg_time:.2f}s/页 | "
|
|
|
+ f"总字符数: {total_chars}"
|
|
|
+ )
|
|
|
|
|
|
# 保存提取后的原始PDF内容到缓存目录
|
|
|
+ from foundation.observability.cachefiles.cache_manager import cache, CacheBaseDir
|
|
|
cache.save(
|
|
|
data=pages,
|
|
|
subdir="document_temp",
|
|
|
@@ -217,50 +368,436 @@ class HybridFullTextExtractor(FullTextExtractor):
|
|
|
|
|
|
return pages
|
|
|
|
|
|
- def _ocr_page(self, page: fitz.Page, page_num: int, original_filename: str) -> str:
|
|
|
+ def _ocr_page_with_glm(self, page: fitz.Page, page_num: int, original_filename: str) -> str:
|
|
|
"""
|
|
|
- 将单页转为图片并调用本地 MinerU OCR。
|
|
|
- 使用 JPEG 格式以减小文件大小,提高传输效率。
|
|
|
+ 将单页转为图片并调用 GLM-OCR 本地 API 识别
|
|
|
+
|
|
|
+ 【优化】优先使用版面分析阶段缓存的图片,避免重复渲染
|
|
|
+
|
|
|
+ 流程:
|
|
|
+ 1. 优先使用缓存图片(如可用)
|
|
|
+ 2. 否则 PyMuPDF 渲染页面为图片(200 DPI)
|
|
|
+ 3. PIL 压缩图片(短边限制 1024px,JPEG 质量 90)
|
|
|
+ 4. Base64 编码
|
|
|
+ 5. POST 请求 GLM-OCR API
|
|
|
+ 6. 解析响应并转换 HTML→Markdown
|
|
|
"""
|
|
|
- # 1. 渲染为图片(保守版优化:220 DPI 提升表格识别精度)
|
|
|
- pix = page.get_pixmap(dpi=self.ocr_dpi)
|
|
|
+ start_time = time.time()
|
|
|
+
|
|
|
+ # 【优化】检查是否有缓存图片
|
|
|
+ cached = self._image_cache.get(page_num)
|
|
|
+ use_cache = cached is not None
|
|
|
+
|
|
|
+ # INFO级别:开始调用GLM-OCR识别(方便查看主要流程)
|
|
|
+ cache_info = "(使用缓存图片)" if use_cache else ""
|
|
|
+ logger.info(f"[GLM-OCR] 开始识别第 {page_num} 页 {cache_info}")
|
|
|
|
|
|
- # 2. 保存为临时 JPEG 文件(比 PNG 更小)
|
|
|
- tmp_path = None
|
|
|
try:
|
|
|
- with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
|
|
|
- tmp_path = tmp_file.name
|
|
|
+ # 1. 获取图片(优先使用缓存)
|
|
|
+ if use_cache:
|
|
|
+ # 【优化】使用版面分析阶段缓存的图片
|
|
|
+ width, height, img_bytes = cached
|
|
|
+ original_kb = len(img_bytes) / 1024
|
|
|
+ logger.debug(f" [GLM-OCR] 第 {page_num} 页使用缓存图片: "
|
|
|
+ f"{original_kb:.1f} KB ({width}x{height})")
|
|
|
+ else:
|
|
|
+ # 兜底:重新渲染(理论上不会发生,因为 table 页都应已缓存)
|
|
|
+ pix = page.get_pixmap(dpi=self.ocr_dpi)
|
|
|
+ img_bytes = pix.tobytes("jpeg")
|
|
|
+ original_kb = len(img_bytes) / 1024
|
|
|
+ logger.warning(f" [GLM-OCR] 第 {page_num} 页无缓存,重新渲染: "
|
|
|
+ f"{original_kb:.1f} KB ({pix.width}x{pix.height})")
|
|
|
|
|
|
- # 保存为 JPEG 格式,质量 90%,几乎无损且文件可控
|
|
|
- pix.save(tmp_path, "jpeg", jpg_quality=self.jpg_quality)
|
|
|
+ # 2. 压缩图片
|
|
|
+ compressed_bytes = self._compress_image(img_bytes)
|
|
|
+ compressed_kb = len(compressed_bytes) / 1024
|
|
|
|
|
|
- # 检查文件是否正确生成
|
|
|
- if not os.path.exists(tmp_path) or os.path.getsize(tmp_path) == 0:
|
|
|
- logger.error(f" [WARN] 无法创建第 {page_num} 页的临时图片")
|
|
|
- return ""
|
|
|
+ # 3. Base64 编码
|
|
|
+ img_base64 = base64.b64encode(compressed_bytes).decode('utf-8').replace('\n', '').replace('\r', '')
|
|
|
+
|
|
|
+ # 4. 构建 OpenAI 兼容格式请求
|
|
|
+ payload = {
|
|
|
+ "model": "GLM-OCR",
|
|
|
+ "messages": [
|
|
|
+ {
|
|
|
+ "role": "user",
|
|
|
+ "content": [
|
|
|
+ {
|
|
|
+ "type": "text",
|
|
|
+ "text": "请详细识别图片中的所有文字内容,保留原始排版格式,以 Markdown 格式输出。"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "type": "image_url",
|
|
|
+ "image_url": {
|
|
|
+ "url": f"data:image/jpeg;base64,{img_base64}"
|
|
|
+ }
|
|
|
+ }
|
|
|
+ ]
|
|
|
+ }
|
|
|
+ ],
|
|
|
+ "max_tokens": 2048,
|
|
|
+ "temperature": 0.1
|
|
|
+ }
|
|
|
+
|
|
|
+ # 5. 调用 GLM-OCR API
|
|
|
+ response = requests.post(
|
|
|
+ self.glm_api_url,
|
|
|
+ headers=self.glm_headers,
|
|
|
+ json=payload,
|
|
|
+ timeout=self.glm_timeout
|
|
|
+ )
|
|
|
+ response.raise_for_status()
|
|
|
+
|
|
|
+ # 6. 解析结果
|
|
|
+ result = response.json()
|
|
|
+ content = self._extract_content(result)
|
|
|
+
|
|
|
+ # 7. 处理 HTML 转 Markdown
|
|
|
+ md_content = self._process_raw_content(content)
|
|
|
+
|
|
|
+ elapsed = time.time() - start_time
|
|
|
+ # INFO级别:识别完成(方便查看主要流程)
|
|
|
+ logger.info(f"[GLM-OCR] 第 {page_num} 页识别完成,耗时: {elapsed:.2f}s,字符数: {len(md_content)}")
|
|
|
+ logger.debug(f" [GLM-OCR] 第 {page_num} 页详细耗时: {elapsed:.2f}s")
|
|
|
+
|
|
|
+ return md_content
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(f" [GLM-OCR] 第 {page_num} 页识别失败: {e}")
|
|
|
+ raise
|
|
|
|
|
|
- # 输出文件大小信息(用于调试)
|
|
|
- file_size_kb = os.path.getsize(tmp_path) / 1024
|
|
|
- logger.debug(f" [INFO] 第 {page_num} 页图片: {file_size_kb:.1f} KB ({pix.width}x{pix.height})")
|
|
|
+ def _ocr_page_with_mineru(self, doc: fitz.Document, page_num: int, original_filename: str) -> str:
|
|
|
+ """
|
|
|
+ 【新增】使用 MinerU 本地 API 识别单页
|
|
|
+
|
|
|
+ 流程:
|
|
|
+ 1. 【优化】优先使用版面分析缓存的图片(JPEG)
|
|
|
+ 2. 无缓存时,提取单页为临时 PDF 文件
|
|
|
+ 3. 调用 MinerU API 上传识别
|
|
|
+ 4. 提取 Markdown 内容
|
|
|
+ 5. 清理临时文件
|
|
|
+
|
|
|
+ Args:
|
|
|
+ doc: 原始 PDF 文档对象
|
|
|
+ page_num: 页码(1-based)
|
|
|
+ original_filename: 原始文件名(用于日志)
|
|
|
|
|
|
- # 3. 构造一个临时的 DocumentSource
|
|
|
- tmp_source = DocumentSource(path=tmp_path)
|
|
|
+ Returns:
|
|
|
+ str: 识别出的 Markdown 文本
|
|
|
+ """
|
|
|
+ import tempfile
|
|
|
+ import os
|
|
|
+
|
|
|
+ start_time = time.time()
|
|
|
+
|
|
|
+ # 【优化】检查是否有缓存图片
|
|
|
+ cached = self._image_cache.get(page_num)
|
|
|
+ use_cache = cached is not None
|
|
|
+
|
|
|
+ # INFO级别:开始识别
|
|
|
+ cache_info = "(使用缓存图片)" if use_cache else ""
|
|
|
+ logger.info(f"[MinerU] 开始识别第 {page_num} 页 {cache_info}")
|
|
|
+
|
|
|
+ tmp_pdf_path = None
|
|
|
+
|
|
|
+ try:
|
|
|
+ # 【优化】优先使用缓存的图片数据
|
|
|
+ if use_cache:
|
|
|
+ width, height, img_bytes = cached
|
|
|
+ logger.debug(f" [MinerU] 第 {page_num} 页使用缓存图片: "
|
|
|
+ f"{len(img_bytes)/1024:.1f} KB ({width}x{height})")
|
|
|
+
|
|
|
+ # 使用图片直接上传(MinerU 支持图片格式)
|
|
|
+ files = {'files': (f"page_{page_num}.jpg", io.BytesIO(img_bytes))}
|
|
|
+ response = requests.post(
|
|
|
+ self.mineru_api_url,
|
|
|
+ files=files,
|
|
|
+ timeout=self.mineru_timeout
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ # 兜底:提取单页为临时 PDF
|
|
|
+ logger.debug(f" [MinerU] 第 {page_num} 页无缓存,创建临时 PDF")
|
|
|
+
|
|
|
+ single_page_doc = fitz.open()
|
|
|
+ single_page_doc.insert_pdf(doc, from_page=page_num-1, to_page=page_num-1)
|
|
|
+
|
|
|
+ # 创建临时文件
|
|
|
+ with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_file:
|
|
|
+ tmp_pdf_path = tmp_file.name
|
|
|
+
|
|
|
+ single_page_doc.save(tmp_pdf_path)
|
|
|
+ single_page_doc.close()
|
|
|
+
|
|
|
+ file_size_kb = os.path.getsize(tmp_pdf_path) / 1024
|
|
|
+ logger.debug(f" [MinerU] 第 {page_num} 页临时文件: {file_size_kb:.1f} KB")
|
|
|
+
|
|
|
+ # 调用 MinerU API
|
|
|
+ with open(tmp_pdf_path, 'rb') as f:
|
|
|
+ files = {'files': (f"page_{page_num}.pdf", f)}
|
|
|
+ response = requests.post(
|
|
|
+ self.mineru_api_url,
|
|
|
+ files=files,
|
|
|
+ timeout=self.mineru_timeout
|
|
|
+ )
|
|
|
|
|
|
- # 4. 调用本地 MinerU
|
|
|
- results = self.mineru_extractor.extract_full_text(tmp_source)
|
|
|
+ if response.status_code != 200:
|
|
|
+ raise RuntimeError(f"MinerU API error: {response.status_code} - {response.text[:200]}")
|
|
|
|
|
|
- if results and len(results) > 0:
|
|
|
- return results[0]["text"]
|
|
|
- return ""
|
|
|
+ # 3. 解析结果
|
|
|
+ result = response.json()
|
|
|
+ content = ""
|
|
|
+
|
|
|
+ if "results" in result and isinstance(result["results"], dict):
|
|
|
+ for filename, file_data in result["results"].items():
|
|
|
+ if isinstance(file_data, dict) and "md_content" in file_data:
|
|
|
+ content = file_data["md_content"]
|
|
|
+ break
|
|
|
+
|
|
|
+ # 4. 处理 HTML 转 Markdown(如果包含 HTML 标签)
|
|
|
+ if "<table" in content.lower() or "<div" in content.lower():
|
|
|
+ logger.debug(f" [MinerU] 检测到 HTML 标签,转换为 Markdown")
|
|
|
+ content = self._process_raw_content(content)
|
|
|
+
|
|
|
+ elapsed = time.time() - start_time
|
|
|
+ logger.info(f"[MinerU] 第 {page_num} 页识别完成,耗时: {elapsed:.2f}s,字符数: {len(content)}")
|
|
|
+
|
|
|
+ return content
|
|
|
|
|
|
except Exception as e:
|
|
|
- logger.error(f" [WARN] 第 {page_num} 页 OCR 失败: {e}")
|
|
|
- return ""
|
|
|
+ logger.error(f" [MinerU] 第 {page_num} 页识别失败: {e}")
|
|
|
+ raise
|
|
|
|
|
|
finally:
|
|
|
# 清理临时文件
|
|
|
- if tmp_path and os.path.exists(tmp_path):
|
|
|
+ if tmp_pdf_path and os.path.exists(tmp_pdf_path):
|
|
|
try:
|
|
|
- os.remove(tmp_path)
|
|
|
+ os.remove(tmp_pdf_path)
|
|
|
+ logger.debug(f" [MinerU] 清理临时文件: {tmp_pdf_path}")
|
|
|
except:
|
|
|
pass
|
|
|
+
|
|
|
+ def _compress_image(self, img_bytes: bytes) -> bytes:
|
|
|
+ """
|
|
|
+ 压缩图片至 GLM-OCR 要求的尺寸限制内
|
|
|
+
|
|
|
+ 【逻辑来源】glm_ocr_api_extractor.py _compress_image 方法
|
|
|
+
|
|
|
+ 压缩规则:
|
|
|
+ - 短边最大 1024px
|
|
|
+ - JPEG 质量 85
|
|
|
+ - 等比缩放
|
|
|
+ """
|
|
|
+ if not PIL_AVAILABLE:
|
|
|
+ logger.debug(" [压缩] PIL 不可用,使用原始图片")
|
|
|
+ return img_bytes
|
|
|
+
|
|
|
+ try:
|
|
|
+ img = Image.open(io.BytesIO(img_bytes))
|
|
|
+
|
|
|
+ # 转为 RGB
|
|
|
+ if img.mode in ('RGBA', 'LA', 'P'):
|
|
|
+ background = Image.new('RGB', img.size, (255, 255, 255))
|
|
|
+ if img.mode == 'P':
|
|
|
+ img = img.convert('RGBA')
|
|
|
+ if img.mode in ('RGBA', 'LA'):
|
|
|
+ background.paste(img, mask=img.split()[-1])
|
|
|
+ img = background
|
|
|
+ elif img.mode != 'RGB':
|
|
|
+ img = img.convert('RGB')
|
|
|
+
|
|
|
+ original_size = img.size
|
|
|
+
|
|
|
+ # 检查是否需要缩放(短边 > 1024px)
|
|
|
+ min_edge = min(img.size)
|
|
|
+ if min_edge > self.MAX_SHORT_EDGE:
|
|
|
+ ratio = self.MAX_SHORT_EDGE / min_edge
|
|
|
+ new_size = (int(img.width * ratio), int(img.height * ratio))
|
|
|
+ img = img.resize(new_size, Image.Resampling.LANCZOS)
|
|
|
+ logger.debug(f" [压缩] 图片缩放: {original_size} -> {img.size}")
|
|
|
+
|
|
|
+ # 压缩为 JPEG
|
|
|
+ buffer = io.BytesIO()
|
|
|
+ img.save(buffer, format='JPEG', quality=self.JPEG_QUALITY, optimize=True)
|
|
|
+
|
|
|
+ compressed_kb = len(buffer.getvalue()) / 1024
|
|
|
+ original_kb = len(img_bytes) / 1024
|
|
|
+ logger.debug(f" [压缩] {original_kb:.1f} KB -> {compressed_kb:.1f} KB")
|
|
|
+
|
|
|
+ return buffer.getvalue()
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f" [压缩] 主流程压缩失败,使用兜底压缩: {e}")
|
|
|
+ # 兜底:简化流程,但保持相同质量
|
|
|
+ try:
|
|
|
+ img = Image.open(io.BytesIO(img_bytes))
|
|
|
+ if img.mode != 'RGB':
|
|
|
+ img = img.convert('RGB')
|
|
|
+ # 确保尺寸符合要求(短边 <= 1024)
|
|
|
+ min_edge = min(img.size)
|
|
|
+ if min_edge > self.MAX_SHORT_EDGE:
|
|
|
+ ratio = self.MAX_SHORT_EDGE / min_edge
|
|
|
+ new_size = (int(img.width * ratio), int(img.height * ratio))
|
|
|
+ img = img.resize(new_size, Image.Resampling.LANCZOS)
|
|
|
+ buffer = io.BytesIO()
|
|
|
+ # 兜底也使用相同质量,确保识别效果
|
|
|
+ img.save(buffer, format='JPEG', quality=self.JPEG_QUALITY, optimize=True)
|
|
|
+ logger.debug(f" [压缩] 兜底压缩成功: {len(buffer.getvalue())/1024:.1f} KB")
|
|
|
+ return buffer.getvalue()
|
|
|
+ except Exception as e2:
|
|
|
+ logger.error(f" [压缩] 兜底压缩也失败: {e2}")
|
|
|
+ # 最后兜底:使用原始图片(可能导致API错误)
|
|
|
+ return img_bytes
|
|
|
+
|
|
|
+ def _extract_content(self, result: Dict[str, Any]) -> str:
|
|
|
+ """
|
|
|
+ 从 OpenAI 兼容响应中提取内容
|
|
|
+
|
|
|
+ 响应格式:
|
|
|
+ {
|
|
|
+ "choices": [{
|
|
|
+ "message": {
|
|
|
+ "content": "识别结果..."
|
|
|
+ }
|
|
|
+ }]
|
|
|
+ }
|
|
|
+ """
|
|
|
+ if "choices" in result and isinstance(result["choices"], list):
|
|
|
+ if len(result["choices"]) > 0:
|
|
|
+ message = result["choices"][0].get("message", {})
|
|
|
+ return message.get("content", "")
|
|
|
+ return ""
|
|
|
+
|
|
|
+ def _process_raw_content(self, raw_content: str) -> str:
|
|
|
+ """
|
|
|
+ 处理原始内容(HTML 转 Markdown)
|
|
|
+
|
|
|
+ 【逻辑来源】glm_ocr_api_extractor.py _process_raw_content 方法
|
|
|
+
|
|
|
+ 处理流程:
|
|
|
+ 1. 检测并转换 HTML 表格
|
|
|
+ 2. 检测 HTML 格式,使用 markdownify 转换
|
|
|
+ 3. 失败则返回原始内容
|
|
|
+ """
|
|
|
+ if not raw_content:
|
|
|
+ return ""
|
|
|
+
|
|
|
+ # 转换 HTML 表格
|
|
|
+ if "<table" in raw_content.lower():
|
|
|
+ raw_content = self._convert_html_tables_to_markdown(raw_content)
|
|
|
+
|
|
|
+ # HTML 转 Markdown
|
|
|
+ if self._is_html_content(raw_content):
|
|
|
+ try:
|
|
|
+ import markdownify
|
|
|
+ return markdownify.markdownify(raw_content, heading_style="ATX").strip()
|
|
|
+ except ImportError:
|
|
|
+ logger.debug(" [转换] markdownify 未安装,跳过 HTML 转换")
|
|
|
+
|
|
|
+ return raw_content.strip()
|
|
|
+
|
|
|
+ def _is_html_content(self, content: str) -> bool:
|
|
|
+ """检查内容是否为 HTML 格式"""
|
|
|
+ if not content:
|
|
|
+ return False
|
|
|
+
|
|
|
+ html_indicators = [
|
|
|
+ "<!DOCTYPE", "<html", "<body", "<div", "<p>", "<table",
|
|
|
+ "<h1", "<h2", "<span", "<br", " ", """
|
|
|
+ ]
|
|
|
+ content_lower = content.lower()
|
|
|
+ html_tag_count = sum(1 for indicator in html_indicators if indicator.lower() in content_lower)
|
|
|
+ return html_tag_count >= 2
|
|
|
+
|
|
|
+ def _convert_html_tables_to_markdown(self, content: str) -> str:
|
|
|
+ """
|
|
|
+ 将 HTML 表格转换为 Markdown 表格格式
|
|
|
+
|
|
|
+ 【逻辑来源】glm_ocr_api_extractor.py _convert_html_tables_to_markdown 方法
|
|
|
+ """
|
|
|
+ import re
|
|
|
+
|
|
|
+ def extract_cell_text(cell_html: str) -> str:
|
|
|
+ text = re.sub(r'<[^>]+>', '', cell_html)
|
|
|
+ text = text.replace(' ', ' ').replace('<', '<').replace('>', '>')
|
|
|
+ text = text.replace('&', '&').replace('"', '"').replace(''', "'")
|
|
|
+ return text.strip()
|
|
|
+
|
|
|
+ def parse_colspan(td_html: str) -> int:
|
|
|
+ match = re.search(r'colspan=["\']?(\d+)["\']?', td_html, re.IGNORECASE)
|
|
|
+ return int(match.group(1)) if match else 1
|
|
|
+
|
|
|
+ def convert_table_match(match):
|
|
|
+ table_html = match.group(0)
|
|
|
+
|
|
|
+ # 提取 thead 和 tbody
|
|
|
+ thead_match = re.search(r'<thead[^>]*>(.*?)</thead>', table_html, re.DOTALL | re.IGNORECASE)
|
|
|
+ tbody_match = re.search(r'<tbody[^>]*>(.*?)</tbody>', table_html, re.DOTALL | re.IGNORECASE)
|
|
|
+
|
|
|
+ all_rows = []
|
|
|
+
|
|
|
+ # 处理 thead 中的行
|
|
|
+ if thead_match:
|
|
|
+ thead_html = thead_match.group(1)
|
|
|
+ tr_matches = re.findall(r'<tr[^>]*>(.*?)</tr>', thead_html, re.DOTALL | re.IGNORECASE)
|
|
|
+ for tr in tr_matches:
|
|
|
+ all_rows.append(tr)
|
|
|
+
|
|
|
+ # 处理 tbody 中的行
|
|
|
+ if tbody_match:
|
|
|
+ tbody_html = tbody_match.group(1)
|
|
|
+ tr_matches = re.findall(r'<tr[^>]*>(.*?)</tr>', tbody_html, re.DOTALL | re.IGNORECASE)
|
|
|
+ for tr in tr_matches:
|
|
|
+ all_rows.append(tr)
|
|
|
+
|
|
|
+ # 如果没有 thead/tbody,直接提取所有 tr
|
|
|
+ if not all_rows:
|
|
|
+ all_rows = re.findall(r'<tr[^>]*>(.*?)</tr>', table_html, re.DOTALL | re.IGNORECASE)
|
|
|
+
|
|
|
+ # 解析所有行
|
|
|
+ parsed_rows = []
|
|
|
+ for tr_html in all_rows:
|
|
|
+ cells = re.findall(r'<(t[dh])[^>]*>(.*?)</\1>', tr_html, re.DOTALL | re.IGNORECASE)
|
|
|
+
|
|
|
+ row_data = []
|
|
|
+ for tag, cell_content in cells:
|
|
|
+ full_cell_match = re.search(rf'<{tag}[^>]*>', tr_html[tr_html.find(cell_content)-50:tr_html.find(cell_content)])
|
|
|
+ cell_start = full_cell_match.group(0) if full_cell_match else f'<{tag}>'
|
|
|
+
|
|
|
+ text = extract_cell_text(cell_content)
|
|
|
+ colspan = parse_colspan(cell_start)
|
|
|
+ row_data.append((text, colspan))
|
|
|
+
|
|
|
+ if row_data:
|
|
|
+ parsed_rows.append(row_data)
|
|
|
+
|
|
|
+ if not parsed_rows:
|
|
|
+ return ""
|
|
|
+
|
|
|
+ # 计算最大列数(考虑 colspan)
|
|
|
+ max_cols = 0
|
|
|
+ for row in parsed_rows:
|
|
|
+ cols = sum(colspan for _, colspan in row)
|
|
|
+ max_cols = max(max_cols, cols)
|
|
|
+
|
|
|
+ # 展开 colspan 并生成 Markdown
|
|
|
+ md_rows = []
|
|
|
+ for row in parsed_rows:
|
|
|
+ expanded_cells = []
|
|
|
+ for text, colspan in row:
|
|
|
+ expanded_cells.append(text)
|
|
|
+ for _ in range(colspan - 1):
|
|
|
+ expanded_cells.append("")
|
|
|
+
|
|
|
+ while len(expanded_cells) < max_cols:
|
|
|
+ expanded_cells.append("")
|
|
|
+
|
|
|
+ md_rows.append("| " + " | ".join(expanded_cells) + " |")
|
|
|
+
|
|
|
+ # 添加分隔行
|
|
|
+ if len(md_rows) > 0:
|
|
|
+ md_rows.insert(1, "| " + " | ".join(["---"] * max_cols) + " |")
|
|
|
+
|
|
|
+ return "\n".join(md_rows)
|
|
|
+
|
|
|
+ return re.sub(r'<table[^>]*>.*?</table>', convert_table_match, content,
|
|
|
+ flags=re.DOTALL | re.IGNORECASE)
|