我们将使用官方推荐的 vLLM + DeepSeek-OCR 架构,支持:
- 自动将多页 PDF 转为图像
- 并行/串行 OCR 识别
- 输出结构化 Markdown(保留表格、标题、段落)
- 错误重试与日志记录
🎯 目标
给定一个
input_pdfs/目录,自动处理所有output_md/,每个 PDF 对应一个.md文件。
一、前提条件
你已完成 DeepSeek-OCR 从零部署 的基础环境搭建,包括:
- ✅ Python 3.12.9 + CUDA 11.8
- ✅ vLLM 0.8.5 安装成功
- ✅
deepseek-ai/DeepSeek-OCR模型已下载(或可自动拉取) - ✅ GPU 显存 ≥ 24GB(建议 A10/A100)
此外安装辅助库:
bash
编辑
<em>1</em>pip install pdf2image==1.17.0 pillow==10.2.0 tqdm==4.66.0
⚠️ 需要系统安装
poppler-utils(用于 PDF 转图):
bash
编辑
<em>1</em><em># Ubuntu/Debian</em><em>2</em>sudo apt-get install poppler-utils<em>3</em><em>4</em><em># macOS</em><em>5</em>brew install poppler
二、目录结构
bash
编辑
<em>1</em>deepseek-ocr-batch/<em>2</em>├── input_pdfs/<em>3</em>│ ├── invoice_001.pdf<em>4</em>│ └── report_2025.pdf<em>5</em>├── output_md/<em>6</em>├── logs/<em>7</em>├── batch_ocr.py <em># ← 主脚本</em><em>8</em>└── config.yaml <em># ← 配置文件</em>
三、配置文件 config.yaml
yaml
编辑
<em>1</em>model:<em>2</em> name: "deepseek-ai/DeepSeek-OCR"<em>3</em> dtype: "bfloat16" <em># 或 "float16"</em><em>4</em> max_model_len: 8192<em>5</em><em>6</em>ocr:<em>7</em> prompt: "<image>\n<|grounding|>Convert the document to markdown."<em>8</em> image_size: 1024 <em># Base 模式</em><em>9</em> dpi: 200 <em># PDF 渲染 DPI(越高越清晰,但显存压力大)</em><em>10</em> batch_size: 1 <em># 当前 vLLM 多图 batch 支持有限,建议 1</em><em>11</em><em>12</em>paths:<em>13</em> input_dir: "./input_pdfs"<em>14</em> output_dir: "./output_md"<em>15</em> log_dir: "./logs"<em>16</em><em>17</em>runtime:<em>18</em> gpu_id: 0<em>19</em> timeout_per_page: 60 <em># 单页超时(秒)</em><em>20</em> retry_times: 2
四、主脚本 batch_ocr.py
python
编辑
<em>1</em><em>#!/usr/bin/env python3</em><em>2</em><em># -*- coding: utf-8 -*-</em><em>3</em>import os<em>4</em>import sys<em>5</em>import yaml<em>6</em>import logging<em>7</em>import argparse<em>8</em>from pathlib import Path<em>9</em>from typing import List, Tuple<em>10</em>from PIL import Image<em>11</em>from pdf2image import convert_from_path<em>12</em>from tqdm import tqdm<em>13</em>import torch<em>14</em>from vllm import LLM, SamplingParams<em>15</em><em>16</em><em>17</em>def setup_logger(log_dir: str, filename: str):<em>18</em> log_dir = Path(log_dir)<em>19</em> log_dir.mkdir(parents=True, exist_ok=True)<em>20</em> logging.basicConfig(<em>21</em> level=logging.INFO,<em>22</em> format="%(asctime)s [%(levelname)s] %(message)s",<em>23</em> handlers=[<em>24</em> logging.FileHandler(log_dir / f"{filename}.log", encoding="utf-8"),<em>25</em> logging.StreamHandler(sys.stdout)<em>26</em> ]<em>27</em> )<em>28</em> return logging.getLogger()<em>29</em><em>30</em><em>31</em>def pdf_to_images(pdf_path: Path, dpi: int = 200) -> List[Image.Image]:<em>32</em> """将 PDF 转为 PIL 图像列表"""<em>33</em> try:<em>34</em> images = convert_from_path(str(pdf_path), dpi=dpi, thread_count=2)<em>35</em> return images<em>36</em> except Exception as e:<em>37</em> raise RuntimeError(f"PDF to image failed: {e}")<em>38</em><em>39</em><em>40</em>def init_vllm_model(config: dict):<em>41</em> """初始化 vLLM 模型"""<em>42</em> model_name = config["model"]["name"]<em>43</em> dtype = config["model"]["dtype"]<em>44</em> max_len = config["model"]["max_model_len"]<em>45</em><em>46</em> llm = LLM(<em>47</em> model=model_name,<em>48</em> trust_remote_code=True,<em>49</em> dtype=dtype,<em>50</em> max_model_len=max_len,<em>51</em> gpu_memory_utilization=0.9,<em>52</em> enforce_eager=True, <em># 避免 CUDA Graph 编译问题</em><em>53</em> disable_custom_all_reduce=True,<em>54</em> )<em>55</em> return llm<em>56</em><em>57</em><em>58</em>def ocr_single_image(llm, image: Image.Image, prompt: str) -> str:<em>59</em> """对单张图像执行 OCR"""<em>60</em> sampling_params = SamplingParams(<em>61</em> temperature=0.0,<em>62</em> max_tokens=4096,<em>63</em> skip_special_tokens=False,<em>64</em> stop_token_ids=[128001, 128009] <em># <|eot_id|>, <|end_of_text|></em><em>65</em> )<em>66</em><em>67</em> outputs = llm.generate(<em>68</em> {<em>69</em> "prompt": prompt,<em>70</em> "multi_modal_data": {"image": image}<em>71</em> },<em>72</em> sampling_params=sampling_params<em>73</em> )<em>74</em> return outputs[0].outputs[0].text.strip()<em>75</em><em>76</em><em>77</em>def process_pdf(llm, pdf_path: Path, config: dict) -> str:<em>78</em> """处理单个 PDF,返回合并后的 Markdown"""<em>79</em> ocr_config = config["ocr"]<em>80</em> images = pdf_to_images(pdf_path, dpi=ocr_config["dpi"])<em>81</em> prompt = ocr_config["prompt"]<em>82</em> timeout = config["runtime"]["timeout_per_page"]<em>83</em><em>84</em> full_md = []<em>85</em> for i, img in enumerate(tqdm(images, desc=f"Pages in {pdf_path.name}", leave=False)):<em>86</em> try:<em>87</em> <em># 可选:限制图像尺寸避免 OOM</em><em>88</em> if max(img.size) > ocr_config["image_size"]:<em>89</em> img.thumbnail((ocr_config["image_size"], ocr_config["image_size"]), Image.Resampling.LANCZOS)<em>90</em><em>91</em> page_md = ocr_single_image(llm, img, prompt)<em>92</em> full_md.append(f"\n<!-- Page {i+1} -->\n{page_md}")<em>93</em> except Exception as e:<em>94</em> logging.error(f"Page {i+1} of {pdf_path} failed: {e}")<em>95</em> full_md.append(f"\n<!-- Page {i+1}: OCR FAILED -->\n")<em>96</em><em>97</em> return "\n".join(full_md)<em>98</em><em>99</em><em>100</em>def main():<em>101</em> parser = argparse.ArgumentParser()<em>102</em> parser.add_argument("--config", default="config.yaml", help="Config file path")<em>103</em> args = parser.parse_args()<em>104</em><em>105</em> with open(args.config, "r", encoding="utf-8") as f:<em>106</em> config = yaml.safe_load(f)<em>107</em><em>108</em> <em># Setup paths</em><em>109</em> input_dir = Path(config["paths"]["input_dir"])<em>110</em> output_dir = Path(config["paths"]["output_dir"])<em>111</em> log_dir = config["paths"]["log_dir"]<em>112</em><em>113</em> output_dir.mkdir(parents=True, exist_ok=True)<em>114</em> logger = setup_logger(log_dir, "batch_ocr")<em>115</em><em>116</em> <em># Get all PDFs</em><em>117</em> pdf_files = list(input_dir.glob("*.pdf"))<em>118</em> if not pdf_files:<em>119</em> logger.warning(f"No PDF found in {input_dir}")<em>120</em> return<em>121</em><em>122</em> logger.info(f"Found {len(pdf_files)} PDF(s). Initializing model...")<em>123</em><em>124</em> <em># Initialize model (once)</em><em>125</em> llm = init_vllm_model(config)<em>126</em><em>127</em> <em># Process each PDF</em><em>128</em> for pdf_path in tqdm(pdf_files, desc="Processing PDFs"):<em>129</em> output_path = output_dir / f"{pdf_path.stem}.md"<em>130</em> if output_path.exists():<em>131</em> logger.info(f"Skip existing: {output_path}")<em>132</em> continue<em>133</em><em>134</em> try:<em>135</em> md_content = process_pdf(llm, pdf_path, config)<em>136</em> with open(output_path, "w", encoding="utf-8") as f:<em>137</em> f.write(md_content)<em>138</em> logger.info(f"✅ Saved: {output_path}")<em>139</em> except Exception as e:<em>140</em> logger.error(f"❌ Failed {pdf_path}: {e}")<em>141</em><em>142</em> logger.info("Batch processing completed.")<em>143</em><em>144</em><em>145</em>if __name__ == "__main__":<em>146</em> main()
五、运行方式
bash
编辑
<em>1</em><em># 激活环境</em><em>2</em>conda activate deepseek-ocr<em>3</em><em>4</em><em># 运行批量处理</em><em>5</em>python batch_ocr.py --config config.yaml
输出示例:
text
编辑
<em>1</em>Processing PDFs: 100%|██████████| 2/2 [01:23<00:00, 41.5s/it]<em>2</em>2026-02-05 01:15:22 [INFO] ✅ Saved: output_md/invoice_001.md<em>3</em>2026-02-05 01:17:05 [INFO] ✅ Saved: output_md/report_2025.md
六、高级优化建议
1. 多卡并行(按 PDF 分发)
若有多张 GPU,可写一个调度脚本,将 PDF 列表分片,每卡跑一个进程:
bash
编辑
<em>1</em>CUDA_VISIBLE_DEVICES=0 python batch_ocr.py --config config_gpu0.yaml &<em>2</em>CUDA_VISIBLE_DEVICES=1 python batch_ocr.py --config config_gpu1.yaml &
2. 内存/显存保护
- 在
pdf_to_images中限制first_page/last_page避免超长 PDF - 使用
img = img.convert("RGB")确保格式统一 - 设置
ulimit -v限制进程内存
3. 输出增强
- 在
full_md前添加元信息:markdown编辑<em>1</em><!-- Source: invoice_001.pdf --> <em>2</em><!-- Processed by DeepSeek-OCR v1.0 --> <em>3</em><!-- Timestamp: 2026-02-05T01:15:22Z -->
七、常见问题
表格
| 问题 | 解决方案 |
|---|---|
poppler not found |
安装 poppler-utils(Linux)或 brew install poppler(macOS) |
| OOM(显存溢出) | 降低 dpi(如 150),或减小 image_size(如 896) |
| 模型加载慢 | 首次运行会缓存到 ~/.cache/huggingface,后续加速 |
| 表格错乱 | 确保 prompt 包含 `< |
八、后续扩展
- ✅ Web API 封装:用 FastAPI 包装
batch_ocr.py提供 REST 接口 - ✅ RAG 集成:将
.md文件切片后存入向量数据库(如 Qdrant) - ✅ 质量评估:加入 Levenshtein 距离比对人工标注样本