ocr_mkcontent.py 19.8 KB
Newer Older
1
2
3
import re

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

6
from magic_pdf.libs.commons import join_path
7
from magic_pdf.libs.language import detect_lang
8
from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
赵小蒙's avatar
赵小蒙 committed
9
from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
10
from magic_pdf.libs.ocr_content_type import BlockType, ContentType
11
from magic_pdf.para.para_split_v3 import ListLineTag
12
13


14
15
16
17
18
19
20
21
22
23
24
25
26
27
def __is_hyphen_at_line_end(line):
    """
    Check if a line ends with one or more letters followed by a hyphen.
    
    Args:
    line (str): The line of text to check.
    
    Returns:
    bool: True if the line ends with one or more letters followed by a hyphen, False otherwise.
    """
    # Use regex to check if the line ends with one or more letters followed by a hyphen
    return bool(re.search(r'[A-Za-z]+-\s*$', line))


28
29
30
def split_long_words(text):
    segments = text.split(' ')
    for i in range(len(segments)):
liukaiwen's avatar
liukaiwen committed
31
        words = re.findall(r'\w+|[^\w]', segments[i], re.UNICODE)
32
        for j in range(len(words)):
33
            if len(words[j]) > 10:
34
35
36
                words[j] = ' '.join(wordninja.split(words[j]))
        segments[i] = ''.join(words)
    return ' '.join(segments)
赵小蒙's avatar
赵小蒙 committed
37
38


赵小蒙's avatar
赵小蒙 committed
39
def ocr_mk_mm_markdown_with_para(pdf_info_list: list, img_buket_path):
xuchao's avatar
xuchao committed
40
    markdown = []
赵小蒙's avatar
赵小蒙 committed
41
    for page_info in pdf_info_list:
42
43
44
        paras_of_layout = page_info.get('para_blocks')
        page_markdown = ocr_mk_markdown_with_para_core_v2(
            paras_of_layout, 'mm', img_buket_path)
45
        markdown.extend(page_markdown)
46
    return '\n\n'.join(markdown)
47
48


赵小蒙's avatar
赵小蒙 committed
49
def ocr_mk_nlp_markdown_with_para(pdf_info_dict: list):
50
    markdown = []
赵小蒙's avatar
赵小蒙 committed
51
    for page_info in pdf_info_dict:
52
53
54
        paras_of_layout = page_info.get('para_blocks')
        page_markdown = ocr_mk_markdown_with_para_core_v2(
            paras_of_layout, 'nlp')
55
56
57
        markdown.extend(page_markdown)
    return '\n\n'.join(markdown)

赵小蒙's avatar
赵小蒙 committed
58

59
60
def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: list,
                                                img_buket_path):
61
    markdown_with_para_and_pagination = []
赵小蒙's avatar
赵小蒙 committed
62
63
    page_no = 0
    for page_info in pdf_info_dict:
64
        paras_of_layout = page_info.get('para_blocks')
65
        if not paras_of_layout:
66
            continue
67
68
        page_markdown = ocr_mk_markdown_with_para_core_v2(
            paras_of_layout, 'mm', img_buket_path)
69
        markdown_with_para_and_pagination.append({
70
71
72
73
            'page_no':
            page_no,
            'md_content':
            '\n\n'.join(page_markdown)
74
        })
赵小蒙's avatar
赵小蒙 committed
75
        page_no += 1
76
77
78
    return markdown_with_para_and_pagination


79
def ocr_mk_markdown_with_para_core(paras_of_layout, mode, img_buket_path=''):
80
81
82
83
84
85
86
    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')
87
                    content = ''
88
                    language = ''
89
                    if span_type == ContentType.Text:
90
91
                        content = span['content']
                        language = detect_lang(content)
92
93
94
                        if (language == 'en'):  # 只对英文长词进行分词处理,中文分词会丢失文本
                            content = ocr_escape_special_markdown_char(
                                split_long_words(content))
95
96
                        else:
                            content = ocr_escape_special_markdown_char(content)
97
                    elif span_type == ContentType.InlineEquation:
98
                        content = f"${span['content']}$"
99
                    elif span_type == ContentType.InterlineEquation:
100
                        content = f"\n$$\n{span['content']}\n$$\n"
101
                    elif span_type in [ContentType.Image, ContentType.Table]:
102
                        if mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
103
                            content = f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
