join_path逻辑修复

# Conflicts:
#	magic_pdf/dict2md/ocr_mkcontent.py

.gitignore

@@ -31,5 +31,6 @@ tmp
 .vscode
 .vscode/
 /tests/
+ocr_demo
 
 /app/common/__init__.py

demo/ocr_demo.py

@@ -4,7 +4,7 @@ import os
 from loguru import logger
 from pathlib import Path
 
-from magic_pdf.dict2md.ocr_mkcontent import mk_nlp_markdown, mk_mm_markdown
+from magic_pdf.dict2md.ocr_mkcontent import mk_mm_markdown2, mk_nlp_markdown, mk_mm_markdown
 from magic_pdf.libs.commons import join_path
 from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
 
@@ -30,12 +30,13 @@ def read_json_file(file_path):
 
 
 if __name__ == '__main__':
-    # ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.pdf"
-    # ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.json"
+    ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.pdf"
+    ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.json"
     # ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.pdf"
     # ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.json"
-    ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_demo\ocr_1_org.pdf"
-    ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\demo_4\ocr_demo\ocr_1.json"
+    
+    ocr_pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.pdf"
+    ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.json"
     try:
         ocr_pdf_model_info = read_json_file(ocr_json_file_path)
         pth = Path(ocr_json_file_path)
@@ -57,7 +58,7 @@ if __name__ == '__main__':
             os.makedirs(parent_dir)
 
         # markdown_content = mk_nlp_markdown(pdf_info_dict)
-        markdown_content = mk_mm_markdown(pdf_info_dict)
+        markdown_content = mk_mm_markdown2(pdf_info_dict)
 
         with open(text_content_save_path, "w", encoding="utf-8") as f:
             f.write(markdown_content)

magic_pdf/dict2md/ocr_mkcontent.py

@@ -1,3 +1,4 @@
+from magic_pdf.libs.markdown_utils import ocr_escape_special_markdown_char
 from magic_pdf.libs.ocr_content_type import ContentType
 
 
@@ -14,7 +15,7 @@ def ocr_mk_nlp_markdown(pdf_info_dict: dict):
                 for span in line['spans']:
                     if not span.get('content'):
                         continue
-                    content = span['content'].replace('$', '\$')  # 转义$
+                    content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                     if span['type'] == ContentType.InlineEquation:
                         content = f"${content}$"
                     elif span['type'] == ContentType.InterlineEquation:
@@ -41,9 +42,9 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
                         if not span.get('image_path'):
                             continue
                         else:
-                            content = f"![]({span['image_path']})"
+                            content = f"![](s3://mllm-raw-media/pdf2md_img/{span['image_path']})"
                     else:
-                        content = span['content'].replace('$', '\$')  # 转义$
+                        content = ocr_escape_special_markdown_char(span['content'])  # 转义特殊符号
                         if span['type'] == ContentType.InlineEquation:
                             content = f"${content}$"
                         elif span['type'] == ContentType.InterlineEquation:
@@ -52,3 +53,36 @@ def ocr_mk_mm_markdown(pdf_info_dict: dict):
                 # 在行末添加两个空格以强制换行
                 markdown.append(line_text.strip() + '  ')
     return '\n'.join(markdown)
+
+
+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
+        for para in paras:
+            para_text = ''
+            for line in para:
+                for span in line['spans']:
+                    span_type = span.get('type')
+                    if span_type == ContentType.Text:
+                        para_text += span['content']
+                    elif span_type == ContentType.InlineEquation:
+                        para_text += f" ${span['content']}$ "
+                    elif span_type == ContentType.InterlineEquation:
+                        para_text += f"$$\n{span['content']}\n$$ "
+                    elif span_type == ContentType.Image:
+                        para_text += f"![](s3://mllm-raw-media/pdf2md_img/{span['image_path']}) "
+            markdown.append(para_text)
+
+    return '\n\n'.join(markdown)
+
+
+def ocr_mk_mm_standard_format():
+    '''
+    content_list
+    type string image/text/table/equation(行间的单独拿出来，行内的和text合并)
+
+    '''
+    pass

magic_pdf/libs/commons.py

@@ -16,7 +16,8 @@ def get_delta_time(input_time):
 
 
 def join_path(*args):
-    return '/'.join(s.rstrip('/') for s in args)
+    return '/'.join(str(s).rstrip('/') for s in args)
+
 
 #配置全局的errlog_path，方便demo同步引用
 error_log_path = "s3://llm-pdf-text/err_logs/"

magic_pdf/libs/markdown_utils.py

