""" 布局处理工具模块 提供布局相关处理功能: - 重叠框检测与去重 - 阅读顺序排序 - OCR Span 与 Layout Block 匹配 - 大面积文本块转换为表格(后处理) 注:底层坐标计算方法(IoU、重叠比例、poly_to_bbox 等)已统一到 coordinate_utils.py """ from typing import Dict, List, Any, Tuple, Optional from loguru import logger import statistics # 导入坐标工具(底层坐标计算方法) try: from .coordinate_utils import CoordinateUtils except ImportError: from coordinate_utils import CoordinateUtils class LayoutUtils: """布局处理工具类""" # ==================== 坐标计算方法(委托给 CoordinateUtils)==================== @staticmethod def calculate_iou(bbox1: List[float], bbox2: List[float]) -> float: """计算两个 bbox 的 IoU(交并比)- 委托给 CoordinateUtils""" return CoordinateUtils.calculate_iou(bbox1, bbox2) @staticmethod def calculate_overlap_ratio(bbox1: List[float], bbox2: List[float]) -> float: """计算重叠面积占小框面积的比例 - 委托给 CoordinateUtils""" return CoordinateUtils.calculate_overlap_ratio(bbox1, bbox2) # ==================== 布局处理方法 ==================== @staticmethod def remove_overlapping_boxes( layout_results: List[Dict[str, Any]], iou_threshold: float = 0.8, overlap_ratio_threshold: float = 0.8 ) -> List[Dict[str, Any]]: """ 处理重叠的布局框(参考 MinerU 的去重策略) 策略: 1. 高 IoU 重叠:保留置信度高的框 2. 包含关系:小框被大框高度包含时,保留大框并扩展边界 3. 同类型优先合并 Args: layout_results: Layout 检测结果列表 iou_threshold: IoU 阈值,超过此值认为高度重叠 overlap_ratio_threshold: 重叠面积占小框面积的比例阈值 Returns: 去重后的布局结果列表 """ if not layout_results or len(layout_results) <= 1: return layout_results # 复制列表避免修改原数据 results = [item.copy() for item in layout_results] need_remove = set() for i in range(len(results)): if i in need_remove: continue for j in range(i + 1, len(results)): if j in need_remove: continue bbox1 = results[i].get('bbox', [0, 0, 0, 0]) bbox2 = results[j].get('bbox', [0, 0, 0, 0]) if len(bbox1) < 4 or len(bbox2) < 4: continue # 计算 IoU iou = LayoutUtils.calculate_iou(bbox1, bbox2) if iou > iou_threshold: # 高度重叠,保留置信度高的 score1 = results[i].get('confidence', results[i].get('score', 0)) score2 = results[j].get('confidence', results[j].get('score', 0)) if score1 >= score2: need_remove.add(j) else: need_remove.add(i) break # i 被移除,跳出内层循环 else: # 检查包含关系 overlap_ratio = LayoutUtils.calculate_overlap_ratio(bbox1, bbox2) if overlap_ratio > overlap_ratio_threshold: # 小框被大框高度包含 area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1]) area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1]) if area1 <= area2: small_idx, large_idx = i, j else: small_idx, large_idx = j, i # 扩展大框的边界 small_bbox = results[small_idx]['bbox'] large_bbox = results[large_idx]['bbox'] results[large_idx]['bbox'] = [ min(small_bbox[0], large_bbox[0]), min(small_bbox[1], large_bbox[1]), max(small_bbox[2], large_bbox[2]), max(small_bbox[3], large_bbox[3]) ] need_remove.add(small_idx) if small_idx == i: break # i 被移除,跳出内层循环 # 返回去重后的结果 return [results[i] for i in range(len(results)) if i not in need_remove] @staticmethod def convert_large_text_to_table( layout_results: List[Dict[str, Any]], image_shape: Tuple[int, int], min_area_ratio: float = 0.25, min_width_ratio: float = 0.4, min_height_ratio: float = 0.3 ) -> List[Dict[str, Any]]: """ 将大面积的文本块转换为表格 判断规则: 1. 面积占比:占页面面积超过 min_area_ratio(默认25%) 2. 尺寸比例:宽度和高度都超过一定比例(避免细长条) 3. 不与其他表格重叠:如果已有表格,不转换 Args: layout_results: Layout 检测结果列表 image_shape: 图像尺寸 (height, width) min_area_ratio: 最小面积占比(0-1),默认0.25(25%) min_width_ratio: 最小宽度占比(0-1),默认0.4(40%) min_height_ratio: 最小高度占比(0-1),默认0.3(30%) Returns: 转换后的布局结果列表 """ if not layout_results: return layout_results img_height, img_width = image_shape img_area = img_height * img_width # 检查是否已有表格 has_table = any( item.get('category', '').lower() in ['table', 'table_body'] for item in layout_results ) # 如果已有表格,不进行转换(避免误判) if has_table: logger.debug("📋 Page already has table elements, skipping text-to-table conversion") return layout_results # 复制列表避免修改原数据 results = [item.copy() for item in layout_results] converted_count = 0 for item in results: category = item.get('category', '').lower() # 只处理文本类型的元素 if category not in ['text', 'ocr_text']: continue bbox = item.get('bbox', [0, 0, 0, 0]) if len(bbox) < 4: continue x1, y1, x2, y2 = bbox[:4] width = x2 - x1 height = y2 - y1 area = width * height # 计算占比 area_ratio = area / img_area if img_area > 0 else 0 width_ratio = width / img_width if img_width > 0 else 0 height_ratio = height / img_height if img_height > 0 else 0 # 判断是否满足转换条件 if (area_ratio >= min_area_ratio and width_ratio >= min_width_ratio and height_ratio >= min_height_ratio): # 转换为表格 item['category'] = 'table' item['original_category'] = category # 保留原始类别 converted_count += 1 logger.info( f"🔄 Converted large text block to table: " f"area={area_ratio:.1%}, size={width_ratio:.1%}×{height_ratio:.1%}, " f"bbox=[{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]" ) if converted_count > 0: logger.info(f"✅ Converted {converted_count} large text block(s) to table(s)") return results @staticmethod def sort_elements_by_reading_order( elements: List[Dict[str, Any]], y_tolerance: float = 15.0 ) -> List[Dict[str, Any]]: """ 根据阅读顺序对元素进行排序,并添加 reading_order 字段 排序规则: 1. 按Y坐标分行(考虑容差,Y坐标相近的元素视为同一行) 2. 同一行内按X坐标从左到右排序 3. 行与行之间按Y坐标从上到下排序 Args: elements: 元素列表(坐标已转换为原始图片坐标系) y_tolerance: Y坐标容差,在此范围内的元素被视为同一行 Returns: 排序后的元素列表,每个元素都添加了 reading_order 字段 """ if not elements: return elements # 为每个元素提取排序用的坐标 elements_with_coords = [] for elem in elements: bbox = elem.get('bbox', [0, 0, 0, 0]) if len(bbox) >= 4: y_top = bbox[1] # 上边界 x_left = bbox[0] # 左边界 else: y_top = 0 x_left = 0 elements_with_coords.append((elem, y_top, x_left)) # 先按Y坐标排序 elements_with_coords.sort(key=lambda x: (x[1], x[2])) # 按Y坐标分行 rows = [] current_row = [] current_row_y = None for elem, y_top, x_left in elements_with_coords: if current_row_y is None: # 第一个元素 current_row.append((elem, x_left)) current_row_y = y_top elif abs(y_top - current_row_y) <= y_tolerance: # 同一行 current_row.append((elem, x_left)) else: # 新的一行 rows.append(current_row) current_row = [(elem, x_left)] current_row_y = y_top # 添加最后一行 if current_row: rows.append(current_row) # 每行内按X坐标排序,然后展平 sorted_elements = [] reading_order = 0 for row in rows: # 行内按X坐标排序 row.sort(key=lambda x: x[1]) for elem, _ in row: # 添加 reading_order 字段 elem['reading_order'] = reading_order sorted_elements.append(elem) reading_order += 1 logger.debug(f"📖 Elements sorted by reading order: {len(sorted_elements)} elements") return sorted_elements class SpanMatcher: """ OCR Span 与 Layout Block 匹配器 参考 MinerU 的处理方式: 1. 整页 OCR 获取所有 spans 2. 将 spans 匹配到对应的 layout blocks 注:底层坐标计算方法已统一到 CoordinateUtils """ # ==================== 坐标计算方法(委托给 CoordinateUtils)==================== @staticmethod def calculate_overlap_area_in_bbox1_ratio( bbox1: List[float], bbox2: List[float] ) -> float: """计算 bbox1 被 bbox2 覆盖的面积比例 - 委托给 CoordinateUtils""" return CoordinateUtils.calculate_overlap_in_bbox1_ratio(bbox1, bbox2) @staticmethod def poly_to_bbox(poly: List) -> List[float]: """将多边形坐标转换为 bbox 格式 - 委托给 CoordinateUtils""" return CoordinateUtils.poly_to_bbox(poly) # ==================== Span 匹配方法 ==================== @staticmethod def match_spans_to_blocks( ocr_spans: List[Dict[str, Any]], layout_blocks: List[Dict[str, Any]], overlap_threshold: float = 0.5 ) -> Dict[int, List[Dict[str, Any]]]: """ 将 OCR spans 匹配到 layout blocks Args: ocr_spans: OCR 识别结果列表,每项包含 'bbox'/'poly', 'text', 'confidence' layout_blocks: Layout 检测结果列表,每项包含 'bbox', 'category' overlap_threshold: 重叠比例阈值,span 被 block 覆盖超过此比例才算匹配 Returns: 匹配结果字典 {block_index: [matched_spans]} """ matched = {i: [] for i in range(len(layout_blocks))} for span in ocr_spans: # 获取 span 的 bbox span_bbox = span.get('bbox', []) if not span_bbox: continue # 转换为标准 bbox 格式 span_bbox = SpanMatcher.poly_to_bbox(span_bbox) # 找到最佳匹配的 block best_match_idx = -1 best_overlap = 0.0 for block_idx, block in enumerate(layout_blocks): block_bbox = block.get('bbox', [0, 0, 0, 0]) overlap = SpanMatcher.calculate_overlap_area_in_bbox1_ratio( span_bbox, block_bbox ) if overlap > overlap_threshold and overlap > best_overlap: best_overlap = overlap best_match_idx = block_idx if best_match_idx >= 0: # 创建带绝对坐标的 span 副本 matched_span = span.copy() matched_span['bbox'] = span_bbox # 确保是标准 bbox 格式 matched[best_match_idx].append(matched_span) return matched @staticmethod def merge_spans_to_text( spans: List[Dict[str, Any]], block_bbox: Optional[List[float]] = None ) -> Tuple[str, List[Dict[str, Any]]]: """ 将多个 spans 合并为单个文本字符串 参考 MinerU 的 span 合并逻辑: 1. 按 Y 坐标分行 2. 同行内按 X 坐标排序 3. 行间添加换行,词间可能添加空格 Args: spans: span 列表 block_bbox: 所属 block 的 bbox(用于参考) Returns: (merged_text, sorted_spans) """ if not spans: return "", [] # 计算 spans 的高度中位数(用于判断同行) heights = [] for span in spans: bbox = span.get('bbox', [0, 0, 0, 0]) if len(bbox) >= 4: h = bbox[3] - bbox[1] if h > 0: heights.append(h) if heights: median_height = statistics.median(heights) y_tolerance = median_height * 0.5 else: y_tolerance = 10 # 为每个 span 添加坐标信息用于排序 spans_with_coords = [] for span in spans: bbox = span.get('bbox', [0, 0, 0, 0]) if len(bbox) >= 4: y_center = (bbox[1] + bbox[3]) / 2 x_left = bbox[0] else: y_center = 0 x_left = 0 spans_with_coords.append((span, y_center, x_left)) # 按 Y 坐标分行 spans_with_coords.sort(key=lambda x: (x[1], x[2])) lines = [] current_line = [] current_line_y = None for span, y_center, x_left in spans_with_coords: if current_line_y is None: current_line.append((span, x_left)) current_line_y = y_center elif abs(y_center - current_line_y) <= y_tolerance: current_line.append((span, x_left)) else: lines.append(current_line) current_line = [(span, x_left)] current_line_y = y_center if current_line: lines.append(current_line) # 合并文本 text_parts = [] sorted_spans = [] for line in lines: # 行内按 X 坐标排序 line.sort(key=lambda x: x[1]) line_texts = [] for span, _ in line: text = span.get('text', '') if text: line_texts.append(text) sorted_spans.append(span) if line_texts: text_parts.append(' '.join(line_texts)) merged_text = '\n'.join(text_parts) return merged_text, sorted_spans @staticmethod def remove_duplicate_spans( spans: List[Dict[str, Any]], iou_threshold: float = 0.9 ) -> List[Dict[str, Any]]: """ 移除重复的 spans(高 IoU 重叠) Args: spans: span 列表 iou_threshold: IoU 阈值 Returns: 去重后的 spans """ if len(spans) <= 1: return spans result = [] removed = set() for i, span1 in enumerate(spans): if i in removed: continue bbox1 = span1.get('bbox', [0, 0, 0, 0]) bbox1 = CoordinateUtils.poly_to_bbox(bbox1) for j in range(i + 1, len(spans)): if j in removed: continue bbox2 = spans[j].get('bbox', [0, 0, 0, 0]) bbox2 = CoordinateUtils.poly_to_bbox(bbox2) iou = CoordinateUtils.calculate_iou(bbox1, bbox2) if iou > iou_threshold: # 保留置信度高的 score1 = span1.get('confidence', span1.get('score', 0)) score2 = spans[j].get('confidence', spans[j].get('score', 0)) if score1 >= score2: removed.add(j) else: removed.add(i) break if i not in removed: result.append(span1) return result