pdf_parse_by_ocr.py 8.24 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
import json
2
3
4
5
6
import os
import time

from loguru import logger

赵小蒙's avatar
赵小蒙 committed
7
from magic_pdf.libs.draw_bbox import draw_layout_bbox, draw_text_bbox
许瑞's avatar
许瑞 committed
8
9
10
11
12
13
14
15
from magic_pdf.libs.commons import (
    read_file,
    join_path,
    fitz,
    get_img_s3_client,
    get_delta_time,
    get_docx_model_output,
)
16
from magic_pdf.libs.coordinate_transform import get_scale_ratio
17
18
19
20
21
from magic_pdf.libs.safe_filename import sanitize_filename
from magic_pdf.pre_proc.detect_footer_by_model import parse_footers
from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model
from magic_pdf.pre_proc.detect_header import parse_headers
from magic_pdf.pre_proc.detect_page_number import parse_pageNos
赵小蒙's avatar
赵小蒙 committed
22
from magic_pdf.pre_proc.ocr_cut_image import cut_image_and_table
23
from magic_pdf.pre_proc.ocr_detect_layout import layout_detect
许瑞's avatar
许瑞 committed
24
25
26
27
from magic_pdf.pre_proc.ocr_dict_merge import (
    remove_overlaps_min_spans,
    merge_spans_to_line_by_layout,
)
28
from magic_pdf.pre_proc.ocr_remove_spans import remove_spans_by_bboxes
许瑞's avatar
许瑞 committed
29
from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
赵小蒙's avatar
赵小蒙 committed
30
31


赵小蒙's avatar
赵小蒙 committed
32
def construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree):
赵小蒙's avatar
赵小蒙 committed
33
    return_dict = {
赵小蒙's avatar
赵小蒙 committed
34
        'preproc_blocks': blocks,
35
        'layout_bboxes': layout_bboxes,
赵小蒙's avatar
赵小蒙 committed
36
37
38
        'page_idx': page_id,
        'page_size': [page_w, page_h],
        '_layout_tree': layout_tree,
赵小蒙's avatar
赵小蒙 committed
39
40
41
42
43
    }
    return return_dict


def parse_pdf_by_ocr(
许瑞's avatar
许瑞 committed
44
45
46
47
48
49
50
51
52
53
    pdf_path,
    s3_pdf_profile,
    pdf_model_output,
    save_path,
    book_name,
    pdf_model_profile=None,
    image_s3_config=None,
    start_page_id=0,
    end_page_id=None,
    debug_mode=False,
赵小蒙's avatar
赵小蒙 committed
54
):
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
    pdf_bytes = read_file(pdf_path, s3_pdf_profile)
    save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
    book_name = sanitize_filename(book_name)
    md_bookname_save_path = ""
    if debug_mode:
        save_path = join_path(save_tmp_path, "md")
        pdf_local_path = join_path(save_tmp_path, "download-pdfs", book_name)

        if not os.path.exists(os.path.dirname(pdf_local_path)):
            # 如果目录不存在,创建它
            os.makedirs(os.path.dirname(pdf_local_path))

        md_bookname_save_path = join_path(save_tmp_path, "md", book_name)
        if not os.path.exists(md_bookname_save_path):
            # 如果目录不存在,创建它
            os.makedirs(md_bookname_save_path)

        with open(pdf_local_path + ".pdf", "wb") as pdf_file:
            pdf_file.write(pdf_bytes)

    pdf_docs = fitz.open("pdf", pdf_bytes)
    # 初始化空的pdf_info_dict
赵小蒙's avatar
赵小蒙 committed
77
    pdf_info_dict = {}
78
79
80
81
82
83
84
    img_s3_client = get_img_s3_client(save_path, image_s3_config)

    start_time = time.time()

    remove_bboxes = []

    end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
赵小蒙's avatar
赵小蒙 committed
85
    for page_id in range(start_page_id, end_page_id + 1):
86
87
88

        # 获取当前页的page对象
        page = pdf_docs[page_id]
