Просмотр исходного кода

feat(layout_detection): enhance layout detection with new methods and post-processing

- Added new methods for layout detection in the BaseLayoutDetector class, including detect, post_process, and utility functions for handling overlapping boxes and converting large text to tables.
- Updated derived classes (DitLayoutDetector, DoclingLayoutDetector, MinerULayoutDetector) to implement the new detection and post-processing logic, improving layout accuracy and flexibility.
- Introduced optional OCR spans in detection methods to support advanced layout processing.
- Refactored import paths for CoordinateUtils to maintain consistency across modules.
zhch158_admin 4 часов назад
Родитель
Сommit
e686a07103

+ 271 - 3
ocr_tools/universal_doc_parser/models/adapters/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Any, List, Union
+from typing import Dict, Any, List, Union, Optional, Tuple
 import numpy as np
 from PIL import Image
 
@@ -44,11 +44,266 @@ class BasePreprocessor(BaseAdapter):
 class BaseLayoutDetector(BaseAdapter):
     """版式检测器基类"""
     
+    def detect(
+        self, 
+        image: Union[np.ndarray, Image.Image],
+        ocr_spans: Optional[List[Dict[str, Any]]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        检测版式(模板方法,自动执行后处理)
+        
+        此方法会:
+        1. 调用子类实现的 _detect_raw() 进行原始检测
+        2. 自动执行后处理(去除重叠框、文本转表格等)
+        
+        Args:
+            image: 输入图像
+            ocr_spans: OCR结果(可选,某些detector可能需要)
+            
+        Returns:
+            后处理后的布局检测结果
+        """
+        # 调用子类实现的原始检测方法
+        layout_results = self._detect_raw(image, ocr_spans)
+        
+        # 自动执行后处理
+        if layout_results:
+            layout_config = self.config.get('post_process', {}) if hasattr(self, 'config') else {}
+            layout_results = self.post_process(layout_results, image, layout_config)
+        
+        return layout_results
+    
     @abstractmethod
-    def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
-        """检测版式"""
+    def _detect_raw(
+        self, 
+        image: Union[np.ndarray, Image.Image],
+        ocr_spans: Optional[List[Dict[str, Any]]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        原始检测方法(子类必须实现)
+        
+        Args:
+            image: 输入图像
+            ocr_spans: OCR结果(可选)
+            
+        Returns:
+            原始检测结果(未后处理)
+        """
         pass
     
+    def post_process(
+        self,
+        layout_results: List[Dict[str, Any]],
+        image: Union[np.ndarray, Image.Image],
+        config: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        后处理布局检测结果
+        
+        默认实现包括:
+        1. 去除重叠框
+        2. 将大面积文本块转换为表格(如果配置启用)
+        
+        子类可以重写此方法以自定义后处理逻辑
+        
+        Args:
+            layout_results: 原始检测结果
+            image: 输入图像
+            config: 后处理配置(可选),如果为None则使用self.config中的post_process配置
+            
+        Returns:
+            后处理后的布局结果
+        """
+        if not layout_results:
+            return layout_results
+        
+        # 获取配置
+        if config is None:
+            config = self.config.get('post_process', {}) if hasattr(self, 'config') else {}
+        
+        # 导入 CoordinateUtils(适配器可以访问)
+        try:
+            from ocr_utils.coordinate_utils import CoordinateUtils
+        except ImportError:
+            try:
+                from ocr_utils import CoordinateUtils
+            except ImportError:
+                # 如果无法导入,返回原始结果
+                return layout_results
+        
+        # 1. 去除重叠框
+        layout_results = self._remove_overlapping_boxes(layout_results, CoordinateUtils)
+        
+        # 2. 将大面积文本块转换为表格(如果配置启用)
+        layout_config = config if config is not None else {}
+        if layout_config.get('convert_large_text_to_table', False):
+            # 获取图像尺寸
+            if isinstance(image, Image.Image):
+                h, w = image.size[1], image.size[0]
+            else:
+                h, w = image.shape[:2] if len(image.shape) >= 2 else (0, 0)
+            
+            layout_results = self._convert_large_text_to_table(
+                layout_results,
+                (h, w),
+                min_area_ratio=layout_config.get('min_text_area_ratio', 0.25),
+                min_width_ratio=layout_config.get('min_text_width_ratio', 0.4),
+                min_height_ratio=layout_config.get('min_text_height_ratio', 0.3)
+            )
+        
+        return layout_results
+    
+    def _remove_overlapping_boxes(
+        self,
+        layout_results: List[Dict[str, Any]],
+        coordinate_utils: Any,
+        iou_threshold: float = 0.8,
+        overlap_ratio_threshold: float = 0.8
+    ) -> List[Dict[str, Any]]:
+        """
+        处理重叠的布局框(参考 MinerU 的去重策略)
+        
+        策略:
+        1. 高 IoU 重叠:保留置信度高的框
+        2. 包含关系:小框被大框高度包含时,保留大框并扩展边界
+        """
+        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 = coordinate_utils.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 = coordinate_utils.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]
+    
+    def _convert_large_text_to_table(
+        self,
+        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. 不与其他表格重叠:如果已有表格,不转换
+        """
+        if not layout_results:
+            return layout_results
+        
+        img_height, img_width = image_shape
+        img_area = img_height * img_width
+        
+        if img_area == 0:
+            return layout_results
+        
+        # 检查是否已有表格
+        has_table = any(
+            item.get('category', '').lower() in ['table', 'table_body']
+            for item in layout_results
+        )
+        
+        # 如果已有表格,不进行转换(避免误判)
+        if has_table:
+            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
+        
+        return results
+    
     def _map_category_id(self, category_id: int) -> str:
         """映射类别ID到字符串"""
         category_map = {
@@ -89,4 +344,17 @@ class BaseOCRRecognizer(BaseAdapter):
     @abstractmethod
     def recognize_text(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
         """识别文本"""
+        pass
+    
+    @abstractmethod
+    def detect_text_boxes(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
+        """
+        只检测文本框(不识别文字内容)
+        
+        子类必须实现此方法。建议使用只运行检测模型的方式(不运行识别模型)以优化性能。
+        如果无法优化,至少实现一个调用 recognize_text() 的版本作为兜底。
+        
+        Returns:
+            文本框列表,每项包含 'bbox', 'poly',可能包含 'confidence'
+        """
         pass

+ 40 - 13
ocr_tools/universal_doc_parser/models/adapters/dit_layout_adapter.py

@@ -512,12 +512,13 @@ class DitLayoutDetector(BaseLayoutDetector):
         self.cfg = None
         self._device = None
     
-    def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
+    def _detect_raw(self, image: Union[np.ndarray, Image.Image], ocr_spans: Optional[List[Dict[str, Any]]] = None) -> List[Dict[str, Any]]:
         """
-        检测布局
+        检测布局(原始检测,不包含后处理)
         
         Args:
             image: 输入图像 (numpy数组或PIL图像)
+            ocr_spans: OCR结果(可选,此detector不使用)
             
         Returns:
             检测结果列表,每个元素包含:
@@ -610,10 +611,35 @@ class DitLayoutDetector(BaseLayoutDetector):
                 }
             })
         
-        # 应用重叠框处理
-        if self._remove_overlap and len(formatted_results) > 1:
-            formatted_results = LayoutUtils.remove_overlapping_boxes(
-                formatted_results,
+        return formatted_results
+    
+    def post_process(
+        self,
+        layout_results: List[Dict[str, Any]],
+        image: Union[np.ndarray, Image.Image],
+        config: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        后处理布局检测结果(DiT 特定实现)
+        
+        包括:
+        1. 应用重叠框处理(使用 DiT 特定的参数)
+        2. 过滤误检的图片框
+        3. 调用基类的后处理(去除重叠框、文本转表格等)
+        """
+        if not layout_results:
+            return layout_results
+        
+        # 获取图像尺寸
+        if isinstance(image, Image.Image):
+            orig_w, orig_h = image.size
+        else:
+            orig_h, orig_w = image.shape[:2] if len(image.shape) >= 2 else (0, 0)
+        
+        # 1. 应用重叠框处理(DiT 特定逻辑)
+        if self._remove_overlap and len(layout_results) > 1:
+            layout_results = LayoutUtils.remove_overlapping_boxes(
+                layout_results,
                 iou_threshold=self._iou_threshold,
                 overlap_ratio_threshold=self._overlap_ratio_threshold,
                 image_size=(orig_w, orig_h),
@@ -622,18 +648,19 @@ class DitLayoutDetector(BaseLayoutDetector):
                 enable_category_priority=self._enable_category_priority
             )
         
-        # 过滤误检的图片框(包含过多文本内容的图片框)
-        if self._filter_false_positive_images and len(formatted_results) > 1:
-            before_count = len(formatted_results)
-            formatted_results = LayoutUtils.filter_false_positive_images(
-                formatted_results,
+        # 2. 过滤误检的图片框(包含过多文本内容的图片框)
+        if self._filter_false_positive_images and len(layout_results) > 1:
+            before_count = len(layout_results)
+            layout_results = LayoutUtils.filter_false_positive_images(
+                layout_results,
                 min_text_area_ratio=self._min_text_area_ratio
             )
-            removed_count = before_count - len(formatted_results)
+            removed_count = before_count - len(layout_results)
             if removed_count > 0:
                 print(f"🔄 Filtered {removed_count} false positive image boxes")
         
-        return formatted_results
+        # 3. 调用基类的后处理(去除重叠框、文本转表格等)
+        return super().post_process(layout_results, image, config)
     
     def detect_batch(
         self, 

+ 3 - 2
ocr_tools/universal_doc_parser/models/adapters/docling_layout_adapter.py

@@ -265,12 +265,13 @@ class DoclingLayoutDetector(BaseLayoutDetector):
         self.image_processor = None
         self._model_path = None
     
-    def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
+    def _detect_raw(self, image: Union[np.ndarray, Image.Image], ocr_spans: Optional[List[Dict[str, Any]]] = None) -> List[Dict[str, Any]]:
         """
-        检测布局
+        检测布局(原始检测,不包含后处理)
         
         Args:
             image: 输入图像 (numpy数组或PIL图像)
+            ocr_spans: OCR结果(可选,此detector不使用)
             
         Returns:
             检测结果列表,每个元素包含:

+ 40 - 3
ocr_tools/universal_doc_parser/models/adapters/mineru_adapter.py

@@ -17,7 +17,7 @@ if str(ocr_platform_root) not in sys.path:
     sys.path.insert(0, str(ocr_platform_root))
 
 from .base import BasePreprocessor, BaseLayoutDetector, BaseVLRecognizer, BaseOCRRecognizer
-from core.coordinate_utils import CoordinateUtils
+from ocr_utils.coordinate_utils import CoordinateUtils
 
 # 导入MinerU组件
 try:
@@ -123,8 +123,8 @@ class MinerULayoutDetector(BaseLayoutDetector):
         """清理资源"""
         pass
         
-    def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
-        """版式检测"""
+    def _detect_raw(self, image: Union[np.ndarray, Image.Image], ocr_spans: Optional[List[Dict[str, Any]]] = None) -> List[Dict[str, Any]]:
+        """版式检测(原始检测,不包含后处理)"""
         if self.layout_model is None:
             raise RuntimeError("Layout model not initialized")
             
@@ -507,6 +507,43 @@ class MinerUOCRRecognizer(BaseOCRRecognizer):
         except Exception as e:
             print(f"❌ OCR recognition failed: {e}")
             return []
+    
+    def detect_text_boxes(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
+        """
+        只检测文本框(不识别文字内容)
+        
+        参考 paddle_ori_cls.py 的做法,使用 rec=False 只运行检测模型
+        这样可以大幅提升性能,因为不需要运行识别模型
+        """
+        if self.ocr_model is None:
+            raise RuntimeError("OCR model not initialized")
+            
+        # 转换为BGR格式
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        
+        try:
+            # 只检测文本框,不识别文字(rec=False)
+            ocr_results = self.ocr_model.ocr(bgr_image, rec=False)
+            
+            # 格式化结果(只有 bbox,没有 text)
+            formatted_results = []
+            if ocr_results and ocr_results[0]:
+                for item in ocr_results[0]:
+                    # item 是文本框坐标(多边形)
+                    if len(item) >= 4:  # 至少4个点
+                        formatted_results.append({
+                            'bbox': CoordinateUtils.poly_to_bbox(item),  # 坐标
+                            'poly': item,  # 多边形坐标
+                            'confidence': 1.0  # 检测阶段没有置信度,设为1.0
+                        })
+                        
+            return formatted_results
+            
+        except Exception as e:
+            print(f"❌ Text box detection failed: {e}")
+            return []
 
 # 导出适配器类
 __all__ = [

+ 1 - 1
ocr_tools/universal_doc_parser/models/adapters/mineru_wired_table.py

@@ -7,7 +7,7 @@ import numpy as np
 from loguru import logger
 
 # 导入坐标转换工具
-from ocr_tools.universal_doc_parser.core.coordinate_utils import CoordinateUtils
+from ocr_utils.coordinate_utils import CoordinateUtils
 
 # 导入子模块
 from .wired_table.debug_utils import WiredTableDebugUtils

+ 4 - 3
ocr_tools/universal_doc_parser/models/adapters/paddle_layout_detector.py

@@ -4,7 +4,7 @@ import cv2
 import numpy as np
 import onnxruntime as ort
 from pathlib import Path
-from typing import Dict, List, Tuple, Union, Any
+from typing import Dict, List, Tuple, Union, Any, Optional
 from PIL import Image
 import sys
 
@@ -111,12 +111,13 @@ class PaddleLayoutDetector(BaseLayoutDetector):
         self.inputs = {}
         self.outputs = {}
     
-    def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
+    def _detect_raw(self, image: Union[np.ndarray, Image.Image], ocr_spans: Optional[List[Dict[str, Any]]] = None) -> List[Dict[str, Any]]:
         """
-        检测布局
+        检测布局(原始检测,不包含后处理)
         
         Args:
             image: 输入图像 (numpy数组或PIL图像)
+            ocr_spans: OCR结果(可选,此detector不使用)
             
         Returns:
             检测结果列表,每个元素包含: