ocr_mkcontent.py 11.4 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
2
from loguru import logger

3
from magic_pdf.libs.commons import join_path
4
from magic_pdf.libs.language import detect_lang
赵小蒙's avatar
赵小蒙 committed
5
from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
赵小蒙's avatar
赵小蒙 committed
6
from magic_pdf.libs.ocr_content_type import ContentType, BlockType
7
8
9
10
11
12
13
14
15
16
17
18
19
import wordninja
import re


def split_long_words(text):
    segments = text.split(' ')
    for i in range(len(segments)):
        words = re.findall(r'\w+|[^\w\s]', segments[i], re.UNICODE)
        for j in range(len(words)):
            if len(words[j]) > 15:
                words[j] = ' '.join(wordninja.split(words[j]))
        segments[i] = ''.join(words)
    return ' '.join(segments)
赵小蒙's avatar
赵小蒙 committed
20
21


赵小蒙's avatar
赵小蒙 committed
22
def ocr_mk_mm_markdown_with_para(pdf_info_list: list, img_buket_path):
xuchao's avatar
xuchao committed
23
    markdown = []
赵小蒙's avatar
赵小蒙 committed
24
    for page_info in pdf_info_list:
25
        paras_of_layout = page_info.get("para_blocks")
赵小蒙's avatar
赵小蒙 committed
26
        page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
27
        markdown.extend(page_markdown)
28
    return '\n\n'.join(markdown)
29
30


赵小蒙's avatar
赵小蒙 committed
31
def ocr_mk_nlp_markdown_with_para(pdf_info_dict: list):
32
    markdown = []
赵小蒙's avatar
赵小蒙 committed
33
    for page_info in pdf_info_dict:
34
        paras_of_layout = page_info.get("para_blocks")
35
        page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "nlp")
36
37
38
        markdown.extend(page_markdown)
    return '\n\n'.join(markdown)

赵小蒙's avatar
赵小蒙 committed
39

赵小蒙's avatar
赵小蒙 committed
40
def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: list, img_buket_path):
41
    markdown_with_para_and_pagination = []
赵小蒙's avatar
赵小蒙 committed
42
43
    page_no = 0
    for page_info in pdf_info_dict:
44
45
        paras_of_layout = page_info.get("para_blocks")
        if not paras_of_layout:
46
            continue
47
        page_markdown = ocr_mk_markdown_with_para_core_v2(paras_of_layout, "mm", img_buket_path)
48
49
        markdown_with_para_and_pagination.append({
            'page_no': page_no,
50
            'md_content': '\n\n'.join(page_markdown)
51
        })
赵小蒙's avatar
赵小蒙 committed
52
        page_no += 1
53
54
55
    return markdown_with_para_and_pagination


赵小蒙's avatar
赵小蒙 committed
56
def ocr_mk_markdown_with_para_core(paras_of_layout, mode, img_buket_path=""):
57
58
59
60
61
62
63
    page_markdown = []
    for paras in paras_of_layout:
        for para in paras:
            para_text = ''
            for line in para:
                for span in line['spans']:
                    span_type = span.get('type')
64
                    content = ''
65
                    language = ''
66
                    if span_type == ContentType.Text:
67
68
69
70
71
72
                        content = span['content']
                        language = detect_lang(content)
                        if language == 'en':  # 只对英文长词进行分词处理,中文分词会丢失文本
                            content = ocr_escape_special_markdown_char(split_long_words(content))
                        else:
                            content = ocr_escape_special_markdown_char(content)
73
                    elif span_type == ContentType.InlineEquation:
74
                        content = f"${span['content']}$"
75
                    elif span_type == ContentType.InterlineEquation:
76
                        content = f"\n$$\n{span['content']}\n$$\n"
77
                    elif span_type in [ContentType.Image, ContentType.Table]:
78
                        if mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
79
                            content = f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
80
81
                        elif mode == 'nlp':
                            pass
82
                    if content != '':
83
84
85
86
                        if language == 'en':  # 英文语境下 content间需要空格分隔
                            para_text += content + ' '
                        else:  # 中文语境下,content间不需要空格分隔
                            para_text += content
