magic-pdf 0.7.0b1__py3-none-any.whl → 0.8.0__py3-none-any.whl

magic_pdf/para/para_split_v2.py

@@ -1,3 +1,5 @@
+import copy
+
 from sklearn.cluster import DBSCAN
 import numpy as np
 from loguru import logger
@@ -100,59 +102,62 @@ def __detect_list_lines(lines, new_layout_bboxes, lang):
 
     if lang != 'en':
         return lines, None
-    else:
-        total_lines = len(lines)
-        line_fea_encode = []
-        """
-        对每一行进行特征编码，编码规则如下：
-        1. 如果行顶格，且大写字母开头或者数字开头，编码为1
-        2. 如果顶格，其他非大写开头编码为4
-        3. 如果非顶格，首字符大写，编码为2
-        4. 如果非顶格，首字符非大写编码为3
-        """
-        if len(lines) > 0:
-            x_map_tag_dict, min_x_tag = cluster_line_x(lines)
-        for l in lines:
-            span_text = __get_span_text(l['spans'][0])
-            first_char = span_text[0]
-            layout = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)
-            if not layout:
-                line_fea_encode.append(0)
+
+    total_lines = len(lines)
+    line_fea_encode = []
+    """
+    对每一行进行特征编码，编码规则如下：
+    1. 如果行顶格，且大写字母开头或者数字开头，编码为1
+    2. 如果顶格，其他非大写开头编码为4
+    3. 如果非顶格，首字符大写，编码为2
+    4. 如果非顶格，首字符非大写编码为3
+    """
+    if len(lines) > 0:
+        x_map_tag_dict, min_x_tag = cluster_line_x(lines)
+    for l in lines:
+        span_text = __get_span_text(l['spans'][0])
+        if not span_text:
+            line_fea_encode.append(0)
+            continue
+        first_char = span_text[0]
+        layout = __find_layout_bbox_by_line(l['bbox'], new_layout_bboxes)
+        if not layout:
+            line_fea_encode.append(0)
+        else:
+            #
+            if x_map_tag_dict[round(l['bbox'][0])] == min_x_tag:
+                # if first_char.isupper() or first_char.isdigit() or not first_char.isalnum():
+                if not first_char.isalnum() or if_match_reference_list(span_text):
+                    line_fea_encode.append(1)
+                else:
+                    line_fea_encode.append(4)
             else:
-                #
-                if x_map_tag_dict[round(l['bbox'][0])] == min_x_tag:
-                    # if first_char.isupper() or first_char.isdigit() or not first_char.isalnum():
-                    if not first_char.isalnum() or if_match_reference_list(span_text):
-                        line_fea_encode.append(1)
-                    else:
-                        line_fea_encode.append(4)
+                if first_char.isupper():
+                    line_fea_encode.append(2)
                 else:
-                    if first_char.isupper():
-                        line_fea_encode.append(2)
-                    else:
-                        line_fea_encode.append(3)
+                    line_fea_encode.append(3)
 
-        # 然后根据编码进行分段, 选出来 1,2,3连续出现至少2次的行，认为是列表。
+    # 然后根据编码进行分段, 选出来 1,2,3连续出现至少2次的行，认为是列表。
 
-        list_indice, list_start_idx = find_repeating_patterns2(line_fea_encode)
-        if len(list_indice) > 0:
+    list_indice, list_start_idx = find_repeating_patterns2(line_fea_encode)
+    if len(list_indice) > 0:
+        if debug_able:
+            logger.info(f"发现了列表，列表行数：{list_indice}， {list_start_idx}")
+
+    # TODO check一下这个特列表里缩进的行左侧是不是对齐的。
+    segments = []
+    for start, end in list_indice:
+        for i in range(start, end + 1):
+            if i > 0:
+                if line_fea_encode[i] == 4:
+                    if debug_able:
+                        logger.info(f"列表行的第{i}行不是顶格的")
+                    break
+        else:
             if debug_able:
