From a49ebd4c52d652bb52892029970f7a7495c40f51 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 25 Aug 2026 20:58:17 +0800 Subject: [PATCH] feat: implement native table priority handling with complex entry marking and internal text removal --- mineru/backend/analysis/pdf/tables.py | 407 +++++++++++++++- mineru/backend/analysis/pdf/window.py | 48 +- .../test_native_pdf_table_pipeline.py | 441 +++++++++++++++++- .../test_pdf_analyze_block_content.py | 7 + 4 files changed, 891 insertions(+), 12 deletions(-) diff --git a/mineru/backend/analysis/pdf/tables.py b/mineru/backend/analysis/pdf/tables.py index b3aa6550..c8ffe020 100644 --- a/mineru/backend/analysis/pdf/tables.py +++ b/mineru/backend/analysis/pdf/tables.py @@ -6,22 +6,29 @@ from __future__ import annotations import html import math from collections import Counter +from dataclasses import dataclass from typing import Any, Literal import cv2 import numpy as np from loguru import logger -from mineru.backend.local_model_runtime import HybridLocalModelContext, run_ocr_inference -from mineru.model.model_types import AtomicModelName -from mineru.types import RAW_FORMULA_NUMBER, BBox, BlockType -from mineru.utils.bbox_utils import ( +from ...local_model_runtime import HybridLocalModelContext, run_ocr_inference +from ....model.model_types import AtomicModelName +from ....types import RAW_ALGORITHM, RAW_FORMULA_NUMBER, RAW_PHONETIC, BBox, BlockType +from ....utils.bbox_utils import ( calculate_overlap_area_in_bbox1_area_ratio, normalize_to_int_bbox, ) -from mineru.utils.ocr_utils import mask_formula_regions_for_ocr_det -from mineru.utils.pdf_document import PDFPage, get_lines_from_chars -from mineru.utils.spatial_text import project_ocr_table_text +from ....utils.native_pdf_table import ( + NativeTableInput, + coerce_native_table_rectangles, + coerce_native_table_rules, + recover_native_pdf_table, +) +from ....utils.ocr_utils import mask_formula_regions_for_ocr_det +from ....utils.pdf_document import PDFPage, get_lines_from_chars +from ....utils.spatial_text import project_ocr_table_text from .constants import ( BATCH_RATIO, @@ -47,6 +54,390 @@ from .geometry import ( from .text.native import _is_supported_rotation +_NATIVE_TABLE_ALWAYS_COMPLEX_BLOCK_TYPES = { + BlockType.IMAGE, + BlockType.CHART, + BlockType.CODE, + RAW_ALGORITHM, +} +_NATIVE_TABLE_FORMULA_BLOCK_TYPES = { + BlockType.EQUATION, + RAW_FORMULA_NUMBER, +} +_NATIVE_TABLE_INTERNAL_TEXT_BLOCK_TYPES = { + BlockType.TEXT, + BlockType.DOC_TITLE, + BlockType.PARAGRAPH_TITLE, + BlockType.ASIDE_TEXT, + BlockType.REF_TEXT, + BlockType.LIST, + BlockType.INDEX, + RAW_PHONETIC, +} +_NATIVE_TABLE_FORMULA_LAYOUT_LABELS = {"inline_formula", "display_formula", "formula_number"} +_NATIVE_TABLE_INTERNAL_BLOCK_COVERAGE = 0.9 +_NATIVE_TABLE_OVERLAP_FALLBACK_COVERAGE = 0.8 +_NATIVE_HIGH_SOURCE_ORDER_KEY = "_native_table_source_order" + + +@dataclass(frozen=True) +class _NativeTablePrioritySummary: + """记录单个窗口内原生表格优先解析的稳定统计。""" + + total: int = 0 + accepted: int = 0 + complex_fallbacks: int = 0 + rejected: int = 0 + errors: int = 0 + removed_internal_text: int = 0 + removed_formula_blocks: int = 0 + removed_formula_layout_items: int = 0 + + +def _has_non_empty_table_content(block: dict[str, Any]) -> bool: + """判断表格块是否已经持有可直接输出的非空内容。""" + + content = block.get("content") + return isinstance(content, str) and bool(content.strip()) + + +def _mark_native_table_complex_entries( + table_entries: list[dict[str, Any]], + page_blocks: list[dict[str, Any]], + layout_res: list[dict[str, Any]], + page_size: tuple[int, int], + *, + formula_is_complex: bool, +) -> None: + """按现有归属几何标记含模型语义对象或重叠表格的候选。""" + + for block in page_blocks: + block_type = block.get("type") + if block_type not in _NATIVE_TABLE_ALWAYS_COMPLEX_BLOCK_TYPES and not ( + formula_is_complex and block_type in _NATIVE_TABLE_FORMULA_BLOCK_TYPES + ): + continue + block_bbox = _bbox_to_pixel_bbox(block.get("bbox"), page_size) + if block_bbox is None: + continue + owner = _select_table_owner(block_bbox, table_entries) + if owner is not None: + owner["complex_reasons"].add(str(block.get("type"))) + + if formula_is_complex: + # inline_formula 不会进入 VL-style block，High 必须从原始 layout 结果单独参与门控。 + for layout_item in layout_res: + label = str(layout_item.get("label") or layout_item.get("type") or "") + if label not in _NATIVE_TABLE_FORMULA_LAYOUT_LABELS: + continue + formula_bbox = _bbox_to_pixel_bbox(layout_item.get("bbox"), page_size) + if formula_bbox is None: + continue + owner = _select_table_owner(formula_bbox, table_entries) + if owner is not None: + owner["complex_reasons"].add(label) + + for entry_index, table_entry in enumerate(table_entries): + table_bbox = table_entry["table_bbox"] + for peer_entry in table_entries[entry_index + 1 :]: + peer_bbox = peer_entry["table_bbox"] + overlap_bbox = _table_bbox_intersection(table_bbox, peer_bbox) + if overlap_bbox is None: + continue + overlap_area = float(overlap_bbox[2] - overlap_bbox[0]) * float(overlap_bbox[3] - overlap_bbox[1]) + table_area = float(table_bbox[2] - table_bbox[0]) * float(table_bbox[3] - table_bbox[1]) + peer_area = float(peer_bbox[2] - peer_bbox[0]) * float(peer_bbox[3] - peer_bbox[1]) + smaller_area = min(table_area, peer_area) + if smaller_area <= 0 or overlap_area / smaller_area < _NATIVE_TABLE_OVERLAP_FALLBACK_COVERAGE: + continue + table_entry["complex_reasons"].add("overlapping_table") + peer_entry["complex_reasons"].add("overlapping_table") + + +def _remove_native_table_internal_text_blocks( + page_blocks: list[dict[str, Any]], + accepted_entries: list[dict[str, Any]], + page_size: tuple[int, int], +) -> int: + """删除已被原生 HTML 完整吸收的重复正文块，并保留视觉注释。""" + + if not accepted_entries: + return 0 + + removed_count = 0 + retained_blocks: list[dict[str, Any]] = [] + for block in page_blocks: + if block.get("type") not in _NATIVE_TABLE_INTERNAL_TEXT_BLOCK_TYPES: + retained_blocks.append(block) + continue + block_bbox = _bbox_to_pixel_bbox(block.get("bbox"), page_size) + if block_bbox is None: + retained_blocks.append(block) + continue + owner = _select_table_owner(block_bbox, accepted_entries) + if owner is None or calculate_overlap_area_in_bbox1_area_ratio(block_bbox, owner["table_bbox"]) < ( + _NATIVE_TABLE_INTERNAL_BLOCK_COVERAGE + ): + retained_blocks.append(block) + continue + removed_count += 1 + + page_blocks[:] = retained_blocks + return removed_count + + +def _remove_medium_native_table_formula_items( + page_blocks: list[dict[str, Any]], + layout_res: list[dict[str, Any]], + accepted_entries: list[dict[str, Any]], + page_size: tuple[int, int], +) -> tuple[int, int]: + """原子删除 Medium 原生命中表内的公式 block 与原始 layout 项。""" + + if not accepted_entries: + return 0, 0 + + removed_block_count = 0 + retained_blocks: list[dict[str, Any]] = [] + for block in page_blocks: + if block.get("type") not in _NATIVE_TABLE_FORMULA_BLOCK_TYPES: + retained_blocks.append(block) + continue + block_bbox = _bbox_to_pixel_bbox(block.get("bbox"), page_size) + if block_bbox is None or _select_table_owner(block_bbox, accepted_entries) is None: + retained_blocks.append(block) + continue + removed_block_count += 1 + page_blocks[:] = retained_blocks + + removed_layout_count = 0 + retained_layout_items: list[dict[str, Any]] = [] + for layout_item in layout_res: + label = str(layout_item.get("label") or layout_item.get("type") or "") + if label not in _NATIVE_TABLE_FORMULA_LAYOUT_LABELS: + retained_layout_items.append(layout_item) + continue + formula_bbox = _bbox_to_pixel_bbox(layout_item.get("bbox"), page_size) + if formula_bbox is None or _select_table_owner(formula_bbox, accepted_entries) is None: + retained_layout_items.append(layout_item) + continue + removed_layout_count += 1 + layout_res[:] = retained_layout_items + return removed_block_count, removed_layout_count + + +def _apply_native_txt_table_priority( + model_list: list[list[dict[str, Any]]], + images_layout_res: list[list[dict[str, Any]]], + pdf_pages: list[PDFPage], + images_list: list[dict[str, Any]], + *, + effort: Literal["medium", "high"], +) -> _NativeTablePrioritySummary: + """在 Medium/High TXT 模型表格识别前回填高置信原生 HTML。""" + + total = 0 + accepted = 0 + complex_fallbacks = 0 + rejected = 0 + errors = 0 + removed_internal_text = 0 + removed_formula_blocks = 0 + removed_formula_layout_items = 0 + + for page_idx, (page_blocks, layout_res, pdf_page, image_dict) in enumerate( + zip(model_list, images_layout_res, pdf_pages, images_list, strict=True) + ): + page_size = _normalize_page_size(image_dict["img_pil"]) + table_entries: list[dict[str, Any]] = [] + for block in page_blocks: + if block.get("type") != BlockType.TABLE: + continue + total += 1 + table_bbox = _bbox_to_pixel_bbox(block.get("bbox"), page_size) + if table_bbox is None: + rejected += 1 + continue + table_entries.append( + { + "table_block": block, + "table_bbox": table_bbox, + "complex_reasons": ( + {"rotated_table"} + if _normalize_visual_block_angle(block.get("angle", 0)) != 0 + else set() + ), + } + ) + + if not table_entries: + continue + + _mark_native_table_complex_entries( + table_entries, + page_blocks, + layout_res, + page_size, + formula_is_complex=effort == "high", + ) + for table_entry in table_entries: + if table_entry["complex_reasons"]: + logger.debug( + "Hybrid native table kept model fallback for complex content: " + f"page_idx={page_idx}, bbox={table_entry['table_block'].get('bbox')}, " + f"reasons={sorted(table_entry['complex_reasons'])}" + ) + eligible_entries = [entry for entry in table_entries if not entry["complex_reasons"]] + complex_fallbacks += len(table_entries) - len(eligible_entries) + if not eligible_entries: + continue + + try: + native_chars = tuple(pdf_page.get_chars()) + native_rules = coerce_native_table_rules(pdf_page.get_drawing_lines()) + native_rectangles = coerce_native_table_rectangles(pdf_page.get_path_infos()) + native_page_size = tuple(float(value) for value in pdf_page.size) + render_scale = float(image_dict.get("scale", 1.0) or 1.0) + except Exception as exc: + errors += len(eligible_entries) + logger.warning( + "Hybrid native table page primitives failed and kept model fallback: " + f"page_idx={page_idx}, tables={len(eligible_entries)}, error={exc}" + ) + continue + + accepted_entries: list[dict[str, Any]] = [] + for table_entry in eligible_entries: + table_block = table_entry["table_block"] + table_bbox = _sidecar_bbox_to_page_bbox( + table_block.get("bbox"), + native_page_size, + render_scale, + ) + if table_bbox is None: + rejected += 1 + continue + try: + result = recover_native_pdf_table( + NativeTableInput( + table_bbox=table_bbox, + page_size=native_page_size, + angle=_normalize_visual_block_angle(table_block.get("angle", 0)), + chars=native_chars, + drawing_lines=native_rules, + rectangles=native_rectangles, + ) + ) + except Exception as exc: + errors += 1 + logger.warning( + "Hybrid native table recovery failed and kept model fallback: " + f"page_idx={page_idx}, bbox={table_block.get('bbox')}, error={exc}" + ) + continue + if result is None or not result.html.strip(): + rejected += 1 + logger.debug( + "Hybrid native table rejected and kept model fallback: " + f"page_idx={page_idx}, bbox={table_block.get('bbox')}" + ) + continue + table_block["content"] = result.html + accepted_entries.append(table_entry) + accepted += 1 + logger.debug( + "Hybrid native table accepted: " + f"page_idx={page_idx}, bbox={table_block.get('bbox')}, " + f"source={result.source}, confidence={result.confidence:.3f}" + ) + + if effort == "medium": + page_removed_formula_blocks, page_removed_formula_layout_items = ( + _remove_medium_native_table_formula_items( + page_blocks, + layout_res, + accepted_entries, + page_size, + ) + ) + removed_formula_blocks += page_removed_formula_blocks + removed_formula_layout_items += page_removed_formula_layout_items + + removed_internal_text += _remove_native_table_internal_text_blocks( + page_blocks, + accepted_entries, + page_size, + ) + + return _NativeTablePrioritySummary( + total=total, + accepted=accepted, + complex_fallbacks=complex_fallbacks, + rejected=rejected, + errors=errors, + removed_internal_text=removed_internal_text, + removed_formula_blocks=removed_formula_blocks, + removed_formula_layout_items=removed_formula_layout_items, + ) + + +def _split_native_high_table_blocks( + model_list: list[list[dict[str, Any]]], +) -> tuple[list[list[dict[str, Any]]], list[list[dict[str, Any]]]]: + """从 High VLM 输入逐表移除原生命中项，并写入临时源顺序。""" + + vlm_blocks_list: list[list[dict[str, Any]]] = [] + accepted_tables_list: list[list[dict[str, Any]]] = [] + for page_blocks in model_list: + accepted_tables = [ + block + for block in page_blocks + if block.get("type") == BlockType.TABLE and _has_non_empty_table_content(block) + ] + if not accepted_tables: + vlm_blocks_list.append(page_blocks) + accepted_tables_list.append([]) + continue + + accepted_ids = {id(block) for block in accepted_tables} + vlm_blocks: list[dict[str, Any]] = [] + for source_order, block in enumerate(page_blocks): + block[_NATIVE_HIGH_SOURCE_ORDER_KEY] = source_order + if id(block) not in accepted_ids: + vlm_blocks.append(block) + vlm_blocks_list.append(vlm_blocks) + accepted_tables_list.append(accepted_tables) + return vlm_blocks_list, accepted_tables_list + + +def _restore_native_high_table_blocks( + vlm_results: list[Any], + accepted_tables_list: list[list[dict[str, Any]]], +) -> list[list[Any]]: + """按临时源顺序把原生表格插回 High VLM 后处理结果。""" + + if len(vlm_results) != len(accepted_tables_list): + raise ValueError("Hybrid high VLM result count does not match native table page count") + + restored_pages: list[list[Any]] = [] + for page_results, accepted_tables in zip(vlm_results, accepted_tables_list, strict=True): + restored_blocks = list(page_results) + for table_block in sorted( + accepted_tables, + key=lambda block: int(block[_NATIVE_HIGH_SOURCE_ORDER_KEY]), + ): + source_order = int(table_block[_NATIVE_HIGH_SOURCE_ORDER_KEY]) + insert_at = len(restored_blocks) + for block_idx, block in enumerate(restored_blocks): + block_order = block.get(_NATIVE_HIGH_SOURCE_ORDER_KEY) if isinstance(block, dict) else None + if isinstance(block_order, int) and block_order > source_order: + insert_at = block_idx + break + restored_blocks.insert(insert_at, table_block) + restored_pages.append(restored_blocks) + return restored_pages + + def _apply_table_rotate_labels( table_items: list[dict[str, Any]], rotate_labels: list[str], @@ -243,7 +634,7 @@ def _collect_medium_table_tasks( page_size = (image_w, image_h) table_entries: list[dict[str, Any]] = [] for block in page_model_list: - if block.get("type") != BlockType.TABLE: + if block.get("type") != BlockType.TABLE or _has_non_empty_table_content(block): continue pixel_bbox = _bbox_to_pixel_bbox(block.get("bbox"), page_size) table_bbox = normalize_to_int_bbox(pixel_bbox, image_size=(image_h, image_w)) diff --git a/mineru/backend/analysis/pdf/window.py b/mineru/backend/analysis/pdf/window.py index ccbaee18..013e44c0 100644 --- a/mineru/backend/analysis/pdf/window.py +++ b/mineru/backend/analysis/pdf/window.py @@ -45,8 +45,11 @@ from .ocr import ( ) from .tables import ( _apply_medium_table_recognition, + _apply_native_txt_table_priority, _apply_table_orientations, _fill_flash_ocr_table_contents, + _restore_native_high_table_blocks, + _split_native_high_table_blocks, ) from .visuals import ( _attach_visual_block_images, @@ -198,7 +201,7 @@ def _process_text_and_formulas( interline_enable = effort == "medium" # medium 识别行内和行间公式；high/xhigh 的 txt 路径只识别行内公式。 - if mfr_enable: + if mfr_enable and any(mfd_res): images_formula_list = local_model_context.mfr_model.batch_predict( mfd_res, np_images, @@ -323,16 +326,55 @@ def process_pdf_windows( vl_style_layout_blocks = _build_vl_style_layout_blocks(images_layout_res, images_pil_list) + if parse_mode == "txt" and effort in {"medium", "high"}: + native_table_summary = _apply_native_txt_table_priority( + vl_style_layout_blocks, + images_layout_res, + window_pages, + images_list, + effort=effort, + ) + if native_table_summary.total: + native_table_stats = { + "effort": effort, + "total": native_table_summary.total, + "accepted": native_table_summary.accepted, + "complex_fallbacks": native_table_summary.complex_fallbacks, + "rejected": native_table_summary.rejected, + "errors": native_table_summary.errors, + "removed_internal_text": native_table_summary.removed_internal_text, + "removed_formula_blocks": native_table_summary.removed_formula_blocks, + "removed_formula_layout_items": native_table_summary.removed_formula_layout_items, + } + logger.bind(native_table_priority=native_table_stats).info( + "Hybrid native table priority. " + f"effort={native_table_stats['effort']}, total={native_table_stats['total']}, " + f"accepted={native_table_stats['accepted']}, " + f"complex_fallbacks={native_table_stats['complex_fallbacks']}, " + f"rejected={native_table_stats['rejected']}, " + f"errors={native_table_stats['errors']}, " + f"removed_internal_text={native_table_stats['removed_internal_text']}, " + f"removed_formula_blocks={native_table_stats['removed_formula_blocks']}, " + f"removed_formula_layout_items={native_table_stats['removed_formula_layout_items']}" + ) + if parse_mode == "txt": if effort == "medium": window_model_list = vl_style_layout_blocks elif effort == "high": - window_model_list = vlm_predictor.batch_extract_with_layout( + high_vlm_blocks, accepted_native_tables = _split_native_high_table_blocks( + vl_style_layout_blocks + ) + high_vlm_results = vlm_predictor.batch_extract_with_layout( images=images_pil_list, - blocks_list=vl_style_layout_blocks, + blocks_list=high_vlm_blocks, not_extract_list=NOT_EXTRACT_TYPES, image_analysis=False, ) + window_model_list = _restore_native_high_table_blocks( + high_vlm_results, + accepted_native_tables, + ) elif effort == "xhigh": window_model_list = vlm_predictor.batch_two_step_extract( images=images_pil_list, diff --git a/tests/unittest/test_native_pdf_table_pipeline.py b/tests/unittest/test_native_pdf_table_pipeline.py index 9edcd5f8..f93cb404 100644 --- a/tests/unittest/test_native_pdf_table_pipeline.py +++ b/tests/unittest/test_native_pdf_table_pipeline.py @@ -6,13 +6,28 @@ from __future__ import annotations from types import SimpleNamespace from unittest.mock import MagicMock +import numpy as np import pytest from PIL import Image +from mineru.backend.analysis.pdf import formulas as pdf_formulas +from mineru.backend.analysis.pdf import layout as pdf_layout from mineru.backend.analysis.pdf import tables as pdf_tables +from mineru.backend.analysis.pdf import window as pdf_window from mineru.model.flash.native_pdf import models as flash_models from mineru.model.flash.native_pdf import tables as flash_tables -from mineru.types import BlockType +from mineru.types import RAW_FORMULA_NUMBER, BlockType + + +def _build_native_pdf_page(*, width: float = 100.0, height: float = 100.0) -> MagicMock: + """构造带完整原生表格页面接口的测试替身。""" + + page = MagicMock() + page.size = (width, height) + page.get_chars.return_value = [] + page.get_drawing_lines.return_value = [] + page.get_path_infos.return_value = [] + return page def test_flash_materialization_prefers_native_html_and_keeps_claims( @@ -80,3 +95,427 @@ def test_flash_ocr_projects_table_text( image.close() assert model_list[0][0]["content"] == "OCR TABLE" + + +def test_medium_native_table_priority_accepts_html_and_removes_internal_text_and_formulas( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """验证 Medium 原生命中后原子清理表内文本和公式并保留表外对象。""" + + html = "
A
" + recover = MagicMock(return_value=SimpleNamespace(html=html, source="vector_grid", confidence=1.0)) + monkeypatch.setattr(pdf_tables, "recover_native_pdf_table", recover) + table_block = { + "type": BlockType.TABLE, + "bbox": [0.1, 0.2, 0.9, 0.8], + "angle": 0, + } + internal_text = { + "type": BlockType.TEXT, + "bbox": [0.2, 0.3, 0.8, 0.4], + } + internal_equation = { + "type": BlockType.EQUATION, + "bbox": [0.2, 0.45, 0.6, 0.55], + } + internal_formula_number = { + "type": RAW_FORMULA_NUMBER, + "bbox": [0.7, 0.45, 0.8, 0.55], + } + outside_equation = { + "type": BlockType.EQUATION, + "bbox": [0.01, 0.01, 0.08, 0.08], + } + caption = { + "type": BlockType.TABLE_CAPTION, + "bbox": [0.1, 0.1, 0.9, 0.18], + "content": "caption", + } + footnote = { + "type": BlockType.TABLE_FOOTNOTE, + "bbox": [0.1, 0.82, 0.9, 0.9], + "content": "footnote", + } + layout_res = [ + {"label": "table", "bbox": [10, 40, 90, 160]}, + {"label": "inline_formula", "bbox": [20, 70, 30, 80]}, + {"label": "display_formula", "bbox": [40, 90, 60, 110]}, + {"label": "formula_number", "bbox": [70, 90, 80, 110]}, + {"label": "inline_formula", "bbox": [1, 1, 5, 5]}, + ] + model_list = [[outside_equation, caption, table_block, internal_text, internal_equation, internal_formula_number, footnote]] + page = _build_native_pdf_page(width=100.0, height=200.0) + image = Image.new("RGB", (100, 200), "white") + try: + summary = pdf_tables._apply_native_txt_table_priority( + model_list, + [layout_res], + [page], + [{"img_pil": image, "scale": 2.0}], + effort="medium", + ) + finally: + image.close() + + assert table_block["content"] == html + assert model_list == [[outside_equation, caption, table_block, footnote]] + assert [item["label"] for item in layout_res] == ["table", "inline_formula"] + assert pdf_formulas._build_formula_inputs([layout_res]) == [ + [ + { + "label": "inline_formula", + "bbox": [1, 1, 5, 5], + "score": 0.0, + "latex": "", + } + ] + ] + assert summary == pdf_tables._NativeTablePrioritySummary( + total=1, + accepted=1, + removed_internal_text=1, + removed_formula_blocks=2, + removed_formula_layout_items=3, + ) + table_input = recover.call_args.args[0] + assert table_input.table_bbox == pytest.approx((10.0, 40.0, 90.0, 160.0)) + assert table_input.page_size == (100.0, 200.0) + assert table_input.angle == 0 + page.get_chars.assert_called_once_with() + page.get_drawing_lines.assert_called_once_with() + page.get_path_infos.assert_called_once_with() + + +@pytest.mark.parametrize( + ("extra_blocks", "layout_res"), + [ + ([{"type": BlockType.IMAGE, "bbox": [0.2, 0.2, 0.4, 0.4]}], []), + ([{"type": BlockType.CODE, "bbox": [0.2, 0.2, 0.4, 0.4]}], []), + ], +) +def test_hybrid_native_table_priority_falls_back_for_complex_content( + monkeypatch: pytest.MonkeyPatch, + extra_blocks: list[dict[str, object]], + layout_res: list[dict[str, object]], +) -> None: + """验证 Medium 图片和代码复杂内容继续保留现有模型回落。""" + + recover = MagicMock() + monkeypatch.setattr(pdf_tables, "recover_native_pdf_table", recover) + table_block = { + "type": BlockType.TABLE, + "bbox": [0.1, 0.1, 0.9, 0.9], + "angle": 0, + } + image = Image.new("RGB", (100, 100), "white") + try: + summary = pdf_tables._apply_native_txt_table_priority( + [[table_block, *extra_blocks]], + [[*layout_res]], + [_build_native_pdf_page()], + [{"img_pil": image, "scale": 1.0}], + effort="medium", + ) + finally: + image.close() + + assert "content" not in table_block + assert summary == pdf_tables._NativeTablePrioritySummary( + total=1, + complex_fallbacks=1, + ) + recover.assert_not_called() + + +def test_high_native_table_priority_falls_back_for_formula_content( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """验证 High 将表内公式视为复杂内容并且不调用原生规则。""" + + recover = MagicMock() + monkeypatch.setattr(pdf_tables, "recover_native_pdf_table", recover) + table_block = {"type": BlockType.TABLE, "bbox": [0.1, 0.1, 0.9, 0.9], "angle": 0} + equation_block = {"type": BlockType.EQUATION, "bbox": [0.2, 0.2, 0.4, 0.4]} + layout_res = [{"label": "inline_formula", "bbox": [20, 20, 40, 40]}] + page_blocks = [table_block, equation_block] + image = Image.new("RGB", (100, 100), "white") + try: + summary = pdf_tables._apply_native_txt_table_priority( + [page_blocks], + [layout_res], + [_build_native_pdf_page()], + [{"img_pil": image, "scale": 1.0}], + effort="high", + ) + finally: + image.close() + + assert page_blocks == [table_block, equation_block] + assert layout_res == [{"label": "inline_formula", "bbox": [20, 20, 40, 40]}] + assert summary == pdf_tables._NativeTablePrioritySummary(total=1, complex_fallbacks=1) + vlm_blocks, accepted_tables = pdf_tables._split_native_high_table_blocks([page_blocks]) + assert vlm_blocks == [page_blocks] + assert accepted_tables == [[]] + recover.assert_not_called() + + +def test_hybrid_native_table_priority_falls_back_for_rotated_table( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """验证真实字符源顺序未覆盖的旋转表继续使用现有模型。""" + + recover = MagicMock() + monkeypatch.setattr(pdf_tables, "recover_native_pdf_table", recover) + table_block = { + "type": BlockType.TABLE, + "bbox": [0.1, 0.1, 0.9, 0.9], + "angle": 270, + } + image = Image.new("RGB", (100, 100), "white") + try: + summary = pdf_tables._apply_native_txt_table_priority( + [[table_block]], + [[]], + [_build_native_pdf_page()], + [{"img_pil": image, "scale": 1.0}], + effort="medium", + ) + finally: + image.close() + + assert "content" not in table_block + assert summary == pdf_tables._NativeTablePrioritySummary( + total=1, + complex_fallbacks=1, + ) + recover.assert_not_called() + + +def test_hybrid_native_table_priority_keeps_none_and_exception_fallbacks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """验证规则拒绝或异常时逐表保留现有模型路径。""" + + recover = MagicMock(side_effect=[None, RuntimeError("broken")]) + monkeypatch.setattr(pdf_tables, "recover_native_pdf_table", recover) + tables = [ + {"type": BlockType.TABLE, "bbox": [0.05, 0.1, 0.45, 0.9], "angle": 0}, + {"type": BlockType.TABLE, "bbox": [0.55, 0.1, 0.95, 0.9], "angle": 0}, + ] + equation_block = {"type": BlockType.EQUATION, "bbox": [0.1, 0.2, 0.3, 0.4]} + page_blocks = [tables[0], equation_block, tables[1]] + layout_res = [{"label": "inline_formula", "bbox": [10, 20, 30, 40]}] + page = _build_native_pdf_page() + image = Image.new("RGB", (100, 100), "white") + try: + summary = pdf_tables._apply_native_txt_table_priority( + [page_blocks], + [layout_res], + [page], + [{"img_pil": image, "scale": 1.0}], + effort="medium", + ) + finally: + image.close() + + assert all("content" not in block for block in tables) + assert page_blocks == [tables[0], equation_block, tables[1]] + assert layout_res == [{"label": "inline_formula", "bbox": [10, 20, 30, 40]}] + assert summary == pdf_tables._NativeTablePrioritySummary( + total=2, + rejected=1, + errors=1, + ) + assert recover.call_count == 2 + page.get_chars.assert_called_once_with() + + +def test_medium_table_tasks_skip_native_html_and_keep_model_fallback() -> None: + """验证 Medium 只为未命中的表格构造 OCR 和结构模型任务。""" + + native_table = { + "type": BlockType.TABLE, + "bbox": [0.0, 0.0, 0.4, 1.0], + "angle": 0, + "content": "
native
", + } + fallback_table = { + "type": BlockType.TABLE, + "bbox": [0.6, 0.0, 1.0, 1.0], + "angle": 0, + } + + tasks = pdf_tables._collect_medium_table_tasks( + [[native_table, fallback_table]], + [[]], + [np.zeros((40, 100, 3), dtype=np.uint8)], + ) + + assert len(tasks) == 1 + assert tasks[0]["table_block"] is fallback_table + assert native_table["content"].startswith("") + + +def test_medium_formula_processing_skips_mfr_after_native_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """验证原生命中清空公式输入后不再加载或调用 Medium MFR。""" + + native_table = { + "type": BlockType.TABLE, + "bbox": [0.1, 0.1, 0.9, 0.9], + "content": "
native formula text
", + } + model_list = [[native_table]] + image = Image.new("RGB", (100, 100), "white") + local_model_context = MagicMock() + local_model_context.mfr_model = MagicMock() + medium_table_recognition = MagicMock() + monkeypatch.setattr(pdf_window, "_apply_medium_table_recognition", medium_table_recognition) + monkeypatch.setattr(pdf_window, "_apply_medium_display_formula_results", MagicMock()) + monkeypatch.setattr(pdf_window, "_apply_medium_formula_number_ocr", MagicMock()) + monkeypatch.setattr(pdf_window, "_ocr_det", MagicMock(return_value=[[]])) + monkeypatch.setattr( + pdf_window, + "_fill_window_block_content_and_lines", + MagicMock(return_value=model_list), + ) + + try: + result = pdf_window._process_text_and_formulas( + [{"img_pil": image, "scale": 1.0}], + [_build_native_pdf_page()], + model_list, + "txt", + "medium", + local_model_context, + [[]], + ) + finally: + image.close() + + assert result is model_list + local_model_context.mfr_model.batch_predict.assert_not_called() + assert medium_table_recognition.call_args.args[2] == [[]] + + +def test_high_native_tables_restore_source_order_and_drop_private_marker() -> None: + """验证 High 原生表格按源顺序回插且临时字段不进入 ModelJson。""" + + blocks = [ + {"type": BlockType.TEXT, "bbox": [0.0, 0.0, 1.0, 0.1]}, + { + "type": BlockType.TABLE, + "bbox": [0.0, 0.1, 1.0, 0.4], + "content": "
native
", + }, + {"type": BlockType.TABLE_CAPTION, "bbox": [0.0, 0.4, 1.0, 0.5]}, + {"type": BlockType.TABLE, "bbox": [0.0, 0.5, 1.0, 0.9]}, + ] + + vlm_blocks, accepted_tables = pdf_tables._split_native_high_table_blocks([blocks]) + assert [block["type"] for block in vlm_blocks[0]] == [ + BlockType.TEXT, + BlockType.TABLE_CAPTION, + BlockType.TABLE, + ] + + vlm_blocks[0][-1]["content"] = "
fallback
" + restored = pdf_tables._restore_native_high_table_blocks(vlm_blocks, accepted_tables) + converted = pdf_layout._convert_vlm_results_to_model_list(restored) + + assert [block["type"] for block in converted[0]] == [ + BlockType.TEXT, + BlockType.TABLE, + BlockType.TABLE_CAPTION, + BlockType.TABLE, + ] + assert converted[0][1]["content"].endswith("native") + assert converted[0][3]["content"].endswith("fallback") + assert all(pdf_tables._NATIVE_HIGH_SOURCE_ORDER_KEY not in block for block in converted[0]) + + +def test_high_txt_window_excludes_native_table_from_vlm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """验证 High TXT 只把规则回落表送入 VLM，并原位保留命中表。""" + + native_html = "
native
" + fallback_html = "
fallback
" + layout_results = [ + [ + {"label": "table", "bbox": [0, 0, 40, 100], "angle": 0}, + {"label": "table", "bbox": [60, 0, 100, 100], "angle": 0}, + ] + ] + page_image = Image.new("RGB", (100, 100), "white") + fake_document = MagicMock(page_count=1) + fake_document.__getitem__.return_value = _build_native_pdf_page() + hybrid_model = MagicMock() + hybrid_model.layout_model.batch_predict.return_value = layout_results + vlm_predictor = MagicMock() + + def fake_native_priority( + model_list: list[list[dict[str, object]]], + _images_layout_res: object, + _pdf_pages: object, + _images_list: object, + *, + effort: object, + ) -> object: + """只命中首个表格，构造同页混合短路场景。""" + + assert effort == "high" + model_list[0][0]["content"] = native_html + return pdf_tables._NativeTablePrioritySummary(total=2, accepted=1, rejected=1) + + def fake_high_extract(*, blocks_list: list[list[dict[str, object]]], **_kwargs: object) -> object: + """校验 VLM 只收到回落表，并模拟现有模型 HTML。""" + + assert len(blocks_list[0]) == 1 + assert blocks_list[0][0]["bbox"] == [0.6, 0.0, 1.0, 1.0] + blocks_list[0][0]["content"] = fallback_html + return blocks_list + + def keep_window_model_list( + _images_list: object, + _pdf_pages: object, + model_list: list[list[dict[str, object]]], + _parse_mode: object, + _effort: object, + _local_model_context: object, + _images_layout_res: object, + ) -> list[list[dict[str, object]]]: + """跳过与本测试无关的正文和公式回填。""" + + return model_list + + vlm_predictor.batch_extract_with_layout.side_effect = fake_high_extract + monkeypatch.setattr(pdf_window, "get_processing_window_size", lambda default: 1) + monkeypatch.setattr( + pdf_window, + "load_images_from_pdf_bytes_range", + MagicMock(return_value=[{"img_pil": page_image, "scale": 1.0}]), + ) + monkeypatch.setattr(pdf_window, "_apply_table_orientations", MagicMock()) + monkeypatch.setattr(pdf_window, "_apply_native_txt_table_priority", fake_native_priority) + monkeypatch.setattr(pdf_window, "_process_text_and_formulas", keep_window_model_list) + monkeypatch.setattr(pdf_window, "_apply_seal_ocr", MagicMock()) + monkeypatch.setattr(pdf_window, "_attach_visual_block_images", MagicMock()) + + result = pdf_window.process_pdf_windows( + b"fake-pdf", + fake_document, + effort="high", + parse_mode="txt", + image_analysis=False, + flash_txt_mode=False, + hybrid_model=hybrid_model, + vlm_predictor=vlm_predictor, + ) + + assert [block["content"] for block in result[0]] == [native_html, fallback_html] + assert all(pdf_tables._NATIVE_HIGH_SOURCE_ORDER_KEY not in block for block in result[0]) + vlm_predictor.batch_extract_with_layout.assert_called_once() + with pytest.raises(ValueError, match="closed image"): + page_image.getpixel((0, 0)) diff --git a/tests/unittest/test_pdf_analyze_block_content.py b/tests/unittest/test_pdf_analyze_block_content.py index 5ae182d0..a75315f4 100644 --- a/tests/unittest/test_pdf_analyze_block_content.py +++ b/tests/unittest/test_pdf_analyze_block_content.py @@ -180,6 +180,7 @@ def test_doc_analyze_converts_vlm_results_before_downstream_processing( return window_model_list xhigh_normalizer = MagicMock(wraps=layout._normalize_xhigh_vlm_blocks) + native_table_priority = MagicMock(return_value=MagicMock(total=0)) monkeypatch.setattr(pipeline, "PDFDocument", MagicMock(return_value=fake_document)) monkeypatch.setattr(pipeline, "HybridLocalModelContextSingleton", MagicMock(return_value=hybrid_singleton)) monkeypatch.setattr( @@ -199,6 +200,7 @@ def test_doc_analyze_converts_vlm_results_before_downstream_processing( MagicMock(return_value=[{"img_pil": page_image, "scale": 1.0}]), ) monkeypatch.setattr(window, "_process_text_and_formulas", fake_process_text_and_formulas) + monkeypatch.setattr(window, "_apply_native_txt_table_priority", native_table_priority) monkeypatch.setattr(window, "_normalize_xhigh_vlm_blocks", xhigh_normalizer) monkeypatch.setattr(window, "_apply_seal_ocr", MagicMock()) monkeypatch.setattr(window, "_supplement_missing_image_block_containers", MagicMock()) @@ -251,6 +253,11 @@ def test_doc_analyze_converts_vlm_results_before_downstream_processing( vlm_predictor.batch_extract_with_layout.assert_called_once() vlm_predictor.batch_two_step_extract.assert_not_called() + if effort == "high" and parse_mode == "txt": + native_table_priority.assert_called_once() + else: + native_table_priority.assert_not_called() + def test_xhigh_vlm_blocks_normalize_visual_annotation_types() -> None: """验证 xhigh VLM 的细分视觉标题和脚注会统一成通用类型。"""