87
88
89
90
            if para_text.strip() == '':
                continue
            else:
                page_markdown.append(para_text.strip() + '  ')
91
92
93
    return page_markdown


赵小蒙's avatar
赵小蒙 committed
94
95
def ocr_mk_markdown_with_para_core_v2(paras_of_layout, mode, img_buket_path=""):
    page_markdown = []
96
97
98
99
100
101
102
103
104
105
    for para_block in paras_of_layout:
        para_type = para_block.get('type')
        if para_type == BlockType.Text:
            para_text = merge_para_with_text(para_block)
        elif para_type == BlockType.Title:
            para_text = f"# {merge_para_with_text(para_block)}"
        elif para_type == BlockType.InterlineEquation:
            para_text = merge_para_with_text(para_block)
        elif para_type == BlockType.Image:
            if mode == 'nlp':
赵小蒙's avatar
赵小蒙 committed
106
                continue
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
            elif mode == 'mm':
                img_blocks = para_block.get('blocks')
                for img_block in img_blocks:
                    if img_block.get('type') == BlockType.ImageBody:
                        for line in img_block.get('lines'):
                            for span in line['spans']:
                                if span.get('type') == ContentType.Image:
                                    para_text = f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
                for img_block in img_blocks:
                    if img_block.get('type') == BlockType.ImageCaption:
                        para_text += merge_para_with_text(img_block)
        elif para_type == BlockType.Table:
            if mode == 'nlp':
                continue
            elif mode == 'mm':
                table_blocks = para_block.get('blocks')
                for table_block in table_blocks:
                    if table_block.get('type') == BlockType.TableBody:
                        for line in table_block.get('lines'):
                            for span in line['spans']:
                                if span.get('type') == ContentType.Table:
                                    para_text = f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
                for table_block in table_blocks:
                    if table_block.get('type') == BlockType.TableCaption:
                        para_text += merge_para_with_text(table_block)
                    elif table_block.get('type') == BlockType.TableFootnote:
                        para_text += merge_para_with_text(table_block)

        if para_text.strip() == '':
            continue
        else:
            page_markdown.append(para_text.strip() + '  ')
赵小蒙's avatar
赵小蒙 committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168

    return page_markdown


def merge_para_with_text(para):
    para_text = ''
    for line in para['lines']:
        for span in line['spans']:
            span_type = span.get('type')
            content = ''
            language = ''
            if span_type == ContentType.Text:
                content = span['content']
                language = detect_lang(content)
                if language == 'en':  # 只对英文长词进行分词处理,中文分词会丢失文本
                    content = ocr_escape_special_markdown_char(split_long_words(content))
                else:
                    content = ocr_escape_special_markdown_char(content)
            elif span_type == ContentType.InlineEquation:
                content = f"${span['content']}$"
            elif span_type == ContentType.InterlineEquation:
                content = f"\n$$\n{span['content']}\n$$\n"
            if content != '':
                if language == 'en':  # 英文语境下 content间需要空格分隔
                    para_text += content + ' '
                else:  # 中文语境下,content间不需要空格分隔
                    para_text += content
    return para_text


169
def para_to_standard_format(para, img_buket_path):
170
171
    para_content = {}
    if len(para) == 1:
172
        para_content = line_to_standard_format(para[0], img_buket_path)
173
174
175
176
177
    elif len(para) > 1:
        para_text = ''
        inline_equation_num = 0
        for line in para:
            for span in line['spans']:
178
                language = ''
179
                span_type = span.get('type')
赵小蒙's avatar
赵小蒙 committed
180
                content = ""
181
                if span_type == ContentType.Text:
182
183
184
185
186
187
                    content = span['content']
                    language = detect_lang(content)
                    if language == 'en':  # 只对英文长词进行分词处理,中文分词会丢失文本
                        content = ocr_escape_special_markdown_char(split_long_words(content))
                    else:
                        content = ocr_escape_special_markdown_char(content)
188
                elif span_type == ContentType.InlineEquation:
189
                    content = f"${span['content']}$"