-                logger.info(f"发现了列表，列表行数：{list_indice}， {list_start_idx}")
-
-        # TODO check一下这个特列表里缩进的行左侧是不是对齐的。
-        segments = []
-        for start, end in list_indice:
-            for i in range(start, end + 1):
-                if i > 0:
-                    if line_fea_encode[i] == 4:
-                        if debug_able:
-                            logger.info(f"列表行的第{i}行不是顶格的")
-                        break
-            else:
-                if debug_able:
-                    logger.info(f"列表行的第{start}到第{end}行是列表")
+                logger.info(f"列表行的第{start}到第{end}行是列表")
 
-        return split_indices(total_lines, list_indice), list_start_idx
+    return split_indices(total_lines, list_indice), list_start_idx
 
 
 def cluster_line_x(lines: list) -> dict:
@@ -164,7 +169,7 @@ def cluster_line_x(lines: list) -> dict:
     x0_lst = np.array([[round(line['bbox'][0]), 0] for line in lines])
     x0_clusters = DBSCAN(eps=min_distance, min_samples=min_sample).fit(x0_lst)
     x0_uniq_label = np.unique(x0_clusters.labels_)
-    #x1_lst = np.array([[line['bbox'][2], 0] for line in lines])
+    # x1_lst = np.array([[line['bbox'][2], 0] for line in lines])
     x0_2_new_val = {}  # 存储旧值对应的新值映射
     min_x0 = round(lines[0]["bbox"][0])
     for label in x0_uniq_label:
@@ -197,7 +202,9 @@ def __valign_lines(blocks, layout_bboxes):
     min_distance = 3
     min_sample = 2
     new_layout_bboxes = []
-
+    # add bbox_fs for para split calculation
+    for block in blocks:
+        block["bbox_fs"] = copy.deepcopy(block["bbox"])
     for layout_box in layout_bboxes:
         blocks_in_layoutbox = [b for b in blocks if
                                b["type"] == BlockType.Text and is_in_layout(b['bbox'], layout_box['layout_bbox'])]
@@ -242,16 +249,15 @@ def __valign_lines(blocks, layout_bboxes):
         # 由于修改了block里的line长度，现在需要重新计算block的bbox
         for block in blocks_in_layoutbox:
             if len(block["lines"]) > 0:
-                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']])]
-
+                block['bbox_fs'] = [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']])]
         """新计算layout的bbox，因为block的bbox变了。"""
-        layout_x0 = min([block['bbox'][0] for block in blocks_in_layoutbox])
-        layout_y0 = min([block['bbox'][1] for block in blocks_in_layoutbox])
-        layout_x1 = max([block['bbox'][2] for block in blocks_in_layoutbox])
-        layout_y1 = max([block['bbox'][3] for block in blocks_in_layoutbox])
+        layout_x0 = min([block['bbox_fs'][0] for block in blocks_in_layoutbox])
+        layout_y0 = min([block['bbox_fs'][1] for block in blocks_in_layoutbox])
+        layout_x1 = max([block['bbox_fs'][2] for block in blocks_in_layoutbox])
+        layout_y1 = max([block['bbox_fs'][3] for block in blocks_in_layoutbox])
         new_layout_bboxes.append([layout_x0, layout_y0, layout_x1, layout_y1])
 
     return new_layout_bboxes
@@ -309,7 +315,7 @@ def __group_line_by_layout(blocks, layout_bboxes):
     # 因为只是一个block一行目前, 一个block就是一个段落
     blocks_group = []
     for lyout in layout_bboxes:
-        blocks_in_layout = [block for block in blocks if is_in_layout(block['bbox'], lyout['layout_bbox'])]
+        blocks_in_layout = [block for block in blocks if is_in_layout(block.get('bbox_fs', None), lyout['layout_bbox'])]
         blocks_group.append(blocks_in_layout)
     return blocks_group
 
@@ -362,7 +368,8 @@ def __split_para_in_layoutbox(blocks_group, new_layout_bbox, lang="en"):
                 for i in range(0, len(list_start)):
                     index = list_start[i] - 1
                     if index >= 0:
-                        if "content" in lines[index]["spans"][-1]:
+                        if "content" in lines[index]["spans"][-1] and lines[index]["spans"][-1].get('type', '') not in [
+                            ContentType.InlineEquation, ContentType.InterlineEquation]:
                             lines[index]["spans"][-1]["content"] += '\n\n'
         layout_list_info = [False, False]  # 这个layout最后是不是列表,记录每一个layout里是不是列表开头，列表结尾
         for content_type, start, end in text_segments:
@@ -474,7 +481,7 @@ def __connect_list_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
                     break
         # 如果这些行的缩进是相等的，那么连到上一个layout的最后一个段落上。
         if len(may_list_lines) > 0 and len(set([x['bbox'][0] for x in may_list_lines])) == 1:
-            #pre_page_paras[-1].append(may_list_lines)
+            # pre_page_paras[-1].append(may_list_lines)
             # 下一页合并到上一页最后一段，打一个cross_page的标签
             for line in may_list_lines:
                 for span in line["spans"]:
@@ -534,7 +541,6 @@ def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
         next_first_line_text = ''.join([__get_span_text(span) for span in next_first_line['spans']])
         next_first_line_type = next_first_line['spans'][0]['type']
         if pre_last_line_type not in [TEXT, INLINE_EQUATION] or next_first_line_type not in [TEXT, INLINE_EQUATION]:
-            #connected_layout_paras.append(layout_paras[i])
             connected_layout_blocks.append(blocks_group[i])
             continue
         pre_layout = __find_layout_bbox_by_line(pre_last_line['bbox'], new_layout_bbox)
@@ -549,10 +555,8 @@ def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
             -1] not in LINE_STOP_FLAG and \
                 next_first_line['bbox'][0] == next_x0_min:  # 前面一行沾满了整个行，并且没有结尾符号.下一行没有空白开头。
             """连接段落条件成立，将前一个layout的段落和后一个layout的段落连接。"""
-            #connected_layout_paras[-1][-1].extend(layout_paras[i][0])
             connected_layout_blocks[-1][-1]["lines"].extend(blocks_group[i][0]["lines"])
-            #layout_paras[i].pop(0)  # 删除后一个layout的第一个段落， 因为他已经被合并到前一个layout的最后一个段落了。
-            blocks_group[i][0]["lines"] = []  #删除后一个layout第一个段落中的lines，因为他已经被合并到前一个layout的最后一个段落了
+            blocks_group[i][0]["lines"] = []  # 删除后一个layout第一个段落中的lines，因为他已经被合并到前一个layout的最后一个段落了
             blocks_group[i][0][LINES_DELETED] = True
             # if len(layout_paras[i]) == 0:
             #     layout_paras.pop(i)
@@ -561,7 +565,6 @@ def __connect_para_inter_layoutbox(blocks_group, new_layout_bbox):
             connected_layout_blocks.append(blocks_group[i])
         else:
             """连接段落条件不成立，将前一个layout的段落加入到结果中。"""
-            #connected_layout_paras.append(layout_paras[i])
             connected_layout_blocks.append(blocks_group[i])
     return connected_layout_blocks
 
