ocr_mkcontent.py 12.4 KB
Newer Older
1
2
import re

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

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


13
def __is_hyphen_at_line_end(line):
14
15
    """Check if a line ends with one or more letters followed by a hyphen.

16
17
    Args:
    line (str): The line of text to check.
18

19
20
21
22
23
24
25
    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))


26
27
def ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict: list,
                                                img_buket_path):
28
    markdown_with_para_and_pagination = []
赵小蒙's avatar
赵小蒙 committed
29
30
    page_no = 0
    for page_info in pdf_info_dict:
31
        paras_of_layout = page_info.get('para_blocks')
32
        if not paras_of_layout:
33
34
35
36
37
38
39
            markdown_with_para_and_pagination.append({
                'page_no':
                    page_no,
                'md_content':
                    '',
            })
            page_no += 1
40
            continue
41
42
        page_markdown = ocr_mk_markdown_with_para_core_v2(
            paras_of_layout, 'mm', img_buket_path)
43
        markdown_with_para_and_pagination.append({
44
            'page_no':
45
                page_no,
46
            'md_content':
47
                '\n\n'.join(page_markdown)
48
        })
赵小蒙's avatar
赵小蒙 committed
49
        page_no += 1
50
51
52
    return markdown_with_para_and_pagination


53
54
def ocr_mk_markdown_with_para_core_v2(paras_of_layout,
                                      mode,
55
56
                                      img_buket_path='',
                                      ):
赵小蒙's avatar
赵小蒙 committed
57
    page_markdown = []
58
    for para_block in paras_of_layout:
赵小蒙's avatar
赵小蒙 committed
59
        para_text = ''
赵小蒙's avatar
赵小蒙 committed
60
        para_type = para_block['type']
61
        if para_type in [BlockType.Text, BlockType.List, BlockType.Index]:
62
            para_text = merge_para_with_text(para_block)
63
        elif para_type == BlockType.Title:
64
            para_text = f'# {merge_para_with_text(para_block)}'
65
        elif para_type == BlockType.InterlineEquation:
66
            para_text = merge_para_with_text(para_block)
67
68
        elif para_type == BlockType.Image:
            if mode == 'nlp':
赵小蒙's avatar
赵小蒙 committed
69
                continue
70
            elif mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
71
                for block in para_block['blocks']:  # 1st.拼image_body
赵小蒙's avatar
赵小蒙 committed
72
73
                    if block['type'] == BlockType.ImageBody:
                        for line in block['lines']:
74
                            for span in line['spans']:
赵小蒙's avatar
赵小蒙 committed
75
                                if span['type'] == ContentType.Image:
76
77
                                    if span.get('image_path', ''):
                                        para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})  \n"
赵小蒙's avatar
赵小蒙 committed
78
                for block in para_block['blocks']:  # 2nd.拼image_caption
赵小蒙's avatar
赵小蒙 committed
79
                    if block['type'] == BlockType.ImageCaption:
80
                        para_text += merge_para_with_text(block) + '  \n'
81
                for block in para_block['blocks']:  # 3rd.拼image_footnote
82
                    if block['type'] == BlockType.ImageFootnote:
83
                        para_text += merge_para_with_text(block) + '  \n'
84
85
86
87
        elif para_type == BlockType.Table:
            if mode == 'nlp':
                continue
            elif mode == 'mm':
赵小蒙's avatar
赵小蒙 committed
88
89
                for block in para_block['blocks']:  # 1st.拼table_caption
                    if block['type'] == BlockType.TableCaption:
90
                        para_text += merge_para_with_text(block) + '  \n'
赵小蒙's avatar
赵小蒙 committed
91
                for block in para_block['blocks']:  # 2nd.拼table_body
赵小蒙's avatar
赵小蒙 committed
92
93
                    if block['type'] == BlockType.TableBody:
                        for line in block['lines']:
94
                            for span in line['spans']:
赵小蒙's avatar
赵小蒙 committed
95
                                if span['type'] == ContentType.Table:
96
97
98
                                    # if processed by table model
                                    if span.get('latex', ''):
                                        para_text += f"\n\n$\n {span['latex']}\n$\n\n"
99
100
                                    elif span.get('html', ''):
                                        para_text += f"\n\n{span['html']}\n\n"
101
                                    elif span.get('image_path', ''):
102
                                        para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})  \n"
赵小蒙's avatar
赵小蒙 committed
103
104
                for block in para_block['blocks']:  # 3rd.拼table_footnote
                    if block['type'] == BlockType.TableFootnote:
105
                        para_text += merge_para_with_text(block) + '  \n'
106
107
108
109
110

        if para_text.strip() == '':
            continue
        else:
            page_markdown.append(para_text.strip() + '  ')
赵小蒙's avatar
赵小蒙 committed
111
112
113
114

    return page_markdown


115
116
117
118
119
120
121
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'
122
        else:
123
124
125
126
            return 'unknown'
    else:
        return 'empty'

127

128
def merge_para_with_text(para_block):
129
130
131
132
133
134
135
    block_text = ''
    for line in para_block['lines']:
        for span in line['spans']:
            if span['type'] in [ContentType.Text]:
                block_text += span['content']
    block_lang = detect_lang(block_text)

赵小蒙's avatar
赵小蒙 committed
136
    para_text = ''
137
138
139
140
141
    for i, line in enumerate(para_block['lines']):

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

142
        for j, span in enumerate(line['spans']):
143

赵小蒙's avatar
赵小蒙 committed
144
            span_type = span['type']
赵小蒙's avatar
赵小蒙 committed
145
146
            content = ''
            if span_type == ContentType.Text:
147
                content = ocr_escape_special_markdown_char(span['content'])
赵小蒙's avatar
赵小蒙 committed
148
            elif span_type == ContentType.InlineEquation:
149
                content = f"${span['content']}$"
赵小蒙's avatar
赵小蒙 committed
150
151
            elif span_type == ContentType.InterlineEquation:
                content = f"\n$$\n{span['content']}\n$$\n"
152

153
            content = content.strip()
154
155
156

            if content:
                langs = ['zh', 'ja', 'ko']
157
                # logger.info(f'block_lang: {block_lang}, content: {content}')
158
159
                if block_lang in langs: # 中文/日语/韩文语境下,换行不需要空格分隔,但是如果是行内公式结尾,还是要加空格
                    if j == len(line['spans']) - 1 and span_type not in [ContentType.InlineEquation]:
160
161
162
                        para_text += content
                    else:
                        para_text += f'{content} '
163
164
165
                else:
                    if span_type in [ContentType.Text, ContentType.InlineEquation]:
                        # 如果span是line的最后一个且末尾带有-连字符,那么末尾不应该加空格,同时应该把-删除
166
                        if j == len(line['spans'])-1 and span_type == ContentType.Text and __is_hyphen_at_line_end(content):
167
168
169
170
171
                            para_text += content[:-1]
                        else:  # 西方文本语境下 content间需要空格分隔
                            para_text += f'{content} '
                    elif span_type == ContentType.InterlineEquation:
                        para_text += content
172
173
            else:
                continue
174
    # 连写字符拆分
175
    # para_text = __replace_ligatures(para_text)
176

赵小蒙's avatar
赵小蒙 committed
177
178
179
    return para_text


180
def para_to_standard_format_v2(para_block, img_buket_path, page_idx, drop_reason=None):
赵小蒙's avatar
赵小蒙 committed
181
    para_type = para_block['type']
182
    para_content = {}
183
    if para_type in [BlockType.Text, BlockType.List, BlockType.Index]:
赵小蒙's avatar
赵小蒙 committed
184
185
        para_content = {
            'type': 'text',
186
            'text': merge_para_with_text(para_block),
赵小蒙's avatar
赵小蒙 committed
187
188
189
190
        }
    elif para_type == BlockType.Title:
        para_content = {
            'type': 'text',
191
            'text': merge_para_with_text(para_block),
192
            'text_level': 1,
赵小蒙's avatar
赵小蒙 committed
193
194
195
196
        }
    elif para_type == BlockType.InterlineEquation:
        para_content = {
            'type': 'equation',
197
            'text': merge_para_with_text(para_block),
198
            'text_format': 'latex',
赵小蒙's avatar
赵小蒙 committed
199
200
        }
    elif para_type == BlockType.Image:
201
        para_content = {'type': 'image', 'img_path': '', 'img_caption': [], 'img_footnote': []}
赵小蒙's avatar
赵小蒙 committed
202
203
        for block in para_block['blocks']:
            if block['type'] == BlockType.ImageBody:
204
205
206
207
208
                for line in block['lines']:
                    for span in line['spans']:
                        if span['type'] == ContentType.Image:
                            if span.get('image_path', ''):
                                para_content['img_path'] = join_path(img_buket_path, span['image_path'])
赵小蒙's avatar
赵小蒙 committed
209
            if block['type'] == BlockType.ImageCaption:
210
                para_content['img_caption'].append(merge_para_with_text(block))
211
            if block['type'] == BlockType.ImageFootnote:
212
                para_content['img_footnote'].append(merge_para_with_text(block))
赵小蒙's avatar
赵小蒙 committed
213
    elif para_type == BlockType.Table:
214
        para_content = {'type': 'table', 'img_path': '', 'table_caption': [], 'table_footnote': []}
赵小蒙's avatar
赵小蒙 committed
215
216
        for block in para_block['blocks']:
            if block['type'] == BlockType.TableBody:
217
218
219
220
221
222
223
224
225
226
227
228
                for line in block['lines']:
                    for span in line['spans']:
                        if span['type'] == ContentType.Table:

                            if span.get('latex', ''):
                                para_content['table_body'] = f"\n\n$\n {span['latex']}\n$\n\n"
                            elif span.get('html', ''):
                                para_content['table_body'] = f"\n\n{span['html']}\n\n"

                            if span.get('image_path', ''):
                                para_content['img_path'] = join_path(img_buket_path, span['image_path'])

赵小蒙's avatar
赵小蒙 committed
229
            if block['type'] == BlockType.TableCaption:
230
                para_content['table_caption'].append(merge_para_with_text(block))
赵小蒙's avatar
赵小蒙 committed
231
            if block['type'] == BlockType.TableFootnote:
232
                para_content['table_footnote'].append(merge_para_with_text(block))
赵小蒙's avatar
赵小蒙 committed
233

234
235
236
237
238
    para_content['page_idx'] = page_idx

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

赵小蒙's avatar
赵小蒙 committed
239
240
241
    return para_content


242
243
244
def union_make(pdf_info_dict: list,
               make_mode: str,
               drop_mode: str,
245
               img_buket_path: str = '',
246
               ):
赵小蒙's avatar
赵小蒙 committed
247
248
    output_content = []
    for page_info in pdf_info_dict:
249
250
        drop_reason_flag = False
        drop_reason = None
251
        if page_info.get('need_drop', False):
252
253
            drop_reason = page_info.get('drop_reason')
            if drop_mode == DropMode.NONE:
赵小蒙's avatar
赵小蒙 committed
254
                pass
255
256
            elif drop_mode == DropMode.NONE_WITH_REASON:
                drop_reason_flag = True
赵小蒙's avatar
赵小蒙 committed
257
            elif drop_mode == DropMode.WHOLE_PDF:
258
259
                raise Exception((f'drop_mode is {DropMode.WHOLE_PDF} ,'
                                 f'drop_reason is {drop_reason}'))
赵小蒙's avatar
赵小蒙 committed
260
            elif drop_mode == DropMode.SINGLE_PAGE:
261
262
                logger.warning((f'drop_mode is {DropMode.SINGLE_PAGE} ,'
                                f'drop_reason is {drop_reason}'))
赵小蒙's avatar
赵小蒙 committed
263
264
                continue
            else:
265
                raise Exception('drop_mode can not be null')
赵小蒙's avatar
赵小蒙 committed
266

267
268
        paras_of_layout = page_info.get('para_blocks')
        page_idx = page_info.get('page_idx')
赵小蒙's avatar
赵小蒙 committed
269
270
271
        if not paras_of_layout:
            continue
        if make_mode == MakeMode.MM_MD:
272
            page_markdown = ocr_mk_markdown_with_para_core_v2(
273
                paras_of_layout, 'mm', img_buket_path)
赵小蒙's avatar
赵小蒙 committed
274
275
            output_content.extend(page_markdown)
        elif make_mode == MakeMode.NLP_MD:
276
            page_markdown = ocr_mk_markdown_with_para_core_v2(
277
                paras_of_layout, 'nlp')
赵小蒙's avatar
赵小蒙 committed
278
279
280
            output_content.extend(page_markdown)
        elif make_mode == MakeMode.STANDARD_FORMAT:
            for para_block in paras_of_layout:
281
                if drop_reason_flag:
282
                    para_content = para_to_standard_format_v2(
283
                        para_block, img_buket_path, page_idx)
284
285
                else:
                    para_content = para_to_standard_format_v2(
286
                        para_block, img_buket_path, page_idx)
赵小蒙's avatar
赵小蒙 committed
287
288
289
290
291
                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