@@ -18,3 +18,14 @@ def escape_special_markdown_char(pymu_blocks):
                         span['text'] = span['text'].replace(char, "\\" + char)
 
     return pymu_blocks
+
+
+def ocr_escape_special_markdown_char(content):
+    """
+    转义正文里对markdown语法有特殊意义的字符
+    """
+    special_chars = ["*", "`", "~", "$"]
+    for char in special_chars:
+        content = content.replace(char, "\\" + char)
+
+    return content

magic_pdf/para/para_split.py

@@ -0,0 +1,206 @@
+from sklearn.cluster import DBSCAN
+import numpy as np
+from loguru import logger
+
+from magic_pdf.libs.boxbase import _is_in
+
+
+LINE_STOP_FLAG = ['.', '!', '?', '。', '！', '？',"：", ":", ")", "）", ";"]
+INLINE_EQUATION = 'inline_equation'
+INTER_EQUATION = "displayed_equation"
+TEXT = "text"
+
+def __add_line_period(blocks, layout_bboxes):
+    """
+    为每行添加句号
+    如果这个行
+    1. 以行内公式结尾，但没有任何标点符号,此时加个句号，认为他就是段落结尾。
+    """
+    for block in blocks:
+        for line in block['lines']:
+            last_span = line['spans'][-1]
+            span_type = last_span['type']
+            if span_type in [TEXT, INLINE_EQUATION]:
+                span_content = last_span['content'].strip()
+                if span_type==INLINE_EQUATION and span_content[-1] not in LINE_STOP_FLAG:
+                    if span_type in [INLINE_EQUATION, INTER_EQUATION]:
+                        last_span['content'] = span_content + '.'
+
+
+
+def __valign_lines(blocks, layout_bboxes):
+    """
+    对齐行的左侧和右侧。
+    扫描行的左侧和右侧，如果x0, x1差距不超过3就强行对齐到所处layout的左右两侧（和layout有一段距离）。
+    3是个经验值，TODO，计算得来
+    
+    """
+    
+    min_distance = 3
+    min_sample = 2
+    
+    for layout_box in layout_bboxes:
+        blocks_in_layoutbox = [b for b in blocks if _is_in(b['bbox'], layout_box['layout_bbox'])]
+        if len(blocks_in_layoutbox)==0:
+            continue
+        
+        x0_lst = np.array([[line['bbox'][0], 0] for block in blocks_in_layoutbox for line in block['lines']])
+        x1_lst = np.array([[line['bbox'][2], 0] for block in blocks_in_layoutbox for line in block['lines']])
+        x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
+        x1_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x1_lst)
+        x0_uniq_label = np.unique(x0_clusters.labels_)
+        x1_uniq_label = np.unique(x1_clusters.labels_)
+        
+        x0_2_new_val = {} # 存储旧值对应的新值映射
+        x1_2_new_val = {}
+        for label in x0_uniq_label:
+            if label==-1:
+                continue
+            x0_index_of_label = np.where(x0_clusters.labels_==label)
+            x0_raw_val = x0_lst[x0_index_of_label][:,0]
+            x0_new_val = np.min(x0_lst[x0_index_of_label][:,0])
+            x0_2_new_val.update({idx: x0_new_val for idx in x0_raw_val})
+        for label in x1_uniq_label:
+            if label==-1:
+                continue
+            x1_index_of_label = np.where(x1_clusters.labels_==label)
+            x1_raw_val = x1_lst[x1_index_of_label][:,0]
+            x1_new_val = np.max(x1_lst[x1_index_of_label][:,0])
+            x1_2_new_val.update({idx: x1_new_val for idx in x1_raw_val})
+        
+        for block in blocks_in_layoutbox:
+            for line in block['lines']:
+                x0, x1 = line['bbox'][0], line['bbox'][2]
+                if x0 in x0_2_new_val:
+                    line['bbox'][0] = int(x0_2_new_val[x0])
+
+                if x1 in x1_2_new_val:
+                    line['bbox'][2] = int(x1_2_new_val[x1])
+            # 其余对不齐的保持不动
+            
+        # 由于修改了block里的line长度，现在需要重新计算block的bbox
+        for block in blocks_in_layoutbox:
+            block['bbox'] = [min([line['bbox'][0] for line in block['lines']]), 
+                            min([line['bbox'][1] for line in block['lines']]), 
+                            max([line['bbox'][2] for line in block['lines']]), 
+                            max([line['bbox'][3] for line in block['lines']])]
+
+
+def __common_pre_proc(blocks, layout_bboxes):
+    """
+    不分语言的，对文本进行预处理
+    """
+    __add_line_period(blocks, layout_bboxes)
+    __valign_lines(blocks, layout_bboxes)
+    
+
+def __pre_proc_zh_blocks(blocks, layout_bboxes):
+    """
+    对中文文本进行分段预处理
+    """
+    pass
+
+
+def __pre_proc_en_blocks(blocks, layout_bboxes):
+    """
+    对英文文本进行分段预处理
+    """
+    pass
+
+
+def __group_line_by_layout(blocks, layout_bboxes, lang="en"):
+    """
+    每个layout内的行进行聚合
+    """
+    # 因为只是一个block一行目前, 一个block就是一个段落
+    lines_group = []
+    
+    for lyout in layout_bboxes:
+        lines = [line for block in blocks if _is_in(block['bbox'], lyout['layout_bbox']) for line in block['lines']]
+        lines_group.append(lines)
+
+    return lines_group
+    
+
+def __split_para_in_layoutbox(lines_group, layout_bboxes, lang="en", char_avg_len=10):
+    """
+    lines_group 进行行分段——layout内部进行分段。
+    1. 先计算每个group的左右边界。
+    2. 然后根据行末尾特征进行分段。
+        末尾特征：以句号等结束符结尾。并且距离右侧边界有一定距离。
+    
+    """
+    def get_span_text(span):
+        c = span.get('content', '')
+        if len(c)==0:
+            c = span.get('image-path', '')
+            
+        return c
+    
+    paras = []
+    right_tail_distance = 1.5 * char_avg_len
+    for lines in lines_group:
+        if len(lines)==0:
+            continue
+        layout_right = max([line['bbox'][2] for line in lines])
+        para = [] # 元素是line
+        for line in lines:
+            line_text = ''.join([get_span_text(span) for span in line['spans']])
+            #logger.info(line_text)
+            last_span_type = line['spans'][-1]['type']
+            if last_span_type in [TEXT, INLINE_EQUATION]:
+                last_char = line['spans'][-1]['content'][-1]
+                if last_char in LINE_STOP_FLAG or line['bbox'][2] < layout_right - right_tail_distance:
+                    para.append(line)
+                    paras.append(para)
+                    # para_text = ''.join([span['content'] for line in para for span in line['spans']])
+                    # logger.info(para_text)
+                    para = []
+                else: 
+                    para.append(line)
+            else: # 其他，图片、表格、行间公式，各自占一段
+                para.append(line)
+                paras.append(para)
+                # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
+                # logger.info(para_text)
+                para = []
+        if len(para)>0:
+            paras.append(para)
+            # para_text = ''.join([get_span_text(span) for line in para for span in line['spans']])
+            # logger.info(para_text)
+            para = []
+                    
+    return paras
+            
+
+def __do_split(blocks, layout_bboxes, lang="en"):
+    """
+    根据line和layout情况进行分段
+    先实现一个根据行末尾特征分段的简单方法。
+    """
+    """
+    算法思路：
+    1. 扫描layout里每一行，找出来行尾距离layout有边界有一定距离的行。
+    2. 从上述行中找到末尾是句号等可作为断行标志的行。
+    3. 参照上述行尾特征进行分段。
+    4. 图、表，目前独占一行，不考虑分段。
+    """
+    lines_group = __group_line_by_layout(blocks, layout_bboxes, lang) # block内分段
+    layout_paras = __split_para_in_layoutbox(lines_group, layout_bboxes, lang) # block间连接分段
+    
+    return layout_paras
+    
+    
+def para_split(blocks, layout_bboxes, lang="en"):
+    """
+    根据line和layout情况进行分段
+    """
+    __common_pre_proc(blocks, layout_bboxes)
+    if lang=='en':
+        __do_split(blocks, layout_bboxes, lang)
+    elif lang=='zh':
+        __do_split(blocks, layout_bboxes, lang)
+    
+    splited_blocks = __do_split(blocks, layout_bboxes, lang)
+    
+    return splited_blocks