@@ -619,7 +622,7 @@ def __connect_para_inter_page(pre_page_paras, next_page_paras, pre_page_layout_b
                 span[CROSS_PAGE] = True
         pre_last_para.extend(next_first_para)
 
-        #next_page_paras[0].pop(0)  # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
+        # next_page_paras[0].pop(0)  # 删除后一个页面的第一个段落， 因为他已经被合并到前一个页面的最后一个段落了。
         next_page_paras[0][0]["lines"] = []
         next_page_paras[0][0][LINES_DELETED] = True
         return True
@@ -663,16 +666,15 @@ def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
         layout_box = new_layout_bbox[layout_i]
         single_line_paras_tag = []
         for i in range(len(layout_para)):
-            #single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
+            # single_line_paras_tag.append(len(layout_para[i]) == 1 and layout_para[i][0]['spans'][0]['type'] == TEXT)
             single_line_paras_tag.append(layout_para[i]['type'] == BlockType.Text and len(layout_para[i]["lines"]) == 1)
         """找出来连续的单行文本，如果连续行高度相同，那么合并为一个段落。"""
         consecutive_single_line_indices = find_consecutive_true_regions(single_line_paras_tag)
         if len(consecutive_single_line_indices) > 0:
-            #index_offset = 0
             """检查这些行是否是高度相同的，居中的"""
             for start, end in consecutive_single_line_indices:
-                #start += index_offset
-                #end += index_offset
+                # start += index_offset
+                # end += index_offset
                 line_hi = np.array([block["lines"][0]['bbox'][3] - block["lines"][0]['bbox'][1] for block in
                                     layout_para[start:end + 1]])
                 first_line_text = ''.join([__get_span_text(span) for span in layout_para[start]["lines"][0]['spans']])
@@ -697,9 +699,9 @@ def __connect_middle_align_text(page_paras, new_layout_bbox, page_num, lang):
                         for i_para in range(start + 1, end + 1):
                             layout_para[i_para]["lines"] = []
                             layout_para[i_para][LINES_DELETED] = True
-                        #layout_para[start:end + 1] = [merge_para]
+                        # layout_para[start:end + 1] = [merge_para]
 
-                        #index_offset -= end - start
+                        # index_offset -= end - start
 
     return
 
@@ -739,7 +741,7 @@ def para_split(pdf_info_dict, debug_mode, lang="en"):
     new_layout_of_pages = []  # 数组的数组，每个元素是一个页面的layoutS
     all_page_list_info = []  # 保存每个页面开头和结尾是否是列表
     for page_num, page in pdf_info_dict.items():
-        blocks = page['preproc_blocks']
+        blocks = copy.deepcopy(page['preproc_blocks'])
         layout_bboxes = page['layout_bboxes']
         new_layout_bbox = __common_pre_proc(blocks, layout_bboxes)
         new_layout_of_pages.append(new_layout_bbox)

magic_pdf/pdf_parse_union_core.py

@@ -41,6 +41,23 @@ def remove_horizontal_overlap_block_which_smaller(all_bboxes):
     return is_useful_block_horz_overlap, all_bboxes
 
 
+def __replace_STX_ETX(text_str:str):
+    """ Replace \u0002 and \u0003, as these characters become garbled when extracted using pymupdf. In fact, they were originally quotation marks.
+Drawback: This issue is only observed in English text; it has not been found in Chinese text so far.
+
+    Args:
+        text_str (str): raw text
+
+    Returns:
+        _type_: replaced text
+    """
+    if text_str:
+        s = text_str.replace('\u0002', "'")
+        s = s.replace("\u0003", "'")
+        return s
+    return text_str
+
+
 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)[
@@ -63,7 +80,7 @@ def txt_spans_extract(pdf_page, inline_equations, interline_equations):
                     spans.append(
                         {
                             "bbox": list(span["bbox"]),
-                            "content": span["text"],
+                            "content": __replace_STX_ETX(span["text"]),
                             "type": ContentType.Text,
                             "score": 1.0,
                         }
@@ -175,7 +192,7 @@ def parse_page_core(pdf_docs, magic_model, page_id, pdf_bytes_md5, imageWriter,
     sorted_blocks = sort_blocks_by_layout(all_bboxes, layout_bboxes)
 
     '''将span填入排好序的blocks中'''
-    block_with_spans, spans = fill_spans_in_blocks(sorted_blocks, spans, 0.6)
+    block_with_spans, spans = fill_spans_in_blocks(sorted_blocks, spans, 0.3)
 
     '''对block进行fix操作'''
     fix_blocks = fix_block_spans(block_with_spans, img_blocks, table_blocks)
@@ -208,13 +225,17 @@ def pdf_parse_union(pdf_bytes,
     magic_model = MagicModel(model_list, pdf_docs)
 
     '''根据输入的起始范围解析pdf'''
-    end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
+    # end_page_id = end_page_id if end_page_id else len(pdf_docs) - 1
+    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(pdf_docs) - 1
+
+    if end_page_id > len(pdf_docs) - 1:
+        logger.warning("end_page_id is out of range, use pdf_docs length")
+        end_page_id = len(pdf_docs) - 1
 
     '''初始化启动时间'''
     start_time = time.time()
 
-    for page_id in range(start_page_id, end_page_id + 1):
-
+    for page_id, page in enumerate(pdf_docs):
         '''debug时输出每页解析的耗时'''
         if debug_mode:
             time_now = time.time()
@@ -224,7 +245,14 @@ def pdf_parse_union(pdf_bytes,
             start_time = time_now
 
         '''解析pdf中的每一页'''
-        page_info = parse_page_core(pdf_docs, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode)
+        if start_page_id <= page_id <= end_page_id:
+            page_info = parse_page_core(pdf_docs, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode)
+        else:
+            page_w = page.rect.width
+            page_h = page.rect.height
+            page_info = ocr_construct_page_component_v2([], [], page_id, page_w, page_h, [],
+                                                [], [], [], [],
+                                                True, "skip page")
         pdf_info_dict[f"page_{page_id}"] = page_info
 
     """分段"""

magic_pdf/pipe/AbsPipe.py

@@ -16,12 +16,15 @@ class AbsPipe(ABC):
     PIP_OCR = "ocr"
     PIP_TXT = "txt"
 
-    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False):
+    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False,
+                 start_page_id=0, end_page_id=None):
         self.pdf_bytes = pdf_bytes
         self.model_list = model_list
         self.image_writer = image_writer
         self.pdf_mid_data = None  # 未压缩
         self.is_debug = is_debug
+        self.start_page_id = start_page_id
+        self.end_page_id = end_page_id
     
     def get_compress_pdf_mid_data(self):
         return JsonCompressor.compress_json(self.pdf_mid_data)

magic_pdf/pipe/OCRPipe.py

@@ -9,17 +9,20 @@ from magic_pdf.user_api import parse_ocr_pdf
 
 class OCRPipe(AbsPipe):
 
-    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False):
-        super().__init__(pdf_bytes, model_list, image_writer, is_debug)
+    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False,
+                 start_page_id=0, end_page_id=None):
+        super().__init__(pdf_bytes, model_list, image_writer, is_debug, start_page_id, end_page_id)
 
     def pipe_classify(self):
         pass
 
     def pipe_analyze(self):
-        self.model_list = doc_analyze(self.pdf_bytes, ocr=True)
+        self.model_list = doc_analyze(self.pdf_bytes, ocr=True,
+                                      start_page_id=self.start_page_id, end_page_id=self.end_page_id)
 
     def pipe_parse(self):
-        self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug)
+        self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug,
+                                          start_page_id=self.start_page_id, end_page_id=self.end_page_id)
 
     def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF):
         result = super().pipe_mk_uni_format(img_parent_path, drop_mode)

