magic-pdf 0.9.3__py3-none-any.whl → 0.10.1__py3-none-any.whl

magic_pdf/pre_proc/ocr_detect_all_bboxes.py

@@ -1,60 +1,181 @@
-from loguru import logger
 
-from magic_pdf.libs.boxbase import get_minbox_if_overlap_by_ratio, calculate_overlap_area_in_bbox1_area_ratio, \
-    calculate_iou, calculate_vertical_projection_overlap_ratio
-from magic_pdf.libs.drop_tag import DropTag
-from magic_pdf.libs.ocr_content_type import BlockType
-from magic_pdf.pre_proc.remove_bbox_overlap import remove_overlap_between_bbox_for_block
-
-
-def ocr_prepare_bboxes_for_layout_split(img_blocks, table_blocks, discarded_blocks, text_blocks,
-                                        title_blocks, interline_equation_blocks, page_w, page_h):
+from magic_pdf.config.ocr_content_type import BlockType
+from magic_pdf.libs.boxbase import (
+    calculate_iou, calculate_overlap_area_in_bbox1_area_ratio,
+    calculate_vertical_projection_overlap_ratio,
+    get_minbox_if_overlap_by_ratio)
+from magic_pdf.pre_proc.remove_bbox_overlap import \
+    remove_overlap_between_bbox_for_block
+
+
+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 = []
     all_discarded_blocks = []
     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, image["score"]])
+        all_bboxes.append(
+            [
+                x0,
+                y0,
+                x1,
+                y1,
+                None,
+                None,
+                None,
+                BlockType.Image,
+                None,
+                None,
+                None,
+                None,
+                image['score'],
+            ]
+        )
 
     for table in table_blocks:
         x0, y0, x1, y1 = table['bbox']
-        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Table, None, None, None, None, table["score"]])
+        all_bboxes.append(
+            [
+                x0,
+                y0,
+                x1,
+                y1,
+                None,
+                None,
+                None,
+                BlockType.Table,
+                None,
+                None,
+                None,
+                None,
+                table['score'],
+            ]
+        )
 
     for text in text_blocks:
         x0, y0, x1, y1 = text['bbox']
-        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Text, None, None, None, None, text["score"]])
+        all_bboxes.append(
+            [
+                x0,
+                y0,
+                x1,
+                y1,
+                None,
+                None,
+                None,
+                BlockType.Text,
+                None,
+                None,
+                None,
+                None,
+                text['score'],
+            ]
+        )
 
     for title in title_blocks:
         x0, y0, x1, y1 = title['bbox']
-        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.Title, None, None, None, None, title["score"]])
+        all_bboxes.append(
+            [
+                x0,
+                y0,
+                x1,
+                y1,
+                None,
+                None,
+                None,
+                BlockType.Title,
+                None,
+                None,
+                None,
+                None,
+                title['score'],
+            ]
+        )
 
     for interline_equation in interline_equation_blocks:
         x0, y0, x1, y1 = interline_equation['bbox']
-        all_bboxes.append([x0, y0, x1, y1, None, None, None, BlockType.InterlineEquation, None, None, None, None, interline_equation["score"]])
-
-    '''block嵌套问题解决'''
-    '''文本框与标题框重叠，优先信任文本框'''
+        all_bboxes.append(
+            [
+                x0,
+                y0,
+                x1,
+                y1,
+                None,
+                None,
+                None,
+                BlockType.InterlineEquation,
+                None,
+                None,
+                None,
+                None,
+                interline_equation['score'],
+            ]
+        )
+
+    """block嵌套问题解决"""
+    """文本框与标题框重叠，优先信任文本框"""
     all_bboxes = fix_text_overlap_title_blocks(all_bboxes)
-    '''任何框体与舍弃框重叠，优先信任舍弃框'''
+    """任何框体与舍弃框重叠，优先信任舍弃框"""
     all_bboxes = remove_need_drop_blocks(all_bboxes, discarded_blocks)
 
     # interline_equation 与title或text框冲突的情况，分两种情况处理
-    '''interline_equation框与文本类型框iou比较接近1的时候，信任行间公式框'''
+    """interline_equation框与文本类型框iou比较接近1的时候，信任行间公式框"""
     all_bboxes = fix_interline_equation_overlap_text_blocks_with_hi_iou(all_bboxes)
-    '''interline_equation框被包含在文本类型框内，且interline_equation比文本区块小很多时信任文本框，这时需要舍弃公式框'''
+    """interline_equation框被包含在文本类型框内，且interline_equation比文本区块小很多时信任文本框，这时需要舍弃公式框"""
     # 通过后续大框套小框逻辑删除
 
-    '''discarded_blocks中只保留宽度超过1/3页面宽度的，高度超过10的，处于页面下半50%区域的（限定footnote）'''
+    """discarded_blocks中只保留宽度超过1/3页面宽度的，高度超过10的，处于页面下半50%区域的（限定footnote）"""
     for discarded in discarded_blocks:
         x0, y0, x1, y1 = discarded['bbox']
