Merge branch 'master' of github.com:papayalove/Magic-PDF

magic_pdf/cli/magicpdf.py

@@ -26,6 +26,8 @@ import json as json_parse
 from datetime import datetime
 import click
 from magic_pdf.pipe.UNIPipe import UNIPipe
+from magic_pdf.pipe.OCRPipe import OCRPipe
+from magic_pdf.pipe.TXTPipe import TXTPipe
 from magic_pdf.libs.config_reader import get_s3_config
 from magic_pdf.libs.path_utils import (
     parse_s3path,
@@ -33,9 +35,9 @@ from magic_pdf.libs.path_utils import (
     remove_non_official_s3_args,
 )
 from magic_pdf.libs.config_reader import get_local_dir
-from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter, MODE_BIN, MODE_TXT
+from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter
 from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
-
+from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
 
 parse_pdf_methods = click.Choice(["ocr", "txt", "auto"])
 
@@ -53,24 +55,34 @@ def prepare_env():
 
 
 def _do_parse(pdf_bytes, model_list, parse_method, image_writer, md_writer, image_dir):
-    uni_pipe = UNIPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
-    jso_useful_key = {
-        "_pdf_type": "txt",
-        "model_list": model_list,
-    }
-    if parse_method == "ocr":
-        jso_useful_key["_pdf_type"] = "ocr"
-
-    uni_pipe.pipe_parse()
-    md_content = uni_pipe.pipe_mk_markdown()
+    if parse_method == "auto":
+        pipe = UNIPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
+    elif parse_method == "txt":
+        pipe = TXTPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
+    elif parse_method == "ocr":
+        pipe = OCRPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
+    else:
+        print("unknow parse method")
+        os.exit(1)
+
+    pipe.pipe_classify()
+    pipe.pipe_parse()
+    md_content = pipe.pipe_mk_markdown()
     part_file_name = datetime.now().strftime("%H-%M-%S")
-    md_writer.write(content=md_content, path=f"{part_file_name}.md", mode=MODE_TXT)
     md_writer.write(
-        content=json_parse.dumps(
-            uni_pipe.pdf_mid_data, ensure_ascii=False, indent=4
-        ),
+        content=md_content, path=f"{part_file_name}.md", mode=AbsReaderWriter.MODE_TXT
+    )
+    md_writer.write(
+        content=json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4),
         path=f"{part_file_name}.json",
-        mode=MODE_TXT,
+        mode=AbsReaderWriter.MODE_TXT,
+    )
+    try:
+        content_list = pipe.pipe_mk_uni_format()
+    except Exception as e:
+        print(e)
+    md_writer.write(
+        str(content_list), f"{part_file_name}.txt", AbsReaderWriter.MODE_TXT
     )
 
 
@@ -106,7 +118,10 @@ def json_command(json, method):
             byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1])
             byte_end += byte_start - 1
         return s3_rw.read_jsonl(
-            remove_non_official_s3_args(s3path), byte_start, byte_end, MODE_BIN
+            remove_non_official_s3_args(s3path),
+            byte_start,
+            byte_end,
+            AbsReaderWriter.MODE_BIN,
         )
 
     jso = json_parse.loads(read_s3_path(json).decode("utf-8"))