104
105
                        elif mode == 'nlp':
                            pass
106
                    if content != '':
107
108
109
110
                        if language == 'en':  # 英文语境下 content间需要空格分隔
                            para_text += content + ' '
                        else:  # 中文语境下,content间不需要空格分隔
                            para_text += content
111
112
113
114
            if para_text.strip() == '':
                continue
            else:
                page_markdown.append(para_text.strip() + '  ')
115
116
117
    return page_markdown


118
119
def ocr_mk_markdown_with_para_core_v2(paras_of_layout,
                                      mode,
120
121
122
123
                                      img_buket_path='',
                                      parse_type="auto",
                                      lang=None
                                      ):
赵小蒙's avatar
赵小蒙 committed
124
    page_markdown = []
125
    for para_block in paras_of_layout:
赵小蒙's avatar
赵小蒙 committed
126
        para_text = ''
赵小蒙's avatar
赵小蒙 committed
127
        para_type = para_block['type']
128
        if para_type in [BlockType.Text, BlockType.List, BlockType.Index]:
129
            para_text = merge_para_with_text(para_block, parse_type=parse_type, lang=lang)
130
        elif para_type == BlockType.Title:
131
            para_text = f'# {merge_para_with_text(para_block, parse_type=parse_type, lang=lang)}'
132
        elif para_type == BlockType.InterlineEquation:
133
            para_text = merge_para_with_text(para_block, parse_type=parse_type, lang=lang)
134
135
        elif para_type == BlockType.Image:
            if mode == 'nlp':
赵小蒙's avatar
赵小蒙 committed
136
                continue
137
            elif mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
138
                for block in para_block['blocks']:  # 1st.拼image_body
赵小蒙's avatar
赵小蒙 committed
139
140
                    if block['type'] == BlockType.ImageBody:
                        for line in block['lines']:
141
                            for span in line['spans']:
赵小蒙's avatar
赵小蒙 committed
142
                                if span['type'] == ContentType.Image:
143
                                    para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})  \n"
赵小蒙's avatar
赵小蒙 committed
144
                for block in para_block['blocks']:  # 2nd.拼image_caption
赵小蒙's avatar
赵小蒙 committed
145
                    if block['type'] == BlockType.ImageCaption:
146
                        para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
147
148
                for block in para_block['blocks']:  # 2nd.拼image_caption
                    if block['type'] == BlockType.ImageFootnote:
149
                        para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
150
151
152
153
        elif para_type == BlockType.Table:
            if mode == 'nlp':
                continue
            elif mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
154
155
                for block in para_block['blocks']:  # 1st.拼table_caption
                    if block['type'] == BlockType.TableCaption:
156
                        para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
157
                for block in para_block['blocks']:  # 2nd.拼table_body
赵小蒙's avatar
赵小蒙 committed
158
159
                    if block['type'] == BlockType.TableBody:
                        for line in block['lines']:
160
                            for span in line['spans']:
赵小蒙's avatar
赵小蒙 committed
161
                                if span['type'] == ContentType.Table:
162
163
164
                                    # if processed by table model
                                    if span.get('latex', ''):
                                        para_text += f"\n\n$\n {span['latex']}\n$\n\n"
165
166
                                    elif span.get('html', ''):
                                        para_text += f"\n\n{span['html']}\n\n"
167
                                    else:
168
                                        para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})  \n"
赵小蒙's avatar
赵小蒙 committed
169
170
                for block in para_block['blocks']:  # 3rd.拼table_footnote
                    if block['type'] == BlockType.TableFootnote:
171
                        para_text += merge_para_with_text(block, parse_type=parse_type, lang=lang)
172
173
174
175
176

        if para_text.strip() == '':
            continue
        else:
            page_markdown.append(para_text.strip() + '  ')
赵小蒙's avatar
赵小蒙 committed
177
178
179
180

    return page_markdown


181
182
183
184
185
186
187
def detect_language(text):
    en_pattern = r'[a-zA-Z]+'
    en_matches = re.findall(en_pattern, text)
    en_length = sum(len(match) for match in en_matches)
    if len(text) > 0:
        if en_length / len(text) >= 0.5:
            return 'en'
188
        else:
189
190
191
192
            return 'unknown'
    else:
        return 'empty'

193

194
def merge_para_with_text(para_block, parse_type="auto", lang=None):
赵小蒙's avatar
赵小蒙 committed
195
    para_text = ''
