Browse Source

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

zhch158_admin 22 hours ago
parent
commit
1172c53f35
1 changed files with 145 additions and 4 deletions
  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
 from typing import List, Dict, Tuple, Optional, Union
 import json
 import json
 from bs4 import BeautifulSoup
 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:
 class TableLineGenerator:
@@ -46,6 +56,9 @@ class TableLineGenerator:
         self.columns = []
         self.columns = []
         self.row_height = 0
         self.row_height = 0
         self.col_widths = []
         self.col_widths = []
+        
+        self.is_skew_corrected = False # 是否已经校正过倾斜(默认 False)
+        self.original_image = None
 
 
 
 
     @staticmethod
     @staticmethod
@@ -125,10 +138,17 @@ class TableLineGenerator:
             'table_bbox': table_bbox,
             'table_bbox': table_bbox,
             'actual_rows': actual_rows,
             'actual_rows': actual_rows,
             'actual_cols': actual_cols,
             '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)} 个文本框")
         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
         return table_bbox, ocr_data
 
 
@@ -221,6 +241,115 @@ class TableLineGenerator:
         else:
         else:
             return self._analyze_by_clustering(y_tolerance, x_tolerance, min_row_height)
             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:
     def _analyze_by_cell_index(self) -> Dict:
         """
         """
         基于单元格的 row/col 索引分析(MinerU 专用)
         基于单元格的 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.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]
         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 {
         return {
             'rows': self.rows,
             'rows': self.rows,
             'columns': self.columns,
             'columns': self.columns,
@@ -340,7 +473,10 @@ class TableLineGenerator:
             'total_cols': actual_cols,
             'total_cols': actual_cols,
             'mode': 'hybrid',  # ✅ 添加 mode 字段
             'mode': 'hybrid',  # ✅ 添加 mode 字段
             'modified_h_lines': [],  # ✅ 添加修改记录字段
             '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:
     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(),
             'table_bbox': self._get_table_bbox(),
             'mode': 'fixed',  # ✅ 添加 mode 字段
             'mode': 'fixed',  # ✅ 添加 mode 字段
             'modified_h_lines': [],  # ✅ 添加修改记录字段
             '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
     @staticmethod
@@ -510,7 +649,9 @@ class TableLineGenerator:
     def _get_table_bbox(self) -> List[int]:
     def _get_table_bbox(self) -> List[int]:
         """获取表格整体边界框"""
         """获取表格整体边界框"""
         if not self.rows or not self.columns:
         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_min = min(row['y_start'] for row in self.rows)
         y_max = max(row['y_end'] for row in self.rows)
         y_max = max(row['y_end'] for row in self.rows)