ocr_mkcontent.py 9.99 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
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_nlp_markdown(pdf_info_dict: list):
赵小蒙's avatar
赵小蒙 committed
23
24
    markdown = []

赵小蒙's avatar
赵小蒙 committed
25
    for page_info in pdf_info_dict:
赵小蒙's avatar
赵小蒙 committed
26
27
28
29
30
31
32
        blocks = page_info.get("preproc_blocks")
        if not blocks:
            continue
        for block in blocks:
            for line in block['lines']:
                line_text = ''
                for span in line['spans']:
33
34
                    if not span.get('content'):
                        continue
赵小蒙's avatar
赵小蒙 committed
35
                    content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
赵小蒙's avatar
赵小蒙 committed
36
                    if span['type'] == ContentType.InlineEquation:
赵小蒙's avatar
赵小蒙 committed
37
                        content = f"${content}$"
赵小蒙's avatar
赵小蒙 committed
38
                    elif span['type'] == ContentType.InterlineEquation:
赵小蒙's avatar
赵小蒙 committed
39
40
41
42
43
                        content = f"$$\n{content}\n$$"
                    line_text += content + ' '
                # 在行末添加两个空格以强制换行
                markdown.append(line_text.strip() + '  ')
    return '\n'.join(markdown)
44

赵小蒙's avatar
赵小蒙 committed
45

赵小蒙's avatar
赵小蒙 committed
46
def ocr_mk_mm_markdown(pdf_info_dict: list):
47
48
    markdown = []

赵小蒙's avatar
赵小蒙 committed
49
    for page_info in pdf_info_dict:
50
51
52
53
54
55
56
57
58
59
60
        blocks = page_info.get("preproc_blocks")
        if not blocks:
            continue
        for block in blocks:
            for line in block['lines']:
                line_text = ''
                for span in line['spans']:
                    if not span.get('content'):
                        if not span.get('image_path'):
                            continue
                        else:
61
                            content = f"![]({span['image_path']})"
62
                    else:
赵小蒙's avatar
赵小蒙 committed
63
                        content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
赵小蒙's avatar
赵小蒙 committed
64
                        if span['type'] == ContentType.InlineEquation:
65
                            content = f"${content}$"
赵小蒙's avatar
赵小蒙 committed
66
                        elif span['type'] == ContentType.InterlineEquation:
67
68
69
70
71
                            content = f"$$\n{content}\n$$"
                    line_text += content + ' '
                # 在行末添加两个空格以强制换行
                markdown.append(line_text.strip() + '  ')
    return '\n'.join(markdown)
72

73

赵小蒙's avatar
赵小蒙 committed
74
def ocr_mk_mm_markdown_with_para(pdf_info_list: list, img_buket_path):
xuchao's avatar
xuchao committed
75
    markdown = []
赵小蒙's avatar
赵小蒙 committed
76
    for page_info in pdf_info_list:
77
        paras_of_layout = page_info.get("para_blocks")
赵小蒙's avatar
赵小蒙 committed
78
        page_markdown = ocr_mk_markdown_with_para_core(paras_of_layout, "mm", img_buket_path)
79
        markdown.extend(page_markdown)
80
    return '\n\n'.join(markdown)
81
82


赵小蒙's avatar
赵小蒙 committed
83
def ocr_mk_nlp_markdown_with_para(pdf_info_dict: list):
84
    markdown = []
赵小蒙's avatar
赵小蒙 committed
85
    for page_info in pdf_info_dict:
86
        paras_of_layout = page_info.get("para_blocks")
87
        page_markdown = ocr_mk_markdown_with_para_core(paras_of_layout, "nlp")
88
89
90
        markdown.extend(page_markdown)
    return '\n\n'.join(markdown)

赵小蒙's avatar
赵小蒙 committed
91
def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: list):
92
    markdown_with_para_and_pagination = []
赵小蒙's avatar
赵小蒙 committed
93
94
    page_no = 0
    for page_info in pdf_info_dict:
95
96
        paras_of_layout = page_info.get("para_blocks")
        if not paras_of_layout:
97
            continue
98
        page_markdown = ocr_mk_markdown_with_para_core(paras_of_layout, "mm")
99
100
        markdown_with_para_and_pagination.append({
            'page_no': page_no,
101
            'md_content': '\n\n'.join(page_markdown)
102
        })
赵小蒙's avatar
赵小蒙 committed
103
        page_no += 1
104
105
106
    return markdown_with_para_and_pagination


赵小蒙's avatar
赵小蒙 committed
107
def ocr_mk_markdown_with_para_core(paras_of_layout, mode, img_buket_path):
108
109
110
111
112
113
114
    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')
115
                    content = ''
116
                    language = ''
117
                    if span_type == ContentType.Text:
