mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-09-01 15:22:18 +08:00
feat: implement native table priority handling with complex entry marking and internal text removal
This commit is contained in:
@@ -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))
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 = "<table><tbody><tr><td>A</td></tr></tbody></table>"
|
||||
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": "<table><tbody><tr><td>native</td></tr></tbody></table>",
|
||||
}
|
||||
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("<table>")
|
||||
|
||||
|
||||
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": "<table><tbody><tr><td>native formula text</td></tr></tbody></table>",
|
||||
}
|
||||
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": "<table><tbody><tr><td>native</td></tr></tbody></table>",
|
||||
},
|
||||
{"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"] = "<table><tbody><tr><td>fallback</td></tr></tbody></table>"
|
||||
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</td></tr></tbody></table>")
|
||||
assert converted[0][3]["content"].endswith("fallback</td></tr></tbody></table>")
|
||||
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 = "<table><tbody><tr><td>native</td></tr></tbody></table>"
|
||||
fallback_html = "<table><tbody><tr><td>fallback</td></tr></tbody></table>"
|
||||
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))
|
||||
|
||||
@@ -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 的细分视觉标题和脚注会统一成通用类型。"""
|
||||
|
||||
Reference in New Issue
Block a user