layout_utils.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518
  1. """
  2. 布局处理工具模块
  3. 提供布局相关处理功能:
  4. - 重叠框检测与去重
  5. - 阅读顺序排序
  6. - OCR Span 与 Layout Block 匹配
  7. - 大面积文本块转换为表格(后处理)
  8. 注:底层坐标计算方法(IoU、重叠比例、poly_to_bbox 等)已统一到 coordinate_utils.py
  9. """
  10. from typing import Dict, List, Any, Tuple, Optional
  11. from loguru import logger
  12. import statistics
  13. # 导入坐标工具(底层坐标计算方法)
  14. try:
  15. from .coordinate_utils import CoordinateUtils
  16. except ImportError:
  17. from coordinate_utils import CoordinateUtils
  18. class LayoutUtils:
  19. """布局处理工具类"""
  20. # ==================== 坐标计算方法(委托给 CoordinateUtils)====================
  21. @staticmethod
  22. def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float:
  23. """计算两个 bbox 的 IoU(交并比)- 委托给 CoordinateUtils"""
  24. return CoordinateUtils.calculate_iou(bbox1, bbox2)
  25. @staticmethod
  26. def calculate_overlap_ratio(bbox1: List[float], bbox2: List[float]) -> float:
  27. """计算重叠面积占小框面积的比例 - 委托给 CoordinateUtils"""
  28. return CoordinateUtils.calculate_overlap_ratio(bbox1, bbox2)
  29. # ==================== 布局处理方法 ====================
  30. @staticmethod
  31. def remove_overlapping_boxes(
  32. layout_results: List[Dict[str, Any]],
  33. iou_threshold: float = 0.8,
  34. overlap_ratio_threshold: float = 0.8
  35. ) -> List[Dict[str, Any]]:
  36. """
  37. 处理重叠的布局框(参考 MinerU 的去重策略)
  38. 策略:
  39. 1. 高 IoU 重叠:保留置信度高的框
  40. 2. 包含关系:小框被大框高度包含时,保留大框并扩展边界
  41. 3. 同类型优先合并
  42. Args:
  43. layout_results: Layout 检测结果列表
  44. iou_threshold: IoU 阈值,超过此值认为高度重叠
  45. overlap_ratio_threshold: 重叠面积占小框面积的比例阈值
  46. Returns:
  47. 去重后的布局结果列表
  48. """
  49. if not layout_results or len(layout_results) <= 1:
  50. return layout_results
  51. # 复制列表避免修改原数据
  52. results = [item.copy() for item in layout_results]
  53. need_remove = set()
  54. for i in range(len(results)):
  55. if i in need_remove:
  56. continue
  57. for j in range(i + 1, len(results)):
  58. if j in need_remove:
  59. continue
  60. bbox1 = results[i].get('bbox', [0, 0, 0, 0])
  61. bbox2 = results[j].get('bbox', [0, 0, 0, 0])
  62. if len(bbox1) < 4 or len(bbox2) < 4:
  63. continue
  64. # 计算 IoU
  65. iou = LayoutUtils.calculate_iou(bbox1, bbox2)
  66. if iou > iou_threshold:
  67. # 高度重叠,保留置信度高的
  68. score1 = results[i].get('confidence', results[i].get('score', 0))
  69. score2 = results[j].get('confidence', results[j].get('score', 0))
  70. if score1 >= score2:
  71. need_remove.add(j)
  72. else:
  73. need_remove.add(i)
  74. break # i 被移除,跳出内层循环
  75. else:
  76. # 检查包含关系
  77. overlap_ratio = LayoutUtils.calculate_overlap_ratio(bbox1, bbox2)
  78. if overlap_ratio > overlap_ratio_threshold:
  79. # 小框被大框高度包含
  80. area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
  81. area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
  82. if area1 <= area2:
  83. small_idx, large_idx = i, j
  84. else:
  85. small_idx, large_idx = j, i
  86. # 扩展大框的边界
  87. small_bbox = results[small_idx]['bbox']
  88. large_bbox = results[large_idx]['bbox']
  89. results[large_idx]['bbox'] = [
  90. min(small_bbox[0], large_bbox[0]),
  91. min(small_bbox[1], large_bbox[1]),
  92. max(small_bbox[2], large_bbox[2]),
  93. max(small_bbox[3], large_bbox[3])
  94. ]
  95. need_remove.add(small_idx)
  96. if small_idx == i:
  97. break # i 被移除,跳出内层循环
  98. # 返回去重后的结果
  99. return [results[i] for i in range(len(results)) if i not in need_remove]
  100. @staticmethod
  101. def convert_large_text_to_table(
  102. layout_results: List[Dict[str, Any]],
  103. image_shape: Tuple[int, int],
  104. min_area_ratio: float = 0.25,
  105. min_width_ratio: float = 0.4,
  106. min_height_ratio: float = 0.3
  107. ) -> List[Dict[str, Any]]:
  108. """
  109. 将大面积的文本块转换为表格
  110. 判断规则:
  111. 1. 面积占比:占页面面积超过 min_area_ratio(默认25%)
  112. 2. 尺寸比例:宽度和高度都超过一定比例(避免细长条)
  113. 3. 不与其他表格重叠:如果已有表格,不转换
  114. Args:
  115. layout_results: Layout 检测结果列表
  116. image_shape: 图像尺寸 (height, width)
  117. min_area_ratio: 最小面积占比(0-1),默认0.25(25%)
  118. min_width_ratio: 最小宽度占比(0-1),默认0.4(40%)
  119. min_height_ratio: 最小高度占比(0-1),默认0.3(30%)
  120. Returns:
  121. 转换后的布局结果列表
  122. """
  123. if not layout_results:
  124. return layout_results
  125. img_height, img_width = image_shape
  126. img_area = img_height * img_width
  127. # 检查是否已有表格
  128. has_table = any(
  129. item.get('category', '').lower() in ['table', 'table_body']
  130. for item in layout_results
  131. )
  132. # 如果已有表格,不进行转换(避免误判)
  133. if has_table:
  134. logger.debug("📋 Page already has table elements, skipping text-to-table conversion")
  135. return layout_results
  136. # 复制列表避免修改原数据
  137. results = [item.copy() for item in layout_results]
  138. converted_count = 0
  139. for item in results:
  140. category = item.get('category', '').lower()
  141. # 只处理文本类型的元素
  142. if category not in ['text', 'ocr_text']:
  143. continue
  144. bbox = item.get('bbox', [0, 0, 0, 0])
  145. if len(bbox) < 4:
  146. continue
  147. x1, y1, x2, y2 = bbox[:4]
  148. width = x2 - x1
  149. height = y2 - y1
  150. area = width * height
  151. # 计算占比
  152. area_ratio = area / img_area if img_area > 0 else 0
  153. width_ratio = width / img_width if img_width > 0 else 0
  154. height_ratio = height / img_height if img_height > 0 else 0
  155. # 判断是否满足转换条件
  156. if (area_ratio >= min_area_ratio and
  157. width_ratio >= min_width_ratio and
  158. height_ratio >= min_height_ratio):
  159. # 转换为表格
  160. item['category'] = 'table'
  161. item['original_category'] = category # 保留原始类别
  162. converted_count += 1
  163. logger.info(
  164. f"🔄 Converted large text block to table: "
  165. f"area={area_ratio:.1%}, size={width_ratio:.1%}×{height_ratio:.1%}, "
  166. f"bbox=[{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]"
  167. )
  168. if converted_count > 0:
  169. logger.info(f"✅ Converted {converted_count} large text block(s) to table(s)")
  170. return results
  171. @staticmethod
  172. def sort_elements_by_reading_order(
  173. elements: List[Dict[str, Any]],
  174. y_tolerance: float = 15.0
  175. ) -> List[Dict[str, Any]]:
  176. """
  177. 根据阅读顺序对元素进行排序,并添加 reading_order 字段
  178. 排序规则:
  179. 1. 按Y坐标分行(考虑容差,Y坐标相近的元素视为同一行)
  180. 2. 同一行内按X坐标从左到右排序
  181. 3. 行与行之间按Y坐标从上到下排序
  182. Args:
  183. elements: 元素列表(坐标已转换为原始图片坐标系)
  184. y_tolerance: Y坐标容差,在此范围内的元素被视为同一行
  185. Returns:
  186. 排序后的元素列表,每个元素都添加了 reading_order 字段
  187. """
  188. if not elements:
  189. return elements
  190. # 为每个元素提取排序用的坐标
  191. elements_with_coords = []
  192. for elem in elements:
  193. bbox = elem.get('bbox', [0, 0, 0, 0])
  194. if len(bbox) >= 4:
  195. y_top = bbox[1] # 上边界
  196. x_left = bbox[0] # 左边界
  197. else:
  198. y_top = 0
  199. x_left = 0
  200. elements_with_coords.append((elem, y_top, x_left))
  201. # 先按Y坐标排序
  202. elements_with_coords.sort(key=lambda x: (x[1], x[2]))
  203. # 按Y坐标分行
  204. rows = []
  205. current_row = []
  206. current_row_y = None
  207. for elem, y_top, x_left in elements_with_coords:
  208. if current_row_y is None:
  209. # 第一个元素
  210. current_row.append((elem, x_left))
  211. current_row_y = y_top
  212. elif abs(y_top - current_row_y) <= y_tolerance:
  213. # 同一行
  214. current_row.append((elem, x_left))
  215. else:
  216. # 新的一行
  217. rows.append(current_row)
  218. current_row = [(elem, x_left)]
  219. current_row_y = y_top
  220. # 添加最后一行
  221. if current_row:
  222. rows.append(current_row)
  223. # 每行内按X坐标排序,然后展平
  224. sorted_elements = []
  225. reading_order = 0
  226. for row in rows:
  227. # 行内按X坐标排序
  228. row.sort(key=lambda x: x[1])
  229. for elem, _ in row:
  230. # 添加 reading_order 字段
  231. elem['reading_order'] = reading_order
  232. sorted_elements.append(elem)
  233. reading_order += 1
  234. logger.debug(f"📖 Elements sorted by reading order: {len(sorted_elements)} elements")
  235. return sorted_elements
  236. class SpanMatcher:
  237. """
  238. OCR Span 与 Layout Block 匹配器
  239. 参考 MinerU 的处理方式:
  240. 1. 整页 OCR 获取所有 spans
  241. 2. 将 spans 匹配到对应的 layout blocks
  242. 注:底层坐标计算方法已统一到 CoordinateUtils
  243. """
  244. # ==================== 坐标计算方法(委托给 CoordinateUtils)====================
  245. @staticmethod
  246. def calculate_overlap_area_in_bbox1_ratio(
  247. bbox1: List[float],
  248. bbox2: List[float]
  249. ) -> float:
  250. """计算 bbox1 被 bbox2 覆盖的面积比例 - 委托给 CoordinateUtils"""
  251. return CoordinateUtils.calculate_overlap_in_bbox1_ratio(bbox1, bbox2)
  252. @staticmethod
  253. def poly_to_bbox(poly: List) -> List[float]:
  254. """将多边形坐标转换为 bbox 格式 - 委托给 CoordinateUtils"""
  255. return CoordinateUtils.poly_to_bbox(poly)
  256. # ==================== Span 匹配方法 ====================
  257. @staticmethod
  258. def match_spans_to_blocks(
  259. ocr_spans: List[Dict[str, Any]],
  260. layout_blocks: List[Dict[str, Any]],
  261. overlap_threshold: float = 0.5
  262. ) -> Dict[int, List[Dict[str, Any]]]:
  263. """
  264. 将 OCR spans 匹配到 layout blocks
  265. Args:
  266. ocr_spans: OCR 识别结果列表,每项包含 'bbox'/'poly', 'text', 'confidence'
  267. layout_blocks: Layout 检测结果列表,每项包含 'bbox', 'category'
  268. overlap_threshold: 重叠比例阈值,span 被 block 覆盖超过此比例才算匹配
  269. Returns:
  270. 匹配结果字典 {block_index: [matched_spans]}
  271. """
  272. matched = {i: [] for i in range(len(layout_blocks))}
  273. for span in ocr_spans:
  274. # 获取 span 的 bbox
  275. span_bbox = span.get('bbox', [])
  276. if not span_bbox:
  277. continue
  278. # 转换为标准 bbox 格式
  279. span_bbox = SpanMatcher.poly_to_bbox(span_bbox)
  280. # 找到最佳匹配的 block
  281. best_match_idx = -1
  282. best_overlap = 0.0
  283. for block_idx, block in enumerate(layout_blocks):
  284. block_bbox = block.get('bbox', [0, 0, 0, 0])
  285. overlap = SpanMatcher.calculate_overlap_area_in_bbox1_ratio(
  286. span_bbox, block_bbox
  287. )
  288. if overlap > overlap_threshold and overlap > best_overlap:
  289. best_overlap = overlap
  290. best_match_idx = block_idx
  291. if best_match_idx >= 0:
  292. # 创建带绝对坐标的 span 副本
  293. matched_span = span.copy()
  294. matched_span['bbox'] = span_bbox # 确保是标准 bbox 格式
  295. matched[best_match_idx].append(matched_span)
  296. return matched
  297. @staticmethod
  298. def merge_spans_to_text(
  299. spans: List[Dict[str, Any]],
  300. block_bbox: Optional[List[float]] = None
  301. ) -> Tuple[str, List[Dict[str, Any]]]:
  302. """
  303. 将多个 spans 合并为单个文本字符串
  304. 参考 MinerU 的 span 合并逻辑:
  305. 1. 按 Y 坐标分行
  306. 2. 同行内按 X 坐标排序
  307. 3. 行间添加换行,词间可能添加空格
  308. Args:
  309. spans: span 列表
  310. block_bbox: 所属 block 的 bbox(用于参考)
  311. Returns:
  312. (merged_text, sorted_spans)
  313. """
  314. if not spans:
  315. return "", []
  316. # 计算 spans 的高度中位数(用于判断同行)
  317. heights = []
  318. for span in spans:
  319. bbox = span.get('bbox', [0, 0, 0, 0])
  320. if len(bbox) >= 4:
  321. h = bbox[3] - bbox[1]
  322. if h > 0:
  323. heights.append(h)
  324. if heights:
  325. median_height = statistics.median(heights)
  326. y_tolerance = median_height * 0.5
  327. else:
  328. y_tolerance = 10
  329. # 为每个 span 添加坐标信息用于排序
  330. spans_with_coords = []
  331. for span in spans:
  332. bbox = span.get('bbox', [0, 0, 0, 0])
  333. if len(bbox) >= 4:
  334. y_center = (bbox[1] + bbox[3]) / 2
  335. x_left = bbox[0]
  336. else:
  337. y_center = 0
  338. x_left = 0
  339. spans_with_coords.append((span, y_center, x_left))
  340. # 按 Y 坐标分行
  341. spans_with_coords.sort(key=lambda x: (x[1], x[2]))
  342. lines = []
  343. current_line = []
  344. current_line_y = None
  345. for span, y_center, x_left in spans_with_coords:
  346. if current_line_y is None:
  347. current_line.append((span, x_left))
  348. current_line_y = y_center
  349. elif abs(y_center - current_line_y) <= y_tolerance:
  350. current_line.append((span, x_left))
  351. else:
  352. lines.append(current_line)
  353. current_line = [(span, x_left)]
  354. current_line_y = y_center
  355. if current_line:
  356. lines.append(current_line)
  357. # 合并文本
  358. text_parts = []
  359. sorted_spans = []
  360. for line in lines:
  361. # 行内按 X 坐标排序
  362. line.sort(key=lambda x: x[1])
  363. line_texts = []
  364. for span, _ in line:
  365. text = span.get('text', '')
  366. if text:
  367. line_texts.append(text)
  368. sorted_spans.append(span)
  369. if line_texts:
  370. text_parts.append(' '.join(line_texts))
  371. merged_text = '\n'.join(text_parts)
  372. return merged_text, sorted_spans
  373. @staticmethod
  374. def remove_duplicate_spans(
  375. spans: List[Dict[str, Any]],
  376. iou_threshold: float = 0.9
  377. ) -> List[Dict[str, Any]]:
  378. """
  379. 移除重复的 spans(高 IoU 重叠)
  380. Args:
  381. spans: span 列表
  382. iou_threshold: IoU 阈值
  383. Returns:
  384. 去重后的 spans
  385. """
  386. if len(spans) <= 1:
  387. return spans
  388. result = []
  389. removed = set()
  390. for i, span1 in enumerate(spans):
  391. if i in removed:
  392. continue
  393. bbox1 = span1.get('bbox', [0, 0, 0, 0])
  394. bbox1 = CoordinateUtils.poly_to_bbox(bbox1)
  395. for j in range(i + 1, len(spans)):
  396. if j in removed:
  397. continue
  398. bbox2 = spans[j].get('bbox', [0, 0, 0, 0])
  399. bbox2 = CoordinateUtils.poly_to_bbox(bbox2)
  400. iou = CoordinateUtils.calculate_iou(bbox1, bbox2)
  401. if iou > iou_threshold:
  402. # 保留置信度高的
  403. score1 = span1.get('confidence', span1.get('score', 0))
  404. score2 = spans[j].get('confidence', spans[j].get('score', 0))
  405. if score1 >= score2:
  406. removed.add(j)
  407. else:
  408. removed.add(i)
  409. break
  410. if i not in removed:
  411. result.append(span1)
  412. return result