LQAgentPlatform/utils_test/MinerU_Test/mineru_test.py

import sys
import os
import json
import base64
import urllib.request
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
from foundation.observability.monitoring.time_statistics import track_execution_time

@track_execution_time
def save_images_from_response(response_data, output_dir="test/mineru_temp"):
    """从API响应中提取并保存图片"""
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 检查是否有图片数据
    if "results" not in response_data or "scan" not in response_data["results"]:
        print("❌ 响应中未找到图片数据")
        return

    scan_result = response_data["results"]["scan"]
    saved_count = 0

    # 保存base64图片
    if "images" in scan_result:
        print(f"发现 {len(scan_result['images'])} 个base64图片")
        for filename, base64_data in scan_result["images"].items():
            try:
                # 解析base64数据格式
                if base64_data.startswith("data:image/"):
                    # 提取实际的base64数据
                    header, base64_string = base64_data.split(",", 1)
                    file_extension = header.split("/")[1].split(";")[0]

                    # 解码base64
                    image_data = base64.b64decode(base64_string)

                    # 保存图片文件
                    output_filename = f"table_{saved_count}.{file_extension}"
                    output_path = os.path.join(output_dir, output_filename)

                    with open(output_path, "wb") as f:
                        f.write(image_data)

                    print(f"✅ 保存图片: {output_filename} ({len(image_data)} bytes)")
                    saved_count += 1

                else:
                    # 直接base64数据
                    image_data = base64.b64decode(base64_data)
                    output_path = os.path.join(output_dir, f"table_{saved_count}.jpg")

                    with open(output_path, "wb") as f:
                        f.write(image_data)

                    print(f"✅ 保存图片: table_{saved_count}.jpg ({len(image_data)} bytes)")
                    saved_count += 1

            except Exception as e:
                print(f"❌ 保存图片 {filename} 失败: {e}")

    # 检查middle_json中的图片信息（如果有额外的）
    if "middle_json" in scan_result:
        try:
            middle_json = json.loads(scan_result["middle_json"])

            # 查找图片路径和HTML信息
            for pdf_info in middle_json.get("pdf_info", []):
                for preproc_block in pdf_info.get("preproc_blocks", []):
                    if preproc_block.get("type") == "table":
                        for block in preproc_block.get("blocks", []):
                            if block.get("type") == "table_body":
                                for line in block.get("lines", []):
                                    for span in line.get("spans", []):
                                        if span.get("type") == "table":
                                            # 保存HTML内容
                                            html_content = span.get('html', 'N/A')
                                            if html_content != 'N/A':
                                                html_file = os.path.join(output_dir, f"table_html_{saved_count}.html")
                                                with open(html_file, "w", encoding="utf-8") as f:
                                                    f.write(html_content)
                                                print(f"✅ 保存表格HTML: table_html_{saved_count}.html")

        except json.JSONDecodeError as e:
            print(f"❌ 解析middle_json失败: {e}")

    print(f"✅ 图片保存完成，共保存 {saved_count} 个图片文件到 {output_dir}")
    return saved_count

@track_execution_time
def save_middle_json(response_data, output_dir="test/mineru_temp"):
    """提取并保存middle_json为单独文件"""
    os.makedirs(output_dir, exist_ok=True)

    if "results" not in response_data or "scan" not in response_data["results"]:
        print("❌ 响应中未找到middle_json数据")
        return False

    scan_result = response_data["results"]["scan"]
    if "middle_json" not in scan_result:
        print("❌ scan结果中未找到middle_json")
        return False

    try:
        # 解析middle_json字符串
        middle_json_data = json.loads(scan_result["middle_json"])

        # 保存格式化的middle_json
        middle_json_file = os.path.join(output_dir, "middle_json_pretty.json")
        with open(middle_json_file, "w", encoding="utf-8") as f:
            json.dump(middle_json_data, f, ensure_ascii=False, indent=4)

        print(f"✅ Middle JSON已保存到 {middle_json_file}")

        # 打印middle_json的基本信息
        if "pdf_info" in middle_json_data:
            pdf_info = middle_json_data["pdf_info"]
            print(f"📄 PDF信息: {len(pdf_info)} 页")

            # 统计各类型块的数量
            total_blocks = 0
            table_count = 0
            text_count = 0
            title_count = 0

            for page_info in pdf_info:
                preproc_blocks = page_info.get("preproc_blocks", [])
                total_blocks += len(preproc_blocks)

                for block in preproc_blocks:
                    block_type = block.get("type", "")
                    if block_type == "table":
                        table_count += 1
                    elif block_type == "text":
                        text_count += 1
                    elif block_type == "title":
                        title_count += 1

            print(f"📊 统计信息:")
            print(f"   总块数: {total_blocks}")
            print(f"   表格块: {table_count}")
            print(f"   文本块: {text_count}")
            print(f"   标题块: {title_count}")

        return True

    except json.JSONDecodeError as e:
        print(f"❌ 解析middle_json失败: {e}")
        return False
    except Exception as e:
        print(f"❌ 保存middle_json失败: {e}")
        return False

def parse_file():
    """使用 API 解析文件"""
    # 配置 API URL
    API_URL = "http://aiclu.small-app.wang:8000/file_parse"

    # 1. 准备文件（本地 PDF/图像）
    files = [
        ("files", ("scan.png", open("test/sgfa_mineru_testimage.png", "rb"), "image/png")),  # 支持多文件
    ]

    # 2. 准备参数
    data = {
        "backend": "pipeline",
        "parse_method": "auto",
        "formula_enable": "true",
        "table_enable": "true",
        "return_md": "true",  # 返回 Markdown 内容
        "start_page_id": 0,
        "end_page_id": None,  # 解析所有页
        "return_middle_json": "true",  # 确保返回middle_json以获取图片信息
        "return_images": "true",
        "return_middle_json": "true",
        "response_format_zip": "true",
    }

    # 3. 发送请求
    response = requests.post(API_URL, files=files, data=data)

    try:
        result = response.json()

        # 打印主要响应内容
        print("=== API响应概要 ===")
        if "results" in result and "scan" in result["results"]:
            scan_result = result["results"]["scan"]
            print(f"MD内容长度: {len(scan_result.get('md_content', ''))}")
            print(f"包含middle_json: {'middle_json' in scan_result}")
            print(f"包含images: {'images' in scan_result}")
            print(f"Backend: {result.get('backend', 'N/A')}")
            print(f"Version: {result.get('version', 'N/A')}")

        # 提取并保存 middle_json
        middle_json_saved = save_middle_json(result)

        # 保存图片信息
        saved_images = save_images_from_response(result)

        # 保存完整响应到JSON文件
        with open("test/mineru_temp/api_response.json", "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print("✅ 完整API响应已保存到 test/mineru_temp/api_response.json")

        # 显示保存结果统计
        print(f"\n=== 保存统计 ===")
        print(f"Middle JSON保存: {'✅ 成功' if middle_json_saved else '❌ 失败'}")
        print(f"图片文件保存: {saved_images} 个")

    except json.JSONDecodeError as e:
        print(f"❌ 解析JSON响应失败: {e}")
        print("原始响应:", response.text)

parse_file()