|
@@ -6,7 +6,7 @@
|
|
|
import json
|
|
import json
|
|
|
from pathlib import Path
|
|
from pathlib import Path
|
|
|
from PIL import Image, ImageDraw
|
|
from PIL import Image, ImageDraw
|
|
|
-from typing import Dict, List, Tuple
|
|
|
|
|
|
|
+from typing import Dict, List, Tuple, Union, Optional
|
|
|
import numpy as np
|
|
import numpy as np
|
|
|
import argparse
|
|
import argparse
|
|
|
import sys
|
|
import sys
|
|
@@ -127,7 +127,7 @@ class TableTemplateApplier:
|
|
|
|
|
|
|
|
def apply_template_fixed(self,
|
|
def apply_template_fixed(self,
|
|
|
image: Image.Image,
|
|
image: Image.Image,
|
|
|
- ocr_data: List[Dict],
|
|
|
|
|
|
|
+ ocr_data: Union[List[Dict], Dict], # 🆕 支持 Dict
|
|
|
anchor_x: int = None,
|
|
anchor_x: int = None,
|
|
|
anchor_y: int = None,
|
|
anchor_y: int = None,
|
|
|
num_rows: int = None,
|
|
num_rows: int = None,
|
|
@@ -138,7 +138,7 @@ class TableTemplateApplier:
|
|
|
|
|
|
|
|
Args:
|
|
Args:
|
|
|
image: 目标图片
|
|
image: 目标图片
|
|
|
- ocr_data: OCR识别结果(用于自动检测锚点)
|
|
|
|
|
|
|
+ ocr_data: OCR识别结果(用于自动检测锚点),可以是列表或完整字典
|
|
|
anchor_x: 表格起始X坐标(None=自动检测)
|
|
anchor_x: 表格起始X坐标(None=自动检测)
|
|
|
anchor_y: 表头起始Y坐标(None=自动检测)
|
|
anchor_y: 表头起始Y坐标(None=自动检测)
|
|
|
num_rows: 总行数(None=自动检测)
|
|
num_rows: 总行数(None=自动检测)
|
|
@@ -148,7 +148,78 @@ class TableTemplateApplier:
|
|
|
Returns:
|
|
Returns:
|
|
|
绘制了表格线的图片
|
|
绘制了表格线的图片
|
|
|
"""
|
|
"""
|
|
|
- img_with_lines = image.copy()
|
|
|
|
|
|
|
+ # 🆕 1. 实例化生成器并进行倾斜校正
|
|
|
|
|
+ ocr_data_dict = {'text_boxes': ocr_data}
|
|
|
|
|
+ # 尝试从 ocr_data 列表中获取角度信息(如果它是从 ocr_data 字典中提取出来的 list)
|
|
|
|
|
+ # 但通常 ocr_data 这里只是 text_boxes 列表。
|
|
|
|
|
+ # 我们需要传递包含 image_rotation_angle 和 skew_angle 的字典。
|
|
|
|
|
+ # 由于调用者可能会传入 list,我们需要检查是否有更多信息。
|
|
|
|
|
+ # 这里假设调用者会在传入 list 前处理好,或者我们在这里无法获取。
|
|
|
|
|
+ # 不过,如果是从 parse_ocr_data 获取的 ocr_data,它应该是个 dict。
|
|
|
|
|
+ # apply_template_fixed 的签名是 ocr_data: List[Dict],这意味着它只接收 text_boxes。
|
|
|
|
|
+ # 这可能是一个问题。我们需要修改调用处或者在这里处理。
|
|
|
|
|
+ # 看看 apply_template_to_single_file 是怎么调用的。
|
|
|
|
|
+
|
|
|
|
|
+ # apply_template_to_single_file:
|
|
|
|
|
+ # text_boxes = ocr_data.get('text_boxes', [])
|
|
|
|
|
+ # applier.apply_template_fixed(image, text_boxes, ...)
|
|
|
|
|
+
|
|
|
|
|
+ # 这样我们就丢失了角度信息。
|
|
|
|
|
+ # 我应该修改 apply_template_fixed 的签名,让它接收 Dict 类型的 ocr_data,或者单独传递角度。
|
|
|
|
|
+ # 为了保持兼容性,我可以修改 apply_template_fixed 内部处理。
|
|
|
|
|
+ # 但最好的方式是让它接收整个 ocr_data 字典,就像 apply_template_hybrid 一样。
|
|
|
|
|
+
|
|
|
|
|
+ # 不过,为了最小化修改,我可以在 apply_template_to_single_file 里把角度传进来?
|
|
|
|
|
+ # 不,那得改很多。
|
|
|
|
|
+
|
|
|
|
|
+ # 让我们看看能不能在 apply_template_fixed 里重新构造 ocr_data_dict。
|
|
|
|
|
+ # 如果传入的 ocr_data 是 list,那我们确实没法知道角度。
|
|
|
|
|
+ # 除非我们改变 apply_template_to_single_file 的调用方式。
|
|
|
|
|
+
|
|
|
|
|
+ # 让我们先修改 apply_template_to_single_file 的调用方式,传整个 ocr_data 进去。
|
|
|
|
|
+ # 但是 apply_template_fixed 的签名明确写了 ocr_data: List[Dict]。
|
|
|
|
|
+
|
|
|
|
|
+ # 既然我正在修改这个文件,我可以改变它的签名。
|
|
|
|
|
+ # 或者,我可以像 apply_template_hybrid 一样,增加一个参数 ocr_data_full: Dict = None
|
|
|
|
|
+
|
|
|
|
|
+ # 实际上,apply_template_hybrid 已经接收 ocr_data_dict: Dict。
|
|
|
|
|
+ # apply_template_fixed 接收 List[Dict]。
|
|
|
|
|
+ # 这是一个不一致的地方。
|
|
|
|
|
+
|
|
|
|
|
+ # 我决定修改 apply_template_fixed 的参数,让它也能利用 TableLineGenerator 进行校正。
|
|
|
|
|
+ # 但是 TableLineGenerator 需要完整的 ocr_data 字典才能读取角度。
|
|
|
|
|
+
|
|
|
|
|
+ # 方案:修改 apply_template_fixed 接收 ocr_data_dict。
|
|
|
|
|
+ # 为了兼容旧代码,如果传入的是 list,就包装一下。
|
|
|
|
|
+
|
|
|
|
|
+ # 但是 Python 类型提示 List[Dict] 和 Dict 是不一样的。
|
|
|
|
|
+ # 我可以把参数名改成 ocr_input,类型 Union[List[Dict], Dict]。
|
|
|
|
|
+
|
|
|
|
|
+ # 或者,既然这是内部使用的工具,我直接修改签名,让它接收 Dict。
|
|
|
|
|
+ # 检查一下是否有其他地方调用这个方法。
|
|
|
|
|
+ # 只在 apply_template_to_single_file 调用了。
|
|
|
|
|
+
|
|
|
|
|
+ # 所以我将修改 apply_template_fixed 接收 ocr_data_dict: Dict。
|
|
|
|
|
+
|
|
|
|
|
+ generator = TableLineGenerator(image, {'text_boxes': ocr_data} if isinstance(ocr_data, list) else ocr_data)
|
|
|
|
|
+ corrected_image, angle = generator.correct_skew()
|
|
|
|
|
+
|
|
|
|
|
+ # 获取角度信息
|
|
|
|
|
+ image_rotation_angle = generator.ocr_data.get('image_rotation_angle', 0.0)
|
|
|
|
|
+ skew_angle = generator.ocr_data.get('skew_angle', 0.0)
|
|
|
|
|
+
|
|
|
|
|
+ if abs(angle) > 0.1 or image_rotation_angle != 0:
|
|
|
|
|
+ print(f"🔄 [TemplateApplier] 自动校正: 旋转={image_rotation_angle}°, 倾斜={skew_angle:.2f}°")
|
|
|
|
|
+ # 更新 OCR 数据(generator 内部已经更新了)
|
|
|
|
|
+ ocr_data = generator.ocr_data.get('text_boxes', [])
|
|
|
|
|
+ # 使用校正后的图片
|
|
|
|
|
+ img_with_lines = corrected_image.copy()
|
|
|
|
|
+ else:
|
|
|
|
|
+ img_with_lines = image.copy()
|
|
|
|
|
+ # 如果是字典,提取 list
|
|
|
|
|
+ if isinstance(ocr_data, dict):
|
|
|
|
|
+ ocr_data = ocr_data.get('text_boxes', [])
|
|
|
|
|
+
|
|
|
draw = ImageDraw.Draw(img_with_lines)
|
|
draw = ImageDraw.Draw(img_with_lines)
|
|
|
|
|
|
|
|
# 🔍 自动检测锚点
|
|
# 🔍 自动检测锚点
|
|
@@ -202,13 +273,15 @@ class TableTemplateApplier:
|
|
|
|
|
|
|
|
print(f"✅ 表格绘制完成: {len(horizontal_lines)}行 × {len(vertical_lines)-1}列")
|
|
print(f"✅ 表格绘制完成: {len(horizontal_lines)}行 × {len(vertical_lines)-1}列")
|
|
|
|
|
|
|
|
- # 🔑 生成结构信息
|
|
|
|
|
|
|
+ # 🔑 生成结构信息
|
|
|
structure = self._build_structure(
|
|
structure = self._build_structure(
|
|
|
horizontal_lines,
|
|
horizontal_lines,
|
|
|
vertical_lines,
|
|
vertical_lines,
|
|
|
anchor_x,
|
|
anchor_x,
|
|
|
anchor_y,
|
|
anchor_y,
|
|
|
- mode='fixed'
|
|
|
|
|
|
|
+ mode='fixed',
|
|
|
|
|
+ image_rotation_angle=image_rotation_angle,
|
|
|
|
|
+ skew_angle=skew_angle
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
return img_with_lines, structure
|
|
return img_with_lines, structure
|
|
@@ -238,7 +311,23 @@ class TableTemplateApplier:
|
|
|
Returns:
|
|
Returns:
|
|
|
绘制了表格线的图片, 结构信息
|
|
绘制了表格线的图片, 结构信息
|
|
|
"""
|
|
"""
|
|
|
- img_with_lines = image.copy()
|
|
|
|
|
|
|
+ # 🆕 1. 实例化生成器并进行倾斜校正
|
|
|
|
|
+ generator = TableLineGenerator(image, ocr_data_dict)
|
|
|
|
|
+ corrected_image, angle = generator.correct_skew()
|
|
|
|
|
+
|
|
|
|
|
+ # 🆕 获取图片旋转角度
|
|
|
|
|
+ image_rotation_angle = ocr_data_dict.get('image_rotation_angle', 0.0)
|
|
|
|
|
+ skew_angle = ocr_data_dict.get('skew_angle', 0.0)
|
|
|
|
|
+
|
|
|
|
|
+ if abs(angle) > 0.1 or image_rotation_angle != 0:
|
|
|
|
|
+ print(f"🔄 [TemplateApplier] 自动校正: 旋转={image_rotation_angle}°, 倾斜={skew_angle:.2f}°")
|
|
|
|
|
+ # 更新 OCR 数据
|
|
|
|
|
+ ocr_data_dict = generator.ocr_data
|
|
|
|
|
+ # 使用校正后的图片
|
|
|
|
|
+ img_with_lines = corrected_image.copy()
|
|
|
|
|
+ else:
|
|
|
|
|
+ img_with_lines = image.copy()
|
|
|
|
|
+
|
|
|
draw = ImageDraw.Draw(img_with_lines)
|
|
draw = ImageDraw.Draw(img_with_lines)
|
|
|
|
|
|
|
|
ocr_data = ocr_data_dict.get('text_boxes', [])
|
|
ocr_data = ocr_data_dict.get('text_boxes', [])
|
|
@@ -286,7 +375,9 @@ class TableTemplateApplier:
|
|
|
vertical_lines,
|
|
vertical_lines,
|
|
|
anchor_x,
|
|
anchor_x,
|
|
|
anchor_y,
|
|
anchor_y,
|
|
|
- mode='hybrid'
|
|
|
|
|
|
|
+ mode='hybrid',
|
|
|
|
|
+ image_rotation_angle=image_rotation_angle,
|
|
|
|
|
+ skew_angle=skew_angle
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
return img_with_lines, structure
|
|
return img_with_lines, structure
|
|
@@ -389,7 +480,9 @@ class TableTemplateApplier:
|
|
|
vertical_lines: List[int],
|
|
vertical_lines: List[int],
|
|
|
anchor_x: int,
|
|
anchor_x: int,
|
|
|
anchor_y: int,
|
|
anchor_y: int,
|
|
|
- mode: str = 'fixed') -> Dict:
|
|
|
|
|
|
|
+ mode: str = 'fixed',
|
|
|
|
|
+ image_rotation_angle: float = 0.0,
|
|
|
|
|
+ skew_angle: float = 0.0) -> Dict:
|
|
|
"""构建表格结构信息(统一)"""
|
|
"""构建表格结构信息(统一)"""
|
|
|
# 生成行区间
|
|
# 生成行区间
|
|
|
rows = []
|
|
rows = []
|
|
@@ -432,7 +525,10 @@ class TableTemplateApplier:
|
|
|
'mode': mode_value, # ✅ 确保有 mode 字段
|
|
'mode': mode_value, # ✅ 确保有 mode 字段
|
|
|
'anchor': {'x': anchor_x, 'y': anchor_y},
|
|
'anchor': {'x': anchor_x, 'y': anchor_y},
|
|
|
'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': abs(skew_angle) > 0.1 or image_rotation_angle != 0
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
def apply_template_to_single_file(
|
|
def apply_template_to_single_file(
|