-        all_discarded_blocks.append([x0, y0, x1, y1, None, None, None, BlockType.Discarded, None, None, None, None, discarded["score"]])
+        all_discarded_blocks.append(
+            [
+                x0,
+                y0,
+                x1,
+                y1,
+                None,
+                None,
+                None,
+                BlockType.Discarded,
+                None,
+                None,
+                None,
+                None,
+                discarded['score'],
+            ]
+        )
         # 将footnote加入到all_bboxes中，用来计算layout
         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, discarded["score"]])
-
-    '''经过以上处理后，还存在大框套小框的情况，则删除小框'''
+            all_bboxes.append(
+                [
+                    x0,
+                    y0,
+                    x1,
+                    y1,
+                    None,
+                    None,
+                    None,
+                    BlockType.Footnote,
+                    None,
+                    None,
+                    None,
+                    None,
+                    discarded['score'],
+                ]
+            )
+
+    """经过以上处理后，还存在大框套小框的情况，则删除小框"""
     all_bboxes = remove_overlaps_min_blocks(all_bboxes)
     all_discarded_blocks = remove_overlaps_min_blocks(all_discarded_blocks)
-    '''将剩余的bbox做分离处理，防止后面分layout时出错'''
+    """将剩余的bbox做分离处理，防止后面分layout时出错"""
     all_bboxes, drop_reasons = remove_overlap_between_bbox_for_block(all_bboxes)
 
     return all_bboxes, all_discarded_blocks, drop_reasons
@@ -64,18 +185,64 @@ def add_bboxes(blocks, block_type, bboxes):
     for block in blocks:
         x0, y0, x1, y1 = block['bbox']
         if block_type in [
-            BlockType.ImageBody, BlockType.ImageCaption, BlockType.ImageFootnote,
-            BlockType.TableBody, BlockType.TableCaption, BlockType.TableFootnote
+            BlockType.ImageBody,
+            BlockType.ImageCaption,
+            BlockType.ImageFootnote,
+            BlockType.TableBody,
+            BlockType.TableCaption,
+            BlockType.TableFootnote,
         ]:
-            bboxes.append([x0, y0, x1, y1, None, None, None, block_type, None, None, None, None, block["score"], block["group_id"]])
+            bboxes.append(
+                [
+                    x0,
+                    y0,
+                    x1,
+                    y1,
+                    None,
+                    None,
+                    None,
+                    block_type,
+                    None,
+                    None,
+                    None,
+                    None,
+                    block['score'],
+                    block['group_id'],
+                ]
+            )
         else:
-            bboxes.append([x0, y0, x1, y1, None, None, None, block_type, None, None, None, None, block["score"]])
+            bboxes.append(
+                [
+                    x0,
+                    y0,
+                    x1,
+                    y1,
+                    None,
+                    None,
+                    None,
+                    block_type,
+                    None,
+                    None,
+                    None,
+                    None,
+                    block['score'],
+                ]
+            )
 
 
 def ocr_prepare_bboxes_for_layout_split_v2(
-        img_body_blocks, img_caption_blocks, img_footnote_blocks,
-        table_body_blocks, table_caption_blocks, table_footnote_blocks,
-        discarded_blocks, text_blocks, title_blocks, interline_equation_blocks, page_w, page_h
+    img_body_blocks,
+    img_caption_blocks,
+    img_footnote_blocks,
+    table_body_blocks,
+    table_caption_blocks,
+    table_footnote_blocks,
+    discarded_blocks,
+    text_blocks,
+    title_blocks,
+    interline_equation_blocks,
+    page_w,
+    page_h,
 ):
     all_bboxes = []
 
@@ -89,40 +256,40 @@ def ocr_prepare_bboxes_for_layout_split_v2(
     add_bboxes(title_blocks, BlockType.Title, all_bboxes)
     add_bboxes(interline_equation_blocks, BlockType.InterlineEquation, all_bboxes)
 
-    '''block嵌套问题解决'''
-    '''文本框与标题框重叠，优先信任文本框'''
+    """block嵌套问题解决"""
+    """文本框与标题框重叠，优先信任文本框"""
     all_bboxes = fix_text_overlap_title_blocks(all_bboxes)
-    '''任何框体与舍弃框重叠，优先信任舍弃框'''
+    """任何框体与舍弃框重叠，优先信任舍弃框"""
     all_bboxes = remove_need_drop_blocks(all_bboxes, discarded_blocks)
 
     # interline_equation 与title或text框冲突的情况，分两种情况处理
-    '''interline_equation框与文本类型框iou比较接近1的时候，信任行间公式框'''
+    """interline_equation框与文本类型框iou比较接近1的时候，信任行间公式框"""
     all_bboxes = fix_interline_equation_overlap_text_blocks_with_hi_iou(all_bboxes)
-    '''interline_equation框被包含在文本类型框内，且interline_equation比文本区块小很多时信任文本框，这时需要舍弃公式框'''
+    """interline_equation框被包含在文本类型框内，且interline_equation比文本区块小很多时信任文本框，这时需要舍弃公式框"""
     # 通过后续大框套小框逻辑删除
 
-    '''discarded_blocks'''
+    """discarded_blocks"""
     all_discarded_blocks = []
     add_bboxes(discarded_blocks, BlockType.Discarded, all_discarded_blocks)
 
-    '''footnote识别：宽度超过1/3页面宽度的，高度超过10的，处于页面下半50%区域的'''
+    """footnote识别：宽度超过1/3页面宽度的，高度超过10的，处于页面下半50%区域的"""
     footnote_blocks = []
     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):
             footnote_blocks.append([x0, y0, x1, y1])
 
