Parcourir la source

feat: 添加线条检测功能以增强表格分类,支持根据线条信息覆盖分类结果

zhch158_admin il y a 3 jours
Parent
commit
2418ba72d7

+ 80 - 4
ocr_tools/universal_doc_parser/models/adapters/paddle_table_classifier.py

@@ -10,6 +10,7 @@ from typing import Dict, Any, Union
 from pathlib import Path
 import numpy as np
 from PIL import Image
+import cv2
 from loguru import logger
 
 from .base import BaseAdapter
@@ -73,27 +74,30 @@ class PaddleTableClassifier(BaseAdapter):
     
     def classify(
         self, 
-        image: Union[np.ndarray, Image.Image]
+        image: Union[np.ndarray, Image.Image],
+        use_line_detection: bool = True
     ) -> Dict[str, Any]:
         """
         分类单个表格图像
         
         Args:
             image: 表格图像(numpy数组或PIL图像)
+            use_line_detection: 是否使用线条检测辅助判断(默认True)
             
         Returns:
             分类结果字典:
             {
                 'table_type': 'wired' | 'wireless',
                 'confidence': float,
-                'raw_label': str  # AtomicModel.WiredTable 或 AtomicModel.WirelessTable
+                'raw_label': str,  # AtomicModel.WiredTable 或 AtomicModel.WirelessTable
+                'line_detection': dict  # 线条检测结果(如果启用)
             }
         """
         if self.model is None:
             raise RuntimeError("Model not initialized. Call initialize() first.")
         
         try:
-            # 调用 MinerU 的预测接口
+            # Step 1: 调用 MinerU 的预测接口
             label, confidence = self.model.predict(image)
             
             # 转换标签为简化形式
@@ -111,7 +115,37 @@ class PaddleTableClassifier(BaseAdapter):
                 'raw_label': str(label)
             }
             
-            logger.debug(f"Table classified as '{table_type}' (confidence: {confidence:.3f})")
+            # Step 2: 使用线条检测辅助判断(覆盖低置信度结果)
+            if use_line_detection:
+                line_info = self._detect_table_lines(image)
+                result['line_detection'] = line_info
+                
+                # 🔑 关键逻辑:只有横线没有竖线 → 强制无线表格
+                has_horizontal = line_info['horizontal_lines'] > 0
+                has_vertical = line_info['vertical_lines'] > 0
+                
+                if (has_horizontal and not has_vertical) or (not has_horizontal and has_vertical):
+                    # 只有横线或者竖线,强制判断为无线表格
+                    if table_type != 'wireless':
+                        logger.info(
+                            f"📊 Line detection override: {table_type}→wireless "
+                            f"(H={line_info['horizontal_lines']}, V={line_info['vertical_lines']})"
+                        )
+                        result['table_type'] = 'wireless'
+                        result['override_reason'] = 'only_horizontal_lines' if has_horizontal else 'only_vertical_lines'
+                elif not has_horizontal and not has_vertical:
+                    # 没有线条,可能是纯文本表格,判断为无线
+                    if table_type != 'wireless':
+                        logger.info(f"📊 Line detection override: {table_type}→wireless (no lines detected)")
+                        result['table_type'] = 'wireless'
+                        result['override_reason'] = 'no_lines_detected'
+            
+            logger.debug(
+                f"Table classified as '{result['table_type']}' "
+                f"(confidence: {confidence:.3f}, "
+                f"H={result.get('line_detection', {}).get('horizontal_lines', 'N/A')}, "
+                f"V={result.get('line_detection', {}).get('vertical_lines', 'N/A')})"
+            )
             return result
             
         except Exception as e:
@@ -124,6 +158,48 @@ class PaddleTableClassifier(BaseAdapter):
                 'error': str(e)
             }
     
+    def _detect_table_lines(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, int]:
+        """
+        检测表格图像中的横线和竖线数量
+        
+        Args:
+            image: 表格图像
+            
+        Returns:
+            {'horizontal_lines': int, 'vertical_lines': int}
+        """
+        # 转换为numpy数组
+        if isinstance(image, Image.Image):
+            img_array = np.array(image)
+        else:
+            img_array = image.copy()
+        
+        # 转换为灰度图
+        if len(img_array.shape) == 3:
+            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+        else:
+            gray = img_array
+        
+        # 二值化
+        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+        
+        h, w = binary.shape
+        
+        # 检测横线
+        horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (max(20, w//30), 1))
+        horizontal_mask = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel)
+        horizontal_lines = cv2.findContours(horizontal_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
+        
+        # 检测竖线
+        vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(20, h//30)))
+        vertical_mask = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel)
+        vertical_lines = cv2.findContours(vertical_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
+        
+        return {
+            'horizontal_lines': len(horizontal_lines),
+            'vertical_lines': len(vertical_lines)
+        }
+    
     def batch_classify(
         self, 
         images: list[Union[np.ndarray, Image.Image]]