ocr_mkcontent.py 3.43 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
from magic_pdf.libs.commons import s3_image_save_path, join_path
赵小蒙's avatar
赵小蒙 committed
2
from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
赵小蒙's avatar
赵小蒙 committed
3
4
5
from magic_pdf.libs.ocr_content_type import ContentType


赵小蒙's avatar
赵小蒙 committed
6
def ocr_mk_nlp_markdown(pdf_info_dict: dict):
赵小蒙's avatar
赵小蒙 committed
7
8
9
10
11
12
13
14
15
16
    markdown = []

    for _, page_info in pdf_info_dict.items():
        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']:
17
18
                    if not span.get('content'):
                        continue
赵小蒙's avatar
赵小蒙 committed
19
                    content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
赵小蒙's avatar
赵小蒙 committed
20
                    if span['type'] == ContentType.InlineEquation:
赵小蒙's avatar
赵小蒙 committed
21
                        content = f"${content}$"
赵小蒙's avatar
赵小蒙 committed
22
                    elif span['type'] == ContentType.InterlineEquation:
赵小蒙's avatar
赵小蒙 committed
23
24
25
26
27
                        content = f"$$\n{content}\n$$"
                    line_text += content + ' '
                # 在行末添加两个空格以强制换行
                markdown.append(line_text.strip() + '  ')
    return '\n'.join(markdown)
28

赵小蒙's avatar
赵小蒙 committed
29

赵小蒙's avatar
赵小蒙 committed
30
def ocr_mk_mm_markdown(pdf_info_dict: dict):
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

    markdown = []

    for _, page_info in pdf_info_dict.items():
        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:
赵小蒙's avatar
赵小蒙 committed
46
                            content = f"![]({join_path(s3_image_save_path, span['image_path'])})"
47
                    else:
赵小蒙's avatar
赵小蒙 committed
48
                        content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
赵小蒙's avatar
赵小蒙 committed
49
                        if span['type'] == ContentType.InlineEquation:
50
                            content = f"${content}$"
赵小蒙's avatar
赵小蒙 committed
51
                        elif span['type'] == ContentType.InterlineEquation:
52
53
54
55
56
                            content = f"$$\n{content}\n$$"
                    line_text += content + ' '
                # 在行末添加两个空格以强制换行
                markdown.append(line_text.strip() + '  ')
    return '\n'.join(markdown)
57

58

xuchao's avatar
xuchao committed
59
60
61
62
63
64
def mk_mm_markdown2(pdf_info_dict:dict):
    markdown = []
    for _, page_info in pdf_info_dict.items():
        paras = page_info.get("para_blocks")
        if not paras:
            continue
65
        for para in paras:
xuchao's avatar
xuchao committed
66
67
68
69
            para_text = ''
            for line in para:
                for span in line['spans']:
                    span_type = span.get('type')
70
                    if span_type == ContentType.Text:
xuchao's avatar
xuchao committed
71
                        para_text += span['content']
72
                    elif span_type == ContentType.InlineEquation:
xuchao's avatar
xuchao committed
73
                        para_text += f" ${span['content']}$ "
74
                    elif span_type == ContentType.InterlineEquation:
xuchao's avatar
xuchao committed
75
                        para_text += f"$$\n{span['content']}\n$$ "
76
                    elif span_type == ContentType.Image:
赵小蒙's avatar
赵小蒙 committed
77
                        para_text += f"![]({join_path(s3_image_save_path, span['image_path'])})"
xuchao's avatar
xuchao committed
78
            markdown.append(para_text)
79

xuchao's avatar
xuchao committed
80
    return '\n\n'.join(markdown)
81
82


83
84
85
86
87
88
89
def ocr_mk_mm_standard_format():
    '''
    content_list
    type string image/text/table/equation(行间的单独拿出来,行内的和text合并)

    '''
    pass