141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
|
import os
|
|||
|
|
from io import BytesIO
|
|||
|
|
|
|||
|
|
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
|
|||
|
|
from magic_pdf.data.dataset import PymuDocDataset
|
|||
|
|
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
|
|||
|
|
from magic_pdf.config.enums import SupportedPdfParseMethod
|
|||
|
|
|
|||
|
|
def process_pdf_with_magic(file_content, callback=None):
|
|||
|
|
"""
|
|||
|
|
使用magic_pdf处理PDF文件
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
file_content: PDF文件内容
|
|||
|
|
callback: 回调函数，用于更新进度
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
解析后的内容列表
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
from magic_pdf.processor import PDFProcessor
|
|||
|
|
from magic_pdf.extractor import TextExtractor, ImageExtractor
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.1, "初始化Magic PDF解析器")
|
|||
|
|
|
|||
|
|
# 创建临时文件
|
|||
|
|
temp_dir = os.path.join(os.getcwd(), "temp")
|
|||
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
temp_pdf_path = os.path.join(temp_dir, "temp.pdf")
|
|||
|
|
with open(temp_pdf_path, "wb") as f:
|
|||
|
|
f.write(file_content)
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.2, "开始解析PDF")
|
|||
|
|
|
|||
|
|
# 初始化处理器
|
|||
|
|
processor = PDFProcessor(temp_pdf_path)
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.3, "提取文本内容")
|
|||
|
|
|
|||
|
|
# 提取文本
|
|||
|
|
text_extractor = TextExtractor(processor)
|
|||
|
|
text_content = text_extractor.extract()
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.5, "提取图片内容")
|
|||
|
|
|
|||
|
|
# 提取图片
|
|||
|
|
image_extractor = ImageExtractor(processor)
|
|||
|
|
images = image_extractor.extract()
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.7, "组织解析结果")
|
|||
|
|
|
|||
|
|
# 组织结果
|
|||
|
|
content_list = []
|
|||
|
|
|
|||
|
|
# 添加文本内容
|
|||
|
|
for page_num, page_text in enumerate(text_content):
|
|||
|
|
content_list.append({
|
|||
|
|
"type": "text",
|
|||
|
|
"page": page_num + 1,
|
|||
|
|
"text": page_text
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 添加图片内容
|
|||
|
|
for i, img in enumerate(images):
|
|||
|
|
content_list.append({
|
|||
|
|
"type": "image",
|
|||
|
|
"page": img.get("page", i + 1),
|
|||
|
|
"image_path": img.get("path", ""),
|
|||
|
|
"caption": img.get("caption", "")
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 清理临时文件
|
|||
|
|
try:
|
|||
|
|
os.remove(temp_pdf_path)
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(1.0, "PDF解析完成")
|
|||
|
|
|
|||
|
|
return content_list
|
|||
|
|
|
|||
|
|
except ImportError:
|
|||
|
|
# 如果magic_pdf未安装，使用简单的文本提取
|
|||
|
|
if callback:
|
|||
|
|
callback(0.2, "Magic PDF未安装，使用备用方法")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import PyPDF2
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.3, "使用PyPDF2提取文本")
|
|||
|
|
|
|||
|
|
pdf_reader = PyPDF2.PdfReader(BytesIO(file_content))
|
|||
|
|
content_list = []
|
|||
|
|
|
|||
|
|
for i, page in enumerate(pdf_reader.pages):
|
|||
|
|
if callback and i % 5 == 0:
|
|||
|
|
progress = 0.3 + (i / len(pdf_reader.pages)) * 0.6
|
|||
|
|
callback(progress, f"正在处理第 {i+1}/{len(pdf_reader.pages)} 页")
|
|||
|
|
|
|||
|
|
text = page.extract_text()
|
|||
|
|
if text:
|
|||
|
|
content_list.append({
|
|||
|
|
"type": "text",
|
|||
|
|
"page": i + 1,
|
|||
|
|
"text": text
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
if callback:
|
|||
|
|
callback(0.9, "文本提取完成")
|
|||
|
|
|
|||
|
|
return content_list
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
if callback:
|
|||
|
|
callback(0.5, f"PDF解析失败: {str(e)}")
|
|||
|
|
|
|||
|
|
# 最简单的备用方案
|
|||
|
|
return [{
|
|||
|
|
"type": "text",
|
|||
|
|
"page": 1,
|
|||
|
|
"text": "无法解析PDF文件内容"
|
|||
|
|
}]
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
if callback:
|
|||
|
|
callback(0.5, f"PDF解析失败: {str(e)}")
|
|||
|
|
|
|||
|
|
# 出错时返回空列表
|
|||
|
|
return [{
|
|||
|
|
"type": "text",
|
|||
|
|
"page": 1,
|
|||
|
|
"text": f"解析失败: {str(e)}"
|
|||
|
|
}]
|