190
                    inline_equation_num += 1
191
192
193
194
195

                if language == 'en':  # 英文语境下 content间需要空格分隔
                    para_text += content + ' '
                else:  # 中文语境下,content间不需要空格分隔
                    para_text += content
196
197
198
199
200
201
202
        para_content = {
            'type': 'text',
            'text': para_text,
            'inline_equation_num': inline_equation_num
        }
    return para_content

赵小蒙's avatar
赵小蒙 committed
203

赵小蒙's avatar
赵小蒙 committed
204
def make_standard_format_with_para(pdf_info_dict: list, img_buket_path: str):
赵小蒙's avatar
赵小蒙 committed
205
    content_list = []
赵小蒙's avatar
赵小蒙 committed
206
    for page_info in pdf_info_dict:
207
208
        paras_of_layout = page_info.get("para_blocks")
        if not paras_of_layout:
赵小蒙's avatar
赵小蒙 committed
209
            continue
210
211
        for paras in paras_of_layout:
            for para in paras:
212
                para_content = para_to_standard_format(para, img_buket_path)
213
                content_list.append(para_content)
赵小蒙's avatar
赵小蒙 committed
214
215
216
    return content_list


217
def line_to_standard_format(line, img_buket_path):
赵小蒙's avatar
赵小蒙 committed
218
219
220
221
222
223
224
225
226
227
    line_text = ""
    inline_equation_num = 0
    for span in line['spans']:
        if not span.get('content'):
            if not span.get('image_path'):
                continue
            else:
                if span['type'] == ContentType.Image:
                    content = {
                        'type': 'image',
228
                        'img_path': join_path(img_buket_path, span['image_path'])
赵小蒙's avatar
赵小蒙 committed
229
230
231
232
233
                    }
                    return content
                elif span['type'] == ContentType.Table:
                    content = {
                        'type': 'table',
234
                        'img_path': join_path(img_buket_path, span['image_path'])
赵小蒙's avatar
赵小蒙 committed
235
236
237
238
                    }
                    return content
        else:
            if span['type'] == ContentType.InterlineEquation:
赵小蒙's avatar
赵小蒙 committed
239
                interline_equation = span['content']
赵小蒙's avatar
赵小蒙 committed
240
241
242
243
244
245
                content = {
                    'type': 'equation',
                    'latex': f"$$\n{interline_equation}\n$$"
                }
                return content
            elif span['type'] == ContentType.InlineEquation:
赵小蒙's avatar
赵小蒙 committed
246
                inline_equation = span['content']
赵小蒙's avatar
赵小蒙 committed
247
248
249
                line_text += f"${inline_equation}$"
                inline_equation_num += 1
            elif span['type'] == ContentType.Text:
250
251
                text_content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                line_text += text_content
赵小蒙's avatar
赵小蒙 committed
252
253
254
255
256
257
258
259
    content = {
        'type': 'text',
        'text': line_text,
        'inline_equation_num': inline_equation_num
    }
    return content


赵小蒙's avatar
赵小蒙 committed
260
def ocr_mk_mm_standard_format(pdf_info_dict: list):
赵小蒙's avatar
update  
赵小蒙 committed
261
    """
262
    content_list
赵小蒙's avatar
赵小蒙 committed
263
264
265
266
267
    type         string      image/text/table/equation(行间的单独拿出来,行内的和text合并)
    latex        string      latex文本字段。
    text         string      纯文本格式的文本数据。
    md           string      markdown格式的文本数据。
    img_path     string      s3://full/path/to/img.jpg
赵小蒙's avatar
update  
赵小蒙 committed
268
    """
赵小蒙's avatar
赵小蒙 committed
269
    content_list = []
赵小蒙's avatar
赵小蒙 committed
270
    for page_info in pdf_info_dict:
赵小蒙's avatar
赵小蒙 committed
271
272
273
274
275
276
277
278
        blocks = page_info.get("preproc_blocks")
        if not blocks:
            continue
        for block in blocks:
            for line in block['lines']:
                content = line_to_standard_format(line)
                content_list.append(content)
    return content_list