magic_pdf/pdf_parse_by_model.py

@@ -257,7 +257,6 @@ def parse_pdf_by_model(
         footnote_bboxes_by_model = parse_footnotes_by_model(page_id, page, model_output_json, md_bookname_save_path, debug_mode=debug_mode)
         # 通过规则识别到的footnote
         footnote_bboxes_by_rule = parse_footnotes_by_rule(remain_text_blocks, page_height, page_id, main_text_font)
-
         """进入pdf过滤器，去掉一些不合理的pdf"""
         is_good_pdf, err = pdf_filter(page, remain_text_blocks, table_bboxes, image_bboxes)
         if not is_good_pdf:

magic_pdf/pdf_parse_by_ocr.py

@@ -16,6 +16,7 @@ from magic_pdf.libs.commons import (
 from magic_pdf.libs.coordinate_transform import get_scale_ratio
 from magic_pdf.libs.ocr_content_type import ContentType
 from magic_pdf.libs.safe_filename import sanitize_filename
+from magic_pdf.para.para_split import para_split
 from magic_pdf.pre_proc.detect_footer_by_model import parse_footers
 from magic_pdf.pre_proc.detect_footnote import parse_footnotes_by_model
 from magic_pdf.pre_proc.detect_header import parse_headers
@@ -31,12 +32,13 @@ from magic_pdf.pre_proc.ocr_span_list_modify import remove_spans_by_bboxes, remo
 from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox
 
 
-def construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
+def construct_page_component(blocks, para_blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
                              images, tables, interline_equations, inline_equations,
                              dropped_text_block, dropped_image_block, dropped_table_block,
                              need_remove_spans_bboxes_dict):
     return_dict = {
         'preproc_blocks': blocks,
+        "para_blocks": para_blocks, # 分好段落的blocks
         'layout_bboxes': layout_bboxes,
         'page_idx': page_id,
         'page_size': [page_w, page_h],
@@ -156,6 +158,9 @@ def parse_pdf_by_ocr(
                     int(x1 / horizontal_scale_ratio),
                     int(y1 / vertical_scale_ratio),
                 ]
+                # 删除高度或者宽度为0的spans
+                if bbox[2] - bbox[0] == 0 or bbox[3] - bbox[1] == 0:
+                    continue
                 """要删除的"""
                 #  3: 'header',      # 页眉
                 #  4: 'page number', # 页码
@@ -193,6 +198,9 @@ def parse_pdf_by_ocr(
             else:
                 continue
 
+
+
+
         # 删除重叠spans中较小的那些
         spans = remove_overlaps_min_spans(spans)
 
@@ -202,7 +210,7 @@ def parse_pdf_by_ocr(
         spans, dropped_text_block, dropped_image_block, dropped_table_block = remove_spans_by_bboxes_dict(spans, need_remove_spans_bboxes_dict)
 
         # 对image和table截图
-        spans = cut_image_and_table(spans, page, page_id, book_name, save_path)
+        spans = cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client)
 
         # 行内公式调整, 高度调整至与同行文字高度一致(优先左侧, 其次右侧)
         displayed_list = []
@@ -228,13 +236,13 @@ def parse_pdf_by_ocr(
         blocks = merge_lines_to_block(lines)
 
         # 根据block合并段落
-
-
+        para_blocks = para_split(blocks, layout_bboxes)
+        
         # 获取QA需要外置的list
         images, tables, interline_equations, inline_equations = get_qa_need_list(blocks)
 
         # 构造pdf_info_dict
-        page_info = construct_page_component(blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
+        page_info = construct_page_component(blocks, para_blocks, layout_bboxes, page_id, page_w, page_h, layout_tree,
                                              images, tables, interline_equations, inline_equations,
                                              dropped_text_block, dropped_image_block, dropped_table_block,
                                              need_remove_spans_bboxes_dict)

magic_pdf/pre_proc/ocr_cut_image.py

@@ -3,7 +3,7 @@ from magic_pdf.libs.ocr_content_type import ContentType
 from magic_pdf.libs.pdf_image_tools import cut_image
 
 
-def cut_image_and_table(spans, page, page_id, book_name, save_path):
+def cut_image_and_table(spans, page, page_id, book_name, save_path, img_s3_client):
     def s3_return_path(type):
         return join_path(book_name, type)
 
@@ -13,8 +13,8 @@ def cut_image_and_table(spans, page, page_id, book_name, save_path):
     for span in spans:
         span_type = span['type']
         if span_type == ContentType.Image:
-            span['image_path'] = cut_image(span['bbox'], page_id, page, img_save_path('images'))
+            span['image_path'] = cut_image(span['bbox'], page_id, page, img_save_path('images'), s3_return_path=s3_return_path('images'), img_s3_client=img_s3_client)
         elif span_type == ContentType.Table:
-            span['image_path'] = cut_image(span['bbox'], page_id, page, img_save_path('tables'))
+            span['image_path'] = cut_image(span['bbox'], page_id, page, img_save_path('tables'), s3_return_path=s3_return_path('tables'), img_s3_client=img_s3_client)
 
     return spans