""" 表格模板应用器 将人工标注的表格结构应用到其他页面 """ import json from pathlib import Path from PIL import Image, ImageDraw from typing import Dict, List, Tuple import numpy as np import argparse try: from table_line_generator import TableLineGenerator except ImportError: from .table_line_generator import TableLineGenerator class TableTemplateApplier: """表格模板应用器""" def __init__(self, template_config_path: str): """ 初始化模板应用器 Args: template_config_path: 模板配置文件路径（人工标注的结果） """ with open(template_config_path, 'r', encoding='utf-8') as f: self.template = json.load(f) # 🎯 从标注结果提取固定参数 self.col_widths = self.template['col_widths'] # 🔧 计算数据行的标准行高（排除表头） rows = self.template['rows'] if len(rows) > 1: # 计算每行的实际高度 row_heights = [row['y_end'] - row['y_start'] for row in rows] # 🎯 假设第一行是表头，从第二行开始计算 data_row_heights = row_heights[1:] if len(row_heights) > 1 else row_heights # 使用中位数作为标准行高（更稳健） self.row_height = int(np.median(data_row_heights)) self.header_height = row_heights[0] if row_heights else self.row_height print(f"📏 表头高度: {self.header_height}px") print(f"📏 数据行高度: {self.row_height}px") print(f" （从 {len(data_row_heights)} 行数据中计算，中位数）") else: # 兜底方案 self.row_height = self.template.get('row_height', 60) self.header_height = self.row_height # 🎯 计算列的相对位置（从第一列开始的偏移量） self.col_offsets = [0] for width in self.col_widths: self.col_offsets.append(self.col_offsets[-1] + width) # 🎯 提取表头的Y坐标（作为参考） self.template_header_y = rows[0]['y_start'] if rows else 0 print(f"\n✅ 加载模板配置:") print(f" 表头高度: {self.header_height}px") print(f" 数据行高度: {self.row_height}px") print(f" 列数: {len(self.col_widths)}") print(f" 列宽: {self.col_widths}") def detect_table_anchor(self, ocr_data: List[Dict]) -> Tuple[int, int]: """ 检测表格的锚点位置（表头左上角） 策略： 1. 找到Y坐标最小的文本框（表头第一行） 2. 找到X坐标最小的文本框（第一列） Args: ocr_data: OCR识别结果 Returns: (anchor_x, anchor_y): 表格左上角坐标 """ if not ocr_data: return (0, 0) # 找到最小的X和Y坐标 min_x = min(item['bbox'][0] for item in ocr_data) min_y = min(item['bbox'][1] for item in ocr_data) return (min_x, min_y) def detect_table_rows(self, ocr_data: List[Dict], header_y: int) -> int: """ 检测表格的行数（包括表头） 策略： 1. 找到Y坐标最大的文本框 2. 根据数据行高计算行数 3. 加上表头行 Args: ocr_data: OCR识别结果 header_y: 表头起始Y坐标 Returns: 总行数（包括表头） """ if not ocr_data: return 1 # 至少有表头 max_y = max(item['bbox'][3] for item in ocr_data) # 🔧 计算数据区的高度（排除表头） data_start_y = header_y + self.header_height data_height = max_y - data_start_y # 计算数据行数 num_data_rows = max(int(data_height / self.row_height), 0) # 总行数 = 1行表头 + n行数据 total_rows = 1 + num_data_rows print(f"📊 行数计算:") print(f" 表头Y: {header_y}, 数据区起始Y: {data_start_y}") print(f" 最大Y: {max_y}, 数据区高度: {data_height}px") print(f" 数据行数: {num_data_rows}, 总行数: {total_rows}") return total_rows def apply_to_image(self, image: Image.Image, ocr_data: List[Dict], anchor_x: int = None, anchor_y: int = None, num_rows: int = None, line_width: int = 2, line_color: Tuple[int, int, int] = (0, 0, 0)) -> Image.Image: """ 将模板应用到图片 Args: image: 目标图片 ocr_data: OCR识别结果（用于自动检测锚点） anchor_x: 表格起始X坐标（None=自动检测） anchor_y: 表头起始Y坐标（None=自动检测） num_rows: 总行数（None=自动检测） line_width: 线条宽度 line_color: 线条颜色 Returns: 绘制了表格线的图片 """ img_with_lines = image.copy() draw = ImageDraw.Draw(img_with_lines) # 🔍 自动检测锚点 if anchor_x is None or anchor_y is None: detected_x, detected_y = self.detect_table_anchor(ocr_data) anchor_x = anchor_x or detected_x anchor_y = anchor_y or detected_y # 🔍 自动检测行数 if num_rows is None: num_rows = self.detect_table_rows(ocr_data, anchor_y) print(f"\n📍 表格锚点: ({anchor_x}, {anchor_y})") print(f"📊 总行数: {num_rows} (1表头 + {num_rows-1}数据)") # 🎨 生成横线坐标 horizontal_lines = [] # 第1条线：表头顶部 horizontal_lines.append(anchor_y) # 第2条线：表头底部/数据区顶部 horizontal_lines.append(anchor_y + self.header_height) # 后续横线：数据行分隔线 current_y = anchor_y + self.header_height for i in range(num_rows - 1): # 减1因为表头已经占了1行 current_y += self.row_height horizontal_lines.append(current_y) # 🎨 生成竖线坐标 vertical_lines = [] for offset in self.col_offsets: x = anchor_x + offset vertical_lines.append(x) print(f"📏 横线坐标: {horizontal_lines[:3]}... (共{len(horizontal_lines)}条)") print(f"📏 竖线坐标: {vertical_lines[:3]}... (共{len(vertical_lines)}条)") # 🖊️ 绘制横线 x_start = vertical_lines[0] x_end = vertical_lines[-1] for y in horizontal_lines: draw.line([(x_start, y), (x_end, y)], fill=line_color, width=line_width) # 🖊️ 绘制竖线 y_start = horizontal_lines[0] y_end = horizontal_lines[-1] for x in vertical_lines: draw.line([(x, y_start), (x, y_end)], fill=line_color, width=line_width) return img_with_lines def generate_structure_for_image(self, ocr_data: List[Dict], anchor_x: int = None, anchor_y: int = None, num_rows: int = None) -> Dict: """ 为新图片生成表格结构配置 Args: ocr_data: OCR识别结果 anchor_x: 表格起始X坐标（None=自动检测） anchor_y: 表头起始Y坐标（None=自动检测） num_rows: 总行数（None=自动检测） Returns: 表格结构配置 """ # 🔍 自动检测锚点 if anchor_x is None or anchor_y is None: detected_x, detected_y = self.detect_table_anchor(ocr_data) anchor_x = anchor_x or detected_x anchor_y = anchor_y or detected_y # 🔍 自动检测行数 if num_rows is None: num_rows = self.detect_table_rows(ocr_data, anchor_y) # 🎨 生成横线坐标 horizontal_lines = [] horizontal_lines.append(anchor_y) horizontal_lines.append(anchor_y + self.header_height) current_y = anchor_y + self.header_height for i in range(num_rows - 1): current_y += self.row_height horizontal_lines.append(current_y) # 🎨 生成竖线坐标 vertical_lines = [] for offset in self.col_offsets: x = anchor_x + offset vertical_lines.append(x) # 🎨 生成行区间 rows = [] for i in range(num_rows): rows.append({ 'y_start': horizontal_lines[i], 'y_end': horizontal_lines[i + 1], 'bboxes': [] }) # 🎨 生成列区间 columns = [] for i in range(len(vertical_lines) - 1): columns.append({ 'x_start': vertical_lines[i], 'x_end': vertical_lines[i + 1] }) return { 'rows': rows, 'columns': columns, 'horizontal_lines': horizontal_lines, 'vertical_lines': vertical_lines, 'header_height': self.header_height, 'row_height': self.row_height, 'col_widths': self.col_widths, 'table_bbox': [ vertical_lines[0], horizontal_lines[0], vertical_lines[-1], horizontal_lines[-1] ], 'anchor': {'x': anchor_x, 'y': anchor_y}, 'num_rows': num_rows } def apply_template_to_single_file( applier: TableTemplateApplier, image_file: Path, json_file: Path, output_dir: Path, line_width: int = 2, line_color: Tuple[int, int, int] = (0, 0, 0) ) -> bool: """ 应用模板到单个文件 Args: applier: 模板应用器实例 image_file: 图片文件路径 json_file: OCR JSON文件路径 output_dir: 输出目录 line_width: 线条宽度 line_color: 线条颜色 Returns: 是否成功 """ print(f"📄 处理: {image_file.name}") try: # 加载OCR数据 with open(json_file, 'r', encoding='utf-8') as f: raw_data = json.load(f) # 🔧 解析OCR数据（支持PPStructure格式） if 'parsing_res_list' in raw_data and 'overall_ocr_res' in raw_data: table_bbox, ocr_data = TableLineGenerator.parse_ppstructure_result(raw_data) else: raise ValueError("不是PPStructure格式的OCR结果") print(f" ✅ 加载OCR数据: {len(ocr_data)} 个文本框") # 加载图片 image = Image.open(image_file) print(f" ✅ 加载图片: {image.size}") # 🎯 应用模板 img_with_lines = applier.apply_to_image( image, ocr_data, line_width=line_width, line_color=line_color ) # 保存图片 output_file = output_dir / f"{image_file.stem}_with_lines.png" img_with_lines.save(output_file) # 🆕 生成并保存结构配置 structure = applier.generate_structure_for_image(ocr_data) structure_file = output_dir / f"{image_file.stem}_structure.json" with open(structure_file, 'w', encoding='utf-8') as f: json.dump(structure, f, indent=2, ensure_ascii=False) print(f" ✅ 保存图片: {output_file.name}") print(f" ✅ 保存配置: {structure_file.name}") print(f" 📊 表格: {structure['num_rows']}行 x {len(structure['columns'])}列") return True except Exception as e: print(f" ❌ 处理失败: {e}") import traceback traceback.print_exc() return False def apply_template_batch( template_config_path: str, image_dir: str, json_dir: str, output_dir: str, line_width: int = 2, line_color: Tuple[int, int, int] = (0, 0, 0) ): """ 批量应用模板到所有图片 Args: template_config_path: 模板配置路径 image_dir: 图片目录 json_dir: OCR JSON目录 output_dir: 输出目录 line_width: 线条宽度 line_color: 线条颜色 """ applier = TableTemplateApplier(template_config_path) image_path = Path(image_dir) json_path = Path(json_dir) output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) # 查找所有图片 image_files = list(image_path.glob("*.jpg")) + list(image_path.glob("*.png")) image_files.sort() print(f"\n🔍 找到 {len(image_files)} 个图片文件") print(f"📂 图片目录: {image_dir}") print(f"📂 JSON目录: {json_dir}") print(f"📂 输出目录: {output_dir}\n") results = [] success_count = 0 failed_count = 0 for idx, image_file in enumerate(image_files, 1): print(f"\n{'='*60}") print(f"[{idx}/{len(image_files)}] 处理: {image_file.name}") print(f"{'='*60}") # 查找对应的JSON文件 json_file = json_path / f"{image_file.stem}.json" if not json_file.exists(): print(f"⚠️ 找不到OCR结果: {json_file.name}") results.append({ 'source': str(image_file), 'status': 'skipped', 'reason': 'no_json' }) failed_count += 1 continue if apply_template_to_single_file( applier, image_file, json_file, output_path, line_width, line_color ): results.append({ 'source': str(image_file), 'json': str(json_file), 'status': 'success' }) success_count += 1 else: results.append({ 'source': str(image_file), 'json': str(json_file), 'status': 'error' }) failed_count += 1 print() # 保存批处理结果 result_file = output_path / "batch_results.json" with open(result_file, 'w', encoding='utf-8') as f: json.dump(results, f, indent=2, ensure_ascii=False) # 统计 skipped_count = sum(1 for r in results if r['status'] == 'skipped') print(f"\n{'='*60}") print(f"🎉 批处理完成！") print(f"{'='*60}") print(f"✅ 成功: {success_count}") print(f"❌ 失败: {failed_count}") print(f"⚠️ 跳过: {skipped_count}") print(f"📊 总计: {len(results)}") print(f"📄 结果保存: {result_file}") def main(): """主函数""" parser = argparse.ArgumentParser( description='应用表格模板到其他页面', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例用法: 1. 批量处理整个目录: python table_template_applier.py \\ --template output/康强_北京农村商业银行_page_001_structure.json \\ --image-dir /path/to/images \\ --json-dir /path/to/jsons \\ --output-dir /path/to/output 2. 处理单个文件: python table_template_applier.py \\ --template output/康强_北京农村商业银行_page_001_structure.json \\ --image-file /path/to/page_002.png \\ --json-file /path/to/page_002.json \\ --output-dir /path/to/output 输出内容: - {name}_with_lines.png: 带表格线的图片 - {name}_structure.json: 表格结构配置 - batch_results.json: 批处理统计结果 """ ) # 模板参数 parser.add_argument( '-t', '--template', type=str, required=True, help='模板配置文件路径（人工标注的第一页结构）' ) # 文件参数组 file_group = parser.add_argument_group('文件参数（单文件模式）') file_group.add_argument( '--image-file', type=str, help='图片文件路径' ) file_group.add_argument( '--json-file', type=str, help='OCR JSON文件路径' ) # 目录参数组 dir_group = parser.add_argument_group('目录参数（批量模式）') dir_group.add_argument( '--image-dir', type=str, help='图片目录' ) dir_group.add_argument( '--json-dir', type=str, help='OCR JSON目录' ) # 输出参数组 output_group = parser.add_argument_group('输出参数') output_group.add_argument( '-o', '--output-dir', type=str, required=True, help='输出目录（必需）' ) # 绘图参数组 draw_group = parser.add_argument_group('绘图参数') draw_group.add_argument( '-w', '--width', type=int, default=2, help='线条宽度（默认: 2）' ) draw_group.add_argument( '-c', '--color', default='black', choices=['black', 'blue', 'red'], help='线条颜色（默认: black）' ) args = parser.parse_args() # 颜色映射 color_map = { 'black': (0, 0, 0), 'blue': (0, 0, 255), 'red': (255, 0, 0) } line_color = color_map[args.color] # 验证模板文件 template_path = Path(args.template) if not template_path.exists(): print(f"❌ 错误: 模板文件不存在: {template_path}") return output_path = Path(args.output_dir) output_path.mkdir(parents=True, exist_ok=True) # 判断模式 if args.image_file and args.json_file: # 单文件模式 image_file = Path(args.image_file) json_file = Path(args.json_file) if not image_file.exists(): print(f"❌ 错误: 图片文件不存在: {image_file}") return if not json_file.exists(): print(f"❌ 错误: JSON文件不存在: {json_file}") return print("\n🔧 单文件处理模式") print(f"📄 模板: {template_path.name}") print(f"📄 图片: {image_file.name}") print(f"📄 JSON: {json_file.name}") print(f"📂 输出: {output_path}\n") applier = TableTemplateApplier(str(template_path)) success = apply_template_to_single_file( applier, image_file, json_file, output_path, args.width, line_color ) if success: print("\n✅ 处理完成!") else: print("\n❌ 处理失败!") elif args.image_dir and args.json_dir: # 批量模式 image_dir = Path(args.image_dir) json_dir = Path(args.json_dir) if not image_dir.exists(): print(f"❌ 错误: 图片目录不存在: {image_dir}") return if not json_dir.exists(): print(f"❌ 错误: JSON目录不存在: {json_dir}") return print("\n🔧 批量处理模式") print(f"📄 模板: {template_path.name}") apply_template_batch( str(template_path), str(image_dir), str(json_dir), str(output_path), args.width, line_color ) else: parser.print_help() print("\n❌ 错误: 请指定单文件模式或批量模式的参数") print("\n提示:") print(" 单文件模式: --image-file + --json-file") print(" 批量模式: --image-dir + --json-dir") if __name__ == "__main__": print("🚀 启动表格模板批量应用程序...") import sys if len(sys.argv) == 1: # 如果没有命令行参数，使用默认配置运行 print("ℹ️ 未提供命令行参数，使用默认配置运行...") # 默认配置 default_config = { "template": "output/table_structures/康强_北京农村商业银行_page_001_structure.json", "image-file": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行/康强_北京农村商业银行_page_002.png", "json-file": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行_page_002.json", "output-dir": "output/batch_results", "width": "2", "color": "black" } print("⚙️ 默认参数:") for key, value in default_config.items(): print(f" --{key}: {value}") # 构造参数 sys.argv = [sys.argv[0]] for key, value in default_config.items(): sys.argv.extend([f"--{key}", str(value)]) sys.exit(main())