Просмотр исходного кода

refactor(para_split_v3): merge list and index block detection

- Combine __is_list_block() and __is_index_block() into a single function __is_list_or_index_block()
- Simplify block type determination logic
- Remove redundant code and improve readability
- Optimize block merging process
myhloli 1 год назад
Родитель
Сommit
fdcb49d327
1 измененных файлов с 43 добавлено и 71 удалено
  1. 43 71
      magic_pdf/para/para_split_v3.py

+ 43 - 71
magic_pdf/para/para_split_v3.py

@@ -15,6 +15,9 @@ class ListLineTag:
 
 
 def __process_blocks(blocks):
+    # 对所有block预处理
+    # 1.通过title和interline_equation将block分组
+    # 2.bbox边界根据line信息重置
 
     result = []
     current_group = []
@@ -47,11 +50,15 @@ def __process_blocks(blocks):
     return result
 
 
-def __is_list_block(block):
+def __is_list_or_index_block(block):
     # 一个block如果是list block 应该同时满足以下特征
     # 1.block内有多个line 2.block 内有多个line左侧顶格写 3.block内有多个line 右侧不顶格(狗牙状)
     # 1.block内有多个line 2.block 内有多个line左侧顶格写 3.多个line以endflag结尾
     # 1.block内有多个line 2.block 内有多个line左侧顶格写 3.block内有多个line 左侧不顶格
+
+    # index block 是一种特殊的list block
+    # 一个block如果是index block 应该同时满足以下特征
+    # 1.block内有多个line 2.block 内有多个line两侧均顶格写 3.line的开头或者结尾均为数字
     if len(block['lines']) >= 3:
         first_line = block['lines'][0]
         line_height = first_line['bbox'][3] - first_line['bbox'][1]
@@ -60,6 +67,7 @@ def __is_list_block(block):
         left_close_num = 0
         left_not_close_num = 0
         right_not_close_num = 0
+        right_close_num = 0
         lines_text_list = []
         for line in block['lines']:
 
@@ -79,68 +87,18 @@ def __is_list_block(block):
                 # logger.info(f"{line_text}, {block['bbox_fs']}, {line['bbox']}")
                 left_not_close_num += 1
 
-            # 计算右侧是否不顶格,拍脑袋用0.3block宽度做阈值
-            closed_area = 0.3 * block_weight
-            # closed_area = 5 * line_height
-            if block['bbox_fs'][2] - line['bbox'][2] > closed_area:
-                right_not_close_num += 1
-
-        # 判断lines_text_list中的元素是否有超过80%都以LIST_END_FLAG结尾
-        line_end_flag = False
-        if len(lines_text_list) > 0:
-            num_end_count = 0
-            for line_text in lines_text_list:
-                if len(line_text) > 0:
-                    if line_text[-1] in LIST_END_FLAG:
-                        num_end_count += 1
-
-            if num_end_count / len(lines_text_list) >= 0.8:
-                line_end_flag = True
-
-        if left_close_num >= 2 and (right_not_close_num >= 2 or line_end_flag or left_not_close_num >= 2):
-            for line in block['lines']:
-                if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
-                    line[ListLineTag.IS_LIST_START_LINE] = True
-                if abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
-                    line[ListLineTag.IS_LIST_END_LINE] = True
-
-            return True
-        else:
-            return False
-    else:
-        return False
-
-
-def __is_index_block(block):
-    # 一个block如果是index block 应该同时满足以下特征
-    # 1.block内有多个line 2.block 内有多个line两侧均顶格写 3.line的开头或者结尾均为数字
-    if len(block['lines']) >= 3:
-        first_line = block['lines'][0]
-        line_height = first_line['bbox'][3] - first_line['bbox'][1]
-
-        left_close_num = 0
-        right_close_num = 0
-
-        lines_text_list = []
-        for line in block['lines']:
-
-            # 计算line左侧顶格数量是否大于2,是否顶格用abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height/2 来判断
-            if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
-                left_close_num += 1
-
-            # 计算右侧是否不顶格
+            # 计算右侧是否顶格
             if abs(block['bbox_fs'][2] - line['bbox'][2]) < line_height / 2:
                 right_close_num += 1