magic_pdf/pipe/TXTPipe.py

@@ -10,17 +10,20 @@ from magic_pdf.user_api import parse_txt_pdf
 
 class TXTPipe(AbsPipe):
 
-    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False):
-        super().__init__(pdf_bytes, model_list, image_writer, is_debug)
+    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False,
+                 start_page_id=0, end_page_id=None):
+        super().__init__(pdf_bytes, model_list, image_writer, is_debug, start_page_id, end_page_id)
 
     def pipe_classify(self):
         pass
 
     def pipe_analyze(self):
-        self.model_list = doc_analyze(self.pdf_bytes, ocr=False)
+        self.model_list = doc_analyze(self.pdf_bytes, ocr=False,
+                                      start_page_id=self.start_page_id, end_page_id=self.end_page_id)
 
     def pipe_parse(self):
-        self.pdf_mid_data = parse_txt_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug)
+        self.pdf_mid_data = parse_txt_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug,
+                                          start_page_id=self.start_page_id, end_page_id=self.end_page_id)
 
     def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF):
         result = super().pipe_mk_uni_format(img_parent_path, drop_mode)

magic_pdf/pipe/UNIPipe.py

@@ -13,9 +13,10 @@ from magic_pdf.user_api import parse_union_pdf, parse_ocr_pdf
 
 class UNIPipe(AbsPipe):
 
