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