赵小蒙's avatar
赵小蒙 committed
89
90
91
        # 获取当前页的宽高
        page_w = page.rect.width
        page_h = page.rect.height
92
93
94

        if debug_mode:
            time_now = time.time()
许瑞's avatar
许瑞 committed
95
96
97
            logger.info(
                f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
            )
98
99
100
            start_time = time_now

        # 获取当前页的模型数据
许瑞's avatar
许瑞 committed
101
102
103
        ocr_page_info = get_docx_model_output(
            pdf_model_output, pdf_model_profile, page_id
        )
104
105
106
107
108

        """从json中获取每页的页码、页眉、页脚的bbox"""
        page_no_bboxes = parse_pageNos(page_id, page, ocr_page_info)
        header_bboxes = parse_headers(page_id, page, ocr_page_info)
        footer_bboxes = parse_footers(page_id, page, ocr_page_info)
许瑞's avatar
许瑞 committed
109
110
111
        footnote_bboxes = parse_footnotes_by_model(
            page_id, page, ocr_page_info, md_bookname_save_path, debug_mode=debug_mode
        )
112
113
114
115
116
117
118
119

        # 构建需要remove的bbox列表
        need_remove_spans_bboxes = []
        need_remove_spans_bboxes.extend(page_no_bboxes)
        need_remove_spans_bboxes.extend(header_bboxes)
        need_remove_spans_bboxes.extend(footer_bboxes)
        need_remove_spans_bboxes.extend(footnote_bboxes)

许瑞's avatar
许瑞 committed
120
        layout_dets = ocr_page_info["layout_dets"]
赵小蒙's avatar
赵小蒙 committed
121
        spans = []
122

123
        # 计算模型坐标和pymu坐标的缩放比例
许瑞's avatar
许瑞 committed
124
125
126
        horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
            ocr_page_info, page
        )
127

赵小蒙's avatar
赵小蒙 committed
128
        for layout_det in layout_dets:
许瑞's avatar
许瑞 committed
129
            category_id = layout_det["category_id"]
赵小蒙's avatar
赵小蒙 committed
130
            allow_category_id_list = [1, 7, 13, 14, 15]
赵小蒙's avatar
赵小蒙 committed
131
            if category_id in allow_category_id_list:
许瑞's avatar
许瑞 committed
132
133
134
135
136
137
138
139
                x0, y0, _, _, x1, y1, _, _ = layout_det["poly"]
                bbox = [
                    int(x0 / horizontal_scale_ratio),
                    int(y0 / vertical_scale_ratio),
                    int(x1 / horizontal_scale_ratio),
                    int(y1 / vertical_scale_ratio),
                ]
                """要删除的"""
赵小蒙's avatar
赵小蒙 committed
140
141
142
143
                #  3: 'header',      # 页眉
                #  4: 'page number', # 页码
                #  5: 'footnote',    # 脚注
                #  6: 'footer',      # 页脚
许瑞's avatar
许瑞 committed
144
                """当成span拼接的"""
赵小蒙's avatar
赵小蒙 committed
145
146
147
148
149
                #  1: 'image', # 图片
                #  7: 'table',       # 表格
                #  13: 'inline_equation',     # 行内公式
                #  14: 'displayed_equation',      # 行间公式
                #  15: 'text',      # ocr识别文本
许瑞's avatar
许瑞 committed
150
                """layout信息"""
赵小蒙's avatar
赵小蒙 committed
151
152
                #  11: 'full column',   # 单栏
                #  12: 'sub column',    # 多栏
赵小蒙's avatar
赵小蒙 committed
153
                span = {
许瑞's avatar
许瑞 committed
154
                    "bbox": bbox,
赵小蒙's avatar
赵小蒙 committed
155
                }
赵小蒙's avatar
赵小蒙 committed
156
                if category_id == 1:
许瑞's avatar
许瑞 committed
157
                    span["type"] = "image"
158

赵小蒙's avatar
赵小蒙 committed
159
                elif category_id == 7:
许瑞's avatar
许瑞 committed
160
                    span["type"] = "table"
161

赵小蒙's avatar
赵小蒙 committed
162
                elif category_id == 13:
许瑞's avatar
许瑞 committed
163
164
                    span["content"] = layout_det["latex"]
                    span["type"] = "inline_equation"
赵小蒙's avatar
赵小蒙 committed
165
                elif category_id == 14:
许瑞's avatar
许瑞 committed
166
167
                    span["content"] = layout_det["latex"]
                    span["type"] = "displayed_equation"
赵小蒙's avatar
赵小蒙 committed
168
                elif category_id == 15:
许瑞's avatar
许瑞 committed
169
170
                    span["content"] = layout_det["text"]
                    span["type"] = "text"
赵小蒙's avatar
赵小蒙 committed
171
172
173
174
175
                # print(span)
                spans.append(span)
            else:
                continue

赵小蒙's avatar
赵小蒙 committed
176
177
178
        # 删除重叠spans中较小的那些
        spans = remove_overlaps_min_spans(spans)

179
180
181
        # 删除remove_span_block_bboxes中的bbox
        spans = remove_spans_by_bboxes(spans, need_remove_spans_bboxes)

赵小蒙's avatar
赵小蒙 committed
182
183
184
        # 对image和table截图
        spans = cut_image_and_table(spans, page, page_id, book_name, save_path)

赵小蒙's avatar
赵小蒙 committed
185
186
187
188
189
        # 行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)

        # 模型识别错误的行间公式, type类型转换成行内公式

        # bbox去除粘连
许瑞's avatar
许瑞 committed
190
        spans = remove_overlap_between_bbox(spans)
赵小蒙's avatar
赵小蒙 committed
191
192

        # 对tpye=["displayed_equation", "image", "table"]进行额外处理,如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
赵小蒙's avatar
赵小蒙 committed
193

赵小蒙's avatar
赵小蒙 committed
194
        # 从ocr_page_info中解析layout信息(按自然阅读方向排序,并修复重叠和交错的bad case)
赵小蒙's avatar
赵小蒙 committed
195
        layout_bboxes, layout_tree = layout_detect(ocr_page_info['subfield_dets'], page, ocr_page_info)
赵小蒙's avatar
赵小蒙 committed
196

赵小蒙's avatar
赵小蒙 committed
197
        # 将spans合并成line(在layout内,从上到下,从左到右)
198
        lines = merge_spans_to_line_by_layout(spans, layout_bboxes)
赵小蒙's avatar
赵小蒙 committed
199

赵小蒙's avatar
赵小蒙 committed
200
201
202
        # 目前不做block拼接,先做个结构,每个block中只有一个line,block的bbox就是line的bbox
        blocks = []
        for line in lines:
许瑞's avatar
许瑞 committed
203
204
205
206
207
208
            blocks.append(
                {
                    "bbox": line["bbox"],
                    "lines": [line],
                }
            )
赵小蒙's avatar
赵小蒙 committed
209
210

        # 构造pdf_info_dict
赵小蒙's avatar
赵小蒙 committed
211
        page_info = construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree)
赵小蒙's avatar
赵小蒙 committed
212
213
        pdf_info_dict[f"page_{page_id}"] = page_info

214
215
    # 在测试时,保存调试信息
    if debug_mode:
许瑞's avatar
许瑞 committed
216
217
218
        params_file_save_path = join_path(
            save_tmp_path, "md", book_name, "preproc_out.json"
        )
219
220
        with open(params_file_save_path, "w", encoding="utf-8") as f:
            json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
赵小蒙's avatar
赵小蒙 committed
221

222
223
224
        # drow_bbox
        draw_layout_bbox(pdf_info_dict, pdf_path, md_bookname_save_path)
        draw_text_bbox(pdf_info_dict, pdf_path, md_bookname_save_path)
赵小蒙's avatar
赵小蒙 committed
225

赵小蒙's avatar
赵小蒙 committed
226
    return pdf_info_dict