252 lines
8.5 KiB
Python
252 lines
8.5 KiB
Python
"""
|
||
包含两个MagicModel类中重复使用的方法和逻辑
|
||
"""
|
||
from typing import List, Dict, Any, Callable
|
||
from mineru.utils.boxbase import bbox_distance, bbox_center_distance, is_in
|
||
|
||
|
||
def reduct_overlap(bboxes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
去除重叠的bbox，保留不被其他bbox包含的bbox
|
||
|
||
Args:
|
||
bboxes: 包含bbox信息的字典列表
|
||
|
||
Returns:
|
||
去重后的bbox列表
|
||
"""
|
||
N = len(bboxes)
|
||
keep = [True] * N
|
||
for i in range(N):
|
||
for j in range(N):
|
||
if i == j:
|
||
continue
|
||
if is_in(bboxes[i]['bbox'], bboxes[j]['bbox']):
|
||
keep[i] = False
|
||
return [bboxes[i] for i in range(N) if keep[i]]
|
||
|
||
|
||
def tie_up_category_by_distance_v3(
|
||
get_subjects_func: Callable,
|
||
get_objects_func: Callable,
|
||
extract_subject_func: Callable = None,
|
||
extract_object_func: Callable = None
|
||
):
|
||
"""
|
||
通用的类别关联方法，用于将主体对象与客体对象进行关联
|
||
|
||
参数:
|
||
get_subjects_func: 函数，提取主体对象
|
||
get_objects_func: 函数，提取客体对象
|
||
extract_subject_func: 函数，自定义提取主体属性（默认使用bbox和其他属性）
|
||
extract_object_func: 函数，自定义提取客体属性（默认使用bbox和其他属性）
|
||
|
||
返回:
|
||
关联后的对象列表
|
||
"""
|
||
subjects = get_subjects_func()
|
||
objects = get_objects_func()
|
||
|
||
# 如果没有提供自定义提取函数，使用默认函数
|
||
if extract_subject_func is None:
|
||
extract_subject_func = lambda x: x
|
||
if extract_object_func is None:
|
||
extract_object_func = lambda x: x
|
||
|
||
ret = []
|
||
N, M = len(subjects), len(objects)
|
||
subjects.sort(key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2)
|
||
objects.sort(key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2)
|
||
|
||
OBJ_IDX_OFFSET = 10000
|
||
SUB_BIT_KIND, OBJ_BIT_KIND = 0, 1
|
||
|
||
all_boxes_with_idx = [(i, SUB_BIT_KIND, sub["bbox"][0], sub["bbox"][1]) for i, sub in enumerate(subjects)] + [
|
||
(i + OBJ_IDX_OFFSET, OBJ_BIT_KIND, obj["bbox"][0], obj["bbox"][1]) for i, obj in enumerate(objects)
|
||
]
|
||
seen_idx = set()
|
||
seen_sub_idx = set()
|
||
|
||
while N > len(seen_sub_idx):
|
||
candidates = []
|
||
for idx, kind, x0, y0 in all_boxes_with_idx:
|
||
if idx in seen_idx:
|
||
continue
|
||
candidates.append((idx, kind, x0, y0))
|
||
|
||
if len(candidates) == 0:
|
||
break
|
||
left_x = min([v[2] for v in candidates])
|
||
top_y = min([v[3] for v in candidates])
|
||
|
||
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_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
|
||
|
||
for i in range(1, len(candidates)):
|
||
if candidates[i][1] ^ fst_kind == 1:
|
||
nxt = candidates[i]
|
||
break
|
||
if nxt is None:
|
||
break
|
||
|
||
if fst_kind == SUB_BIT_KIND:
|
||
sub_idx, obj_idx = fst_idx, nxt[0] - OBJ_IDX_OFFSET
|
||
else:
|
||
sub_idx, obj_idx = nxt[0], fst_idx - OBJ_IDX_OFFSET
|
||
|
||
pair_dis = bbox_distance(subjects[sub_idx]["bbox"], objects[obj_idx]["bbox"])
|
||
nearest_dis = float("inf")
|
||
for i in range(N):
|
||
# 取消原先算法中 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"]))
|
||
|
||
if pair_dis >= 3 * nearest_dis:
|
||
seen_idx.add(sub_idx)
|
||
continue
|
||
|
||
seen_idx.add(sub_idx)
|
||
seen_idx.add(obj_idx + OBJ_IDX_OFFSET)
|
||
seen_sub_idx.add(sub_idx)
|
||
|
||
ret.append(
|
||
{
|
||
"sub_bbox": extract_subject_func(subjects[sub_idx]),
|
||
"obj_bboxes": [extract_object_func(objects[obj_idx])],
|
||
"sub_idx": sub_idx,
|
||
}
|
||
)
|
||
|
||
for i in range(len(objects)):
|
||
j = i + OBJ_IDX_OFFSET
|
||
if j in seen_idx:
|
||
continue
|
||
seen_idx.add(j)
|
||
nearest_dis, nearest_sub_idx = float("inf"), -1
|
||
for k in range(len(subjects)):
|
||
dis = bbox_distance(objects[i]["bbox"], subjects[k]["bbox"])
|
||
if dis < nearest_dis:
|
||
nearest_dis = dis
|
||
nearest_sub_idx = k
|
||
|
||
for k in range(len(subjects)):
|
||
if k != nearest_sub_idx:
|
||
continue
|
||
if k in seen_sub_idx:
|
||
for kk in range(len(ret)):
|
||
if ret[kk]["sub_idx"] == k:
|
||
ret[kk]["obj_bboxes"].append(extract_object_func(objects[i]))
|
||
break
|
||
else:
|
||
ret.append(
|
||
{
|
||
"sub_bbox": extract_subject_func(subjects[k]),
|
||
"obj_bboxes": [extract_object_func(objects[i])],
|
||
"sub_idx": k,
|
||
}
|
||
)
|
||
seen_sub_idx.add(k)
|
||
seen_idx.add(k)
|
||
|
||
for i in range(len(subjects)):
|
||
if i in seen_sub_idx:
|
||
continue
|
||
ret.append(
|
||
{
|
||
"sub_bbox": extract_subject_func(subjects[i]),
|
||
"obj_bboxes": [],
|
||
"sub_idx": i,
|
||
}
|
||
)
|
||
|
||
return ret
|
||
|
||
|
||
def tie_up_category_by_index(
|
||
get_subjects_func: Callable,
|
||
get_objects_func: Callable,
|
||
extract_subject_func: Callable = None,
|
||
extract_object_func: Callable = None
|
||
):
|
||
"""
|
||
基于index的类别关联方法，用于将主体对象与客体对象进行关联
|
||
客体优先匹配给index最接近的主体，index差值相同时使用bbox中心点距离作为tiebreaker
|
||
|
||
参数:
|
||
get_subjects_func: 函数，提取主体对象
|
||
get_objects_func: 函数，提取客体对象
|
||
extract_subject_func: 函数，自定义提取主体属性（默认使用bbox和其他属性）
|
||
extract_object_func: 函数，自定义提取客体属性（默认使用bbox和其他属性）
|
||
|
||
返回:
|
||
关联后的对象列表，按主体index升序排列
|
||
"""
|
||
subjects = get_subjects_func()
|
||
objects = get_objects_func()
|
||
|
||
# 如果没有提供自定义提取函数，使用默认函数
|
||
if extract_subject_func is None:
|
||
extract_subject_func = lambda x: x
|
||
if extract_object_func is None:
|
||
extract_object_func = lambda x: x
|
||
|
||
# 初始化结果字典，key为主体索引，value为关联信息
|
||
result_dict = {}
|
||
|
||
# 初始化所有主体
|
||
for i, subject in enumerate(subjects):
|
||
result_dict[i] = {
|
||
"sub_bbox": extract_subject_func(subject),
|
||
"obj_bboxes": [],
|
||
"sub_idx": i,
|
||
}
|
||
|
||
# 为每个客体找到最匹配的主体
|
||
for obj in objects:
|
||
if len(subjects) == 0:
|
||
# 如果没有主体，跳过客体
|
||
continue
|
||
|
||
obj_index = obj["index"]
|
||
min_index_diff = float("inf")
|
||
best_subject_indices = []
|
||
|
||
# 找出index差值最小的所有主体
|
||
for i, subject in enumerate(subjects):
|
||
sub_index = subject["index"]
|
||
index_diff = abs(obj_index - sub_index)
|
||
|
||
if index_diff < min_index_diff:
|
||
min_index_diff = index_diff
|
||
best_subject_indices = [i]
|
||
elif index_diff == min_index_diff:
|
||
best_subject_indices.append(i)
|
||
|
||
# 如果有多个主体的index差值相同，使用中心点距离作为tiebreaker
|
||
if len(best_subject_indices) > 1:
|
||
min_center_dist = float("inf")
|
||
best_subject_idx = best_subject_indices[0]
|
||
|
||
for idx in best_subject_indices:
|
||
center_dist = bbox_center_distance(obj["bbox"], subjects[idx]["bbox"])
|
||
if center_dist < min_center_dist:
|
||
min_center_dist = center_dist
|
||
best_subject_idx = idx
|
||
else:
|
||
best_subject_idx = best_subject_indices[0]
|
||
|
||
# 将客体添加到最佳主体的obj_bboxes中
|
||
result_dict[best_subject_idx]["obj_bboxes"].append(extract_object_func(obj))
|
||
|
||
# 转换为列表并按主体index排序
|
||
ret = list(result_dict.values())
|
||
ret.sort(key=lambda x: x["sub_idx"])
|
||
|
||
return ret
|