瀏覽代碼

feat: Implement skew correction functionality in TableLineGenerator, enhancing image processing and OCR data handling

zhch158_admin 1 天之前
父節點
當前提交
1172c53f35
共有 1 個文件被更改,包括 145 次插入4 次删除
  1. 145 4
      table_line_generator/table_line_generator.py

+ 145 - 4
table_line_generator/table_line_generator.py

@@ -10,6 +10,16 @@ from pathlib import Path
 from typing import List, Dict, Tuple, Optional, Union
 import json
 from bs4 import BeautifulSoup
+import sys
+
+# 添加父目录到路径,以便导入 merger 模块
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+try:
+    from merger.bbox_extractor import BBoxExtractor
+except ImportError:
+    # 尝试相对导入 (当作为包安装时)
+    from ..merger.bbox_extractor import BBoxExtractor
 
 
 class TableLineGenerator:
@@ -46,6 +56,9 @@ class TableLineGenerator:
         self.columns = []
         self.row_height = 0
         self.col_widths = []
+        
+        self.is_skew_corrected = False # 是否已经校正过倾斜(默认 False)
+        self.original_image = None
 
 
     @staticmethod
@@ -125,10 +138,17 @@ class TableLineGenerator:
             'table_bbox': table_bbox,
             'actual_rows': actual_rows,
             'actual_cols': actual_cols,
-            'text_boxes': table_cells
+            'text_boxes': table_cells,
+            'image_rotation_angle': table_data.get('image_rotation_angle', 0.0),
+            'skew_angle': table_data.get('skew_angle', 0.0),
+            'original_skew_angle': table_data.get('skew_angle', 0.0)
         }
         
         print(f"📊 MinerU 数据解析完成: {len(table_cells)} 个文本框")
+        if ocr_data['image_rotation_angle'] != 0:
+            print(f"   🔄 读取到图片旋转角度: {ocr_data['image_rotation_angle']}°")
+        if ocr_data['skew_angle'] != 0:
+            print(f"   📐 读取到倾斜角度: {ocr_data['skew_angle']:.2f}°")
         
         return table_bbox, ocr_data
 
@@ -221,6 +241,115 @@ class TableLineGenerator:
         else:
             return self._analyze_by_clustering(y_tolerance, x_tolerance, min_row_height)
 
