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

赵小蒙's avatar
赵小蒙 committed
5
from loguru import logger
赵小蒙's avatar
赵小蒙 committed
6

赵小蒙'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
from magic_pdf.libs.drop_tag import DropTag
赵小蒙's avatar
赵小蒙 committed
18
from magic_pdf.libs.ocr_content_type import ContentType
19
from magic_pdf.libs.safe_filename import sanitize_filename
xuchao's avatar
xuchao committed
20
from magic_pdf.para.para_split import para_split
21
22
23
24
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
25
from magic_pdf.pre_proc.ocr_cut_image import cut_image_and_table
26
from magic_pdf.pre_proc.ocr_detect_layout import layout_detect
许瑞's avatar
许瑞 committed
27
from magic_pdf.pre_proc.ocr_dict_merge import (
赵小蒙's avatar
赵小蒙 committed
28
    merge_spans_to_line_by_layout, merge_lines_to_block,
许瑞's avatar
许瑞 committed
29
)
30
from magic_pdf.pre_proc.ocr_span_list_modify import remove_spans_by_bboxes, remove_overlaps_min_spans, \
赵小蒙's avatar
赵小蒙 committed
31
32
    adjust_bbox_for_standalone_block, modify_y_axis, modify_inline_equation, get_qa_need_list, \
    remove_spans_by_bboxes_dict
许瑞's avatar
许瑞 committed
33
from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
赵小蒙's avatar
赵小蒙 committed
34

赵小蒙's avatar
赵小蒙 committed
35

36
def construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
赵小蒙's avatar
赵小蒙 committed
37
                             images, tables, interline_equations, inline_equations,
38
                             dropped_text_block, dropped_image_block, dropped_table_block, dropped_equation_block,
赵小蒙's avatar
赵小蒙 committed
39
                             need_remove_spans_bboxes_dict):
赵小蒙's avatar
赵小蒙 committed
40
    return_dict = {
赵小蒙's avatar
赵小蒙 committed
41
        'preproc_blocks': blocks,
42
        'layout_bboxes': layout_bboxes,
赵小蒙's avatar
赵小蒙 committed
43
44
45
        'page_idx': page_id,
        'page_size': [page_w, page_h],
        '_layout_tree': layout_tree,
赵小蒙's avatar
赵小蒙 committed
46
47
48
49
        'images': images,
        'tables': tables,
        'interline_equations': interline_equations,
        'inline_equations': inline_equations,
赵小蒙's avatar
赵小蒙 committed
50
51
52
        'droped_text_block': dropped_text_block,
        'droped_image_block': dropped_image_block,
        'droped_table_block': dropped_table_block,
53
        'dropped_equation_block': dropped_equation_block,
赵小蒙's avatar
赵小蒙 committed
54
        'droped_bboxes': need_remove_spans_bboxes_dict,
赵小蒙's avatar
赵小蒙 committed
55
56
57
58
59
    }
    return return_dict


def parse_pdf_by_ocr(
赵小蒙's avatar
赵小蒙 committed
60
61
62
63
64
65
66
67
68
69
        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
70
):
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    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)
赵小蒙's avatar
赵小蒙 committed
90

91
92
    pdf_docs = fitz.open("pdf", pdf_bytes)
    # 初始化空的pdf_info_dict
赵小蒙's avatar
赵小蒙 committed
93
    pdf_info_dict = {}
94
95
96
97
98
    img_s3_client = get_img_s3_client(save_path, image_s3_config)

    start_time = time.time()

    end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
赵小蒙's avatar
赵小蒙 committed
99
    for page_id in range(start_page_id, end_page_id + 1):
100
101
102

        # 获取当前页的page对象
        page = pdf_docs[page_id]
赵小蒙's avatar
赵小蒙 committed
103
104
105
        # 获取当前页的宽高
        page_w = page.rect.width
        page_h = page.rect.height
106
107
108

        if debug_mode:
            time_now = time.time()
许瑞's avatar
许瑞 committed
109
110
111
            logger.info(
                f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
            )
112
113
114
            start_time = time_now

        # 获取当前页的模型数据
许瑞's avatar
许瑞 committed
115
116
117
        ocr_page_info = get_docx_model_output(
            pdf_model_output, pdf_model_profile, page_id
        )