196
197
198
199
200
    for i, line in enumerate(para_block['lines']):

        if i >= 1 and line.get(ListLineTag.IS_LIST_START_LINE, False):
            para_text += '  \n'

201
202
        line_text = ''
        line_lang = ''
203
204
205
206
        for span in line['spans']:
            span_type = span['type']
            if span_type == ContentType.Text:
                line_text += span['content'].strip()
207
        if line_text != '':
208
            line_lang = detect_lang(line_text)
赵小蒙's avatar
赵小蒙 committed
209
        for span in line['spans']:
赵小蒙's avatar
赵小蒙 committed
210
            span_type = span['type']
赵小蒙's avatar
赵小蒙 committed
211
212
213
            content = ''
            if span_type == ContentType.Text:
                content = span['content']
214
215
                # language = detect_lang(content)
                language = detect_language(content)
216
217
                # 判断是否小语种
                if lang is not None and lang != 'en':
赵小蒙's avatar
赵小蒙 committed
218
                    content = ocr_escape_special_markdown_char(content)
219
220
221
222
223
224
                else:  # 非小语种逻辑
                    if language == 'en' and parse_type == 'ocr':  # 只对英文长词进行分词处理,中文分词会丢失文本
                        content = ocr_escape_special_markdown_char(
                            split_long_words(content))
                    else:
                        content = ocr_escape_special_markdown_char(content)
赵小蒙's avatar
赵小蒙 committed
225
            elif span_type == ContentType.InlineEquation:
226
                content = f" ${span['content']}$ "
赵小蒙's avatar
赵小蒙 committed
227
228
            elif span_type == ContentType.InterlineEquation:
                content = f"\n$$\n{span['content']}\n$$\n"
229

赵小蒙's avatar
赵小蒙 committed
230
            if content != '':
231
232
233
                langs = ['zh', 'ja', 'ko']
                if line_lang in langs:  # 遇到一些一个字一个span的文档,这种单字语言判断不准,需要用整行文本判断
                    para_text += content  # 中文/日语/韩文语境下,content间不需要空格分隔
234
235
                elif line_lang == 'en':
                    # 如果是前一行带有-连字符,那么末尾不应该加空格
drunkpig's avatar
drunkpig committed
236
237
                    if __is_hyphen_at_line_end(content):
                        para_text += content[:-1]
238
239
                    else:
                        para_text += content + ' '
240
                else:
241
                    para_text += content + ' '  # 西方文本语境下 content间需要空格分隔
赵小蒙's avatar
赵小蒙 committed
242
243
244
    return para_text


245
def para_to_standard_format(para, img_buket_path):
246
247
    para_content = {}
    if len(para) == 1:
248
        para_content = line_to_standard_format(para[0], img_buket_path)
249
250
251
252
253
    elif len(para) > 1:
        para_text = ''
        inline_equation_num = 0
        for line in para:
            for span in line['spans']:
254
                language = ''
255
                span_type = span.get('type')
256
                content = ''
257
                if span_type == ContentType.Text:
258
259
260
                    content = span['content']
                    language = detect_lang(content)
                    if language == 'en':  # 只对英文长词进行分词处理,中文分词会丢失文本
261
262
                        content = ocr_escape_special_markdown_char(
                            split_long_words(content))
263
264
                    else:
                        content = ocr_escape_special_markdown_char(content)
265
                elif span_type == ContentType.InlineEquation:
266
                    content = f"${span['content']}$"
267
                    inline_equation_num += 1
268
269
270
271
                if language == 'en':  # 英文语境下 content间需要空格分隔
                    para_text += content + ' '
                else:  # 中文语境下,content间不需要空格分隔
                    para_text += content
272
273
274
        para_content = {
            'type': 'text',
            'text': para_text,
275
            'inline_equation_num': inline_equation_num,
276
277
278
        }
    return para_content

赵小蒙's avatar
赵小蒙 committed
279

280
def para_to_standard_format_v2(para_block, img_buket_path, page_idx, parse_type="auto", lang=None, drop_reason=None):
赵小蒙's avatar
赵小蒙 committed
281
    para_type = para_block['type']
282
    para_content = {}
