PP-StructureV3 如何批处理PDF?

[tao-xiaoxin]

1000 多页的PDF 处理,GPU 推理一直卡死,如何批处理指定参数? 代码如下

import uvicorn
import os
import tempfile
import time
import json
import sys
import uuid
import base64
import requests
from io import BytesIO
import argparse
import re

from datetime import datetime
from typing import Optional, List, Union
from pathlib import Path

from loguru import logger
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Request
from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from paddleocr import PPStructureV3, PaddleOCR

base_dir = os.path.dirname(os.path.abspath(__file__))
results_dir = os.path.join(base_dir, "results")
log_dir = os.path.join(base_dir, "logs")

# 创建必要的目录
os.makedirs(log_dir, exist_ok=True)
os.makedirs(results_dir, exist_ok=True)

# 控制台日志 - 彩色输出，包含所有级别的详细信息
logger.add(
    sys.stdout,
    format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <magenta>{process}</magenta>:<magenta>{thread}</magenta> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | <blue>{extra}</blue> - <level>{message}</level>",
    level="DEBUG",
    colorize=True,
    backtrace=True,
    diagnose=True
)

# 详细文件日志 - 包含所有级别和完整上下文信息
logger.add(
    f"{log_dir}/debug_{{time:YYYY-MM-DD}}.log",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | PID:{process} | Thread:{thread} | {name}:{function}:{line} | Extra:{extra} | {message}",
    level="DEBUG",
    rotation="1 day",
    retention="30 days",
    compression="zip",
    backtrace=True,
    diagnose=True
)

# 普通文件日志 - INFO及以上级别
logger.add(
    f"{log_dir}/info_{{time:YYYY-MM-DD}}.log",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} | {message}",
    level="INFO",
    rotation="1 day",
    retention="30 days",
    compression="zip"
)

# 错误日志单独文件 - 包含完整的错误追踪信息
logger.add(
    f"{log_dir}/error_{{time:YYYY-MM-DD}}.log",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | PID:{process} | Thread:{thread} | {name}:{function}:{line} | {message}",
    level="ERROR",
    rotation="1 day",
    retention="30 days",
    compression="zip",
    backtrace=True,
    diagnose=True
)

app = FastAPI(
    title="PaddleOCR V3 API",
    description="基于 PaddleOCR V3 的文档结构化识别服务",
    version="1.0.0",
    docs_url="/api/v1/docs",
    redoc_url="/api/v1/redoc",
    openapi_url="/api/v1/openapi.json"
)

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# 挂载 results 目录为静态资源
app.mount("/results", StaticFiles(directory=results_dir), name="results")


def initialize_pipeline():
    logger.info("🚀 Initializing PaddleOCR...")
    start_load = time.time()
    pipeline = PPStructureV3(
        device="gpu:0",
    )
    end_load = time.time()
    logger.info(f"🚀 PaddleOCR initialized  PPStructureV3 in {end_load - start_load:.2f} seconds")
    logger.info("✅ PaddleOCR initialized successfully!")
    return pipeline, ocr


pipeline, ocr = initialize_pipeline()


@app.get("/")
def root():
    """根路径 - 服务状态检查"""
    logger.info("📊 Root endpoint accessed")
    return {
        "service": "PaddleOCR API",
        "status": "running",
        "version": "1.0.0",
        "timestamp": datetime.now().isoformat(),
        "current_time_utc": datetime.now(),
        "documentation": {
            "swagger_ui": "/api/v1/docs",
            "redoc": "/api/v1/redoc",
            "openapi_json": "/api/v1/openapi.json"
        },
        "endpoints": {
            "layout_parsing": "/api/v1/layout-parsing",
        }
    }


def process_parsing_result(res, _output_dir,use_ocr =False):
    json_dir = os.path.join(_output_dir, "json")
    md_dir = os.path.join(_output_dir, "markdown")
    res.save_to_json(save_path=json_dir)  ## 保存当前图像的结构化json结果
    if not use_ocr:
        res.save_to_markdown(save_path=md_dir)  ## 保存当前图像的markdown结果
        md_info = res.markdown
        md_data = md_info.copy()
        images_list = list(md_data.get("markdown_images").keys())
        md_data.pop("markdown_images")
    else:
        md_data = ""
        images_list = []
        md_info = ""
    logger.info(res.print())
    json_data = res.json.get("res").copy()
    json_data.pop("model_settings")
    json_data.pop("input_path")

    return json_data, md_data, images_list, md_info




def save_base64_to_file(input_str, output_dir, suffix, request_id):
    """
    将base64字符串保存为文件，写入到指定目录output_dir，文件名使用request_id，返回文件路径。
    """
    import os
    import base64
    import re
    os.makedirs(output_dir, exist_ok=True)
    base64_pattern = r"^data:.*?;base64,(.*)"
    match = re.match(base64_pattern, input_str)
    if match:
        file_data = base64.b64decode(match.group(1))
    else:
        try:
            file_data = base64.b64decode(input_str)
        except Exception:
            file_data = None
    if file_data:
        file_name = f"{request_id}.{suffix}"
        file_path = os.path.join(output_dir, file_name)
        with open(file_path, "wb") as f:
            f.write(file_data)
        return file_path
    return None

@app.post("/api/v1/layout-parsing")
async def layout_parsing(request: Request):
    # 先尝试解析 form
    form = await request.form()
    if "input_file" in form and "file_type" in form:
        input_file = form["input_file"]
        file_type = int(form["file_type"])
    else:
        # 再尝试解析 json
        body = await request.json()
        input_file = body.get("input_file")
        file_type = int(body.get("file_type", 2))
    # 后续逻辑不变
    start_time = time.time()
    request_id = uuid.uuid4().hex
    logger.info(f"📝 layout-parsing请求 - RequestID: {request_id}")
    suffix = 'pdf' if file_type == 1 else 'png'
    _output_dir = os.path.join(results_dir, request_id)
    data = list()
    markdown_texts = ""
    temp_file_path = None
    try:
        # 处理 base64 字符串输入
        if isinstance(input_file, str):
            temp_file_path = save_base64_to_file(input_file, _output_dir, suffix, request_id)
            if temp_file_path:
                input_file = temp_file_path
        output = pipeline.predict(input_file)
        if suffix == "image":
            res = output[0]
            json_data, md_data, images_list, _ = process_parsing_result(res, _output_dir)
            markdown_texts = md_data
            detail = {
                "json": json_data,
                "markdown": md_data,
                "images": images_list,
                "page_num": 0,
            }
            data.append(detail)
        else:
            markdown_list = []
            for res in output:
                json_data, md_data, images_list, md_info = process_parsing_result(res, _output_dir)
                page_num = json_data.get("page_index")
                markdown_list.append(md_info)
                detail = {
                    "json": json_data,
                    "markdown": md_data,
                    "images": images_list,
                    "page_num": page_num,
                }
                data.append(detail)
                markdown_texts = pipeline.concatenate_markdown_pages(markdown_list)
                mkd_file_path = os.path.join(_output_dir, "markdown", f"{request_id}.md")
                with open(mkd_file_path, "w", encoding="utf-8") as f:
                    f.write(markdown_texts)
        total_time = time.time() - start_time
        logger.info(f"✅ 版面解析请求完成 - RequestID: {request_id},总耗时: {total_time:.2f}秒")
        return JSONResponse(content={
            "code": 200,
            "message": "解析成功!",
            "request_id": request_id,
            "file_type": file_type,
            "markdown_texts": markdown_texts,
            "data": data
        })
    except Exception as e:
        error_time = time.time() - start_time
        logger.error(f"❌ 版面解析失败 - RequestID: {request_id}, 错误: {str(e)}, 耗时: {error_time:.2f}秒")
        return JSONResponse(
            status_code=500,
            content={
                "code": 500,
                "message": f"版面解析失败: {str(e)}",
                "request_id": request_id,
                "data": None
            }
        )
    finally:
        # 清理临时文件
        if temp_file_path and os.path.exists(temp_file_path):
            os.remove(temp_file_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="启动 PaddleOCR 服务器")
    parser.add_argument("--port", type=int, default=int(os.environ.get("PADDLEOCR_PORT", 2024)), help="服务端口号")
    args = parser.parse_args()

    logger.info(f"🚀 启动 PaddleOCR 服务器... 端口: {args.port}")
    uvicorn.run(app, host="0.0.0.0", port=args.port)