Jelajahi Sumber

feat(ocr_tools): add OCRBasedLayoutEvaluator for layout quality assessment

- Introduced OCRBasedLayoutEvaluator class to evaluate layout detection quality using OCR results, specifically for wireless tables.
- Implemented methods for assessing boundary accuracy, column alignment, and content purity, providing a comprehensive evaluation report.
- Added functionality to compare models based on their evaluation scores, enhancing the model selection process for layout detection.
zhch158_admin 4 jam lalu
induk
melakukan
ced05f5449

+ 467 - 0
ocr_tools/universal_doc_parser/core/ocr_based_layout_evaluator.py

@@ -0,0 +1,467 @@
+"""
+结合 OCR 结果的 Layout 检测质量评估器
+专门用于评估无线表格的检测准确性
+"""
+from typing import Dict, List, Any, Tuple, Optional, Union
+import numpy as np
+from PIL import Image
+import cv2
+from loguru import logger
+from collections import defaultdict
+
+try:
+    from ocr_utils.coordinate_utils import CoordinateUtils
+except ImportError:
+    from ocr_utils import CoordinateUtils
+
+
+class OCRBasedLayoutEvaluator:
+    """基于 OCR 结果的 Layout 检测质量评估器
+    
+    核心思路:
+    1. 利用 OCR 文本块的分布特征来判断表格边界
+    2. 无线表格的文本块应该呈现列状对齐特征
+    3. 表格区域外的文本块不应该被包含在表格中
+    """
+    
+    def __init__(self):
+        pass
+    
+    def evaluate_with_ocr(
+        self,
+        layout_results: List[Dict[str, Any]],
+        ocr_spans: List[Dict[str, Any]],
+        image: Union[np.ndarray, Image.Image],
+        table_category: str = 'table_body'
+    ) -> Dict[str, Any]:
+        """
+        结合 OCR 结果评估 layout 检测质量
+        
+        Args:
+            layout_results: Layout 检测结果
+            ocr_spans: OCR 文本块列表,每项包含 'bbox', 'text', 'confidence'
+            image: 输入图像
+            table_category: 表格类别名称
+            
+        Returns:
+            评估结果字典
+        """
+        if isinstance(image, Image.Image):
+            img_array = np.array(image)
+        else:
+            img_array = image.copy()
+        
+        h, w = img_array.shape[:2] if len(img_array.shape) == 2 else img_array.shape[:2]
+        
+        # 提取表格区域
+        table_boxes = [
+            r for r in layout_results 
+            if r.get('category', '').lower() in [table_category.lower(), 'table']
+        ]
+        
+        if not table_boxes:
+            return {
+                'has_table': False,
+                'table_count': 0,
+                'overall_score': 0.0
+            }
+        
+        all_scores = []
+        all_issues = []
+        
+        for i, table_box in enumerate(table_boxes):
+            bbox = table_box.get('bbox', [])
+            if len(bbox) < 4:
+                continue
+            
+            # 评估单个表格
+            table_eval = self._evaluate_single_table_with_ocr(
+                bbox, ocr_spans, layout_results, (h, w), i
+            )
+            
+            all_scores.append(table_eval['score'])
+            all_issues.extend(table_eval['issues'])
+        
+        avg_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
+        
+        return {
+            'has_table': True,
+            'table_count': len(table_boxes),
+            'overall_score': float(avg_score),
+            'table_scores': all_scores,
+            'issues': all_issues
+        }
+    
+    def _evaluate_single_table_with_ocr(
+        self,
+        table_bbox: List[float],
+        ocr_spans: List[Dict[str, Any]],
+        all_layout_results: List[Dict[str, Any]],
+        image_shape: Tuple[int, int],
+        table_idx: int
+    ) -> Dict[str, Any]:
+        """评估单个表格的质量"""
+        x1, y1, x2, y2 = table_bbox[0], table_bbox[1], table_bbox[2], table_bbox[3]
+        h, w = image_shape
+        
+        # 1. 分析表格区域内的 OCR spans
+        spans_inside = self._get_spans_in_region(ocr_spans, table_bbox)
+        
+        # 2. 分析表格区域上方的 OCR spans(可能被错误包含的内容)
+        top_region = [0, max(0, y1 - int(h * 0.15)), w, y1]  # 表格上方15%的区域
+        spans_above = self._get_spans_in_region(ocr_spans, top_region)
+        
+        # 3. 评估表格边界准确性
+        boundary_score = self._evaluate_boundary_accuracy(
+            spans_inside, spans_above, table_bbox, image_shape
+        )
+        
+        # 4. 评估表格内容特征(列状对齐)
+        alignment_score = self._evaluate_column_alignment(spans_inside, table_bbox)
+        
+        # 5. 评估内容纯度(是否包含非表格文本块)
+        purity_score = self._evaluate_content_purity(
+            spans_inside, all_layout_results, table_bbox
+        )
+        
+        # 综合评分
+        overall_score = (
+            boundary_score * 0.4 +      # 边界准确性 40%
+            alignment_score * 0.3 +    # 列对齐特征 30%
+            purity_score * 0.3          # 内容纯度 30%
+        )
+        
+        # 生成问题报告
+        issues = []
+        if boundary_score < 0.6:
+            issues.append(f"表格{table_idx+1}: 边界可能包含了上方非表格内容")
+        if alignment_score < 0.5:
+            issues.append(f"表格{table_idx+1}: 文本块缺少列状对齐特征(可能不是表格)")
+        if purity_score < 0.6:
+            issues.append(f"表格{table_idx+1}: 包含非表格类型的文本块")
+        
+        return {
+            'score': float(overall_score),
+            'boundary_score': float(boundary_score),
+            'alignment_score': float(alignment_score),
+            'purity_score': float(purity_score),
+            'spans_inside_count': len(spans_inside),
+            'spans_above_count': len(spans_above),
+            'issues': issues
+        }
+    
+    def _get_spans_in_region(
+        self,
+        ocr_spans: List[Dict[str, Any]],
+        region_bbox: List[float]
+    ) -> List[Dict[str, Any]]:
+        """获取区域内的 OCR spans"""
+        x1, y1, x2, y2 = region_bbox[0], region_bbox[1], region_bbox[2], region_bbox[3]
+        
+        spans_in_region = []
+        for span in ocr_spans:
+            span_bbox = span.get('bbox', [])
+            if len(span_bbox) < 4:
+                continue
+            
+            sx1, sy1, sx2, sy2 = span_bbox[0], span_bbox[1], span_bbox[2], span_bbox[3]
+            
+            # 计算重叠
+            overlap_x1 = max(x1, sx1)
+            overlap_y1 = max(y1, sy1)
+            overlap_x2 = min(x2, sx2)
+            overlap_y2 = min(y2, sy2)
+            
+            if overlap_x2 > overlap_x1 and overlap_y2 > overlap_y1:
+                # 计算重叠比例
+                overlap_area = (overlap_x2 - overlap_x1) * (overlap_y2 - overlap_y1)
+                span_area = (sx2 - sx1) * (sy2 - sy1)
+                
+                if span_area > 0 and overlap_area / span_area > 0.5:  # 超过50%重叠
+                    spans_in_region.append(span)
+        
+        return spans_in_region
+    
+    def _evaluate_boundary_accuracy(
+        self,
+        spans_inside: List[Dict[str, Any]],
+        spans_above: List[Dict[str, Any]],
+        table_bbox: List[float],
+        image_shape: Tuple[int, int]
+    ) -> float:
+        """
+        评估表格边界准确性
+        
+        核心逻辑:
+        1. 如果表格上方有大量文本块,而表格区域内也包含这些文本块,说明边界设置错误
+        2. 表格应该从表头开始,不应该包含上方的账户信息
+        """
+        if not spans_inside:
+            return 0.0
+        
+        # 如果上方没有文本块,说明边界可能是合理的
+        if not spans_above:
+            return 1.0
+        
+        # 检查上方文本块是否被错误包含在表格中
+        # 通过比较上方文本块和表格内文本块的位置关系
+        
+        # 计算上方文本块的平均Y坐标
+        above_y_coords = []
+        for span in spans_above:
+            bbox = span.get('bbox', [])
+            if len(bbox) >= 4:
+                above_y_coords.append((bbox[1] + bbox[3]) / 2)  # 中心Y坐标
+        
+        if not above_y_coords:
+            return 1.0
+        
+        avg_above_y = sum(above_y_coords) / len(above_y_coords)
+        table_top = table_bbox[1]
+        
+        # 如果上方文本块的平均位置接近表格顶部,说明可能被错误包含
+        # 计算距离比例
+        h = image_shape[0]
+        distance_ratio = abs(avg_above_y - table_top) / h if h > 0 else 1.0
+        
+        # 距离越近,得分越低
+        if distance_ratio < 0.05:  # 距离小于5%图像高度
+            return 0.3  # 很可能错误包含
+        elif distance_ratio < 0.1:  # 距离小于10%
+            return 0.6
+        else:
+            return 1.0
+    
+    def _evaluate_column_alignment(
+        self,
+        spans_inside: List[Dict[str, Any]],
+        table_bbox: List[float]
+    ) -> float:
+        """
+        评估表格列对齐特征
+        
+        无线表格的文本块应该呈现列状对齐特征:
+        1. 多个文本块的X坐标应该聚集在几个列位置
+        2. 同一列的文本块应该有相似的X坐标
+        """
+        if len(spans_inside) < 3:
+            return 0.5  # 文本块太少,无法判断
+        
+        # 提取所有文本块的X坐标(左边界)
+        x_coords = []
+        for span in spans_inside:
+            bbox = span.get('bbox', [])
+            if len(bbox) >= 4:
+                x_coords.append(bbox[0])
+        
+        if len(x_coords) < 3:
+            return 0.5
+        
+        # 使用K-means聚类或简单的直方图分析来检测列
+        # 简化版:计算X坐标的分布特征
+        
+        x_coords_sorted = sorted(x_coords)
+        table_width = table_bbox[2] - table_bbox[0]
+        
+        # 将表格宽度分成若干列(假设至少3列)
+        num_bins = max(3, min(10, len(x_coords) // 5))  # 动态确定列数
+        
+        # 计算每个bin中的文本块数量
+        bin_width = table_width / num_bins
+        bin_counts = [0] * num_bins
+        
+        for x in x_coords:
+            bin_idx = min(int((x - table_bbox[0]) / bin_width), num_bins - 1)
+            bin_counts[bin_idx] += 1
+        
+        # 计算分布的集中度(好的表格应该有明显的列聚集)
+        # 使用熵来衡量分布的集中度
+        total = sum(bin_counts)
+        if total == 0:
+            return 0.5
+        
+        # 计算归一化熵(0-1,越低表示越集中)
+        probs = [count / total for count in bin_counts if count > 0]
+        entropy = -sum(p * np.log2(p) for p in probs) if probs else 0
+        max_entropy = np.log2(num_bins)
+        normalized_entropy = entropy / max_entropy if max_entropy > 0 else 1.0
+        
+        # 集中度得分(熵越低,得分越高)
+        alignment_score = 1.0 - normalized_entropy
+        
+        # 检查是否有明显的列(至少3列有明显聚集)
+        significant_columns = sum(1 for count in bin_counts if count >= len(x_coords) * 0.1)
+        
+        if significant_columns < 3:
+            alignment_score *= 0.7  # 列数不足,降低得分
+        
+        return max(0.0, min(1.0, alignment_score))
+    
+    def _evaluate_content_purity(
+        self,
+        spans_inside: List[Dict[str, Any]],
+        all_layout_results: List[Dict[str, Any]],
+        table_bbox: List[float]
+    ) -> float:
+        """
+        评估表格内容纯度
+        
+        检查表格区域内的文本块是否应该属于表格
+        如果表格区域内有很多其他layout类别的元素,说明可能错误包含了非表格内容
+        """
+        if not spans_inside:
+            return 0.0
+        
+        # 统计表格区域内其他layout类别的元素
+        other_category_elements = 0
+        
+        for layout_result in all_layout_results:
+            layout_bbox = layout_result.get('bbox', [])
+            if len(layout_bbox) < 4:
+                continue
+            
+            category = layout_result.get('category', '').lower()
+            
+            # 跳过表格类别本身
+            if 'table' in category:
+                continue
+            
+            # 检查是否在表格区域内
+            overlap = CoordinateUtils.calculate_iou(layout_bbox, table_bbox)
+            
+            if overlap > 0.3:  # 重叠超过30%
+                # 检查这个元素是否包含文本块
+                element_spans = [s for s in spans_inside 
+                               if self._span_in_bbox(s, layout_bbox)]
+                
+                if element_spans:
+                    other_category_elements += len(element_spans)
+        
+        # 计算纯度
+        total_spans = len(spans_inside)
+        if total_spans == 0:
+            return 1.0
+        
+        purity = 1.0 - (other_category_elements / total_spans)
+        return max(0.0, min(1.0, purity))
+    
+    def _span_in_bbox(
+        self,
+        span: Dict[str, Any],
+        bbox: List[float]
+    ) -> bool:
+        """检查span是否在bbox内"""
+        span_bbox = span.get('bbox', [])
+        if len(span_bbox) < 4:
+            return False
+        
+        sx1, sy1, sx2, sy2 = span_bbox[0], span_bbox[1], span_bbox[2], span_bbox[3]
+        bx1, by1, bx2, by2 = bbox[0], bbox[1], bbox[2], bbox[3]
+        
+        # 检查中心点是否在bbox内
+        center_x = (sx1 + sx2) / 2
+        center_y = (sy1 + sy2) / 2
+        
+        return bx1 <= center_x <= bx2 and by1 <= center_y <= by2
+    
+    def compare_models_with_ocr(
+        self,
+        results1: List[Dict[str, Any]],
+        results2: List[Dict[str, Any]],
+        ocr_spans: List[Dict[str, Any]],
+        image: Union[np.ndarray, Image.Image],
+        model1_name: str = "Model1",
+        model2_name: str = "Model2"
+    ) -> Dict[str, Any]:
+        """比较两个模型的检测质量(结合OCR)"""
+        eval1 = self.evaluate_with_ocr(results1, ocr_spans, image)
+        eval2 = self.evaluate_with_ocr(results2, ocr_spans, image)
+        
+        score1 = eval1.get('overall_score', 0)
+        score2 = eval2.get('overall_score', 0)
+        
+        better_model = model1_name if score1 > score2 else model2_name
+        if abs(score1 - score2) < 0.05:
+            better_model = "tie"
+        
+        return {
+            'model1': {
+                'name': model1_name,
+                'metrics': eval1
+            },
+            'model2': {
+                'name': model2_name,
+                'metrics': eval2
+            },
+            'better_model': better_model,
+            'recommendation': self._generate_recommendation(eval1, eval2, model1_name, model2_name)
+        }
+    
+    def _generate_recommendation(
+        self,
+        eval1: Dict[str, Any],
+        eval2: Dict[str, Any],
+        name1: str,
+        name2: str
+    ) -> str:
+        """生成推荐建议"""
+        score1 = eval1.get('overall_score', 0)
+        score2 = eval2.get('overall_score', 0)
+        
+        better = name1 if score1 > score2 else name2
+        better_eval = eval1 if score1 > score2 else eval2
+        
+        if abs(score1 - score2) < 0.05:
+            return "两个模型的表格检测质量相近"
+        
+        reasons = []
+        if better_eval.get('issues'):
+            worse_eval = eval2 if score1 > score2 else eval1
+            if len(better_eval['issues']) < len(worse_eval.get('issues', [])):
+                reasons.append(f"问题更少({len(better_eval['issues'])} vs {len(worse_eval.get('issues', []))})")
+        
+        reason_str = "、" .join(reasons) if reasons else "表格边界更准确"
+        
+        return f"推荐使用 {better}:{reason_str}"
+    
+    def print_evaluation_report(self, evaluation: Dict[str, Any]):
+        """打印评估报告"""
+        print("\n" + "="*60)
+        print("基于 OCR 的表格检测质量评估报告")
+        print("="*60)
+        
+        if 'model1' in evaluation:
+            # 比较报告
+            print(f"\n【{evaluation['model1']['name']}】")
+            print(f"  综合得分: {evaluation['model1']['metrics']['overall_score']:.3f}")
+            if evaluation['model1']['metrics'].get('table_scores'):
+                for i, score in enumerate(evaluation['model1']['metrics']['table_scores']):
+                    print(f"  表格{i+1}得分: {score:.3f}")
+            if evaluation['model1']['metrics'].get('issues'):
+                print(f"  问题:")
+                for issue in evaluation['model1']['metrics']['issues']:
+                    print(f"    - {issue}")
+            
+            print(f"\n【{evaluation['model2']['name']}】")
+            print(f"  综合得分: {evaluation['model2']['metrics']['overall_score']:.3f}")
+            if evaluation['model2']['metrics'].get('table_scores'):
+                for i, score in enumerate(evaluation['model2']['metrics']['table_scores']):
+                    print(f"  表格{i+1}得分: {score:.3f}")
+            if evaluation['model2']['metrics'].get('issues'):
+                print(f"  问题:")
+                for issue in evaluation['model2']['metrics']['issues']:
+                    print(f"    - {issue}")
+            
+            print(f"\n【推荐】{evaluation['recommendation']}")
+        else:
+            # 单模型报告
+            print(f"\n综合得分: {evaluation['overall_score']:.3f}")
+            if evaluation.get('table_scores'):
+                for i, score in enumerate(evaluation['table_scores']):
+                    print(f"表格{i+1}得分: {score:.3f}")
+            if evaluation.get('issues'):
+                print(f"\n发现的问题:")
+                for issue in evaluation['issues']:
+                    print(f"  - {issue}")
+        
+        print("="*60 + "\n")