RagFlow_Plugs/services/pdf_parser/pdf_splitter.py

import fitz
from PIL import Image
import io
from typing import List, Dict, Tuple

class PDFSplitter:
    """PDF扫描件按页拆分工具"""
    
    @staticmethod
    def split_pdf(pdf_path: str) -> List[Dict[str, any]]:
        """
        将PDF按页拆分，转换为图像并记录页码，同时保存图片到指定目录
        
        Args:
            pdf_path: PDF文件路径
            
        Returns:
            List[Dict]: 包含每一页信息的列表，每个字典包含:
                - page_number: 页码
                - image: PIL图像对象
                - image_bytes: 图像字节流
                - image_path: 保存的图片路径
        """
        import os
        from pathlib import Path
        
        try:
            # 打开PDF文件
            pdf_document = fitz.open(pdf_path)
            
            # 获取PDF文件名（不含扩展名）
            pdf_filename = os.path.splitext(os.path.basename(pdf_path))[0]
            
            # 输出目录
            output_dir = r"d:\project\work\ragflow_plugs\book\output\temp"
            
            # 确保输出目录存在
            os.makedirs(output_dir, exist_ok=True)
            
            result = []
            for page_num in range(len(pdf_document)):
                # 获取页面
                page = pdf_document[page_num]
                # 页码从1开始
                page_number = page_num + 1
                
                # 将页面转换为图像
                # 使用较高分辨率，DPI=300
                pix = page.get_pixmap(dpi=300)
                
                # 将fitz pixmap转换为PIL图像
                image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                
                # 将图像转换为字节流，便于后续处理
                image_bytes = io.BytesIO()
                image.save(image_bytes, format='PNG')
                image_bytes.seek(0)
                
                # 保存图片到指定目录
                image_filename = f"{pdf_filename}_{page_number}.png"
                image_path = os.path.join(output_dir, image_filename)
                image.save(image_path, format='PNG')
                
                result.append({
                    "page_number": page_number,
                    "image": image,
                    "image_bytes": image_bytes,
                    "image_path": image_path
                })
            
            # 关闭PDF文件
            pdf_document.close()
            
            return result
        except Exception as e:
            raise Exception(f"PDF拆分失败: {str(e)}")
    
    @staticmethod
    def save_page_image(image: Image.Image, output_path: str) -> None:
        """
        保存单页图像
        
        Args:
            image: PIL图像对象
            output_path: 输出文件路径
        """
        image.save(output_path, format='PNG')