Explorar o código

Merge pull request #3325 from myhloli/dev

fix: improve bounding box update logic based on score comparison
Xiaomeng Zhao hai 3 meses
pai
achega
bb933ff9f6
Modificáronse 1 ficheiros con 22 adicións e 18 borrados
  1. 22 18
      mineru/utils/model_utils.py

+ 22 - 18
mineru/utils/model_utils.py

@@ -219,27 +219,31 @@ def remove_overlaps_min_blocks(res_list):
             )
 
             if overlap_box is not None:
-                res_to_remove = None
-                large_res = None
 
-                # 确定哪个是小块(要移除的)
+                # 根据重叠框确定哪个是小块,哪个是大块
                 if overlap_box == res_list[i]['bbox']:
-                    res_to_remove = res_list[i]
-                    large_res = res_list[j]
+                    small_res, large_res = res_list[i], res_list[j]
                 elif overlap_box == res_list[j]['bbox']:
-                    res_to_remove = res_list[j]
-                    large_res = res_list[i]
-
-                if res_to_remove is not None and res_to_remove not in need_remove:
-                    # 更新大块的边界为两者的并集
-                    x1, y1, x2, y2 = large_res['bbox']
-                    sx1, sy1, sx2, sy2 = res_to_remove['bbox']
-                    x1 = min(x1, sx1)
-                    y1 = min(y1, sy1)
-                    x2 = max(x2, sx2)
-                    y2 = max(y2, sy2)
-                    large_res['bbox'] = [x1, y1, x2, y2]
-                    need_remove.append(res_to_remove)
+                    small_res, large_res = res_list[j], res_list[i]
+                else:
+                    continue  # 如果重叠框与任一块都不匹配,跳过处理
+
+                if small_res['score'] <= large_res['score']:
+                    # 如果小块的分数低于大块,则小块为需要移除的块
+                    if small_res is not None and small_res not in need_remove:
+                        # 更新大块的边界为两者的并集
+                        x1, y1, x2, y2 = large_res['bbox']
+                        sx1, sy1, sx2, sy2 = small_res['bbox']
+                        x1 = min(x1, sx1)
+                        y1 = min(y1, sy1)
+                        x2 = max(x2, sx2)
+                        y2 = max(y2, sy2)
+                        large_res['bbox'] = [x1, y1, x2, y2]
+                        need_remove.append(small_res)
+                else:
+                    # 如果大块的分数低于小块,则大块为需要移除的块, 这时不需要更新小块的边界
+                    if large_res is not None and large_res not in need_remove:
+                        need_remove.append(large_res)
 
     # 从列表中移除标记的元素
     for res in need_remove: