|
|
@@ -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")
|