赵小蒙's avatar
赵小蒙 committed
283
284
285
    if para_type == BlockType.Text:
        para_content = {
            'type': 'text',
286
            'text': merge_para_with_text(para_block, parse_type=parse_type, lang=lang),
赵小蒙's avatar
赵小蒙 committed
287
288
289
290
        }
    elif para_type == BlockType.Title:
        para_content = {
            'type': 'text',
291
            'text': merge_para_with_text(para_block, parse_type=parse_type, lang=lang),
292
            'text_level': 1,
赵小蒙's avatar
赵小蒙 committed
293
294
295
296
        }
    elif para_type == BlockType.InterlineEquation:
        para_content = {
            'type': 'equation',
297
            'text': merge_para_with_text(para_block, parse_type=parse_type, lang=lang),
298
            'text_format': 'latex',
赵小蒙's avatar
赵小蒙 committed
299
300
        }
    elif para_type == BlockType.Image:
301
        para_content = {'type': 'image'}
赵小蒙's avatar
赵小蒙 committed
302
303
        for block in para_block['blocks']:
            if block['type'] == BlockType.ImageBody:
304
305
306
                para_content['img_path'] = join_path(
                    img_buket_path,
                    block['lines'][0]['spans'][0]['image_path'])
赵小蒙's avatar
赵小蒙 committed
307
            if block['type'] == BlockType.ImageCaption:
