浏览代码

fix: improve candidate sorting logic in vlm_magic_model.py

myhloli 4 月之前
父节点
当前提交
4f612cbc1d
共有 1 个文件被更改,包括 7 次插入3 次删除
  1. 7 3
      mineru/backend/vlm/vlm_magic_model.py

+ 7 - 3
mineru/backend/vlm/vlm_magic_model.py

@@ -320,7 +320,10 @@ def __tie_up_category_by_distance_v3(
         candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y) ** 2)
         candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y) ** 2)
 
 
         fst_idx, fst_kind, left_x, top_y = candidates[0]
         fst_idx, fst_kind, left_x, top_y = candidates[0]
-        candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y) ** 2)
+        fst_bbox = subjects[fst_idx]['bbox'] if fst_kind == SUB_BIT_KIND else objects[fst_idx - OBJ_IDX_OFFSET]['bbox']
+        candidates.sort(
+            key=lambda x: bbox_distance(fst_bbox, subjects[x[0]]['bbox']) if x[1] == SUB_BIT_KIND else bbox_distance(
+                fst_bbox, objects[x[0] - OBJ_IDX_OFFSET]['bbox']))
         nxt = None
         nxt = None
 
 
         for i in range(1, len(candidates)):
         for i in range(1, len(candidates)):
@@ -339,8 +342,9 @@ def __tie_up_category_by_distance_v3(
         pair_dis = bbox_distance(subjects[sub_idx]["bbox"], objects[obj_idx]["bbox"])
         pair_dis = bbox_distance(subjects[sub_idx]["bbox"], objects[obj_idx]["bbox"])
         nearest_dis = float("inf")
         nearest_dis = float("inf")
         for i in range(N):
         for i in range(N):
-            if i in seen_idx or i == sub_idx:
-                continue
+            # 取消原先算法中 1对1 匹配的偏置
+            # if i in seen_idx or i == sub_idx:
+            #     continue
             nearest_dis = min(nearest_dis, bbox_distance(subjects[i]["bbox"], objects[obj_idx]["bbox"]))
             nearest_dis = min(nearest_dis, bbox_distance(subjects[i]["bbox"], objects[obj_idx]["bbox"]))
 
 
         if pair_dis >= 3 * nearest_dis:
         if pair_dis >= 3 * nearest_dis: