pdf2markdown/helper/content_recognition/utils.py

import os
import tempfile
import cv2
from marker.converters.table import TableConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
from .rapid_table_pipeline.main import table2md_pipeline
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.data.read_api import read_local_images
from markdownify import markdownify as md
import re
from ..image_helper import text_rec
from ..constants import PageDetectionEnum as E


def html2md(html_content):
    md_content = md(html_content)
    md_content = re.sub(r'\\([#*_`])', r'\1', md_content)
    return md_content


def markdown_rec(image):
    # TODO 可以传入文件夹
    image_path = f'{tempfile.mktemp()}.jpg'
    cv2.imwrite(image_path, image)

    try:
        ds = read_local_images(image_path)[0]
        x = ds.apply(doc_analyze, ocr=True)
        x = x.pipe_ocr_mode(None)
        html = x.get_markdown(None)
    finally:
        os.remove(image_path)
    return html2md(html)


def table_rec(image):
    boxes, texts, conficences = text_rec(image)
    if not texts:
        return None
    ocr_result = list(zip(boxes, texts, conficences))
    return table2md_pipeline(image, ocr_result)


table_converter = TableConverter(artifact_dict=create_model_dict())


def scanning_document_rec(scanning_document_image):
    tmp_image_path = f'{tempfile.mktemp()}.jpg'
    try:
        cv2.imwrite(tmp_image_path, scanning_document_image)
        rendered = table_converter(tmp_image_path)
        text, _, _ = text_from_rendered(rendered)
    finally:
        os.remove(tmp_image_path)
    return text


def compute_box_distance(box1, box2):
    x11, y11, x12, y12 = box1
    x21, y21, x22, y22 = box2

    # 计算水平和垂直方向的重叠量
    x_overlap = max(0, min(x12, x22) - max(x11, x21))
    y_overlap = max(0, min(y12, y22) - max(y11, y21))

    # 如果有重叠（x和y都重叠），返回负的重叠深度（取 min 表示最小穿透）
    if x_overlap > 0 and y_overlap > 0:
        return -min(x_overlap, y_overlap)

    distances = []

    # 如果 x 方向有投影重叠，计算上下边的距离
    if x12 > x21 and x11 < x22:
        dist_top = y21 - y12  # box1下边到box2上边
        dist_bottom = y11 - y22  # box1上边到box2下边
        if dist_top > 0:
            distances.append(dist_top)
        if dist_bottom > 0:
            distances.append(dist_bottom)

    # 如果 y 方向有投影重叠，计算左右边的距离
    if y12 > y21 and y11 < y22:
        dist_left = x11 - x22  # box1左边到box2右边
        dist_right = x21 - x12  # box1右边到box2左边
        if dist_left > 0:
            distances.append(dist_left)
        if dist_right > 0:
            distances.append(dist_right)

    # 如果有合法的距离，返回最小值，否则说明边无法对齐，返回 None
    return min(distances) if distances else None


def assign_tables_to_titles(layout_results, max_distance=200):
    tables = [_ for _ in layout_results if _.clsid in (E.TABLE.value, E.SCANNED_DOCUMENT.value)]
    titles = [_ for _ in layout_results if _.clsid == E.TABLE_CAPTION.value]

    table_to_title = {}
    title_to_table = {}

    changed = True
    while changed:
        changed = False
        for title in titles:
            title_id = id(title)

            best_table = None
            min_dist = float('inf')

            for table in tables:
                table_id = id(table)

                dist = compute_box_distance(title.box, table.box)
                if dist is None or dist > max_distance:
                    continue

                if dist < min_dist:
                    min_dist = dist
                    best_table = table

            if best_table is None:
                continue

            table_id = id(best_table)

            current_table = title_to_table.get(title_id)
            if current_table is best_table:
                continue  # 已是最优，无需更新

            prev_title = table_to_title.get(table_id)
            if prev_title:
                prev_title_id = id(prev_title)
                prev_dist = compute_box_distance(prev_title.box, best_table.box)
                if prev_dist is not None and prev_dist <= min_dist:
                    continue  # 原标题绑定得更近，跳过

                # 解绑旧标题
                title_to_table.pop(prev_title_id, None)

            # 更新新绑定
            title_to_table[title_id] = best_table
            table_to_title[table_id] = title
            changed = True  # 有更新

    # 最终写回绑定结果
    for table in tables:
        table_id = id(table)
        title = table_to_title.get(table_id)
        if title:
            table.table_title = title.content
        else:
            table.table_title = None