@@ -119,7 +134,7 @@ def json_command(json, method):
 
     _do_parse(
         pdf_data,
-        jso['doc_layout_result'],
+        jso["doc_layout_result"],
         method,
         local_image_rw,
         local_md_rw,
@@ -148,7 +163,7 @@ def pdf_command(pdf, model, method):
 
     def read_fn(path):
         disk_rw = DiskReaderWriter(os.path.dirname(path))
-        return disk_rw.read(os.path.basename(path), MODE_BIN)
+        return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
 
     pdf_data = read_fn(pdf)
     jso = json_parse.loads(read_fn(model).decode("utf-8"))

magic_pdf/model/magic_model.py

@@ -48,6 +48,7 @@ class MagicModel:
         self.__model_list = model_list
         self.__docs = docs
         self.__fix_axis()
+        #@todo 移除置信度小于0.6的所有block
 
     def __reduct_overlap(self, bboxes):
         N = len(bboxes)

magic_pdf/pdf_parse_by_txt_v2.py

+import time
+
+from loguru import logger
+
+from magic_pdf.layout.layout_sort import get_bboxes_layout
+from magic_pdf.libs.convert_utils import dict_to_list
+from magic_pdf.libs.hash_utils import compute_md5
+from magic_pdf.libs.commons import fitz, get_delta_time
+from magic_pdf.model.magic_model import MagicModel
+from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2
+from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table
+from magic_pdf.pre_proc.ocr_detect_all_bboxes import ocr_prepare_bboxes_for_layout_split
+from magic_pdf.pre_proc.ocr_dict_merge import (
+    sort_blocks_by_layout,
+    fill_spans_in_blocks,
+    fix_block_spans,
+)
+from magic_pdf.libs.ocr_content_type import ContentType
+from magic_pdf.pre_proc.ocr_span_list_modify import (
+    remove_overlaps_min_spans,
+    get_qa_need_list_v2,
+)
+from magic_pdf.pre_proc.equations_replace import (
+    combine_chars_to_pymudict,
+    remove_chars_in_text_blocks,
+    replace_equations_in_textblock,
+)
+from magic_pdf.pre_proc.equations_replace import (
+    combine_chars_to_pymudict,
+    remove_chars_in_text_blocks,
+    replace_equations_in_textblock,
+)
+from magic_pdf.pre_proc.citationmarker_remove import remove_citation_marker
+
+
+def txt_spans_extract(pdf_page, inline_equations, interline_equations):
+    text_raw_blocks = pdf_page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"]
+    char_level_text_blocks = pdf_page.get_text("rawdict", flags=fitz.TEXTFLAGS_TEXT)[
+        "blocks"
+    ]
+    text_blocks = combine_chars_to_pymudict(text_raw_blocks, char_level_text_blocks)
+    text_blocks = replace_equations_in_textblock(
+        text_blocks, inline_equations, interline_equations
+    )
+    text_blocks = remove_citation_marker(text_blocks)
+    text_blocks = remove_chars_in_text_blocks(text_blocks)
+    spans = []
+    for v in text_blocks:
+        for line in v["lines"]:
+            for span in line["spans"]:
+                spans.append(
+                    {
+                        "bbox": list(span["bbox"]),
+                        "content": span["text"],
+                        "type": ContentType.Text,
+                    }
+                )
+    return spans
+
+
+def replace_text_span(pymu_spans, ocr_spans):
+    return list(filter(lambda x: x["type"] != ContentType.Text, ocr_spans)) + pymu_spans
+
+
+def parse_pdf_by_txt(
+    pdf_bytes,
+    model_list,
+    imageWriter,
+    start_page_id=0,
+    end_page_id=None,
+    debug_mode=False,
+):
+    pdf_bytes_md5 = compute_md5(pdf_bytes)
+    pdf_docs = fitz.open("pdf", pdf_bytes)
+
+    """初始化空的pdf_info_dict"""
+    pdf_info_dict = {}
+
+    """用model_list和docs对象初始化magic_model"""
+    magic_model = MagicModel(model_list, pdf_docs)
+
+    """根据输入的起始范围解析pdf"""
+    end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
+
+    """初始化启动时间"""
+    start_time = time.time()
+
+    for page_id in range(start_page_id, end_page_id + 1):
+
+        """debug时输出每页解析的耗时"""
+        if debug_mode:
+            time_now = time.time()
+            logger.info(
+                f"page_id: {page_id}, last_page_cost_time: {get_delta_time(start_time)}"
+            )
+            start_time = time_now
+
+        """从magic_model对象中获取后面会用到的区块信息"""
+        img_blocks = magic_model.get_imgs(page_id)
+        table_blocks = magic_model.get_tables(page_id)
+        discarded_blocks = magic_model.get_discarded(page_id)
+        text_blocks = magic_model.get_text_blocks(page_id)
+        title_blocks = magic_model.get_title_blocks(page_id)
+        inline_equations, interline_equations, interline_equation_blocks = (
+            magic_model.get_equations(page_id)
+        )
+
+        page_w, page_h = magic_model.get_page_size(page_id)
+
+        """将所有区块的bbox整理到一起"""
+        all_bboxes = ocr_prepare_bboxes_for_layout_split(
+            img_blocks,
+            table_blocks,
+            discarded_blocks,
+            text_blocks,
+            title_blocks,
+            interline_equation_blocks,
+            page_w,
+            page_h,
+        )
+
+        """根据区块信息计算layout"""
+        page_boundry = [0, 0, page_w, page_h]
+        layout_bboxes, layout_tree = get_bboxes_layout(
+            all_bboxes, page_boundry, page_id
+        )
+
+        """根据layout顺序，对当前页面所有需要留下的block进行排序"""
+        sorted_blocks = sort_blocks_by_layout(all_bboxes, layout_bboxes)
+
+        """ocr 中文本类的 span 用 pymu spans 替换！"""
+        ocr_spans = magic_model.get_all_spans(page_id)
+        pymu_spans = txt_spans_extract(
+            pdf_docs[page_id], inline_equations, interline_equations
+        )
+        spans = replace_text_span(pymu_spans, ocr_spans)
+
+        """删除重叠spans中较小的那些"""
+        spans, dropped_spans_by_span_overlap = remove_overlaps_min_spans(spans)
+        """对image和table截图"""
+        spans = ocr_cut_image_and_table(
+            spans, pdf_docs[page_id], page_id, pdf_bytes_md5, imageWriter
+        )
+
+        """将span填入排好序的blocks中"""
+        block_with_spans = fill_spans_in_blocks(sorted_blocks, spans)
+
+        """对block进行fix操作"""
+        fix_blocks = fix_block_spans(block_with_spans, img_blocks, table_blocks)
+
+        """获取QA需要外置的list"""
+        images, tables, interline_equations = get_qa_need_list_v2(fix_blocks)
+
+        """构造pdf_info_dict"""
+        page_info = ocr_construct_page_component_v2(
+            fix_blocks,
+            layout_bboxes,
+            page_id,
+            page_w,
+            page_h,
+            layout_tree,
+            images,
+            tables,
+            interline_equations,
+            discarded_blocks,
+        )
+        pdf_info_dict[f"page_{page_id}"] = page_info
+
+    """分段"""
+    pass
+
+    """dict转list"""
+    pdf_info_list = dict_to_list(pdf_info_dict)
+    new_pdf_info_dict = {
+        "pdf_info": pdf_info_list,
+    }
+
+    return new_pdf_info_dict
+
+
+if __name__ == "__main__":
+    if 1:
+        import fitz
+        import json
+
+        with open("/opt/data/pdf/20240418/25536-00.pdf", "rb") as f:
+            pdf_bytes = f.read()
+        pdf_docs = fitz.open("pdf", pdf_bytes)
+
+        with open("/opt/data/pdf/20240418/25536-00.json") as f:
+            model_list = json.loads(f.readline())
+
+        magic_model = MagicModel(model_list, pdf_docs)
+        for i in range(7):
+            print(magic_model.get_imgs(i))
+
+        for page_no, page in enumerate(pdf_docs):
+            inline_equations, interline_equations, interline_equation_blocks = (
+                magic_model.get_equations(page_no)
+            )
+
+            text_raw_blocks = page.get_text("dict", flags=fitz.TEXTFLAGS_TEXT)["blocks"]
+            char_level_text_blocks = page.get_text(
+                "rawdict", flags=fitz.TEXTFLAGS_TEXT
+            )["blocks"]
+            text_blocks = combine_chars_to_pymudict(
+                text_raw_blocks, char_level_text_blocks
+            )
+            text_blocks = replace_equations_in_textblock(
+                text_blocks, inline_equations, interline_equations
+            )
+            text_blocks = remove_citation_marker(text_blocks)
+
+            text_blocks = remove_chars_in_text_blocks(text_blocks)

magic_pdf/pre_proc/ocr_detect_all_bboxes.py

-from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio
+from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio, calculate_overlap_area_in_bbox1_area_ratio, \
+    calculate_iou
 from magic_pdf.libs.drop_tag import DropTag
 from magic_pdf.libs.ocr_content_type import BlockType
 
@@ -6,7 +7,6 @@ from magic_pdf.libs.ocr_content_type import BlockType
 def ocr_prepare_bboxes_for_layout_split(img_blocks, table_blocks, discarded_blocks, text_blocks,
                                         title_blocks, interline_equation_blocks, page_w, page_h):
     all_bboxes = []
-
     for image in img_blocks:
         x0, y0, x1, y1 = image['bbox']
         all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Image, None, None, None, None])
