paddleocr笔记
这段代码是一个基于 PaddleOCR 和 PaddleOCR-VL 模型的 PDF 文档处理脚本，能够提取 PDF 文档中的文本和图片，并根据指定的任务（如 OCR、表格、公式、图表识别）进行进一步处理。处理完成后，可以将结果以结构化的格式（JSON 或 Markdown）输出。下面我将详细讲解这段代码，并解释每个部分的功能。
1. 导入必要的库
import pdfplumber
from paddleocr import PaddleOCR
from PIL import Image
import torch
from transformers import AutoModelForCausalLM, AutoProcessor
import io
import numpy as np
import json
import warnings
pdfplumber：用于提取 PDF 文件中的文本和图像。PaddleOCR：用于执行 OCR（光学字符识别），将图像中的文本转换为机器可读的文本。PIL(Pillow)：用于处理图像数据。torch和transformers：用于加载并运行预训练的 PaddleOCR-VL 模型，处理图像和文本。io：用于处理字节数据流，帮助处理 PDF 中的图片。numpy：用于操作数组，图像数据以 NumPy 数组的形式存储。json：用于将结构化数据转换为 JSON 格式。warnings：用于忽略不必要的警告。
2. 警告过滤
# 忽略可能的 FutureWarnings
warnings.filterwarnings("ignore", category=FutureWarning)
这里是为了避免 Python 执行过程中出现 FutureWarning 警告，保持代码的干净输出。
3. PaddleOCR 初始化
ocr = PaddleOCR(use_textline_orientation=True, lang='en')
PaddleOCR：加载 PaddleOCR 模型。use_textline_orientation=True：替代过时的use_angle_cls参数，用于文本方向识别。lang='en'：指定使用英语模型进行 OCR 识别。
4. 设备和数据类型配置
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
if DEVICE == "cuda":
model_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
else:
model_dtype = torch.float32
- 判断是否有可用的 GPU（
cuda），如果有则使用 GPU 加速，否则使用 CPU。 - 根据设备类型选择合适的数值数据类型（如
float16或float32）。
5. 任务配置
CHOSEN_TASK = "table" # 可选项: 'ocr' | 'table' | 'chart' | 'formula'
PROMPTS = {
"ocr": "OCR:",
"table": "Describe the table structure and content in Markdown format:",
"formula": "Formula Recognition:",
"chart": "Chart Recognition:",
}
CHOSEN_TASK：定义要执行的任务类型，如ocr（基本 OCR）、table（表格识别）、formula（公式识别）等。PROMPTS：根据不同的任务定义不同的提示语，用于指导模型进行相应的识别。
6. 加载 PaddleOCR-VL 模型
try:
print("Loading PaddleOCR-VL model...")
model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=model_dtype
).to(DEVICE).eval()
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading PaddleOCR-VL model: {e}. Advanced tasks will be skipped.")
model = None
processor = None
AutoModelForCausalLM和AutoProcessor：从预训练模型路径加载模型和处理器。trust_remote_code=True：允许加载远程的代码（即模型的实现）。torch_dtype=model_dtype：根据设备类型选择模型的数据类型。- 如果加载模型失败，则跳过高级任务（如表格、图表等任务），只执行基本 OCR 任务。
7. 核心函数：OCR 和 VLLM 处理
infer_with_paddleocr_vl 函数
def infer_with_paddleocr_vl(model, processor, image_pil, prompt, device):
"""使用 PaddleOCR-VL 模型进行高级结构化识别。"""
if model is None or processor is None:
return f"VLLM model not loaded. Task: {prompt}"
try:
inputs = processor(images=image_pil, text=prompt, return_tensors="pt").to(device)
# 确保输入张量与模型 dtype 一致
for k, v in inputs.items():
if v.dtype == torch.float:
inputs[k] = v.to(model_dtype)
# 推理
generated_ids = model.generate(**inputs, max_length=1024, num_beams=5)
response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return response
except Exception as e:
return f"VLLM Inference Error: {e}"
infer_with_paddleocr_vl：该函数用于执行高级的结构化识别任务（如表格、图表、公式等），通过调用 PaddleOCR-VL 模型进行推理。model.generate：生成模型的预测结果。processor.batch_decode：解码模型生成的结果，得到最终的输出文本。
extract_content_with_coords 函数
def extract_content_with_coords(file_path):
"""使用 pdfplumber 提取带有坐标的文本行和图像字节，以块为单位。"""
content_blocks = []
with pdfplumber.open(file_path) as pdf:
for page in pdf.pages:
page_number = page.page_number
# 提取文本行
text_blocks = page.extract_text_lines(x_tolerance=2, y_tolerance=2)
for line in text_blocks:
content_blocks.append({
'type': 'text',
'content': line['text'],
'bbox': (line['x0'], line['top'], line['x1'], line['bottom']),
'page': page_number
})
# 提取图像块
for img in page.images:
bbox = (img['x0'], img['top'], img['x1'], img['bottom'])
try:
cropped_image_pil = page.crop(bbox).to_image(resolution=200).original
img_bytes = io.BytesIO()
cropped_image_pil.save(img_bytes, format='PNG')
content_blocks.append({
'type': 'image',
'content': img_bytes.getvalue(),
'bbox': bbox,
'page': page_number
})
except Exception as e:
print(f"Page {page_number} Image Extraction Error: {e}")
continue
content_blocks.sort(key=lambda x: (x['page'], x['bbox'][1]))
return content_blocks
extract_content_with_coords：此函数使用pdfplumber提取 PDF 中的文本和图像，保留每个元素的坐标信息（bbox）。- 文本块和图像块将被存储在一个列表
content_blocks中。
8. 处理提取的内容
def process_content_blocks(content_blocks, ocr_tool, model, processor, device, chosen_task):
"""处理内容块，对图像执行 OCR 和 VLLM 结构化识别。"""
processed_output = []
for block in content_blocks:
if block['type'] == 'text':
processed_output.append(block)
elif block['type'] == 'image':
image_bytes = block['content']
img_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
img_np = np.array(img_pil)
# OCR 识别
try:
ocr_result_raw = ocr_tool.ocr(img_np, cls=True)
ocr_text = " ".join([line[1][0] for line in ocr_result_raw[0]])
except Exception as e:
ocr_text = f"[OCR Failed: {e}]"
processed_output.append({
'type': 'image_ocr',
'content': ocr_text.strip(),
'bbox': block['bbox'],
'page': block['page']
})
# VLLM 识别
if chosen_task in PROMPTS and chosen_task != 'ocr':
prompt = PROMPTS[chosen_task]
vllm_result = infer_with_paddleocr_vl(model, processor, img_pil, prompt, device)
processed_output.append({
'type': f'{chosen_task}_struct',
'content': vllm_result,
'bbox': block['bbox'],
'page': block['page']
})
return processed_output
- 处理每个内容块，如果是文本，则直接保留；如果是图像，则先进行 OCR 识别，然后根据任务类型进行 VLLM 结构化识别。
9. 输出格式化
def generate_structured_output(processed_blocks, output_format='json'):
"""将处理后的块列表转换为指定的结构化输出格式。"""
if output_format.lower() == 'json':
json_ready_blocks = [
{k: v for k, v in block.items() if k != 'content' or block['type'] != 'image'}
for block in processed_blocks
]
return json.dumps(json_ready_blocks, indent=4, ensure_ascii=False)
elif output_format.lower() == 'markdown':
markdown_output = []
current_page = 0
for block in processed_blocks:
if block['page'] != current_page:
markdown_output.append(f"\n# Page {block['page']}")
current_page = block['page']
if block['type'] == 'text':
markdown_output.append(block['content'])
elif block['type'] == 'image_ocr':
markdown_output.append(f"\n> OCR Result: {block['content']}")
elif '_struct' in block['type']:
markdown_output.append(f"\n\n--- {block['type'].replace('_struct', '').upper()} STRUCTURE START ---")
markdown_output.append(block['content'])
markdown_output.append("--- END ---\n")
return "\n".join(markdown_output)
return f"Unsupported Output Format: {output_format}. Please choose 'json' or 'markdown'."
generate_structured_output：根据用户选择的输出格式（JSON 或 Markdown），将处理后的内容块转换为结构化格式。
10. 主执行流程
def run_pdf_processor(pdf_path, output_format='json', chosen_task='ocr'):
print(f"--- Starting PDF Processing for: {pdf_path} ---")
# 1. 提取内容
content_blocks = extract_content_with_coords(pdf_path)
print(f"Extracted {len(content_blocks)} blocks.")
# 2. 处理内容（OCR 和 VLLM）
processed_blocks = process_content_blocks(
content_blocks,
ocr,
model,
processor,
DEVICE,
chosen_task
)
print(f"Processed {len(processed_blocks)} final blocks (including derived results).")
# 3. 生成结构化输出
final_output = generate_structured_output(processed_blocks, output_format)
return final_output
-
run_pdf_processor：主函数，用于执行整个流程：- 提取 PDF 中的内容。
- 对提取的内容执行 OCR 和 VLLM 识别。
- 生成并输出结构化结果。
11. 运行示例
OUTPUT_FORMAT = 'json'
try:
final_structured_result = run_pdf_processor(pdf_file_path, OUTPUT_FORMAT, CHOSEN_TASK)
# 输出结果
print("\n" * 2)
print("=" * 50)
print(f"FINAL STRUCTURED OUTPUT ({OUTPUT_FORMAT.upper()}):")
print("=" * 50)
print(final_structured_result)
except FileNotFoundError:
print(f"\nERROR: File not found at path: {pdf_file_path}")
print("Please ensure 'Fan_Bi-level_Learning_of_Task-Specific_Decoders_for_Joint_Registration_and_One-Shot_CVPR_2024_paper_dual.pdf' exists in the current directory.")
except Exception as e:
print(f"\nAn unexpected error occurred during processing: {e}")
run_pdf_processor执行任务并打印输出结果。如果文件未找到或处理出现错误，将显示错误信息。
总结：
- OCR 识别：通过 PaddleOCR 识别图像中的文本。
- VLLM 结构化识别：通过 PaddleOCR-VL 模型进行更高级的任务，如表格、公式和图表识别。
- 输出格式化：将处理结果输出为 JSON 或 Markdown 格式，方便进一步使用。
这个脚本为你提供了一个功能强大的 PDF 文档解析和结构化输出工具，适用于需要从文档中提取复杂内容并以结构化格式输出的场景。如果你遇到任何问题或有进一步的疑问，随时告诉我！
浙公网安备 33010602011771号