118
119
120
121
122
123
                        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)
124
                    elif span_type == ContentType.InlineEquation:
125
                        content = f"${span['content']}$"
126
                    elif span_type == ContentType.InterlineEquation:
127
                        content = f"\n$$\n{span['content']}\n$$\n"
128
                    elif span_type in [ContentType.Image, ContentType.Table]:
129
                        if mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
130
                            content = f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
131
132
                        elif mode == 'nlp':
                            pass
133
                    if content != '':
134
135
136
137
                        if language == 'en':  # 英文语境下 content间需要空格分隔
                            para_text += content + ' '
                        else:  # 中文语境下,content间不需要空格分隔
                            para_text += content
138
139
140
141
            if para_text.strip() == '':
                continue
            else:
                page_markdown.append(para_text.strip() + '  ')
142
143
144
    return page_markdown


145
def para_to_standard_format(para, img_buket_path):
146
147
    para_content = {}
    if len(para) == 1:
148
        para_content = line_to_standard_format(para[0], img_buket_path)
149
150
151
152
153
    elif len(para) > 1:
        para_text = ''
        inline_equation_num = 0
        for line in para:
            for span in line['spans']:
154
                language = ''
155
                span_type = span.get('type')
赵小蒙's avatar
赵小蒙 committed
156
                content = ""
157
                if span_type == ContentType.Text:
158
159
160
161
162
163
                    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)
164
                elif span_type == ContentType.InlineEquation:
165
                    content = f"${span['content']}$"
166
                    inline_equation_num += 1
167
168
169
170
171

                if language == 'en':  # 英文语境下 content间需要空格分隔
                    para_text += content + ' '
                else:  # 中文语境下,content间不需要空格分隔
                    para_text += content
172
173
174
175
176
177
178
        para_content = {
            'type': 'text',
            'text': para_text,
            'inline_equation_num': inline_equation_num
        }
    return para_content

赵小蒙's avatar
赵小蒙 committed
179
def make_standard_format_with_para(pdf_info_dict: list, img_buket_path: str):
赵小蒙's avatar
赵小蒙 committed
180
    content_list = []
赵小蒙's avatar
赵小蒙 committed
181
    for page_info in pdf_info_dict:
182
183
        paras_of_layout = page_info.get("para_blocks")
        if not paras_of_layout:
赵小蒙's avatar
赵小蒙 committed
184
            continue
185
186
        for paras in paras_of_layout:
            for para in paras:
187
                para_content = para_to_standard_format(para, img_buket_path)
188
                content_list.append(para_content)
赵小蒙's avatar
赵小蒙 committed
189
190
191
    return content_list


192
def line_to_standard_format(line, img_buket_path):
赵小蒙's avatar
赵小蒙 committed
193
194
195
196
197
198
199
200
201
202
    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',
203
                        'img_path': join_path(img_buket_path, span['image_path'])
赵小蒙's avatar
赵小蒙 committed
204
205
206
207
208
                    }
                    return content
                elif span['type'] == ContentType.Table:
                    content = {
                        'type': 'table',
209
                        'img_path': join_path(img_buket_path, span['image_path'])
赵小蒙's avatar
赵小蒙 committed
210
211
212
213
                    }
                    return content
        else:
            if span['type'] == ContentType.InterlineEquation:
赵小蒙's avatar
赵小蒙 committed
214
                interline_equation = span['content']
赵小蒙's avatar
赵小蒙 committed
215
216
217
218
219
220
                content = {
                    'type': 'equation',
                    'latex': f"$$\n{interline_equation}\n$$"
                }
                return content
            elif span['type'] == ContentType.InlineEquation:
赵小蒙's avatar
赵小蒙 committed
221
                inline_equation = span['content']
赵小蒙's avatar
赵小蒙 committed
222
223
224
                line_text += f"${inline_equation}$"
                inline_equation_num += 1
            elif span['type'] == ContentType.Text:
225
226
                text_content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                line_text += text_content
赵小蒙's avatar
赵小蒙 committed
227
228
229
230
231
232
233
234
    content = {
        'type': 'text',
        'text': line_text,
        'inline_equation_num': inline_equation_num
    }
    return content


赵小蒙's avatar
赵小蒙 committed
235
def ocr_mk_mm_standard_format(pdf_info_dict: list):
赵小蒙's avatar
update  
赵小蒙 committed
236
    """
237
    content_list
赵小蒙's avatar
赵小蒙 committed
238
239
240
241
242
    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
243
    """
赵小蒙's avatar
赵小蒙 committed
244
    content_list = []
赵小蒙's avatar
赵小蒙 committed
245
    for page_info in pdf_info_dict:
赵小蒙's avatar
赵小蒙 committed
246
247
248
249
250
251
252
253
        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