@@ -27,34 +27,66 @@ def ocr_prepare_bboxes_for_layout_split(img_blocks, table_blocks, discarded_bloc
         x0, y0, x1, y1 = interline_equation['bbox']
         all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.InterlineEquation, None, None, None, None])
 
+    '''block嵌套问题解决'''
+    '''文本框与标题框重叠，优先信任标题框'''
+    all_bboxes = fix_text_overlap_title_blocks(all_bboxes)
+    '''任何框体与舍弃框重叠，优先信任舍弃框'''
+    all_bboxes = remove_need_drop_blocks(all_bboxes, discarded_blocks)
+    '''经过以上处理后，还存在大框套小框的情况，则删除小框'''
+    all_bboxes = remove_overlaps_min_blocks(all_bboxes)
+
     '''discarded_blocks中只保留宽度超过1/3页面宽度的，高度超过10的，处于页面下半50%区域的（限定footnote）'''
     for discarded in discarded_blocks:
         x0, y0, x1, y1 = discarded['bbox']
         if (x1 - x0) > (page_w / 3) and (y1 - y0) > 10 and y0 > (page_h / 2):
             all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Footnote, None, None, None, None])
 
-    '''block嵌套问题解决'''
-    # @todo 1. text block大框套小框，删除小框 2. 图片或文本框与舍弃框重叠，优先信任舍弃框 3. 文本框与标题框重叠，优先信任文本框
-    all_bboxes, dropped_blocks = remove_overlaps_min_blocks(all_bboxes)
+    return all_bboxes
+
+
+def fix_text_overlap_title_blocks(all_bboxes):
+    # 先提取所有text和title block
+    text_blocks = []
+    for block in all_bboxes:
+        if block[7] == BlockType.Text:
+            text_blocks.append(block)
+    title_blocks = []
+    for block in all_bboxes:
+        if block[7] == BlockType.Title:
+            title_blocks.append(block)
 
