|
|
@@ -0,0 +1,446 @@
|
|
|
+"""
|
|
|
+extractor.py - PDF 表格提取与 DOCX 生成主模块
|
|
|
+
|
|
|
+负责:
|
|
|
+- 打开 PDF,逐页提取表格
|
|
|
+- 检测跨页表格并自动合并
|
|
|
+- 协调 table_parser 和 docx_writer 完成完整流程
|
|
|
+- 支持单文件输出和目录式输出(JSON + 每表独立 DOCX)
|
|
|
+"""
|
|
|
+
|
|
|
+import json
|
|
|
+import logging
|
|
|
+import os
|
|
|
+import re
|
|
|
+
|
|
|
+import pdfplumber
|
|
|
+
|
|
|
+from .table_parser import TableInfo, tables_on_page
|
|
|
+from .docx_writer import build_document, build_single_table_document
|
|
|
+from .text_scorer import get_scorer
|
|
|
+
|
|
|
+logger = logging.getLogger(__name__)
|
|
|
+
|
|
|
+
|
|
|
+# 跨页表格检测用容差
|
|
|
+WIDTH_MATCH_EPSILON = 0.1 # 列宽比例差异小于此值认为匹配
|
|
|
+PAGE_EDGE_EPSILON = 80 # pt,距页边距小于此值认为"靠近边缘"
|
|
|
+
|
|
|
+
|
|
|
+def _table_first_row_preview(table: TableInfo, max_chars: int = 20) -> list:
|
|
|
+ """提取表格首行前几列作为预览。"""
|
|
|
+ preview = []
|
|
|
+ if table.rows == 0:
|
|
|
+ return preview
|
|
|
+ for c in range(min(table.cols, 5)):
|
|
|
+ cell = table.cells[0][c]
|
|
|
+ if cell:
|
|
|
+ txt = cell.text.strip().replace("\n", " ")[:max_chars]
|
|
|
+ preview.append(txt)
|
|
|
+ else:
|
|
|
+ preview.append("")
|
|
|
+ return preview
|
|
|
+
|
|
|
+
|
|
|
+def _has_continuity_signal(t1: TableInfo, t2: TableInfo) -> bool:
|
|
|
+ """检查跨页表格之间的内容连续性信号。
|
|
|
+
|
|
|
+ 当列宽不完全匹配时,用以下信号判断是否为同表跨页:
|
|
|
+
|
|
|
+ 信号 A——空首列延续:t2 的首行第 0 列为空,但 t1 对应位置有内容
|
|
|
+ (考虑纵向合并,从下往上取 t1 末列第一个非空单元格)。
|
|
|
+
|
|
|
+ 信号 B——NLP 文本延续:使用语言模型评分,判断 t1 与 t2 交界处
|
|
|
+ 的文本是否可以拼接成完整语义单元(如「资源投入」+「和管理」)。
|
|
|
+
|
|
|
+ 信号 C——全表空首列:t2 在首列几乎全部为空,是 t1 分类标签的续表。
|
|
|
+ """
|
|
|
+ if t1.rows == 0 or t2.rows == 0:
|
|
|
+ return False
|
|
|
+
|
|
|
+ signals = []
|
|
|
+
|
|
|
+ # ---- 辅助:获取 t1 末行指定列的第一个非空文本 ----
|
|
|
+ def _last_nonempty(tbl: TableInfo, col: int) -> str:
|
|
|
+ for r in range(tbl.rows - 1, -1, -1):
|
|
|
+ txt = tbl.cells[r][col].text.strip()
|
|
|
+ if txt:
|
|
|
+ return txt
|
|
|
+ return ""
|
|
|
+
|
|
|
+ # ---- 信号 A:空首列延续 ----
|
|
|
+ last_nonempty_t1_col0 = _last_nonempty(t1, 0)
|
|
|
+ first_cell_t2 = t2.cells[0][0].text.strip()
|
|
|
+ if not first_cell_t2 and last_nonempty_t1_col0:
|
|
|
+ signals.append("A")
|
|
|
+
|
|
|
+ # ---- 信号 B:NLP 文本延续(基于语言模型评分) ----
|
|
|
+ try:
|
|
|
+ scorer = get_scorer()
|
|
|
+ for c in range(min(t1.cols, t2.cols)):
|
|
|
+ t1_last = _last_nonempty(t1, c)
|
|
|
+ t2_first = t2.cells[0][c].text.strip()
|
|
|
+ if not t1_last or not t2_first:
|
|
|
+ continue
|
|
|
+ if len(t1_last) < 2 or len(t2_first) < 2:
|
|
|
+ continue
|
|
|
+ if scorer.is_continuation(t1_last, t2_first, threshold=1.5):
|
|
|
+ logger.debug(f"NLP 文本延续信号: '{t1_last[-20:]}' + '{t2_first[:20]}'")
|
|
|
+ signals.append("B")
|
|
|
+ break
|
|
|
+ except FileNotFoundError as e:
|
|
|
+ logger.warning(f"语言模型未加载,跳过 NLP 检测: {e}")
|
|
|
+ except Exception as e:
|
|
|
+ logger.debug(f"NLP 评分异常(不影响主流程): {e}")
|
|
|
+
|
|
|
+ # ---- 信号 C:全表空首列 -> 可能是续表 ----
|
|
|
+ all_first_col_empty = all(
|
|
|
+ not t2.cells[r][0].text.strip()
|
|
|
+ for r in range(min(t2.rows, 5))
|
|
|
+ )
|
|
|
+ if all_first_col_empty and last_nonempty_t1_col0:
|
|
|
+ signals.append("C")
|
|
|
+
|
|
|
+ return len(signals) > 0
|
|
|
+
|
|
|
+
|
|
|
+def _tables_span_pages(t1: TableInfo, t2: TableInfo,
|
|
|
+ page1_height: float = None,
|
|
|
+ page2_height: float = None) -> bool:
|
|
|
+ """判断两个表格是否属于跨页的同一个表格。
|
|
|
+
|
|
|
+ 两阶段判断:
|
|
|
+ 阶段一(严格匹配):列数相同 + 列宽匹配 + 页面位置 + 非重复表头
|
|
|
+ 阶段二(宽松匹配):列数相同 + 页面位置 + 内容连续性信号
|
|
|
+
|
|
|
+ 宽松匹配用于处理 pdfplumber 列宽解析不一致的情况
|
|
|
+ (如跨页后因合并单元格导致列边界不同)。
|
|
|
+ """
|
|
|
+ if t1.cols != t2.cols:
|
|
|
+ return False
|
|
|
+
|
|
|
+ # ---- 公共前提:页面位置靠近 ----
|
|
|
+ if t1.bbox and page1_height:
|
|
|
+ if page1_height - t1.bbox[3] > PAGE_EDGE_EPSILON:
|
|
|
+ return False
|
|
|
+ if t2.bbox and page2_height:
|
|
|
+ if t2.bbox[1] > PAGE_EDGE_EPSILON:
|
|
|
+ return False
|
|
|
+
|
|
|
+ # ---- 公共前提:首行不是重复表头 ----
|
|
|
+ if t1.rows > 0 and t2.rows > 0:
|
|
|
+ if _table_first_row_preview(t1) == _table_first_row_preview(t2):
|
|
|
+ return False
|
|
|
+
|
|
|
+ # ---- 阶段一:列宽严格匹配 ----
|
|
|
+ if len(t1.col_widths) == len(t2.col_widths) and len(t1.col_widths) > 0:
|
|
|
+ widths_match = True
|
|
|
+ for i in range(len(t1.col_widths)):
|
|
|
+ if abs(t1.col_widths[i] - t2.col_widths[i]) > WIDTH_MATCH_EPSILON:
|
|
|
+ widths_match = False
|
|
|
+ break
|
|
|
+ if widths_match:
|
|
|
+ return True
|
|
|
+
|
|
|
+ # ---- 阶段二:内容连续性宽松匹配 ----
|
|
|
+ if _has_continuity_signal(t1, t2):
|
|
|
+ logger.debug(
|
|
|
+ f"列宽不匹配但检测到延续信号,视为跨页:"
|
|
|
+ f"第 {t1.page_num+1} 页 → 第 {t2.page_num+1} 页"
|
|
|
+ )
|
|
|
+ return True
|
|
|
+
|
|
|
+ return False
|
|
|
+
|
|
|
+
|
|
|
+def _merge_tables(t1: TableInfo, t2: TableInfo) -> TableInfo:
|
|
|
+ """合并两个跨页表格,并修复纵向合并的语义断层。"""
|
|
|
+ skip_header = False
|
|
|
+ if t1.rows > 0 and t2.rows > 0:
|
|
|
+ header_same = True
|
|
|
+ for c in range(min(t1.cols, t2.cols)):
|
|
|
+ if t1.cells[0][c].text.strip() != t2.cells[0][c].text.strip():
|
|
|
+ header_same = False
|
|
|
+ break
|
|
|
+ if header_same:
|
|
|
+ skip_header = True
|
|
|
+
|
|
|
+ t2_start = 1 if skip_header else 0
|
|
|
+ merged_rows = t1.rows + (t2.rows - t2_start)
|
|
|
+ merged_cells = []
|
|
|
+ for r in range(t1.rows):
|
|
|
+ merged_cells.append(list(t1.cells[r]))
|
|
|
+ for r in range(t2_start, t2.rows):
|
|
|
+ merged_cells.append(list(t2.cells[r]))
|
|
|
+
|
|
|
+ # ---- 修复纵向合并的语义断层 ----
|
|
|
+ # t1 的纵向合并标签(如"分项服务方案")被分页截断,
|
|
|
+ # t2 对应列首行为空。此时应扩展 t1 的合并覆盖到 t2 区域。
|
|
|
+ for c in range(t1.cols):
|
|
|
+ # t1 部分:从下往上找最后一个非空文本的合并起始单元格
|
|
|
+ t1_origin_r = None
|
|
|
+ t1_span = 0
|
|
|
+ for r in range(t1.rows - 1, -1, -1):
|
|
|
+ cell = merged_cells[r][c]
|
|
|
+ if cell.row_span > 0 and cell.text.strip():
|
|
|
+ t1_origin_r = r
|
|
|
+ t1_span = cell.row_span
|
|
|
+ break
|
|
|
+
|
|
|
+ # t2 部分:第一个合并起始单元格
|
|
|
+ t2_start_r = None
|
|
|
+ t2_span = 0
|
|
|
+ for r in range(t1.rows, merged_rows):
|
|
|
+ cell = merged_cells[r][c]
|
|
|
+ if cell.row_span > 0:
|
|
|
+ t2_start_r = r
|
|
|
+ t2_span = max(cell.row_span, 1)
|
|
|
+ break
|
|
|
+
|
|
|
+ if (t1_origin_r is not None and t2_start_r is not None
|
|
|
+ and t1_span > 0 and t2_span > 0):
|
|
|
+ t1_merge_end = t1_origin_r + t1_span - 1
|
|
|
+ reaches_bottom = (t1_merge_end >= t1.rows - 1)
|
|
|
+ t2_text = merged_cells[t2_start_r][c].text.strip()
|
|
|
+
|
|
|
+ if not t2_text and reaches_bottom:
|
|
|
+ new_end = t2_start_r + t2_span - 1
|
|
|
+ new_span = new_end - t1_origin_r + 1
|
|
|
+ merged_cells[t1_origin_r][c].row_span = new_span
|
|
|
+ for rr in range(t2_start_r, min(t2_start_r + t2_span, merged_rows)):
|
|
|
+ merged_cells[rr][c].row_span = 0
|
|
|
+ merged_cells[rr][c].col_span = 0
|
|
|
+
|
|
|
+ all_pages = list(t1.original_pages or [t1.page_num + 1])
|
|
|
+ all_pages.extend(t2.original_pages or [t2.page_num + 1])
|
|
|
+
|
|
|
+ return TableInfo(
|
|
|
+ rows=merged_rows,
|
|
|
+ cols=t1.cols,
|
|
|
+ cells=merged_cells,
|
|
|
+ col_widths=t1.col_widths,
|
|
|
+ bbox=None,
|
|
|
+ page_num=t1.page_num,
|
|
|
+ cross_page=True,
|
|
|
+ original_pages=sorted(set(all_pages)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def extract_tables_from_pdf(pdf_path: str) -> list:
|
|
|
+ """从 PDF 中提取所有表格,包括跨页表格的合并。
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ list[TableInfo]: 所有合并后的完整表格列表
|
|
|
+ """
|
|
|
+ if not os.path.exists(pdf_path):
|
|
|
+ raise FileNotFoundError(f"PDF 文件不存在: {pdf_path}")
|
|
|
+
|
|
|
+ all_tables = []
|
|
|
+
|
|
|
+ with pdfplumber.open(pdf_path) as pdf:
|
|
|
+ total_pages = len(pdf.pages)
|
|
|
+ logger.info(f"PDF 共 {total_pages} 页,开始提取表格...")
|
|
|
+
|
|
|
+ # ---- 阶段 1:逐页提取表格 ----
|
|
|
+ page_tables = []
|
|
|
+ for page_num, page in enumerate(pdf.pages):
|
|
|
+ tables = tables_on_page(page, page_num)
|
|
|
+ if tables:
|
|
|
+ page_tables.append((page_num, tables))
|
|
|
+ logger.info(f"第 {page_num+1} 页: 发现 {len(tables)} 个表格")
|
|
|
+
|
|
|
+ if not page_tables:
|
|
|
+ logger.info("未检测到任何表格")
|
|
|
+ return []
|
|
|
+
|
|
|
+ # ---- 阶段 2:展平并合并跨页表格 ----
|
|
|
+ flat_tables = []
|
|
|
+ for page_num, tables in page_tables:
|
|
|
+ for table in tables:
|
|
|
+ table.page_num = page_num
|
|
|
+ table.original_pages = [page_num + 1]
|
|
|
+ flat_tables.append(table)
|
|
|
+
|
|
|
+ merged = []
|
|
|
+ i = 0
|
|
|
+ while i < len(flat_tables):
|
|
|
+ current = flat_tables[i]
|
|
|
+
|
|
|
+ if i + 1 < len(flat_tables):
|
|
|
+ next_table = flat_tables[i + 1]
|
|
|
+ if next_table.page_num == current.page_num + 1:
|
|
|
+ # 获取两页的高度用于位置判断
|
|
|
+ p1_h = pdf.pages[current.page_num].height
|
|
|
+ p2_h = pdf.pages[next_table.page_num].height
|
|
|
+ if _tables_span_pages(
|
|
|
+ current, next_table,
|
|
|
+ page1_height=p1_h, page2_height=p2_h,
|
|
|
+ ):
|
|
|
+ logger.info(
|
|
|
+ f"检测到跨页表格:第 {current.page_num+1} 页 → "
|
|
|
+ f"第 {next_table.page_num+1} 页"
|
|
|
+ )
|
|
|
+ merged_table = _merge_tables(current, next_table)
|
|
|
+ merged.append(merged_table)
|
|
|
+ i += 2
|
|
|
+ continue
|
|
|
+
|
|
|
+ merged.append(current)
|
|
|
+ i += 1
|
|
|
+
|
|
|
+ all_tables = merged
|
|
|
+
|
|
|
+ logger.info(f"表格提取完成,共 {len(all_tables)} 个完整表格")
|
|
|
+ for idx, t in enumerate(all_tables):
|
|
|
+ pages_str = (
|
|
|
+ f"(跨页: {t.original_pages})" if t.cross_page else ""
|
|
|
+ )
|
|
|
+ logger.info(f" 表格 {idx+1}: {t.rows} 行 x {t.cols} 列 {pages_str}")
|
|
|
+
|
|
|
+ return all_tables
|
|
|
+
|
|
|
+
|
|
|
+def _sanitize_filename(name: str) -> str:
|
|
|
+ """清理文件名中的非法字符。"""
|
|
|
+ return re.sub(r'[\\/:*?"<>|]', "_", name)
|
|
|
+
|
|
|
+
|
|
|
+def _build_json_metadata(tables: list, pdf_path: str, output_dir: str) -> dict:
|
|
|
+ """构建表格元数据 JSON。"""
|
|
|
+ pdf_name = os.path.basename(pdf_path)
|
|
|
+ metadata = {
|
|
|
+ "pdf_file": pdf_name,
|
|
|
+ "pdf_path": os.path.abspath(pdf_path),
|
|
|
+ "total_tables": len(tables),
|
|
|
+ "tables": [],
|
|
|
+ }
|
|
|
+
|
|
|
+ for idx, t in enumerate(tables):
|
|
|
+ preview = _table_first_row_preview(t)
|
|
|
+ entry = {
|
|
|
+ "id": idx + 1,
|
|
|
+ "file": f"tables/table_{idx+1:03d}.docx",
|
|
|
+ "pages": t.original_pages or [t.page_num + 1],
|
|
|
+ "rows": t.rows,
|
|
|
+ "cols": t.cols,
|
|
|
+ "cross_page": t.cross_page,
|
|
|
+ "preview": preview,
|
|
|
+ }
|
|
|
+ metadata["tables"].append(entry)
|
|
|
+
|
|
|
+ return metadata
|
|
|
+
|
|
|
+
|
|
|
+class PDFTableExtractor:
|
|
|
+ """PDF 表格提取器,封装完整流程。"""
|
|
|
+
|
|
|
+ def __init__(self, pdf_path: str):
|
|
|
+ self.pdf_path = pdf_path
|
|
|
+ self.tables = []
|
|
|
+
|
|
|
+ def extract(self) -> list:
|
|
|
+ """提取 PDF 中的表格。"""
|
|
|
+ self.tables = extract_tables_from_pdf(self.pdf_path)
|
|
|
+ return self.tables
|
|
|
+
|
|
|
+ def to_docx(
|
|
|
+ self,
|
|
|
+ output_path: str = None,
|
|
|
+ page_width_cm: float = 21.0,
|
|
|
+ page_height_cm: float = 29.7,
|
|
|
+ margin_cm: float = 2.0,
|
|
|
+ ) -> str:
|
|
|
+ """提取表格并生成单个 DOCX 文件(所有表格依次排列)。
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ str: DOCX 文件路径
|
|
|
+ """
|
|
|
+ if not self.tables:
|
|
|
+ self.extract()
|
|
|
+ if not self.tables:
|
|
|
+ logger.warning("未提取到任何表格,无法生成 DOCX")
|
|
|
+ return ""
|
|
|
+
|
|
|
+ if not output_path:
|
|
|
+ base = os.path.splitext(self.pdf_path)[0]
|
|
|
+ output_path = f"{base}_表格输出.docx"
|
|
|
+
|
|
|
+ return build_document(
|
|
|
+ self.tables,
|
|
|
+ output_path,
|
|
|
+ page_width_cm=page_width_cm,
|
|
|
+ page_height_cm=page_height_cm,
|
|
|
+ margin_cm=margin_cm,
|
|
|
+ )
|
|
|
+
|
|
|
+ def to_docx_directory(
|
|
|
+ self,
|
|
|
+ output_dir: str = None,
|
|
|
+ page_width_cm: float = 21.0,
|
|
|
+ page_height_cm: float = 29.7,
|
|
|
+ margin_cm: float = 2.0,
|
|
|
+ ) -> str:
|
|
|
+ """提取表格并生成目录式输出:
|
|
|
+ - output_dir/tables.json 全局元数据
|
|
|
+ - output_dir/tables/table_NNN.docx 每个表格独立文件
|
|
|
+
|
|
|
+ Args:
|
|
|
+ output_dir: 输出目录,默认与 PDF 同目录 + "_表格输出"
|
|
|
+ page_width_cm: 页面宽度
|
|
|
+ page_height_cm: 页面高度
|
|
|
+ margin_cm: 页边距
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ str: 输出目录路径
|
|
|
+ """
|
|
|
+ if not self.tables:
|
|
|
+ self.extract()
|
|
|
+ if not self.tables:
|
|
|
+ logger.warning("未提取到任何表格,无法生成输出目录")
|
|
|
+ return ""
|
|
|
+
|
|
|
+ if not output_dir:
|
|
|
+ base = os.path.splitext(self.pdf_path)[0]
|
|
|
+ output_dir = f"{base}_表格输出"
|
|
|
+
|
|
|
+ os.makedirs(output_dir, exist_ok=True)
|
|
|
+
|
|
|
+ # ---- 创建表格子目录 ----
|
|
|
+ tables_dir = os.path.join(output_dir, "tables")
|
|
|
+ os.makedirs(tables_dir, exist_ok=True)
|
|
|
+
|
|
|
+ # ---- 写入每个表格的独立 DOCX ----
|
|
|
+ for idx, table_info in enumerate(self.tables):
|
|
|
+ if table_info.rows == 0 or table_info.cols == 0:
|
|
|
+ continue
|
|
|
+
|
|
|
+ docx_name = f"table_{idx+1:03d}.docx"
|
|
|
+ docx_path = os.path.join(tables_dir, docx_name)
|
|
|
+
|
|
|
+ label = f"表格 {idx+1}"
|
|
|
+ pages = table_info.original_pages or [table_info.page_num + 1]
|
|
|
+ if table_info.cross_page:
|
|
|
+ label += f"(跨页: 第 {pages[0]}-{pages[-1]} 页)"
|
|
|
+ else:
|
|
|
+ label += f"(第 {pages[0]} 页)"
|
|
|
+
|
|
|
+ build_single_table_document(
|
|
|
+ table_info,
|
|
|
+ docx_path,
|
|
|
+ page_width_cm=page_width_cm,
|
|
|
+ page_height_cm=page_height_cm,
|
|
|
+ margin_cm=margin_cm,
|
|
|
+ table_label=label,
|
|
|
+ )
|
|
|
+
|
|
|
+ # ---- 写入 JSON 元数据 ----
|
|
|
+ meta = _build_json_metadata(self.tables, self.pdf_path, output_dir)
|
|
|
+ json_path = os.path.join(output_dir, "tables.json")
|
|
|
+ with open(json_path, "w", encoding="utf-8") as f:
|
|
|
+ json.dump(meta, f, ensure_ascii=False, indent=2)
|
|
|
+
|
|
|
+ logger.info(f"目录输出完成,共 {len(self.tables)} 个表格")
|
|
|
+ logger.info(f" 元数据: {json_path}")
|
|
|
+ logger.info(f" 表格目录: {tables_dir}/")
|
|
|
+
|
|
|
+ return output_dir
|