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")
|