Parcourir la source

fix: enhance block coverage detection by merging OCR and table results

myhloli il y a 3 mois
Parent
commit
60d32b8ac1
1 fichiers modifiés avec 65 ajouts et 0 suppressions
  1. 65 0
      mineru/utils/model_utils.py

+ 65 - 0
mineru/utils/model_utils.py

@@ -311,6 +311,71 @@ def get_res_list_from_layout_res(layout_res, iou_threshold=0.7, overlap_threshol
             del res['bbox']
             layout_res.remove(res)
 
+    # 新增:检测大block内部是否包含多个小block
+    # 合并ocr和table列表进行检测
+    combined_res_list = ocr_res_list + filtered_table_res_list
+
+    # 计算每个block的坐标和面积
+    block_info = []
+    for block in combined_res_list:
+        xmin, ymin = int(block['poly'][0]), int(block['poly'][1])
+        xmax, ymax = int(block['poly'][4]), int(block['poly'][5])
+        area = (xmax - xmin) * (ymax - ymin)
+        score = block.get('score', 0.5)  # 如果没有score字段,默认为0.5
+        block_info.append((xmin, ymin, xmax, ymax, area, score, block))
+
+    blocks_to_remove = []
+
+    # 检查每个block内部是否有3个及以上的小block
+    for i, (xmin, ymin, xmax, ymax, area, score, block) in enumerate(block_info):
+        # 查找内部的小block
+        blocks_inside = [(j, j_score, j_block) for j, (xj_min, yj_min, xj_max, yj_max, j_area, j_score, j_block) in
+                         enumerate(block_info)
+                         if i != j and is_inside(block_info[j], block_info[i])]
+
+        # 如果内部有3个及以上的小block
+        if len(blocks_inside) >= 3:
+            # 计算小block的平均分数
+            avg_score = sum(s for _, s, _ in blocks_inside) / len(blocks_inside)
+
+            # 比较大block的分数和小block的平均分数
+            if score > avg_score:
+                # 保留大block,扩展其边界
+                # 首先将所有小block标记为要删除
+                for j, _, j_block in blocks_inside:
+                    if j_block not in blocks_to_remove:
+                        blocks_to_remove.append(j_block)
+
+                # 扩展大block的边界以包含所有小block
+                new_xmin, new_ymin, new_xmax, new_ymax = xmin, ymin, xmax, ymax
+                for _, _, j_block in blocks_inside:
+                    j_xmin, j_ymin = int(j_block['poly'][0]), int(j_block['poly'][1])
+                    j_xmax, j_ymax = int(j_block['poly'][4]), int(j_block['poly'][5])
+                    new_xmin = min(new_xmin, j_xmin)
+                    new_ymin = min(new_ymin, j_ymin)
+                    new_xmax = max(new_xmax, j_xmax)
+                    new_ymax = max(new_ymax, j_ymax)
+
+                # 更新大block的边界
+                block['poly'][0] = block['poly'][6] = new_xmin
+                block['poly'][1] = block['poly'][3] = new_ymin
+                block['poly'][2] = block['poly'][4] = new_xmax
+                block['poly'][5] = block['poly'][7] = new_ymax
+            else:
+                # 保留小blocks,删除大block
+                blocks_to_remove.append(block)
+
+    # 移除需要删除的blocks
+    for block in blocks_to_remove:
+        if block in ocr_res_list:
+            ocr_res_list.remove(block)
+        elif block in filtered_table_res_list:
+            filtered_table_res_list.remove(block)
+
+        # 同时从layout_res中删除
+        if block in layout_res:
+            layout_res.remove(block)
+
     return ocr_res_list, filtered_table_res_list, single_page_mfdetrec_res