308
                para_content['img_caption'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
309
            if block['type'] == BlockType.ImageFootnote:
310
                para_content['img_footnote'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
311
    elif para_type == BlockType.Table:
312
        para_content = {'type': 'table'}
赵小蒙's avatar
赵小蒙 committed
313
314
        for block in para_block['blocks']:
            if block['type'] == BlockType.TableBody:
315
                if block["lines"][0]["spans"][0].get('latex', ''):
liukaiwen's avatar
liukaiwen committed
316
                    para_content['table_body'] = f"\n\n$\n {block['lines'][0]['spans'][0]['latex']}\n$\n\n"
317
318
                elif block["lines"][0]["spans"][0].get('html', ''):
                    para_content['table_body'] = f"\n\n{block['lines'][0]['spans'][0]['html']}\n\n"
赵小蒙's avatar
赵小蒙 committed
319
320
                para_content['img_path'] = join_path(img_buket_path, block["lines"][0]["spans"][0]['image_path'])
            if block['type'] == BlockType.TableCaption:
321
                para_content['table_caption'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
322
            if block['type'] == BlockType.TableFootnote:
323
                para_content['table_footnote'] = merge_para_with_text(block, parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
324

325
326
327
328
329
    para_content['page_idx'] = page_idx

    if drop_reason is not None:
        para_content['drop_reason'] = drop_reason

赵小蒙's avatar
赵小蒙 committed
330
331
332
    return para_content


赵小蒙's avatar
赵小蒙 committed
333
def make_standard_format_with_para(pdf_info_dict: list, img_buket_path: str):
赵小蒙's avatar
赵小蒙 committed
334
    content_list = []
赵小蒙's avatar
赵小蒙 committed
335
    for page_info in pdf_info_dict:
336
        paras_of_layout = page_info.get('para_blocks')
337
        if not paras_of_layout:
赵小蒙's avatar
赵小蒙 committed
338
            continue
赵小蒙's avatar
赵小蒙 committed
339
        for para_block in paras_of_layout:
340
341
            para_content = para_to_standard_format_v2(para_block,
                                                      img_buket_path)
赵小蒙's avatar
赵小蒙 committed
342
            content_list.append(para_content)
赵小蒙's avatar
赵小蒙 committed
343
344
345
    return content_list


346
def line_to_standard_format(line, img_buket_path):
347
    line_text = ''
赵小蒙's avatar
赵小蒙 committed
348
349
350
351
352
353
354
355
356
    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',
357
358
                        'img_path': join_path(img_buket_path,
                                              span['image_path']),
赵小蒙's avatar
赵小蒙 committed
359
360
361
362
363
                    }
                    return content
                elif span['type'] == ContentType.Table:
                    content = {
                        'type': 'table',
364
365
                        'img_path': join_path(img_buket_path,
                                              span['image_path']),
赵小蒙's avatar
赵小蒙 committed
366
367
368
369
                    }
                    return content
        else:
            if span['type'] == ContentType.InterlineEquation:
赵小蒙's avatar
赵小蒙 committed
370
                interline_equation = span['content']
赵小蒙's avatar
赵小蒙 committed
371
372
                content = {
                    'type': 'equation',
373
                    'latex': f'$$\n{interline_equation}\n$$'
赵小蒙's avatar
赵小蒙 committed
374
375
376
                }
                return content
            elif span['type'] == ContentType.InlineEquation:
赵小蒙's avatar
赵小蒙 committed
377
                inline_equation = span['content']
378
                line_text += f'${inline_equation}$'
赵小蒙's avatar
赵小蒙 committed
379
380
                inline_equation_num += 1
            elif span['type'] == ContentType.Text:
381
382
                text_content = ocr_escape_special_markdown_char(
                    span['content'])  # 转义特殊符号
383
                line_text += text_content
赵小蒙's avatar
赵小蒙 committed
384
385
386
    content = {
        'type': 'text',
        'text': line_text,
387
        'inline_equation_num': inline_equation_num,
赵小蒙's avatar
赵小蒙 committed
388
389
390
391
    }
    return content


赵小蒙's avatar
赵小蒙 committed
392
def ocr_mk_mm_standard_format(pdf_info_dict: list):
393
394
395
396
    """content_list 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
赵小蒙 committed
397
    content_list = []
赵小蒙's avatar
赵小蒙 committed
398
    for page_info in pdf_info_dict:
399
        blocks = page_info.get('preproc_blocks')
赵小蒙's avatar
赵小蒙 committed
400
401
402
403
404
405
406
        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
赵小蒙's avatar
赵小蒙 committed
407
408


409
410
411
def union_make(pdf_info_dict: list,
               make_mode: str,
               drop_mode: str,
412
413
414
               img_buket_path: str = '',
               parse_type: str = "auto",
               lang=None):
赵小蒙's avatar
赵小蒙 committed
415
416
    output_content = []
    for page_info in pdf_info_dict:
417
418
        drop_reason_flag = False
        drop_reason = None
419
        if page_info.get('need_drop', False):
420
421
            drop_reason = page_info.get('drop_reason')
            if drop_mode == DropMode.NONE:
赵小蒙's avatar
赵小蒙 committed
422
                pass
423
424
            elif drop_mode == DropMode.NONE_WITH_REASON:
                drop_reason_flag = True
赵小蒙's avatar
赵小蒙 committed
425
            elif drop_mode == DropMode.WHOLE_PDF:
426
427
                raise Exception((f'drop_mode is {DropMode.WHOLE_PDF} ,'
                                 f'drop_reason is {drop_reason}'))
赵小蒙's avatar
赵小蒙 committed
428
            elif drop_mode == DropMode.SINGLE_PAGE:
429
430
                logger.warning((f'drop_mode is {DropMode.SINGLE_PAGE} ,'
                                f'drop_reason is {drop_reason}'))
赵小蒙's avatar
赵小蒙 committed
431
432
                continue
            else:
433
                raise Exception('drop_mode can not be null')
赵小蒙's avatar
赵小蒙 committed
434

435
436
        paras_of_layout = page_info.get('para_blocks')
        page_idx = page_info.get('page_idx')
赵小蒙's avatar
赵小蒙 committed
437
438
439
        if not paras_of_layout:
            continue
        if make_mode == MakeMode.MM_MD:
440
            page_markdown = ocr_mk_markdown_with_para_core_v2(
441
                paras_of_layout, 'mm', img_buket_path, parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
442
443
            output_content.extend(page_markdown)
        elif make_mode == MakeMode.NLP_MD:
444
            page_markdown = ocr_mk_markdown_with_para_core_v2(
445
                paras_of_layout, 'nlp', parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
446
447
448
            output_content.extend(page_markdown)
        elif make_mode == MakeMode.STANDARD_FORMAT:
            for para_block in paras_of_layout:
449
                if drop_reason_flag:
450
                    para_content = para_to_standard_format_v2(
451
                        para_block, img_buket_path, page_idx, parse_type=parse_type, lang=lang, drop_reason=drop_reason)
452
453
                else:
                    para_content = para_to_standard_format_v2(
454
                        para_block, img_buket_path, page_idx, parse_type=parse_type, lang=lang)
赵小蒙's avatar
赵小蒙 committed
455
456
457
458
459
                output_content.append(para_content)
    if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
        return '\n\n'.join(output_content)
    elif make_mode == MakeMode.STANDARD_FORMAT:
        return output_content