magic_model_utils.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262
  1. """
  2. 布局处理的公共工具类
  3. 包含两个MagicModel类中重复使用的方法和逻辑
  4. """
  5. from typing import List, Dict, Any, Union
  6. from mineru.utils.boxbase import bbox_relative_pos, calculate_iou, bbox_distance, is_in
  7. def reduct_overlap(bboxes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
  8. """
  9. 去除重叠的bbox,保留不被其他bbox包含的bbox
  10. Args:
  11. bboxes: 包含bbox信息的字典列表
  12. Returns:
  13. 去重后的bbox列表
  14. """
  15. N = len(bboxes)
  16. keep = [True] * N
  17. for i in range(N):
  18. for j in range(N):
  19. if i == j:
  20. continue
  21. if is_in(bboxes[i]['bbox'], bboxes[j]['bbox']):
  22. keep[i] = False
  23. return [bboxes[i] for i in range(N) if keep[i]]
  24. def bbox_distance_with_relative_check(bbox1: List[int], bbox2: List[int]) -> float:
  25. """
  26. 计算两个bbox之间的距离,考虑相对位置约束
  27. Args:
  28. bbox1: 第一个bbox [x1, y1, x2, y2]
  29. bbox2: 第二个bbox [x1, y1, x2, y2]
  30. Returns:
  31. 距离值,如果不满足条件返回无穷大
  32. """
  33. left, right, bottom, top = bbox_relative_pos(bbox1, bbox2)
  34. flags = [left, right, bottom, top]
  35. count = sum([1 if v else 0 for v in flags])
  36. if count > 1:
  37. return float('inf')
  38. if left or right:
  39. l1 = bbox1[3] - bbox1[1]
  40. l2 = bbox2[3] - bbox2[1]
  41. else:
  42. l1 = bbox1[2] - bbox1[0]
  43. l2 = bbox2[2] - bbox2[0]
  44. if l2 > l1 and (l2 - l1) / l1 > 0.3:
  45. return float('inf')
  46. return bbox_distance(bbox1, bbox2)
  47. def tie_up_category_by_distance_v3(
  48. data_source: Union[List[Dict], Dict],
  49. subject_category_filter,
  50. object_category_filter,
  51. extract_bbox_func=None,
  52. extract_score_func=None,
  53. create_item_func=None
  54. ) -> List[Dict[str, Any]]:
  55. """
  56. 基于距离关联不同类型的区块/元素
  57. Args:
  58. data_source: 数据源,可以是列表或包含layout_dets的字典
  59. subject_category_filter: 主体类别过滤函数或值
  60. object_category_filter: 对象类别过滤函数或值
  61. extract_bbox_func: 提取bbox的函数,默认使用'bbox'键
  62. extract_score_func: 提取score的函数,默认使用'score'键
  63. create_item_func: 创建返回项的函数
  64. Returns:
  65. 关联结果列表
  66. """
  67. # 默认函数
  68. if extract_bbox_func is None:
  69. extract_bbox_func = lambda x: x['bbox']
  70. if extract_score_func is None:
  71. extract_score_func = lambda x: x['score']
  72. if create_item_func is None:
  73. create_item_func = lambda x: {'bbox': extract_bbox_func(x), 'score': extract_score_func(x)}
  74. # 处理数据源
  75. if isinstance(data_source, dict) and 'layout_dets' in data_source:
  76. items = data_source['layout_dets']
  77. else:
  78. items = data_source
  79. # 过滤主体和对象
  80. if callable(subject_category_filter):
  81. subjects = list(filter(subject_category_filter, items))
  82. else:
  83. subjects = list(filter(lambda x: x.get('category_id') == subject_category_filter or x.get('type') == subject_category_filter, items))
  84. if callable(object_category_filter):
  85. objects = list(filter(object_category_filter, items))
  86. else:
  87. objects = list(filter(lambda x: x.get('category_id') == object_category_filter or x.get('type') == object_category_filter, items))
  88. # 转换为标准格式并去重
  89. subjects = reduct_overlap([create_item_func(x) for x in subjects])
  90. objects = reduct_overlap([create_item_func(x) for x in objects])
  91. ret = []
  92. N, M = len(subjects), len(objects)
  93. subjects.sort(key=lambda x: extract_bbox_func(x)[0] ** 2 + extract_bbox_func(x)[1] ** 2)
  94. objects.sort(key=lambda x: extract_bbox_func(x)[0] ** 2 + extract_bbox_func(x)[1] ** 2)
  95. OBJ_IDX_OFFSET = 10000
  96. SUB_BIT_KIND, OBJ_BIT_KIND = 0, 1
  97. all_boxes_with_idx = [(i, SUB_BIT_KIND, extract_bbox_func(sub)[0], extract_bbox_func(sub)[1]) for i, sub in enumerate(subjects)] + \
  98. [(i + OBJ_IDX_OFFSET, OBJ_BIT_KIND, extract_bbox_func(obj)[0], extract_bbox_func(obj)[1]) for i, obj in enumerate(objects)]
  99. seen_idx = set()
  100. seen_sub_idx = set()
  101. seen_sub_idx_len = len(seen_sub_idx)
  102. while N > seen_sub_idx_len:
  103. candidates = []
  104. for idx, kind, x0, y0 in all_boxes_with_idx:
  105. if idx in seen_idx:
  106. continue
  107. candidates.append((idx, kind, x0, y0))
  108. if len(candidates) == 0:
  109. break
  110. left_x = min([v[2] for v in candidates])
  111. top_y = min([v[3] for v in candidates])
  112. candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y) ** 2)
  113. fst_idx, fst_kind, left_x, top_y = candidates[0]
  114. fst_bbox = extract_bbox_func(subjects[fst_idx]) if fst_kind == SUB_BIT_KIND else extract_bbox_func(objects[fst_idx - OBJ_IDX_OFFSET])
  115. candidates.sort(
  116. key=lambda x: bbox_distance(fst_bbox, extract_bbox_func(subjects[x[0]])) if x[1] == SUB_BIT_KIND else bbox_distance(
  117. fst_bbox, extract_bbox_func(objects[x[0] - OBJ_IDX_OFFSET])))
  118. nxt = None
  119. for i in range(1, len(candidates)):
  120. if candidates[i][1] ^ fst_kind == 1:
  121. nxt = candidates[i]
  122. break
  123. if nxt is None:
  124. break
  125. if fst_kind == SUB_BIT_KIND:
  126. sub_idx, obj_idx = fst_idx, nxt[0] - OBJ_IDX_OFFSET
  127. else:
  128. sub_idx, obj_idx = nxt[0], fst_idx - OBJ_IDX_OFFSET
  129. pair_dis = bbox_distance(extract_bbox_func(subjects[sub_idx]), extract_bbox_func(objects[obj_idx]))
  130. nearest_dis = float('inf')
  131. for i in range(N):
  132. # 取消原先算法中 1对1 匹配的偏置
  133. # if i in seen_idx or i == sub_idx:continue
  134. nearest_dis = min(nearest_dis, bbox_distance(extract_bbox_func(subjects[i]), extract_bbox_func(objects[obj_idx])))
  135. if pair_dis >= 3 * nearest_dis:
  136. seen_idx.add(sub_idx)
  137. continue
  138. seen_idx.add(sub_idx)
  139. seen_idx.add(obj_idx + OBJ_IDX_OFFSET)
  140. seen_sub_idx.add(sub_idx)
  141. ret.append({
  142. 'sub_bbox': subjects[sub_idx],
  143. 'obj_bboxes': [objects[obj_idx]],
  144. 'sub_idx': sub_idx,
  145. })
  146. # 处理剩余的对象
  147. for i in range(len(objects)):
  148. j = i + OBJ_IDX_OFFSET
  149. if j in seen_idx:
  150. continue
  151. seen_idx.add(j)
  152. nearest_dis, nearest_sub_idx = float('inf'), -1
  153. for k in range(len(subjects)):
  154. dis = bbox_distance(extract_bbox_func(objects[i]), extract_bbox_func(subjects[k]))
  155. if dis < nearest_dis:
  156. nearest_dis = dis
  157. nearest_sub_idx = k
  158. for k in range(len(subjects)):
  159. if k != nearest_sub_idx:
  160. continue
  161. if k in seen_sub_idx:
  162. for kk in range(len(ret)):
  163. if ret[kk]['sub_idx'] == k:
  164. ret[kk]['obj_bboxes'].append(objects[i])
  165. break
  166. else:
  167. ret.append({
  168. 'sub_bbox': subjects[k],
  169. 'obj_bboxes': [objects[i]],
  170. 'sub_idx': k,
  171. })
  172. seen_sub_idx.add(k)
  173. seen_idx.add(k)
  174. # 处理剩余的主体
  175. for i in range(len(subjects)):
  176. if i in seen_sub_idx:
  177. continue
  178. ret.append({
  179. 'sub_bbox': subjects[i],
  180. 'obj_bboxes': [],
  181. 'sub_idx': i,
  182. })
  183. return ret
  184. def remove_high_iou_low_confidence(layout_dets: List[Dict], iou_threshold: float = 0.9):
  185. """
  186. 删除高IOU且置信度较低的检测结果
  187. Args:
  188. layout_dets: 布局检测结果列表
  189. iou_threshold: IOU阈值
  190. """
  191. need_remove_list = []
  192. for i in range(len(layout_dets)):
  193. for j in range(i + 1, len(layout_dets)):
  194. layout_det1 = layout_dets[i]
  195. layout_det2 = layout_dets[j]
  196. if calculate_iou(layout_det1['bbox'], layout_det2['bbox']) > iou_threshold:
  197. layout_det_need_remove = layout_det1 if layout_det1['score'] < layout_det2['score'] else layout_det2
  198. if layout_det_need_remove not in need_remove_list:
  199. need_remove_list.append(layout_det_need_remove)
  200. for need_remove in need_remove_list:
  201. if need_remove in layout_dets:
  202. layout_dets.remove(need_remove)
  203. def remove_low_confidence(layout_dets: List[Dict], confidence_threshold: float = 0.05):
  204. """
  205. 删除置信度特别低的检测结果
  206. Args:
  207. layout_dets: 布局检测结果列表
  208. confidence_threshold: 置信度阈值
  209. """
  210. need_remove_list = []
  211. for layout_det in layout_dets:
  212. if layout_det['score'] <= confidence_threshold:
  213. need_remove_list.append(layout_det)
  214. for need_remove in need_remove_list:
  215. if need_remove in layout_dets:
  216. layout_dets.remove(need_remove)