+    for text_block in text_blocks:
+        for title_block in title_blocks:
+            text_block_bbox = text_block[0], text_block[1], text_block[2], text_block[3]
+            title_block_bbox = title_block[0], title_block[1], title_block[2], title_block[3]
+            if calculate_iou(text_block_bbox, title_block_bbox) > 0.8:
+                all_bboxes.remove(text_block)
+
+    return all_bboxes
+
+
+def remove_need_drop_blocks(all_bboxes, discarded_blocks):
+    for block in all_bboxes.copy():
+        for discarded_block in discarded_blocks:
+            block_bbox = block[0], block[1], block[2], block[3]
+            if calculate_overlap_area_in_bbox1_area_ratio(block_bbox, discarded_block['bbox']) > 0.6:
+                all_bboxes.remove(block)
     return all_bboxes
 
 
 def remove_overlaps_min_blocks(all_bboxes):
-    dropped_blocks = []
     #  删除重叠blocks中较小的那些
     for block1 in all_bboxes.copy():
         for block2 in all_bboxes.copy():
             if block1 != block2:
-                block1_box = block1[0], block1[1], block1[2], block1[3]
-                block2_box = block2[0], block2[1], block2[2], block2[3]
-                overlap_box = get_minbox_if_overlap_by_ratio(block1_box, block2_box, 0.8)
+                block1_bbox = [block1[0], block1[1], block1[2], block1[3]]
+                block2_bbox = [block2[0], block2[1], block2[2], block2[3]]
+                overlap_box = get_minbox_if_overlap_by_ratio(block1_bbox, block2_bbox, 0.8)
                 if overlap_box is not None:
                     bbox_to_remove = next(
                         (block for block in all_bboxes if [block[0], block[1], block[2], block[3]] == overlap_box),
                         None)
                     if bbox_to_remove is not None:
                         all_bboxes.remove(bbox_to_remove)
-                        bbox_to_remove['tag'] = DropTag.BLOCK_OVERLAP
-                        dropped_blocks.append(bbox_to_remove)
-    return all_bboxes, dropped_blocks
+
+    return all_bboxes

magic_pdf/user_api.py

@@ -16,7 +16,7 @@ from loguru import logger
 
 from magic_pdf.rw import AbsReaderWriter
 from magic_pdf.pdf_parse_by_ocr_v2 import parse_pdf_by_ocr
-from magic_pdf.pdf_parse_by_txt import parse_pdf_by_txt
+from magic_pdf.pdf_parse_by_txt_v2 import parse_pdf_by_txt
 
 
 PARSE_TYPE_TXT = "txt"