deepdata/my_ops/mineru_pdf_parser/process.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import sys
import json
import time
import base64
import asyncio
import aiohttp
import re
import requests
import mimetypes
from io import BytesIO
from pathlib import Path
from urllib.parse import urlparse
from typing import Dict, Any
from loguru import logger

try:
    from deepdata.core.base_op import Mapper
except ImportError:
    print("⚠️ 本地调试模式：未找到 'deepdata' 库，正在启用模拟 (Mock) Mapper...")
    class Mapper:
        def __init__(self, *args, **kwargs): pass

# 匹配绝对地址图片链接
ABSOLUTE_IMG_REGEX = re.compile(r'!\[([^\]]*)\]\((https?://[^\s)]+)\)')
# 匹配相对地址图片链接
RELATIVE_IMG_REGEX = re.compile(r'!\[([^\]]*)\]\((images/[^)]+)\)')

class MinerUPdfParser(Mapper):
    """基于远端 MinerU FastAPI 的 PDF 结构化解析算子"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        
        self.mineru_api = kwargs.get("mineruApi", "http://mineru.ailab.ccoop.cc:18585").rstrip("/")
        self.server_url = kwargs.get("server_url", "")
        self.backend = kwargs.get("backend", "pipeline")
        
        lang_str = kwargs.get("lang_list", "ch")
        self.lang_list = [lang.strip() for lang in lang_str.split(",") if lang.strip()]
        
        self.parse_method = kwargs.get("parse_method", "auto")
        self.formula_enable = kwargs.get("formula_enable", "true") == "true"
        self.table_enable = kwargs.get("table_enable", "true") == "true"
        self.return_images = kwargs.get("return_images", "true") == "true"
        
        # 🔑 Base64 嵌入开关
        self.embed_images_base64 = kwargs.get("embed_images_base64", "true") == "true"
        
        # 🆕 下载远程图片开关
        self.download_remote_images = kwargs.get("download_remote_images", "true") == "true"
        
        # 🆕 上传到外部图床配置
        self.upload_images = kwargs.get("upload_images", "false") == "true"
        self.upload_api_url = kwargs.get("upload_api_url", "").strip()
        # 注意：这里保留 rstrip("/")，但在拼接函数中会做更严格的清洗
        self.preview_base_url = kwargs.get("preview_base_url", "").rstrip("/")
        
        try:
            self.timeout = int(kwargs.get("timeout", 900))
        except (ValueError, TypeError):
            self.timeout = 900
            
        self.images_subdir = "images"

    def _is_same_origin(self, url1: str, url2: str) -> bool:
        """判断两个 URL 是否同源（协议 + 域名 + 端口）"""
        if not url1 or not url2:
            return False
        try:
            parsed1 = urlparse(url1)
            parsed2 = urlparse(url2)
            return (parsed1.scheme == parsed2.scheme and 
                    parsed1.hostname == parsed2.hostname and 
                    parsed1.port == parsed2.port)
        except Exception:
            return False

    async def _call_api_async(self, file_path: str) -> dict:
        url = f"{self.mineru_api}/file_parse"
        file_obj = open(file_path, "rb")
        
        try:
            form = aiohttp.FormData()
            form.add_field("files", file_obj, filename=Path(file_path).name, content_type="application/pdf")

            def add_field(key: str, value: Any):
                if value is None: return
                if isinstance(value, list):
                    form.add_field(key, ",".join(map(str, value)))
                elif isinstance(value, bool):
                    form.add_field(key, str(value).lower())
                else:
                    form.add_field(key, str(value))

            add_field("backend", self.backend)
            add_field("lang_list", self.lang_list)
            add_field("parse_method", self.parse_method)
            add_field("formula_enable", self.formula_enable)
            add_field("table_enable", self.table_enable)
            add_field("server_url", self.server_url)
            add_field("return_md", True)
            add_field("return_images", self.return_images)
            add_field("response_format_zip", False)
            add_field("start_page_id", 0)
            add_field("end_page_id", 99999)
            add_field("output_dir", "./output")

            timeout = aiohttp.ClientTimeout(total=self.timeout)
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(url, data=form) as resp:
                    resp_text = await resp.text()
                    if resp.status != 200:
                        raise Exception(f"MinerU API HTTP {resp.status}: {resp_text}")
                    return json.loads(resp_text)
        finally:
            file_obj.close()

    def _upload_image_to_api(self, img_bytes: bytes, filename: str) -> str:
        """上传图片到外部 API 并返回完整预览 URL"""
        if not self.upload_images or not self.upload_api_url:
            return ""
            
        ct = mimetypes.guess_type(filename)[0] or "image/jpeg"
        files = {"file": (filename, BytesIO(img_bytes), ct)}
        
        try:
            resp = requests.post(self.upload_api_url, files=files, timeout=30)
            resp.raise_for_status()
            
            try:
                data = resp.json()
                # 尝试获取 path, url, data 字段，兼容不同返回格式
                raw_path = data.get("path") or data.get("url") or data.get("data")
                path = str(raw_path) if raw_path else ""
            except:
                path = resp.text.strip()
            
            # 🛡️ 核心修复：强力清洗路径，防止双斜杠 //
            # 1. 去除首尾的引号、空格
            clean_path = path.strip('\"\' ')
            # 2. 去除开头的斜杠 (因为我们会手动拼接一个分隔符)
            clean_path = clean_path.lstrip('/')
            
            # 3. 去除 Base URL 末尾的斜杠和空格
            clean_base = self.preview_base_url.strip(' /')
            
            if not clean_base or not clean_path:
                 raise ValueError("Base URL 或 路径 解析为空")
            
            # 4. 标准拼接：base/path
            return f"{clean_base}/{clean_path}"
            
        except Exception as e:
            logger.warning(f"⚠️ 图片上传失败 [{filename}]: {e}")
            return ""

    def _decode_image_data(self, img_data: Any) -> bytes:
        """将 Base64 字符串或原始数据转为 bytes"""
        if isinstance(img_data, bytes):
            return img_data
        if isinstance(img_data, str):
            try:
                if img_data.startswith("data:image"):
                    _, b64_data = img_data.split(",", 1)
                else:
                    b64_data = img_data
                return base64.b64decode(b64_data)
            except Exception:
                return None
        return None

    def _download_image(self, url: str) -> bytes:
        """下载图片"""
        headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
        try:
            response = requests.get(url, headers=headers, timeout=30)
            response.raise_for_status()
            return response.content
        except Exception as e:
            logger.warning(f"⚠️ 下载图片失败 [{url}]: {e}")
            return None

    def _process_all_images(self, md_content: str, images_dict: Dict[str, Any], export_path: str) -> str:
        """统一处理所有类型图片（绝对地址 + 相对地址）"""
        
        # ==================== 1. 处理相对地址图片 ====================
        def replace_relative(match):
            alt = match.group(1)
            img_filename = match.group(2)  # 如 "images/xxx.jpg"
            img_data = images_dict.get(Path(img_filename).name)  # 取文件名部分
            
            if not img_data:
                return match.group(0)
                
            img_bytes = self._decode_image_data(img_data)
            if not img_bytes:
                return match.group(0)
            
            # 🟢 优先上传
            if self.upload_images and self.upload_api_url:
                upload_url = self._upload_image_to_api(img_bytes, Path(img_filename).name)
                if upload_url:
                    return f"![{alt}]({upload_url})"
            
            # 🟡 Base64 内联
            if self.embed_images_base64:
                if isinstance(img_data, str) and not img_data.startswith("data:"):
                    ext = Path(img_filename).suffix.lstrip('.') or 'jpg'
                    img_data = f"data:image/{ext};base64,{img_data}"
                return f"![{alt}]({img_data})"
            
            # 🔴 本地存储
            final_export_path = os.path.abspath(export_path)
            images_dir = os.path.join(final_export_path, self.images_subdir)
            os.makedirs(images_dir, exist_ok=True)
            
            save_path = os.path.join(images_dir, Path(img_filename).name)
            with open(save_path, "wb") as f:
                f.write(img_bytes)
                
            return f"![{alt}]({os.path.join(self.images_subdir, Path(img_filename).name)})"
        
        md_content = RELATIVE_IMG_REGEX.sub(replace_relative, md_content)
        
        # ==================== 2. 处理绝对地址图片 ====================
        def replace_absolute(match):
            alt = match.group(1)
            img_url = match.group(2)
            
            # 判断是否与预览地址同源
            if self.upload_images and self.preview_base_url and self._is_same_origin(img_url, self.preview_base_url):
                return match.group(0) # 同源则保留原链接
            
            # 不同源则下载后上传
            img_bytes = self._download_image(img_url)
            if not img_bytes:
                return match.group(0)  # 下载失败保留原链接
            
            filename = Path(urlparse(img_url).path).name or "remote_image.jpg"
            
            # 🟢 优先上传
            if self.upload_images and self.upload_api_url:
                upload_url = self._upload_image_to_api(img_bytes, filename)
                if upload_url:
                    return f"![{alt}]({upload_url})"
            
            # 🟡 Base64 内联
            if self.embed_images_base64:
                mime_type = mimetypes.guess_type(img_url)[0] or "image/png"
                b64_str = base64.b64encode(img_bytes).decode("utf-8")
                return f"![{alt}](data:{mime_type};base64,{b64_str})"
            
            # 🔴 本地存储
            import hashlib
            mime_type = mimetypes.guess_type(img_url)[0] or "image/png"
            ext = mimetypes.guess_extension(mime_type) or '.jpg'
            new_filename = hashlib.md5(img_url.encode()).hexdigest()[:12] + ext
            
            final_export_path = os.path.abspath(export_path)
            images_dir = os.path.join(final_export_path, self.images_subdir)
            os.makedirs(images_dir, exist_ok=True)
            
            save_path = os.path.join(images_dir, new_filename)
            with open(save_path, "wb") as f:
                f.write(img_bytes)
            
            return f"![{alt}]({os.path.join(self.images_subdir, new_filename)})"
        
        if self.download_remote_images:
            md_content = ABSOLUTE_IMG_REGEX.sub(replace_absolute, md_content)
        
        return md_content

    def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        filename = sample.get("fileName", "")
        filepath = sample.get("filePath", "")
        export_path = sample.get("export_path", "./output")

        # 支持多种图片格式和 PDF
        if not filename.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif", ".pdf")):
            return sample
        
        start_time = time.time()
        try:
            logger.info(f"🚀 开始解析: {filename}")
            response = asyncio.run(self._call_api_async(filepath))

            if "results" not in response:
                raise ValueError("Invalid response structure: missing 'results'")
            
            results = response["results"]
            file_stem = Path(filepath).stem
            task_data = results.get(file_stem)
            if not task_data and results:
                task_data = next(iter(results.values()))

            md_content = task_data.get("md_content", "")
            images_dict = task_data.get("images", {})

            # 统一处理所有图片（相对地址 + 绝对地址）
            md_content = self._process_all_images(md_content, images_dict, export_path)
            logger.info("✅ 所有图片处理完成")

            sample["text"] = md_content
            sample["target_type"] = "md"
            logger.info(f"✅ 解析成功: {filename} | 耗时 {time.time() - start_time:.2f}s | MD 长度 {len(md_content)}")

        except Exception as e:
            logger.exception(f"❌ 解析失败 {filename}: {e}")
            raise

        return sample

# ==========================================================
# 🧪 本地调试与测试入口
# ==========================================================
if __name__ == "__main__":
    print("="*60)
    print("🧪 MinerUPdfParser 本地测试")
    print("="*60)

    current_dir = os.path.dirname(os.path.abspath(__file__))
    test_pdf = os.path.join(current_dir, "test.pdf")
    
    if not os.path.exists(test_pdf):
        print(f"❌ 找不到测试文件: {test_pdf}")
        sys.exit(1)

    output_dir = os.path.join(current_dir, "test_output")
    os.makedirs(output_dir, exist_ok=True)

    mock_sample = {
        "fileName": "test.pdf",
        "filePath": test_pdf,
        "export_path": output_dir,
        "text": "",
        "target_type": ""
    }

    mock_kwargs = {
        "mineruApi": "http://mineru.ailab.ccoop.cc:18585",
        "embed_images_base64": "true",
        "download_remote_images": "true",
        "upload_images": "true",
        "upload_api_url": "http://127.0.0.1:8000/upload/form",
        "preview_base_url": "http://127.0.0.1:8000/file"  # 即使这里没带斜杠，代码也会处理带斜杠的情况
    }

    try:
        parser = MinerUPdfParser(**mock_kwargs)
        result = parser.execute(mock_sample)
        
        # 保存生成的 Markdown
        md_content = result.get("text", "")
        md_file = os.path.join(output_dir, "result.md")
        with open(md_file, "w", encoding="utf-8") as f:
            f.write(md_content)
        
        print("\n" + "="*60)
        print("✅ 测试通过!")
        print(f"📂 输出目录: {result.get('export_path')}")
        print(f"📄 结果文件: {md_file}")
        print("="*60)
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)