-    '''移除在footnote下面的任何框'''
+    """移除在footnote下面的任何框"""
     need_remove_blocks = find_blocks_under_footnote(all_bboxes, footnote_blocks)
     if len(need_remove_blocks) > 0:
         for block in need_remove_blocks:
             all_bboxes.remove(block)
             all_discarded_blocks.append(block)
 
-    '''经过以上处理后，还存在大框套小框的情况，则删除小框'''
+    """经过以上处理后，还存在大框套小框的情况，则删除小框"""
     all_bboxes = remove_overlaps_min_blocks(all_bboxes)
     all_discarded_blocks = remove_overlaps_min_blocks(all_discarded_blocks)
-    '''将剩余的bbox做分离处理，防止后面分layout时出错'''
+    """将剩余的bbox做分离处理，防止后面分layout时出错"""
     all_bboxes, drop_reasons = remove_overlap_between_bbox_for_block(all_bboxes)
 
     return all_bboxes, all_discarded_blocks
@@ -135,7 +302,13 @@ def find_blocks_under_footnote(all_bboxes, footnote_blocks):
         for footnote_bbox in footnote_blocks:
             footnote_x0, footnote_y0, footnote_x1, footnote_y1 = footnote_bbox
             # 如果footnote的纵向投影覆盖了block的纵向投影的80%且block的y0大于等于footnote的y1
-            if block_y0 >= footnote_y1 and calculate_vertical_projection_overlap_ratio((block_x0, block_y0, block_x1, block_y1), footnote_bbox) >= 0.8:
+            if (
+                block_y0 >= footnote_y1
+                and calculate_vertical_projection_overlap_ratio(
+                    (block_x0, block_y0, block_x1, block_y1), footnote_bbox
+                )
+                >= 0.8
+            ):
                 if block not in need_remove_blocks:
                     need_remove_blocks.append(block)
                     break
@@ -203,7 +376,12 @@ def remove_need_drop_blocks(all_bboxes, discarded_blocks):
     for block in all_bboxes:
         for discarded_block in discarded_blocks:
             block_bbox = block[:4]
-            if calculate_overlap_area_in_bbox1_area_ratio(block_bbox, discarded_block['bbox']) > 0.6:
+            if (
+                calculate_overlap_area_in_bbox1_area_ratio(
+                    block_bbox, discarded_block['bbox']
+                )
+                > 0.6
+            ):
                 if block not in need_remove:
                     need_remove.append(block)
                     break
@@ -223,10 +401,18 @@ def remove_overlaps_min_blocks(all_bboxes):
             if block1 != block2:
                 block1_bbox = block1[:4]
                 block2_bbox = block2[:4]
-                overlap_box = get_minbox_if_overlap_by_ratio(block1_bbox, block2_bbox, 0.8)
+                overlap_box = get_minbox_if_overlap_by_ratio(
+                    block1_bbox, block2_bbox, 0.8
+                )
                 if overlap_box is not None:
-                    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:
+                    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]

magic_pdf/pre_proc/ocr_dict_merge.py

@@ -1,8 +1,8 @@
+from magic_pdf.config.drop_tag import DropTag
+from magic_pdf.config.ocr_content_type import BlockType, ContentType
 from magic_pdf.libs.boxbase import (__is_overlaps_y_exceeds_threshold,
                                     _is_in_or_part_overlap_with_area_ratio,
                                     calculate_overlap_area_in_bbox1_area_ratio)
-from magic_pdf.libs.drop_tag import DropTag
-from magic_pdf.libs.ocr_content_type import BlockType, ContentType
 
 
 # 将每一个line中的span从左到右排序
@@ -24,7 +24,7 @@ def line_sort_spans_by_left_to_right(lines):
     return line_objects
 
 
-def merge_spans_to_line(spans):
+def merge_spans_to_line(spans, threshold=0.6):
     if len(spans) == 0:
         return []
     else:
@@ -49,7 +49,7 @@ def merge_spans_to_line(spans):
                 continue
 
             # 如果当前的span与当前行的最后一个span在y轴上重叠，则添加到当前行
-            if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox'], 0.5):
+            if __is_overlaps_y_exceeds_threshold(span['bbox'], current_line[-1]['bbox'], threshold):
                 current_line.append(span)
             else:
                 # 否则，开始新行
@@ -157,7 +157,7 @@ def fill_spans_in_blocks(blocks, spans, radio):
             BlockType.ImageBody, BlockType.ImageCaption, BlockType.ImageFootnote,
             BlockType.TableBody, BlockType.TableCaption, BlockType.TableFootnote
         ]:
-            block_dict["group_id"] = block[-1]
+            block_dict['group_id'] = block[-1]
         block_spans = []
         for span in spans:
             span_bbox = span['bbox']