|
|
@@ -15,9 +15,9 @@ import statistics
|
|
|
|
|
|
# 导入坐标工具(底层坐标计算方法)
|
|
|
try:
|
|
|
- from .coordinate_utils import CoordinateUtils
|
|
|
+ from ocr_utils.coordinate_utils import CoordinateUtils
|
|
|
except ImportError:
|
|
|
- from coordinate_utils import CoordinateUtils
|
|
|
+ from ocr_utils import CoordinateUtils
|
|
|
|
|
|
|
|
|
class LayoutUtils:
|
|
|
@@ -38,182 +38,6 @@ class LayoutUtils:
|
|
|
# ==================== 布局处理方法 ====================
|
|
|
|
|
|
@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
|