-    def __init__(self, pdf_bytes: bytes, jso_useful_key: dict, image_writer: AbsReaderWriter, is_debug: bool = False):
+    def __init__(self, pdf_bytes: bytes, jso_useful_key: dict, image_writer: AbsReaderWriter, is_debug: bool = False,
+                 start_page_id=0, end_page_id=None):
         self.pdf_type = jso_useful_key["_pdf_type"]
-        super().__init__(pdf_bytes, jso_useful_key["model_list"], image_writer, is_debug)
+        super().__init__(pdf_bytes, jso_useful_key["model_list"], image_writer, is_debug, start_page_id, end_page_id)
         if len(self.model_list) == 0:
             self.input_model_is_empty = True
         else:
@@ -26,17 +27,21 @@ class UNIPipe(AbsPipe):
 
     def pipe_analyze(self):
         if self.pdf_type == self.PIP_TXT:
-            self.model_list = doc_analyze(self.pdf_bytes, ocr=False)
+            self.model_list = doc_analyze(self.pdf_bytes, ocr=False,
+                                          start_page_id=self.start_page_id, end_page_id=self.end_page_id)
         elif self.pdf_type == self.PIP_OCR:
-            self.model_list = doc_analyze(self.pdf_bytes, ocr=True)
+            self.model_list = doc_analyze(self.pdf_bytes, ocr=True,
+                                          start_page_id=self.start_page_id, end_page_id=self.end_page_id)
 
     def pipe_parse(self):
         if self.pdf_type == self.PIP_TXT:
             self.pdf_mid_data = parse_union_pdf(self.pdf_bytes, self.model_list, self.image_writer,
-                                                is_debug=self.is_debug, input_model_is_empty=self.input_model_is_empty)
+                                                is_debug=self.is_debug, input_model_is_empty=self.input_model_is_empty,
+                                                start_page_id=self.start_page_id, end_page_id=self.end_page_id)
         elif self.pdf_type == self.PIP_OCR:
             self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer,
-                                              is_debug=self.is_debug)
+                                              is_debug=self.is_debug,
+                                              start_page_id=self.start_page_id, end_page_id=self.end_page_id)
 
     def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF):
         result = super().pipe_mk_uni_format(img_parent_path, drop_mode)

magic_pdf/pre_proc/ocr_detect_all_bboxes.py

@@ -133,6 +133,7 @@ def remove_need_drop_blocks(all_bboxes, discarded_blocks):
 
 
 def remove_overlaps_min_blocks(all_bboxes):
+    #  重叠block，小的不能直接删除，需要和大的那个合并成一个更大的。
     #  删除重叠blocks中较小的那些
     need_remove = []
     for block1 in all_bboxes:
@@ -142,9 +143,17 @@ def remove_overlaps_min_blocks(all_bboxes):
                 block2_bbox = block2[:4]
                 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[:4] == overlap_box), None)
-                    if bbox_to_remove is not None and bbox_to_remove not in need_remove:
-                        need_remove.append(bbox_to_remove)
+                    block_to_remove = next((block for block in all_bboxes if block[:4] == overlap_box), None)
+                    if block_to_remove is not None and block_to_remove not in need_remove:
+                        large_block = block1 if block1 != block_to_remove else block2
+                        x1, y1, x2, y2 = large_block[:4]
+                        sx1, sy1, sx2, sy2 = block_to_remove[:4]
+                        x1 = min(x1, sx1)
+                        y1 = min(y1, sy1)
+                        x2 = max(x2, sx2)
+                        y2 = max(y2, sy2)
+                        large_block[:4] = [x1, y1, x2, y2]
+                        need_remove.append(block_to_remove)
 
     if len(need_remove) > 0:
         for block in need_remove: