scripts/pdf_parse_structured.py
10.5 KB| 1 | #!/usr/bin/env python3 |
| 2 | """ |
| 3 | pdf_parse_structured.py — 结构化文档解析工具 |
| 4 | |
| 5 | 利用 docling 进行版面分析, 输出结构化的 Markdown 和 JSON, |
| 6 | 保留文档的标题层级、表格、列表等逻辑结构。 |
| 7 | |
| 8 | 用法： |
| 9 | pip install docling |
| 10 | |
| 11 | # 结构化解析为 Markdown + JSON |
| 12 | python pdf_parse_structured.py input.pdf -o ./output |
| 13 | |
| 14 | # 输出包含: output/input.md + output/input.json + output/input_images/ |
| 15 | """ |
| 16 | |
| 17 | import argparse |
| 18 | import json |
| 19 | import os |
| 20 | import sys |
| 21 | import time |
| 22 | from pathlib import Path |
| 23 | from typing import Any |
| 24 | |
| 25 | |
| 26 | # ── 文档解析引擎 ────────────────────────────────────────────── |
| 27 | |
| 28 | def parse_document(input_path: str, do_ocr: bool = True, |
| 29 | do_table: bool = True, pages: str | None = None) -> dict[str, Any]: |
| 30 | """ |
| 31 | 使用 docling 解析文档, 返回结构化结果。 |
| 32 | |
| 33 | 参数: |
| 34 | input_path: PDF 或图片路径 |
| 35 | do_ocr: 是否启用 OCR (扫描件需要) |
| 36 | do_table: 是否检测表格结构 |
| 37 | pages: 页码范围, 如 "1-5,10" (None = 全部) |
| 38 | """ |
| 39 | from docling.document_converter import DocumentConverter |
| 40 | from docling.datamodel.pipeline_options import PdfPipelineOptions |
| 41 | |
| 42 | pipeline_opts = PdfPipelineOptions() |
| 43 | pipeline_opts.do_ocr = do_ocr |
| 44 | pipeline_opts.do_table_structure = do_table |
| 45 | |
| 46 | # 启用图片生成 (用于可视化) |
| 47 | pipeline_opts.generate_page_images = True |
| 48 | pipeline_opts.generate_table_images = True |
| 49 | |
| 50 | # 高级版面分析 |
| 51 | pipeline_opts.do_picture_classification = True |
| 52 | pipeline_opts.do_formula_enrichment = True |
| 53 | |
| 54 | converter = DocumentConverter() |
| 55 | result = converter.convert(input_path, pipeline_options=pipeline_opts) |
| 56 | |
| 57 | doc = result.document |
| 58 | |
| 59 | # ── 提取元数据 ── |
| 60 | metadata = { |
| 61 | "file": os.path.basename(input_path), |
| 62 | "pages": len(doc.pages), |
| 63 | "has_ocr": do_ocr, |
| 64 | } |
| 65 | |
| 66 | # ── 提取各页内容 ── |
| 67 | pages_data = [] |
| 68 | for page_num, page in doc.pages.items(): |
| 69 | page_info = { |
| 70 | "page": page_num, |
| 71 | "width": page.size.width if page.size else 0, |
| 72 | "height": page.size.height if page.size else 0, |
| 73 | } |
| 74 | pages_data.append(page_info) |
| 75 | |
| 76 | # ── 提取文本内容（保留版面顺序） ── |
| 77 | items = [] |
| 78 | # 遍历 docling 的 item 层级 |
| 79 | for item in doc.iterate_items(): |
| 80 | item_data = { |
| 81 | "type": item.label.value if hasattr(item.label, 'value') else str(item.label), |
| 82 | "text": item.text.strip() if item.text else "", |
| 83 | "bbox": { |
| 84 | "x": item.bbox.x if item.bbox else 0, |
| 85 | "y": item.bbox.y if item.bbox else 0, |
| 86 | "w": item.bbox.width if item.bbox else 0, |
| 87 | "h": item.bbox.height if item.bbox else 0, |
| 88 | } if item.bbox else None, |
| 89 | "page": item.page_num if hasattr(item, 'page_num') else 0, |
| 90 | } |
| 91 | items.append(item_data) |
| 92 | |
| 93 | # ── 提取表格 ── |
| 94 | tables = [] |
| 95 | for table in doc.tables: |
| 96 | table_data = { |
| 97 | "page": table.prov[0].page_no if table.prov else 0, |
| 98 | "bbox": { |
| 99 | "x": table.bbox.x if table.bbox else 0, |
| 100 | "y": table.bbox.y if table.bbox else 0, |
| 101 | "w": table.bbox.width if table.bbox else 0, |
| 102 | "h": table.bbox.height if table.bbox else 0, |
| 103 | } if table.bbox else None, |
| 104 | } |
| 105 | |
| 106 | # 提取表格数据为二维数组 |
| 107 | grid = [] |
| 108 | for row in table.data: |
| 109 | grid.append([cell.text if cell.text else "" for cell in row]) |
| 110 | table_data["grid"] = grid |
| 111 | |
| 112 | # 生成 Markdown 表格 |
| 113 | md_lines = [] |
| 114 | if grid: |
| 115 | # 表头 |
| 116 | header = "| " + " | ".join(grid[0]) + " |" |
| 117 | md_lines.append(header) |
| 118 | # 分隔线 |
| 119 | sep = "| " + " | ".join(["---"] * len(grid[0])) + " |" |
| 120 | md_lines.append(sep) |
| 121 | # 数据行 |
| 122 | for row in grid[1:]: |
| 123 | md_lines.append("| " + " | ".join(row) + " |") |
| 124 | |
| 125 | table_data["markdown"] = "\n".join(md_lines) |
| 126 | tables.append(table_data) |
| 127 | |
| 128 | return { |
| 129 | "metadata": metadata, |
| 130 | "pages": pages_data, |
| 131 | "items": items, |
| 132 | "tables": tables, |
| 133 | "text": doc.text if hasattr(doc, 'text') else "", |
| 134 | } |
| 135 | |
| 136 | |
| 137 | # ── Markdown 生成 ────────────────────────────────────────────── |
| 138 | |
| 139 | def generate_markdown(result: dict[str, Any], include_tables: bool = True) -> str: |
| 140 | """将 docling 解析结果转为带结构的 Markdown""" |
| 141 | lines = [] |
| 142 | |
| 143 | # 文档信息 |
| 144 | meta = result["metadata"] |
| 145 | lines.append(f"# 文档解析结果 — {meta['file']}") |
| 146 | lines.append(f"> 页数: {meta['pages']} | OCR: {'是' if meta['has_ocr'] else '否'}") |
| 147 | lines.append("") |
| 148 | |
| 149 | # 按页组织 |
| 150 | current_page = 0 |
| 151 | page_counter = 0 |
| 152 | |
| 153 | for item in result["items"]: |
| 154 | item_page = item.get("page", 0) |
| 155 | |
| 156 | if item_page != current_page: |
| 157 | if current_page > 0: |
| 158 | lines.append("---\n") |
| 159 | page_counter += 1 |
| 160 | lines.append(f"\n## 第 {page_counter} 页\n") |
| 161 | current_page = item_page |
| 162 | |
| 163 | item_type = item["type"] |
| 164 | text = item["text"] |
| 165 | |
| 166 | if not text: |
| 167 | continue |
| 168 | |
| 169 | # 根据类型输出 |
| 170 | if item_type in ("heading", "title", "section-heading"): |
| 171 | # 检测标题层级 (基于字体大小/位置) |
| 172 | level = 3 # 默认 h3 |
| 173 | lines.append(f"{'#' * level} {text}\n") |
| 174 | elif item_type == "paragraph": |
| 175 | lines.append(f"{text}\n") |
| 176 | elif item_type == "list": |
| 177 | lines.append(f"- {text}") |
| 178 | elif item_type == "caption": |
| 179 | lines.append(f"*{text}*\n") |
| 180 | elif item_type == "page-header": |
| 181 | lines.append(f"*{text}* \n") |
| 182 | elif item_type == "page-footer": |
| 183 | lines.append(f"--- *{text}*\n") |
| 184 | else: |
| 185 | lines.append(f"{text}\n") |
| 186 | |
| 187 | # 表格 |
| 188 | if include_tables and result["tables"]: |
| 189 | lines.append("\n## 表格\n") |
| 190 | for i, table in enumerate(result["tables"]): |
| 191 | if table["markdown"]: |
| 192 | lines.append(f"### 表格 {i+1} (第 {table['page']} 页)\n") |
| 193 | lines.append(table["markdown"]) |
| 194 | lines.append("") |
| 195 | |
| 196 | return "\n".join(lines) |
| 197 | |
| 198 | |
| 199 | # ── 工具函数 ─────────────────────────────────────────────────── |
| 200 | |
| 201 | def save_json(data: dict, output_path: str): |
| 202 | """保存 JSON 结果""" |
| 203 | os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True) |
| 204 | with open(output_path, 'w', encoding='utf-8') as f: |
| 205 | json.dump(data, f, ensure_ascii=False, indent=2) |
| 206 | |
| 207 | |
| 208 | def save_markdown(md: str, output_path: str): |
| 209 | """保存 Markdown 结果""" |
| 210 | os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True) |
| 211 | with open(output_path, 'w', encoding='utf-8') as f: |
| 212 | f.write(md) |
| 213 | |
| 214 | |
| 215 | def format_duration(seconds: float) -> str: |
| 216 | if seconds < 1: |
| 217 | return f"{seconds*1000:.0f}ms" |
| 218 | elif seconds < 60: |
| 219 | return f"{seconds:.1f}s" |
| 220 | else: |
| 221 | m, s = divmod(seconds, 60) |
| 222 | return f"{int(m)}m{int(s)}s" |
| 223 | |
| 224 | |
| 225 | # ── CLI ──────────────────────────────────────────────────────── |
| 226 | |
| 227 | def main(): |
| 228 | parser = argparse.ArgumentParser( |
| 229 | description="结构化文档解析 — 版面分析 + 表格提取 + 结构化输出", |
| 230 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 231 | epilog=""" |
| 232 | 示例: |
| 233 | %(prog)s input.pdf -o ./output # 默认: OCR + 表格 |
| 234 | %(prog)s input.pdf -o ./output --no-ocr # 文本层 PDF (更快) |
| 235 | %(prog)s input.pdf -o ./output --no-table # 跳过表格提取 |
| 236 | %(prog)s input.pdf -o ./output --pages "1-3" # 仅前3页 |
| 237 | %(prog)s input.pdf -o ./output --json-only # 仅输出 JSON |
| 238 | %(prog)s input.pdf -o ./output --md-only # 仅输出 Markdown |
| 239 | |
| 240 | 工作流集成: |
| 241 | # 预处理 → 结构化解析 |
| 242 | python scripts/pdf_preprocess.py input.pdf -o ./preprocessed --all |
| 243 | python scripts/pdf_parse_structured.py ./preprocessed/input_p001.jpg -o ./output |
| 244 | |
| 245 | # 直接解析 (平整扫描件) |
| 246 | python scripts/pdf_parse_structured.py input.pdf -o ./output |
| 247 | """ |
| 248 | ) |
| 249 | parser.add_argument('input', help='输入文件 (PDF / 图片)') |
| 250 | parser.add_argument('-o', '--output', default='./output', |
| 251 | help='输出目录 (默认: ./output)') |
| 252 | |
| 253 | parser.add_argument('--no-ocr', action='store_true', |
| 254 | help='禁用 OCR (适用于文本层 PDF)') |
| 255 | parser.add_argument('--no-table', action='store_true', |
| 256 | help='跳过表格结构检测') |
| 257 | parser.add_argument('--pages', type=str, default=None, |
| 258 | help='页码范围, 如 "1-5,10" (默认: 全部)') |
| 259 | parser.add_argument('--json-only', action='store_true', |
| 260 | help='仅输出 JSON') |
| 261 | parser.add_argument('--md-only', action='store_true', |
| 262 | help='仅输出 Markdown') |
| 263 | |
| 264 | args = parser.parse_args() |
| 265 | |
| 266 | input_path = args.input |
| 267 | if not os.path.exists(input_path): |
| 268 | print(f"❌ 文件不存在: {input_path}") |
| 269 | sys.exit(1) |
| 270 | |
| 271 | output_dir = args.output |
| 272 | basename = os.path.splitext(os.path.basename(input_path))[0] |
| 273 | |
| 274 | print(f"📄 输入: {input_path}") |
| 275 | print(f" OCR: {'禁用' if args.no_ocr else '启用'}") |
| 276 | print(f" 表格: {'跳过' if args.no_table else '检测'}") |
| 277 | |
| 278 | t0 = time.time() |
| 279 | |
| 280 | result = parse_document( |
| 281 | input_path, |
| 282 | do_ocr=not args.no_ocr, |
| 283 | do_table=not args.no_table, |
| 284 | pages=args.pages, |
| 285 | ) |
| 286 | |
| 287 | elapsed = time.time() - t0 |
| 288 | |
| 289 | # 输出统计 |
| 290 | print(f"\n📊 解析统计:") |
| 291 | print(f" 页数: {result['metadata']['pages']}") |
| 292 | print(f" 元素: {len(result['items'])} 个") |
| 293 | print(f" 表格: {len(result['tables'])} 个") |
| 294 | print(f" 耗时: {format_duration(elapsed)}") |
| 295 | |
| 296 | # 保存结果 |
| 297 | os.makedirs(output_dir, exist_ok=True) |
| 298 | |
| 299 | if not args.md_only: |
| 300 | json_path = os.path.join(output_dir, f"{basename}.json") |
| 301 | save_json(result, json_path) |
| 302 | print(f"\n ✅ JSON: {json_path}") |
| 303 | |
| 304 | if not args.json_only: |
| 305 | md = generate_markdown(result, include_tables=not args.no_table) |
| 306 | md_path = os.path.join(output_dir, f"{basename}.md") |
| 307 | save_markdown(md, md_path) |
| 308 | print(f" ✅ Markdown: {md_path}") |
| 309 | |
| 310 | print(f"\n✅ 结构化解析完成!") |
| 311 | |
| 312 | |
| 313 | if __name__ == '__main__': |
| 314 | main() |
文件内容是平台检查时读出的文本，只供查看。平台不会运行包里的任何脚本或程序。