118
119
120
121
122

        """从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
123
124
125
        footnote_bboxes = parse_footnotes_by_model(
            page_id, page, ocr_page_info, md_bookname_save_path, debug_mode=debug_mode
        )
126

赵小蒙's avatar
赵小蒙 committed
127
128
        # 构建需要remove的bbox字典
        need_remove_spans_bboxes_dict = {
129
130
131
132
            DropTag.PAGE_NUMBER: page_no_bboxes,
            DropTag.HEADER: header_bboxes,
            DropTag.FOOTER: footer_bboxes,
            DropTag.FOOTNOTE: footnote_bboxes,
赵小蒙's avatar
赵小蒙 committed
133
        }
134

许瑞's avatar
许瑞 committed
135
        layout_dets = ocr_page_info["layout_dets"]
赵小蒙's avatar
赵小蒙 committed
136
        spans = []
137

138
        # 计算模型坐标和pymu坐标的缩放比例
许瑞's avatar
许瑞 committed
139
140
141
        horizontal_scale_ratio, vertical_scale_ratio = get_scale_ratio(
            ocr_page_info, page
        )
142

赵小蒙's avatar
赵小蒙 committed
143
        for layout_det in layout_dets:
许瑞's avatar
许瑞 committed
144
            category_id = layout_det["category_id"]
赵小蒙's avatar
赵小蒙 committed
145
            allow_category_id_list = [1, 7, 13, 14, 15]
赵小蒙's avatar
赵小蒙 committed
146
            if category_id in allow_category_id_list:
许瑞's avatar
许瑞 committed
147
148
149
150
151
152
153
                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),
                ]
154
155
156
                # 删除高度或者宽度为0的spans
                if bbox[2] - bbox[0] == 0 or bbox[3] - bbox[1] == 0:
                    continue
许瑞's avatar
许瑞 committed
157
                """要删除的"""
赵小蒙's avatar
赵小蒙 committed
158
159
160
161
                #  3: 'header',      # 页眉
                #  4: 'page number', # 页码
                #  5: 'footnote',    # 脚注
                #  6: 'footer',      # 页脚
许瑞's avatar
许瑞 committed
162
                """当成span拼接的"""
赵小蒙's avatar
赵小蒙 committed
163
164
165
                #  1: 'image', # 图片
                #  7: 'table',       # 表格
                #  13: 'inline_equation',     # 行内公式
赵小蒙's avatar
赵小蒙 committed
166
                #  14: 'interline_equation',      # 行间公式
赵小蒙's avatar
赵小蒙 committed
167
                #  15: 'text',      # ocr识别文本
许瑞's avatar
许瑞 committed
168
                """layout信息"""
赵小蒙's avatar
赵小蒙 committed
169
170
                #  11: 'full column',   # 单栏
                #  12: 'sub column',    # 多栏
赵小蒙's avatar
赵小蒙 committed
171
                span = {
许瑞's avatar
许瑞 committed
172
                    "bbox": bbox,
赵小蒙's avatar
赵小蒙 committed
173
                }
赵小蒙's avatar
赵小蒙 committed
174
                if category_id == 1:
赵小蒙's avatar
赵小蒙 committed
175
                    span["type"] = ContentType.Image
176

赵小蒙's avatar
赵小蒙 committed
177
                elif category_id == 7:
赵小蒙's avatar
赵小蒙 committed
178
                    span["type"] = ContentType.Table
179

赵小蒙's avatar
赵小蒙 committed
180
                elif category_id == 13:
许瑞's avatar
许瑞 committed
181
                    span["content"] = layout_det["latex"]
赵小蒙's avatar
赵小蒙 committed
182
                    span["type"] = ContentType.InlineEquation
赵小蒙's avatar
赵小蒙 committed
183
                elif category_id == 14:
许瑞's avatar
许瑞 committed
184
                    span["content"] = layout_det["latex"]
赵小蒙's avatar
赵小蒙 committed
185
                    span["type"] = ContentType.InterlineEquation
赵小蒙's avatar
赵小蒙 committed
186
                elif category_id == 15:
许瑞's avatar
许瑞 committed
187
                    span["content"] = layout_det["text"]
赵小蒙's avatar
赵小蒙 committed
188
                    span["type"] = ContentType.Text
赵小蒙's avatar
赵小蒙 committed
189
190
191
192
193
                # print(span)
                spans.append(span)
            else:
                continue

赵小蒙's avatar
赵小蒙 committed
194
        '''删除重叠spans中较小的那些'''
195
        spans, dropped_spans_by_span_overlap = remove_overlaps_min_spans(spans)
赵小蒙's avatar
赵小蒙 committed
196

赵小蒙's avatar
赵小蒙 committed
197
198
199
200
        '''
        删除remove_span_block_bboxes中的bbox
        并增加drop相关数据
        '''
201
        spans, dropped_spans_by_removed_bboxes = remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict)
202

赵小蒙's avatar
赵小蒙 committed
203
        '''对image和table截图'''
204
        spans = cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client)
赵小蒙's avatar
赵小蒙 committed
205

赵小蒙's avatar
赵小蒙 committed
206
        '''行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)'''
liukaiwen's avatar
liukaiwen committed
207
208
209
        displayed_list = []
        text_inline_lines = []
        modify_y_axis(spans, displayed_list, text_inline_lines)
赵小蒙's avatar
赵小蒙 committed
210
211

        '''模型识别错误的行间公式, type类型转换成行内公式'''
liukaiwen's avatar
liukaiwen committed
212
        spans = modify_inline_equation(spans, displayed_list, text_inline_lines)
赵小蒙's avatar
赵小蒙 committed
213

赵小蒙's avatar
赵小蒙 committed
214
        '''bbox去除粘连'''
许瑞's avatar
许瑞 committed
215
        spans = remove_overlap_between_bbox(spans)
赵小蒙's avatar
赵小蒙 committed
216

赵小蒙's avatar
赵小蒙 committed
217
218
219
220
        '''
        对tpye=["interline_equation", "image", "table"]进行额外处理,
        如果左边有字的话,将该span的bbox中y0调整至不高于文字的y0
        '''
221
222
        spans = adjust_bbox_for_standalone_block(spans)

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

赵小蒙's avatar
赵小蒙 committed
226
        '''将spans合并成line(在layout内,从上到下,从左到右)'''
227
        lines, dropped_spans_by_layout = merge_spans_to_line_by_layout(spans, layout_bboxes)
赵小蒙's avatar
赵小蒙 committed
228

赵小蒙's avatar
赵小蒙 committed
229
        '''将lines合并成block'''
赵小蒙's avatar
赵小蒙 committed
230
231
        blocks = merge_lines_to_block(lines)

赵小蒙's avatar
赵小蒙 committed
232
        '''获取QA需要外置的list'''
赵小蒙's avatar
赵小蒙 committed
233
        images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
赵小蒙's avatar
赵小蒙 committed
234

赵小蒙's avatar
赵小蒙 committed
235
        '''drop的span_list合并'''
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
        dropped_spans = []
        dropped_spans.extend(dropped_spans_by_span_overlap)
        dropped_spans.extend(dropped_spans_by_removed_bboxes)
        dropped_spans.extend(dropped_spans_by_layout)

        dropped_text_block = []
        dropped_image_block = []
        dropped_table_block = []
        dropped_equation_block = []
        for span in dropped_spans:
            # drop出的spans进行分类
            if span['type'] == ContentType.Text:
                dropped_text_block.append(span)
            elif span['type'] == ContentType.Image:
                dropped_image_block.append(span)
            elif span['type'] == ContentType.Table:
                dropped_table_block.append(span)
            elif span['type'] in [ContentType.InlineEquation, ContentType.InterlineEquation]:
                dropped_equation_block.append(span)

赵小蒙's avatar
赵小蒙 committed
256
        '''构造pdf_info_dict'''
257
        page_info = construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
赵小蒙's avatar
赵小蒙 committed
258
                                             images, tables, interline_equations, inline_equations,
赵小蒙's avatar
赵小蒙 committed
259
260
                                             dropped_text_block, dropped_image_block, dropped_table_block,
                                             dropped_equation_block,
赵小蒙's avatar
赵小蒙 committed
261
                                             need_remove_spans_bboxes_dict)
赵小蒙's avatar
赵小蒙 committed
262
263
        pdf_info_dict[f"page_{page_id}"] = page_info

264
    """分段"""
265
    para_split(pdf_info_dict, debug_mode=debug_mode)
赵小蒙's avatar
赵小蒙 committed
266
267

    '''在测试时,保存调试信息'''
268
    if debug_mode:
许瑞's avatar
许瑞 committed
269
270
271
        params_file_save_path = join_path(
            save_tmp_path, "md", book_name, "preproc_out.json"
        )
272
273
        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
274

275
        # drow_bbox
276
277
        draw_layout_bbox(pdf_info_dict, pdf_bytes, md_bookname_save_path)
        draw_text_bbox(pdf_info_dict, pdf_bytes, md_bookname_save_path)
赵小蒙's avatar
赵小蒙 committed
278

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