Explorar el Código

feat: add functionality to convert large text blocks into tables during post-processing, enhancing layout organization and improving document structure recognition

zhch158_admin hace 1 semana
padre
commit
a11d9c639c
Se han modificado 1 ficheros con 90 adiciones y 0 borrados
  1. 90 0
      zhch/universal_doc_parser/core/layout_utils.py

+ 90 - 0
zhch/universal_doc_parser/core/layout_utils.py

@@ -5,6 +5,7 @@
 - 重叠框检测与去重
 - 阅读顺序排序
 - OCR Span 与 Layout Block 匹配
+- 大面积文本块转换为表格(后处理)
 
 注:底层坐标计算方法(IoU、重叠比例、poly_to_bbox 等)已统一到 coordinate_utils.py
 """
@@ -124,6 +125,95 @@ class LayoutUtils:
         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