+            else:
+                # 右侧不顶格情况下是否有一段距离,拍脑袋用0.3block宽度做阈值
+                closed_area = 0.3 * block_weight
+                # closed_area = 5 * line_height
+                if block['bbox_fs'][2] - line['bbox'][2] > closed_area:
+                    right_not_close_num += 1
 
-            line_text = ""
-
-            for span in line['spans']:
-                span_type = span['type']
-                if span_type == ContentType.Text:
-                    line_text += span['content'].strip()
-
-            lines_text_list.append(line_text)
-
+        # 判断lines_text_list中的元素是否有超过80%都以LIST_END_FLAG结尾
+        line_end_flag = False
         # 判断lines_text_list中的元素是否有超过80%都以数字开头或都以数字结尾
         line_num_flag = False
         if len(lines_text_list) > 0:
@@ -148,23 +106,36 @@ def __is_index_block(block):
             num_end_count = 0
             for line_text in lines_text_list:
                 if len(line_text) > 0:
+                    if line_text[-1] in LIST_END_FLAG:
+                        num_end_count += 1
                     if line_text[0].isdigit():
                         num_start_count += 1
                     if line_text[-1].isdigit():
                         num_end_count += 1
 
+            if num_end_count / len(lines_text_list) >= 0.8:
+                line_end_flag = True
+
             if num_start_count / len(lines_text_list) >= 0.8 or num_end_count / len(lines_text_list) >= 0.8:
                 line_num_flag = True
 
         if left_close_num >= 2 and right_close_num >= 2 and line_num_flag:
             for line in block['lines']:
                 line[ListLineTag.IS_LIST_START_LINE] = True
+            return BlockType.Index
 
-            return True
+        elif left_close_num >= 2 and (right_not_close_num >= 2 or line_end_flag or left_not_close_num > 2):
+            for line in block['lines']:
+                if abs(block['bbox_fs'][0] - line['bbox'][0]) < line_height / 2:
+                    line[ListLineTag.IS_LIST_START_LINE] = True
+                if abs(block['bbox_fs'][2] - line['bbox'][2]) > line_height:
+                    line[ListLineTag.IS_LIST_END_LINE] = True
+
+            return BlockType.List
         else:
-            return False
+            return BlockType.Text
     else:
-        return False
+        return BlockType.Text
 
 
 def __merge_2_text_blocks(block1, block2):
@@ -206,12 +177,11 @@ def __para_merge_page(blocks):
     for text_blocks_group in page_text_blocks_groups:
 
         if len(text_blocks_group) > 0:
-            # 需要先在合并前对所有block判断是否为list block
+            # 需要先在合并前对所有block判断是否为list or index block
             for block in text_blocks_group:
-                if __is_list_block(block):
-                    block['type'] = BlockType.List
-                elif __is_index_block(block):
-                    block['type'] = BlockType.Index
+                block_type = __is_list_or_index_block(block)
+                block['type'] = block_type
+                # logger.info(f"{block['type']}:{block}")
 
         if len(text_blocks_group) > 1:
             # 倒序遍历
@@ -224,10 +194,12 @@ def __para_merge_page(blocks):
 
                     if current_block['type'] == 'text' and prev_block['type'] == 'text':
                         __merge_2_text_blocks(current_block, prev_block)
-                    if current_block['type'] == BlockType.List and prev_block['type'] == BlockType.List:
-                        __merge_2_list_blocks(current_block, prev_block)
-                    if current_block['type'] == BlockType.Index and prev_block['type'] == BlockType.Index:
+                    elif (
+                            (current_block['type'] == BlockType.List and prev_block['type'] == BlockType.List) or
+                            (current_block['type'] == BlockType.Index and prev_block['type'] == BlockType.Index)
+                    ):
                         __merge_2_list_blocks(current_block, prev_block)
+
         else:
             continue