wehub-resource-sync
275 lines
10 KiB
Python
275 lines
10 KiB
Python
"""
|
|
PP-Structure-V3 文档解析器
|
|
|
|
使用 PP-Structure-V3 进行文档版面解析和内容提取
|
|
"""
|
|
|
|
import base64
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import requests
|
|
|
|
from yuxi.knowledge.parser.base import BaseDocumentProcessor, DocumentParserException
|
|
from yuxi.utils import logger
|
|
|
|
|
|
class PPStructureV3Parser(BaseDocumentProcessor):
|
|
"""PP-Structure-V3 文档解析器 - 使用 PP-Structure-V3 进行版面解析"""
|
|
|
|
def __init__(self, server_url: str | None = None):
|
|
self.server_url = server_url or os.getenv("PADDLEX_URI") or "http://localhost:8080"
|
|
self.base_url = self.server_url.rstrip("/")
|
|
self.endpoint = f"{self.base_url}/layout-parsing"
|
|
|
|
def get_service_name(self) -> str:
|
|
return "pp_structure_v3_ocr"
|
|
|
|
def get_supported_extensions(self) -> list[str]:
|
|
"""PP-Structure-V3 支持 PDF 和多种图像格式"""
|
|
return [".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]
|
|
|
|
def _encode_file_to_base64(self, file_path: str) -> str:
|
|
"""将文件编码为Base64"""
|
|
with open(file_path, "rb") as file:
|
|
encoded = base64.b64encode(file.read()).decode("utf-8")
|
|
return encoded
|
|
|
|
def _process_file_input(self, file_input: str) -> str:
|
|
"""处理文件输入：本地文件路径、URL或Base64内容"""
|
|
# 检查是否为本地文件路径
|
|
if os.path.exists(file_input):
|
|
logger.info(f"📁 检测到本地文件: {file_input}")
|
|
logger.info(f"📏 文件大小: {os.path.getsize(file_input) / 1024 / 1024:.2f} MB")
|
|
return self._encode_file_to_base64(file_input)
|
|
|
|
# 检查是否为URL
|
|
elif file_input.startswith(("http://", "https://")):
|
|
logger.info(f"🌐 检测到URL: {file_input}")
|
|
return file_input
|
|
|
|
# 否则假设为Base64编码内容
|
|
else:
|
|
logger.info(f"📝 假设为Base64编码内容，长度: {len(file_input)} 字符")
|
|
return file_input
|
|
|
|
def _call_layout_api(
|
|
self,
|
|
file_input: str,
|
|
file_type: int | None = None,
|
|
use_table_recognition: bool = True,
|
|
use_formula_recognition: bool = True,
|
|
use_seal_recognition: bool = False,
|
|
**kwargs,
|
|
) -> dict[str, Any]:
|
|
"""调用PP-Structure-V3版面解析API"""
|
|
# 处理文件输入
|
|
processed_file_input = self._process_file_input(file_input)
|
|
payload = {"file": processed_file_input}
|
|
|
|
# 添加核心参数
|
|
optional_params = {
|
|
"fileType": file_type,
|
|
"useTableRecognition": use_table_recognition,
|
|
"useFormulaRecognition": use_formula_recognition,
|
|
"useSealRecognition": use_seal_recognition,
|
|
}
|
|
|
|
# 添加非空参数
|
|
for key, value in optional_params.items():
|
|
if value is not None:
|
|
payload[key] = value
|
|
|
|
# 添加其他kwargs参数
|
|
for key, value in kwargs.items():
|
|
if value is not None:
|
|
payload[key] = value
|
|
|
|
response = requests.post(self.endpoint, json=payload, headers={"Content-Type": "application/json"}, timeout=300)
|
|
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
else:
|
|
error_msg = f"PP-Structure-V3 API请求失败: {response.status_code}"
|
|
try:
|
|
error_result = response.json()
|
|
raise DocumentParserException(f"{error_msg}: {error_result}", self.get_service_name(), "api_error")
|
|
except Exception:
|
|
raise DocumentParserException(f"{error_msg}: {response.text}", self.get_service_name(), "api_error")
|
|
|
|
def _parse_api_result(self, api_result: dict[str, Any], file_path: str) -> dict[str, Any]:
|
|
"""解析API返回结果"""
|
|
# 基本信息
|
|
parsed_result = {
|
|
"success": True,
|
|
"file_path": file_path,
|
|
"file_name": os.path.basename(file_path),
|
|
"log_id": api_result.get("logId"),
|
|
"total_pages": 0,
|
|
"pages": [],
|
|
"full_text": "",
|
|
"summary": {},
|
|
}
|
|
|
|
result_data = api_result.get("result", {})
|
|
layout_results = result_data.get("layoutParsingResults", [])
|
|
|
|
# 数据信息
|
|
parsed_result["total_pages"] = len(layout_results)
|
|
|
|
# 统计信息
|
|
total_tables = 0
|
|
total_formulas = 0
|
|
all_text_content = []
|
|
|
|
# 解析每页结果
|
|
for page_result in layout_results:
|
|
# Markdown内容
|
|
if "markdown" in page_result:
|
|
markdown = page_result["markdown"]
|
|
if markdown.get("text"):
|
|
all_text_content.append(markdown["text"])
|
|
|
|
# 详细识别结果
|
|
if "prunedResult" in page_result:
|
|
pruned = page_result["prunedResult"]
|
|
|
|
# 表格识别
|
|
table_result = pruned.get("table_result", [])
|
|
total_tables += len(table_result)
|
|
|
|
# 公式识别
|
|
formula_result = pruned.get("formula_result", [])
|
|
total_formulas += len(formula_result)
|
|
|
|
# 汇总全文内容
|
|
parsed_result["full_text"] = "\n\n".join(all_text_content)
|
|
|
|
# 汇总统计信息
|
|
parsed_result["summary"] = {
|
|
"total_tables": total_tables,
|
|
"total_formulas": total_formulas,
|
|
"total_characters": len(parsed_result["full_text"]),
|
|
}
|
|
|
|
return parsed_result
|
|
|
|
def check_health(self) -> dict:
|
|
"""检查 PP-Structure-V3 服务健康状态"""
|
|
try:
|
|
response = requests.get(f"{self.base_url}/health", timeout=5)
|
|
|
|
if response.status_code == 200:
|
|
return {
|
|
"status": "healthy",
|
|
"message": "PP-Structure-V3 服务运行正常",
|
|
"details": {"server_url": self.server_url},
|
|
}
|
|
else:
|
|
return {
|
|
"status": "unhealthy",
|
|
"message": f"PP-Structure-V3 服务响应异常: {response.status_code}",
|
|
"details": {"server_url": self.server_url},
|
|
}
|
|
|
|
except requests.exceptions.ConnectionError:
|
|
return {
|
|
"status": "unavailable",
|
|
"message": "PP-Structure-V3 服务无法连接,请检查服务是否启动",
|
|
"details": {"server_url": self.server_url},
|
|
}
|
|
except requests.exceptions.Timeout:
|
|
return {
|
|
"status": "timeout",
|
|
"message": "PP-Structure-V3 服务连接超时",
|
|
"details": {"server_url": self.server_url},
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"status": "error",
|
|
"message": f"PP-Structure-V3 健康检查失败: {str(e)}",
|
|
"details": {"server_url": self.server_url, "error": str(e)},
|
|
}
|
|
|
|
def process_file(self, file_path: str, params: dict | None = None) -> str:
|
|
"""
|
|
使用 PP-Structure-V3 处理文档
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
params: 处理参数
|
|
- use_table_recognition: 启用表格识别 (默认: True)
|
|
- use_formula_recognition: 启用公式识别 (默认: True)
|
|
- use_seal_recognition: 启用印章识别 (默认: False)
|
|
|
|
Returns:
|
|
str: 提取的 Markdown 文本
|
|
"""
|
|
if not os.path.exists(file_path):
|
|
raise DocumentParserException(f"文件不存在: {file_path}", self.get_service_name(), "file_not_found")
|
|
|
|
file_ext = Path(file_path).suffix.lower()
|
|
if not self.supports_file_type(file_ext):
|
|
raise DocumentParserException(
|
|
f"不支持的文件类型: {file_ext}", self.get_service_name(), "unsupported_file_type"
|
|
)
|
|
|
|
# 先检查服务健康状态
|
|
health = self.check_health()
|
|
if health["status"] != "healthy":
|
|
raise DocumentParserException(
|
|
f"PP-Structure-V3 服务不可用: {health['message']}", self.get_service_name(), health["status"]
|
|
)
|
|
|
|
try:
|
|
start_time = time.time()
|
|
params = params or {}
|
|
|
|
# 判断文件类型
|
|
file_type = 0 if file_ext == ".pdf" else 1
|
|
|
|
logger.info(f"PP-Structure-V3 开始处理: {os.path.basename(file_path)}")
|
|
|
|
# 调用API
|
|
api_result = self._call_layout_api(
|
|
file_input=file_path,
|
|
file_type=file_type,
|
|
use_table_recognition=params.get("use_table_recognition", True),
|
|
use_formula_recognition=params.get("use_formula_recognition", True),
|
|
use_seal_recognition=params.get("use_seal_recognition", False),
|
|
)
|
|
|
|
# 检查API调用是否成功
|
|
if api_result.get("errorCode") != 0:
|
|
raise DocumentParserException(
|
|
f"PP-Structure-V3 API错误: {api_result.get('errorMsg', '未知错误')}",
|
|
self.get_service_name(),
|
|
"api_error",
|
|
)
|
|
|
|
# 解析结果
|
|
result = self._parse_api_result(api_result, file_path)
|
|
text = result.get("full_text", "")
|
|
|
|
processing_time = time.time() - start_time
|
|
logger.info(
|
|
f"PP-Structure-V3 处理成功: {os.path.basename(file_path)} - {len(text)} 字符 ({processing_time:.2f}s)"
|
|
)
|
|
|
|
# 记录统计信息
|
|
summary = result.get("summary", {})
|
|
if summary:
|
|
logger.info(f" 统计: {summary.get('total_tables', 0)} 表格, {summary.get('total_formulas', 0)} 公式")
|
|
|
|
return text
|
|
|
|
except DocumentParserException:
|
|
raise
|
|
except Exception as e:
|
|
processing_time = time.time() - start_time
|
|
error_msg = f"PP-Structure-V3 处理失败: {str(e)}"
|
|
logger.error(f"{error_msg} ({processing_time:.2f}s)")
|
|
raise DocumentParserException(error_msg, self.get_service_name(), "processing_failed")
|