137 lines
5.4 KiB
Python
137 lines
5.4 KiB
Python
import fitz # PyMuPDF
|
||
from paddleocr import PaddleOCR
|
||
from PIL import Image
|
||
import io
|
||
import numpy as np
|
||
from typing import List, Dict, Any
|
||
import traceback
|
||
import json
|
||
import traceback
|
||
import os
|
||
|
||
class PaddlePdfProcessor:
|
||
"""
|
||
一个基于 PaddleOCR 的健壮PDF处理器。
|
||
(最终版：适配 PyMuPDF, PaddlePaddle 3.x, 和 PaddleOCR 3.x API)
|
||
"""
|
||
|
||
def __init__(self, pdf_path: str):
|
||
self.pdf_path = pdf_path
|
||
self.doc = fitz.open(pdf_path)
|
||
print("正在初始化 PaddleOCR 模型...")
|
||
self.ocr_engine = PaddleOCR(use_textline_orientation=True, lang='ch')
|
||
print("PaddleOCR 初始化完成。")
|
||
|
||
def process_page(self, page_num: int) -> List[Dict[str, Any]]:
|
||
page = self.doc[page_num - 1]
|
||
pix = page.get_pixmap(dpi=200) # type: ignore
|
||
img_bytes = pix.tobytes("png")
|
||
image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
||
img_array = np.array(image)
|
||
|
||
result = self.ocr_engine.predict(input=img_array)
|
||
# result[0] 是一个 OCRResult 对象
|
||
structured_blocks = self._parse_paddle_result(result[0] if result else None, page_num)
|
||
return structured_blocks
|
||
|
||
def _parse_paddle_result(self, ocr_result, page_num: int) -> List[Dict[str, Any]]:
|
||
"""
|
||
解析 PaddleOCR 3.x 返回的 OCRResult 对象。
|
||
"""
|
||
blocks = []
|
||
# 如果 ocr_result 为 None 或为空，直接返回
|
||
if not ocr_result:
|
||
return blocks
|
||
|
||
# 从 OCRResult 对象中，通过键名提取三个平行的列表
|
||
boxes = ocr_result['dt_polys']
|
||
texts = ocr_result['rec_texts']
|
||
scores = ocr_result['rec_scores']
|
||
|
||
# 使用 zip() 函数同时遍历这三个列表
|
||
for box, text, score in zip(boxes, texts, scores):
|
||
# 将多边形坐标(box)转换为简单的矩形边界框(bbox)
|
||
# box 是一个 numpy 数组，形状类似 [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
|
||
x_coords = [p[0] for p in box]
|
||
y_coords = [p[1] for p in box]
|
||
bbox = [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
|
||
|
||
blocks.append({
|
||
"type": "text",
|
||
"content": text,
|
||
"bbox": bbox,
|
||
"page_num": page_num,
|
||
"confidence": score
|
||
})
|
||
|
||
# 按从上到下、从左到右的阅读顺序对文本块进行排序
|
||
return sorted(blocks, key=lambda b: (b['bbox'][1], b['bbox'][0]))
|
||
|
||
def process_all_pages(self) -> List[Dict[str, Any]]:
|
||
all_content_blocks = []
|
||
for i in range(len(self.doc)):
|
||
page_number = i + 1
|
||
print(f"--- 正在处理第 {page_number}/{len(self.doc)} 页 ---")
|
||
blocks = self.process_page(page_number)
|
||
all_content_blocks.extend(blocks)
|
||
return all_content_blocks
|
||
|
||
def close(self):
|
||
self.doc.close()
|
||
|
||
# --- 主程序入口 ---
|
||
# --- 主程序入口 (包含保存和加载逻辑) ---
|
||
if __name__ == '__main__':
|
||
# 1. 定义输入和输出文件路径
|
||
pdf_path = "../document/sample_bidding_document.pdf"
|
||
# 我们将把解析结果保存到这个JSON文件中
|
||
parsed_output_path = "parsed_output.json"
|
||
|
||
all_blocks = []
|
||
|
||
# 2. 检查解析结果是否已存在
|
||
if os.path.exists(parsed_output_path):
|
||
print(f"✅ 找到了已保存的解析结果 '{parsed_output_path}'，直接加载...")
|
||
with open(parsed_output_path, 'r', encoding='utf-8') as f:
|
||
all_blocks = json.load(f)
|
||
else:
|
||
print(f"未找到已保存的结果，启动完整的PDF解析流程...")
|
||
print(f"准备处理文件: {pdf_path}")
|
||
|
||
try:
|
||
# 启动耗时的解析过程
|
||
processor = PaddlePdfProcessor(pdf_path)
|
||
all_blocks = processor.process_all_pages()
|
||
processor.close()
|
||
|
||
# 3. 将新解析的结果保存到文件中，供下次使用
|
||
print(f"\n💾 正在将解析结果保存到 '{parsed_output_path}'...")
|
||
with open(parsed_output_path, 'w', encoding='utf-8') as f:
|
||
# indent=4 让JSON文件格式化，更易读
|
||
# ensure_ascii=False 确保中文字符能被正确写入
|
||
json.dump(all_blocks, f, indent=4, ensure_ascii=False)
|
||
print("保存成功！")
|
||
|
||
except FileNotFoundError:
|
||
print(f"❌ 错误: 文件 '{pdf_path}' 未找到。")
|
||
except Exception as e:
|
||
print(f"❌ 程序在解析过程中发生严重错误: {e}")
|
||
traceback.print_exc()
|
||
|
||
# 4. 无论数据是加载的还是新生成的，都可以在这里进行后续处理
|
||
if all_blocks:
|
||
print("\n" + "="*50)
|
||
print(f"数据加载/处理完成，共获得 {len(all_blocks)} 个文本块。")
|
||
print("="*50)
|
||
|
||
print("\n--- 处理结果预览 (前 5 个文本块) ---")
|
||
for block in all_blocks[:5]:
|
||
# 注意：json会把元组(tuple)存成列表(list)，所以bbox现在是列表
|
||
print(f"P{block['page_num']}: {block['content']} (置信度: {block['confidence']:.4f})")
|
||
|
||
# 在这里调用您后续的分析函数，例如段落合并
|
||
# print("\n--- 开始进行段落合并 ---")
|
||
# paragraphs = assemble_paragraphs(all_blocks)
|
||
# ...
|
||
else:
|
||
print("未能获取任何文本块，程序结束。") |