瀏覽代碼

feat: Update apply_template_fixed to support both List and Dict for OCR data, enabling skew correction and angle retrieval

zhch158_admin 1 天之前
父節點
當前提交
5c8b27eb5c
共有 1 個文件被更改,包括 106 次插入10 次删除
  1. 106 10
      table_line_generator/table_template_applier.py

+ 106 - 10
table_line_generator/table_template_applier.py

@@ -6,7 +6,7 @@
 import json
 from pathlib import Path
 from PIL import Image, ImageDraw
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union, Optional
 import numpy as np
 import argparse
 import sys
@@ -127,7 +127,7 @@ class TableTemplateApplier:
     
     def apply_template_fixed(self, 
                        image: Image.Image,
-                       ocr_data: List[Dict],
+                       ocr_data: Union[List[Dict], Dict],  # 🆕 支持 Dict
                        anchor_x: int = None,
                        anchor_y: int = None,
                        num_rows: int = None,
@@ -138,7 +138,7 @@ class TableTemplateApplier:
         
         Args:
             image: 目标图片
-            ocr_data: OCR识别结果(用于自动检测锚点)
+            ocr_data: OCR识别结果(用于自动检测锚点),可以是列表或完整字典
             anchor_x: 表格起始X坐标(None=自动检测)
             anchor_y: 表头起始Y坐标(None=自动检测)
             num_rows: 总行数(None=自动检测)
@@ -148,7 +148,78 @@ class TableTemplateApplier:
         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)
         
         # 🔍 自动检测锚点
@@ -202,13 +273,15 @@ class TableTemplateApplier:
 
         print(f"✅ 表格绘制完成: {len(horizontal_lines)}行 × {len(vertical_lines)-1}列")
 
-                # 🔑 生成结构信息
+        # 🔑 生成结构信息
         structure = self._build_structure(
             horizontal_lines, 
             vertical_lines, 
             anchor_x, 
             anchor_y,
-            mode='fixed'
+            mode='fixed',
+            image_rotation_angle=image_rotation_angle,
+            skew_angle=skew_angle
         )
         
         return img_with_lines, structure
@@ -238,7 +311,23 @@ class TableTemplateApplier:
         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)
         
         ocr_data = ocr_data_dict.get('text_boxes', [])
@@ -286,7 +375,9 @@ class TableTemplateApplier:
             vertical_lines, 
             anchor_x, 
             anchor_y,
-            mode='hybrid'
+            mode='hybrid',
+            image_rotation_angle=image_rotation_angle,
+            skew_angle=skew_angle
         )
         
         return img_with_lines, structure
@@ -389,7 +480,9 @@ class TableTemplateApplier:
                         vertical_lines: List[int],
                         anchor_x: int,
                         anchor_y: int,
-                        mode: str = 'fixed') -> Dict:
+                        mode: str = 'fixed',
+                        image_rotation_angle: float = 0.0,
+                        skew_angle: float = 0.0) -> Dict:
         """构建表格结构信息(统一)"""
         # 生成行区间
         rows = []
@@ -432,7 +525,10 @@ class TableTemplateApplier:
             'mode': mode_value,  # ✅ 确保有 mode 字段
             'anchor': {'x': anchor_x, 'y': anchor_y},
             '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(