+    def correct_skew(self, force: bool = False) -> Tuple[Optional[Image.Image], float]:
+        """
+        检测并校正图片倾斜(包含整图旋转和微小倾斜校正)
+        同时会更新 self.ocr_data 中的 bbox 坐标以匹配新图片
+        
+        Args:
+            force: 是否强制重新校正
+            
+        Returns:
+            (corrected_image, total_angle): 校正后的图片和总旋转角度
+        """
+        if self.is_skew_corrected and not force:
+            # 如果已经校正过且不强制更新,直接返回当前状态
+            return self.image, 0.0
+            
+        if not self.ocr_data or 'text_boxes' not in self.ocr_data:
+            return self.image, 0.0
+            
+        text_boxes = self.ocr_data['text_boxes']
+        
+        # 1. 获取旋转和倾斜角度
+        image_rotation_angle = self.ocr_data.get('image_rotation_angle', 0.0)
+        skew_angle = self.ocr_data.get('skew_angle', 0.0)
+        
+        # 如果没有角度需要调整,且没有原始图片备份(说明没做过调整),则直接返回
+        if image_rotation_angle == 0 and abs(skew_angle) < 0.1 and not self.original_image:
+            return self.image, 0.0
+
+        # 准备源图片
+        if self.original_image:
+            # 如果有原始图片备份,从原始图片开始
+            current_image = self.original_image.copy()
+            # 恢复 text_boxes 到原始状态 (这里假设 original_bbox 存储了最初的坐标)
+            # 但实际上我们在 rotate_box_coordinates 时并没有保存 original_bbox 到 list 中
+            # 这是一个问题。如果是多次旋转,坐标会乱。
+            # 简单的做法:如果不复杂的逻辑,我们假设 self.ocr_data['text_boxes'] 里的 bbox 是相对于 self.image 的。
+            # 如果我们要重做,我们需要原始的 bbox。
+            # 在第一次 correct_skew 时,我们应该保存原始 bbox。
+            
+            # 让我们检查一下第一次 correct_skew 的逻辑。
+            # 如果是第一次,我们用 self.image。
+            pass
+        elif self.image:
+            self.original_image = self.image.copy()
+            current_image = self.image
+        else:
+            return None, 0.0
+
+        # 为了支持重做,我们需要保存原始的 OCR 数据。
+        if 'original_text_boxes' not in self.ocr_data:
+             # 深拷贝 text_boxes
+             import copy
+             self.ocr_data['original_text_boxes'] = copy.deepcopy(text_boxes)
+             # 同时保存原始 table_bbox
+             if 'table_bbox' in self.ocr_data:
+                 self.ocr_data['original_table_bbox'] = list(self.ocr_data['table_bbox'])
+        
+        # 使用原始数据进行计算
+        working_text_boxes = [box.copy() for box in self.ocr_data['original_text_boxes']]
+        original_size = self.original_image.size
+        
+        # 2. 执行图片旋转 (image_rotation_angle)
+        if image_rotation_angle != 0:
+            print(f"   🔄 执行图片旋转: {image_rotation_angle}°")
+            current_image = current_image.rotate(image_rotation_angle, expand=True)
+            
+            # 更新 bbox 坐标 (原图坐标 -> 旋转后坐标)
+            for box in working_text_boxes:
+                if 'bbox' in box:
+                    box['bbox'] = BBoxExtractor.rotate_box_coordinates(
+                        box['bbox'], image_rotation_angle, original_size
+                    )
+            
+            # 更新 table_bbox
+            if 'original_table_bbox' in self.ocr_data:
+                self.ocr_data['table_bbox'] = BBoxExtractor.rotate_box_coordinates(
+                    self.ocr_data['original_table_bbox'], image_rotation_angle, original_size
+                )
+        else:
+            # 如果没有旋转,恢复 table_bbox
+             if 'original_table_bbox' in self.ocr_data:
+                 self.ocr_data['table_bbox'] = list(self.ocr_data['original_table_bbox'])
+        
+        # 3. 执行倾斜校正 (skew_angle)
+        if abs(skew_angle) > 0.1:
+            print(f"   📐 执行倾斜校正: {skew_angle:.2f}°")
+            # 图片逆时针歪了 skew_angle 度,需要顺时针转 skew_angle 度校正
+            correction_angle = -skew_angle
+            current_image = current_image.rotate(correction_angle, expand=False, fillcolor='white')
+            
+            # 更新 bbox 坐标
+            working_text_boxes = BBoxExtractor.correct_boxes_skew(
+                working_text_boxes, 
+                correction_angle, 
+                current_image.size
+            )
+            
+            # 更新 table_bbox
+            if 'table_bbox' in self.ocr_data:
+                dummy_box = [{'bbox': self.ocr_data['table_bbox'], 'poly': BBoxExtractor._bbox_to_poly(self.ocr_data['table_bbox'])}]
+                corrected_dummy = BBoxExtractor.correct_boxes_skew(dummy_box, correction_angle, current_image.size)
+                self.ocr_data['table_bbox'] = corrected_dummy[0]['bbox']
+
+        self.image = current_image
+        self.ocr_data['text_boxes'] = working_text_boxes
+            
+        self.is_skew_corrected = True
+        return self.image, image_rotation_angle + skew_angle
+
     def _analyze_by_cell_index(self) -> Dict:
         """
         基于单元格的 row/col 索引分析(MinerU 专用)
@@ -328,6 +457,10 @@ class TableLineGenerator:
         self.row_height = int(np.median([r['y_end'] - r['y_start'] for r in self.rows])) if self.rows else 0
         self.col_widths = [c['x_end'] - c['x_start'] for c in self.columns]
         
+        # 获取角度信息
+        image_rotation_angle = self.ocr_data.get('image_rotation_angle', 0.0)
+        skew_angle = self.ocr_data.get('skew_angle', 0.0)
+        
         return {
             'rows': self.rows,
             'columns': self.columns,
@@ -340,7 +473,10 @@ class TableLineGenerator:
             'total_cols': actual_cols,
             'mode': 'hybrid',  # ✅ 添加 mode 字段
             'modified_h_lines': [],  # ✅ 添加修改记录字段
-            'modified_v_lines': []   # ✅ 添加修改记录字段
+            'modified_v_lines': [],   # ✅ 添加修改记录字段
+            'image_rotation_angle': image_rotation_angle,
+            'skew_angle': skew_angle,
+            'is_skew_corrected': self.is_skew_corrected
         }
     
     def _analyze_by_clustering(self, y_tolerance: int, x_tolerance: int, min_row_height: int) -> Dict:
@@ -415,7 +551,10 @@ class TableLineGenerator:
             'table_bbox': self._get_table_bbox(),
             'mode': 'fixed',  # ✅ 添加 mode 字段
             'modified_h_lines': [],  # ✅ 添加修改记录字段
-            'modified_v_lines': []   # ✅ 添加修改记录字段
+            'modified_v_lines': [],   # ✅ 添加修改记录字段
+            'image_rotation_angle': self.ocr_data.get('image_rotation_angle', 0.0),
+            'skew_angle': self.ocr_data.get('skew_angle', 0.0),
+            'is_skew_corrected': self.is_skew_corrected
         }
 
     @staticmethod
@@ -510,7 +649,9 @@ class TableLineGenerator:
     def _get_table_bbox(self) -> List[int]:
         """获取表格整体边界框"""
         if not self.rows or not self.columns:
-            return [0, 0, self.image.width, self.image.height]
+            if self.image:
+                return [0, 0, self.image.width, self.image.height]
+            return [0, 0, 0, 0]
         
         y_min = min(row['y_start'] for row in self.rows)
         y_max = max(row['y_end'] for row in self.rows)