Bladeren bron

feat: 增加OCR补偿功能以恢复缺失单元格,优化网格结构重建逻辑

zhch158_admin 4 dagen geleden
bovenliggende
commit
f6c8b03d9c
1 gewijzigde bestanden met toevoegingen van 223 en 1 verwijderingen
  1. 223 1
      ocr_tools/universal_doc_parser/models/adapters/wired_table/grid_recovery.py

+ 223 - 1
ocr_tools/universal_doc_parser/models/adapters/wired_table/grid_recovery.py

@@ -531,7 +531,11 @@ class GridRecovery:
         return grid_lines
     
     @staticmethod
-    def recover_grid_structure(bboxes: List[List[float]]) -> List[Dict]:
+    def recover_grid_structure(
+        bboxes: List[List[float]], 
+        ocr_bboxes: Optional[List[Dict]] = None,
+        enable_ocr_compensation: bool = True
+    ) -> List[Dict]:
         """
         从散乱的单元格 bbox 恢复表格的行列结构 (row, col, rowspan, colspan)
         重构版:基于投影网格线 (Projected Grid Lines) 的算法
@@ -539,6 +543,8 @@ class GridRecovery:
         
         Args:
             bboxes: 单元格bbox列表
+            ocr_bboxes: 整页OCR结果 [{'bbox': [x1,y1,x2,y2], 'text': '...'}, ...](可选)
+            enable_ocr_compensation: 是否启用OCR补偿缺失单元格
             
         Returns:
             结构化单元格列表,包含 row, col, rowspan, colspan
@@ -561,6 +567,15 @@ class GridRecovery:
             x_coords.append(b[2])
         col_dividers = GridRecovery.find_grid_lines(x_coords, tolerance=5, min_support=2)
         
+        # 2.5. OCR补偿缺失单元格(在分配row/col之前)
+        if enable_ocr_compensation and ocr_bboxes:
+            compensated_bboxes = GridRecovery._compensate_with_ocr(
+                bboxes, ocr_bboxes, row_dividers, col_dividers
+            )
+            if compensated_bboxes:
+                logger.info(f"🔧 OCR补偿: +{len(compensated_bboxes)} 个缺失单元格")
+                bboxes = bboxes + compensated_bboxes
+        
         # 3. 构建网格结构
         structured_cells = []
         
@@ -722,3 +737,210 @@ class GridRecovery:
             new_cells.append(new_cell)
 
         return new_cells
+    
+    @staticmethod
+    def _compensate_with_ocr(
+        existing_bboxes: List[List[float]],
+        ocr_bboxes: List[Dict],
+        row_dividers: List[float],
+        col_dividers: List[float],
+        min_overlap_ratio: float = 0.3
+    ) -> List[List[float]]:
+        """
+        利用整页OCR信息补偿缺失的单元格
+        
+        策略:
+        1. 计算所有理论单元格位置(基于网格线)
+        2. 检查哪些理论位置有OCR内容但没有检测到单元格
+        3. 根据OCR bbox跨越的网格数量自动判断是否为合并单元格
+        4. 只补偿有相邻单元格的位置(避免孤立补偿)
+        
+        Args:
+            existing_bboxes: 已检测到的单元格bbox
+            ocr_bboxes: 整页OCR结果 [{'bbox': [x1,y1,x2,y2], 'text': '...'}, ...]
+            row_dividers: 行分割线
+            col_dividers: 列分割线
+            min_overlap_ratio: OCR bbox与理论单元格的最小重叠率
+        
+        Returns:
+            补偿的bbox列表
+        """
+        if not ocr_bboxes or len(row_dividers) < 2 or len(col_dividers) < 2:
+            return []
+        
+        # 1. 构建已存在单元格的覆盖区域(快速查找)
+        existing_coverage = set()
+        for bbox in existing_bboxes:
+            # 计算该bbox覆盖的理论网格区域
+            covered_rows = []
+            covered_cols = []
+            
+            for i in range(len(row_dividers) - 1):
+                if GridRecovery._has_overlap_1d(bbox[1], bbox[3], row_dividers[i], row_dividers[i+1]):
+                    covered_rows.append(i)
+            
+            for j in range(len(col_dividers) - 1):
+                if GridRecovery._has_overlap_1d(bbox[0], bbox[2], col_dividers[j], col_dividers[j+1]):
+                    covered_cols.append(j)
+            
+            # 标记覆盖的所有理论单元格
+            for r in covered_rows:
+                for c in covered_cols:
+                    existing_coverage.add((r, c))
+        
+        logger.debug(f"📊 理论网格: {len(row_dividers)-1}行 × {len(col_dividers)-1}列, 已覆盖: {len(existing_coverage)} 个单元格")
+        
+        # 2. 遍历OCR结果,查找缺失的单元格
+        compensated_bboxes = []
+        ocr_processed = set()  # 避免重复补偿
+        
+        for ocr in ocr_bboxes:
+            ocr_bbox = ocr['bbox']
+            ocr_text = ocr.get('text', '')
+            
+            # 计算OCR bbox覆盖的理论网格区域
+            covered_rows = []
+            covered_cols = []
+            
+            for i in range(len(row_dividers) - 1):
+                theoretical_bbox = [col_dividers[0], row_dividers[i], col_dividers[-1], row_dividers[i+1]]
+                overlap = GridRecovery._compute_overlap_ratio(ocr_bbox, theoretical_bbox)
+                if overlap > min_overlap_ratio * 0.5:  # 行方向用更宽松的阈值
+                    covered_rows.append(i)
+            
+            for j in range(len(col_dividers) - 1):
+                theoretical_bbox = [col_dividers[j], row_dividers[0], col_dividers[j+1], row_dividers[-1]]
+                overlap = GridRecovery._compute_overlap_ratio(ocr_bbox, theoretical_bbox)
+                if overlap > min_overlap_ratio * 0.5:  # 列方向用更宽松的阈值
+                    covered_cols.append(j)
+            
+            if not covered_rows or not covered_cols:
+                continue
+            
+            # 找出缺失的单元格(逐个检查,而不是要求全部缺失)
+            missing_cells = []
+            for r in covered_rows:
+                for c in covered_cols:
+                    if (r, c) not in existing_coverage:
+                        missing_cells.append((r, c))
+            
+            if not missing_cells:
+                continue  # 该OCR覆盖的区域没有缺失单元格
+            
+            # 检查缺失单元格是否有相邻单元格(避免孤立补偿)
+            valid_missing_cells = []
+            for r, c in missing_cells:
+                # 检查上下左右是否有相邻单元格(已存在或待补偿)
+                if ((r-1, c) in existing_coverage or
+                    (r+1, c) in existing_coverage or
+                    (r, c-1) in existing_coverage or
+                    (r, c+1) in existing_coverage or
+                    (r-1, c) in missing_cells or
+                    (r+1, c) in missing_cells or
+                    (r, c-1) in missing_cells or
+                    (r, c+1) in missing_cells):
+                    valid_missing_cells.append((r, c))
+            
+            if not valid_missing_cells:
+                logger.debug(f"⚠️ 跳过孤立OCR: '{ocr_text[:20]}' at (R{covered_rows}C{covered_cols})")
+                continue
+            
+            # 新策略:通过观察相邻单元格的分布推断合并尺寸
+            # 为每个缺失单元格分析其应该占据的网格范围
+            for r, c in valid_missing_cells:
+                # 避免重复补偿
+                if (r, c) in ocr_processed:
+                    continue
+                
+                # 分析该位置的行列跨度
+                # 新策略:不仅检查当前OCR内的缺失单元格,还检查整个网格的覆盖情况
+                
+                # 1. 向下探测:检查同列(c)中有多少连续的缺失单元格(不限于当前OCR)
+                rowspan_candidate = 1
+                r_check = r + 1
+                while r_check < len(row_dividers) - 1:
+                    # 关键改变:检查existing_coverage而不只是valid_missing_cells
+                    if (r_check, c) not in existing_coverage and (r_check, c) not in ocr_processed:
+                        rowspan_candidate += 1
+                        r_check += 1
+                    else:
+                        break
+                
+                # 2. 向右探测:检查同行(r)中有多少连续的缺失单元格(不限于当前OCR)
+                colspan_candidate = 1
+                c_check = c + 1
+                while c_check < len(col_dividers) - 1:
+                    # 关键改变:检查existing_coverage而不只是valid_missing_cells
+                    if (r, c_check) not in existing_coverage and (r, c_check) not in ocr_processed:
+                        colspan_candidate += 1
+                        c_check += 1
+                    else:
+                        break
+                
+                # 3. 验证:检查推断出的矩形区域内的所有单元格是否都缺失
+                is_valid_merge = True
+                cells_to_process = []
+                if rowspan_candidate > 1 or colspan_candidate > 1:
+                    for rr in range(r, r + rowspan_candidate):
+                        for cc in range(c, c + colspan_candidate):
+                            if (rr, cc) in existing_coverage or (rr, cc) in ocr_processed:
+                                # 区域内有单元格已存在或已处理
+                                is_valid_merge = False
+                                break
+                            cells_to_process.append((rr, cc))
+                        if not is_valid_merge:
+                            break
+                
+                # 4. 如果不是有效的合并,降级为1×1
+                if not is_valid_merge or not cells_to_process:
+                    rowspan_candidate = 1
+                    colspan_candidate = 1
+                    cells_to_process = [(r, c)]
+                
+                # 生成补偿bbox
+                compensated_bbox = [
+                    col_dividers[c],
+                    row_dividers[r],
+                    col_dividers[c + colspan_candidate],
+                    row_dividers[r + rowspan_candidate]
+                ]
+                
+                # 标记已处理的区域
+                region_key = (r, r + rowspan_candidate - 1, c, c + colspan_candidate - 1)
+                if region_key in ocr_processed:
+                    continue
+                    
+                # 标记所有涉及的单元格为已处理和已覆盖
+                for rr, cc in cells_to_process:
+                    ocr_processed.add((rr, cc))
+                    existing_coverage.add((rr, cc))
+                
+                compensated_bboxes.append(compensated_bbox)
+                
+                merge_info = f"({rowspan_candidate}×{colspan_candidate}合并)" if rowspan_candidate > 1 or colspan_candidate > 1 else "(1×1)"
+                logger.info(
+                    f"✨ 补偿单元格: '{ocr_text[:30]}' at R{r}C{c} {merge_info}"
+                )
+        
+        return compensated_bboxes
+    
+    @staticmethod
+    def _has_overlap_1d(a1: float, a2: float, b1: float, b2: float) -> bool:
+        """判断两个1维区间是否有重叠"""
+        return max(a1, b1) < min(a2, b2)
+    
+    @staticmethod
+    def _compute_overlap_ratio(bbox1: List[float], bbox2: List[float]) -> float:
+        """计算bbox1与bbox2的重叠率(相对于bbox1的面积)"""
+        x1 = max(bbox1[0], bbox2[0])
+        y1 = max(bbox1[1], bbox2[1])
+        x2 = min(bbox1[2], bbox2[2])
+        y2 = min(bbox1[3], bbox2[3])
+        
+        if x2 <= x1 or y2 <= y1:
+            return 0.0
+        
+        overlap_area = (x2 - x1) * (y2 - y1)
+        bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
+        
+        return overlap_area / bbox1_area if bbox1_area > 0 else 0.0