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