import os import sys import json import logging import time import glob import re import argparse import shutil import traceback from datetime import datetime from pdf_to_table import process_pdf_to_table, post_process_money_texts # 配置日志 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) logger = logging.getLogger(__name__) def parse_regions_file(txt_file): """从regions文本文件提取name和money信息""" name_data = [] money_data = [] current_section = None try: logger.info(f"开始解析文件: {txt_file}") # 使用更高效的文件读取方式 with open(txt_file, 'r', encoding='utf-8') as f: # 一次性读取所有行，减少IO操作 lines = f.readlines() # 预先定义正则表达式进行匹配，提高效率 item_pattern = re.compile(r'^\s+(\d+)\.\s+(.+)$') # 直接遍历处理，避免不必要的日志记录 for line in lines: line = line.strip() if not line: continue if 'NAME 区域' in line: current_section = 'name' continue elif 'MONEY 区域' in line: current_section = 'money' continue elif '===' in line: # 区域分隔符，重置当前部分 current_section = None continue # 使用正则表达式匹配，更高效 if current_section: match = item_pattern.match(line) if match: content = match.group(2).strip() if current_section == 'name': name_data.append(content) elif current_section == 'money': money_data.append(content) # 只在处理完成后记录一次结果 logger.info(f"从{txt_file}解析结果: {len(name_data)}个名称, {len(money_data)}个金额") # 只有当实际有内容时才记录示例 if name_data and len(name_data) > 0: examples = name_data[:min(3, len(name_data))] logger.info(f"名称示例: {examples}") if money_data and len(money_data) > 0: examples = money_data[:min(3, len(money_data))] logger.info(f"金额示例: {examples}") except Exception as e: logger.error(f"解析结果文件时出错: {str(e)}") logger.error(f"详细错误: {traceback.format_exc()}") return { 'name': name_data, 'money': money_data } def extract_texts_from_directory(directory, image_prefix): """ 直接从目录中的文本文件提取name和money文本 作为parse_regions_file的备用方法 Args: directory: 包含文本文件的目录路径 image_prefix: 图像前缀(如"5"表示提取5_name_*.txt和5_money_*.txt文件) Returns: 包含name和money列表的字典 """ name_data = [] money_data = [] try: # 使用glob模式直接查找匹配的文件 name_patterns = [ os.path.join(directory, f"{image_prefix}_name_*_texts.txt"), os.path.join(directory, f"{image_prefix}_rotated_name_*_texts.txt") ] money_patterns = [ os.path.join(directory, f"{image_prefix}_money_*_texts.txt"), os.path.join(directory, f"{image_prefix}_rotated_money_*_texts.txt") ] name_files = [] for pattern in name_patterns: name_files.extend(glob.glob(pattern)) money_files = [] for pattern in money_patterns: money_files.extend(glob.glob(pattern)) # 如果在当前目录找不到，检查子目录 if not name_files and not money_files: image_dir = os.path.join(directory, image_prefix) if os.path.isdir(image_dir): # 同样支持旋转与非旋转的文件格式 name_patterns = [ os.path.join(image_dir, f"{image_prefix}_name_*_texts.txt"), os.path.join(image_dir, f"{image_prefix}_rotated_name_*_texts.txt") ] money_patterns = [ os.path.join(image_dir, f"{image_prefix}_money_*_texts.txt"), os.path.join(image_dir, f"{image_prefix}_rotated_money_*_texts.txt") ] name_files = [] for pattern in name_patterns: name_files.extend(glob.glob(pattern)) money_files = [] for pattern in money_patterns: money_files.extend(glob.glob(pattern)) # 定义排序函数 def extract_index(file_path): try: # 使用split+basename更高效 parts = os.path.basename(file_path).split('_') # 处理旋转图片的情况 if 'rotated' in parts: # 如果是rotated文件，则索引位置会不同 idx_pos = 3 if len(parts) >= 4 else 0 else: # 非旋转图片的正常情况 idx_pos = 2 if len(parts) >= 3 else 0 return int(parts[idx_pos]) if idx_pos < len(parts) else 0 except (IndexError, ValueError): return float('inf') # 对文件列表进行排序 name_files.sort(key=extract_index) money_files.sort(key=extract_index) # 只有在找到文件时才记录 if name_files or money_files: logger.info(f"找到 {len(name_files)} 个name文本文件和 {len(money_files)} 个money文本文件") # 使用更高效的读取方式 for file_path in name_files: with open(file_path, 'r', encoding='utf-8') as f: # 一次性读取并处理 lines = [line.strip() for line in f if line.strip()] name_data.extend(lines) logger.info(f"从 {os.path.basename(file_path)} 提取了 {len(lines)} 个名称") for file_path in money_files: with open(file_path, 'r', encoding='utf-8') as f: # 一次性读取并处理 lines = [line.strip() for line in f if line.strip()] money_data.extend(lines) logger.info(f"从 {os.path.basename(file_path)} 提取了 {len(lines)} 个金额") if name_data or money_data: logger.info(f"直接从文本文件提取: {len(name_data)}个名称, {len(money_data)}个金额") except Exception as e: logger.error(f"从目录提取文本时出错: {str(e)}") logger.error(f"详细错误: {traceback.format_exc()}") return { 'name': name_data, 'money': money_data } def process_pdf(pdf_path, output_dir, confidence_threshold=0.6): """ 处理单个PDF文件 Args: pdf_path: PDF文件路径 output_dir: 输出目录 confidence_threshold: 置信度阈值 Returns: 处理结果字典 """ try: # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 获取PDF文件名 pdf_filename = os.path.basename(pdf_path) pdf_name = os.path.splitext(pdf_filename)[0] logger.info(f"开始处理PDF: {pdf_path}") # 获取PDF文件大小 file_size = os.path.getsize(pdf_path) / (1024 * 1024) # 转换为MB logger.info(f"文件大小: {file_size:.2f} MB") # 创建临时目录存放PDF pdf_folder = os.path.join(output_dir, "temp_pdf") os.makedirs(pdf_folder, exist_ok=True) # 复制PDF到临时目录 temp_pdf_path = os.path.join(pdf_folder, pdf_filename) shutil.copy2(pdf_path, temp_pdf_path) # 记录处理开始时间 start_time = time.time() # 处理PDF文件 stats = process_pdf_to_table( pdf_folder=pdf_folder, output_folder=output_dir, confidence_threshold=confidence_threshold ) # 计算处理时间 processing_time = time.time() - start_time logger.info(f"处理耗时: {processing_time:.2f} 秒") # 记录性能数据 stats['performance'] = { 'file_size_mb': round(file_size, 2), 'processing_time_seconds': round(processing_time, 2), 'processing_speed_mb_per_second': round(file_size / processing_time, 4) if processing_time > 0 else 0 } # 解析结果 all_results = [] results_folder = os.path.join(output_dir, "results") if os.path.exists(results_folder): # 使用glob快速查找所有regions.txt文件 region_files = glob.glob(os.path.join(results_folder, "*_regions.txt")) # 按文件名中的数字进行排序 def extract_number(file_path): try: filename = os.path.basename(file_path) number_str = filename.split('_')[0] return int(number_str) except (IndexError, ValueError): return float('inf') # 按文件名中的数字从小到大排序 region_files.sort(key=extract_number) if region_files: logger.info(f"找到 {len(region_files)} 个区域文件用于解析") # 顺序处理区域文件 for region_file in region_files: image_prefix = os.path.basename(region_file).split('_')[0] result = parse_regions_file(region_file) # 只有当解析到内容时才添加结果 if result['name'] or result['money']: # 对当前图片的数据进行处理和匹配 result = process_single_image_data(result, image_prefix) all_results.append(result) logger.info(f"从 {image_prefix} 解析到内容，并完成了匹配处理") else: # 备用方法：直接从texts.txt文件读取 backup_result = extract_texts_from_directory(results_folder, image_prefix) if backup_result['name'] or backup_result['money']: # 对备用方法提取的数据也进行处理和匹配 backup_result = process_single_image_data(backup_result, image_prefix) all_results.append(backup_result) logger.info(f"使用备用方法从 {image_prefix} 成功提取内容，并完成了匹配处理") else: logger.warning("未找到任何regions.txt文件") # 如果没有从regions.txt解析到结果，尝试直接从子目录提取 if not all_results: logger.warning("未从regions.txt文件解析到内容，尝试从子目录直接提取") # 获取所有子目录并排序 subdirs = [] if os.path.exists(results_folder): subdirs = [d for d in os.listdir(results_folder) if os.path.isdir(os.path.join(results_folder, d))] # 按目录名中的数字排序 def extract_dir_number(dirname): try: return int(dirname) except ValueError: return float('inf') subdirs.sort(key=extract_dir_number) if subdirs: logger.info(f"找到 {len(subdirs)} 个子目录，准备提取") # 使用顺序处理子目录 ordered_results = [] for dirname in subdirs: result_dir = os.path.join(results_folder, dirname) backup_result = extract_texts_from_directory(result_dir, dirname) if backup_result['name'] or backup_result['money']: # 对子目录数据进行处理和匹配 backup_result = process_single_image_data(backup_result, dirname) ordered_results.append((int(dirname) if dirname.isdigit() else float('inf'), backup_result)) # 按子目录编号排序后添加到结果 ordered_results.sort(key=lambda x: x[0]) for _, result in ordered_results: all_results.append(result) # 合并所有结果 merged_result = {'name': [], 'money': [], 'name_money': []} if all_results: # 合并所有图片的处理结果，保持每个图片内的匹配关系 for result in all_results: merged_result['name'].extend(result.get('name', [])) merged_result['money'].extend(result.get('money', [])) merged_result['name_money'].extend(result.get('name_money', [])) # 注意：由于每个图片已经单独处理并匹配，此处不再需要全局匹配逻辑 # 记录处理结果摘要 logger.info(f"PDF {pdf_filename} 处理完成，找到 {len(merged_result['name'])} 个名称, {len(merged_result['money'])} 个金额, {len(merged_result['name_money'])} 个名称-金额对") # 保存最终结果到JSON文件 result_json_path = os.path.join(output_dir, f"{pdf_name}_result.json") with open(result_json_path, 'w', encoding='utf-8') as f: json.dump({ 'pdf_file': pdf_filename, 'name': merged_result['name'], 'money': merged_result['money'], 'name_money': merged_result['name_money'], 'stats': stats }, f, ensure_ascii=False, indent=2) logger.info(f"JSON结果已保存到: {result_json_path}") # 保存最终结果到TXT文件 result_txt_path = os.path.join(output_dir, f"{pdf_name}_result.txt") with open(result_txt_path, 'w', encoding='utf-8') as f: f.write(f"PDF文件：{pdf_filename} 处理结果\n") f.write(f"处理时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") f.write("=" * 80 + "\n\n") # 写入name-money对 f.write(f"共找到 {len(merged_result['name_money'])} 个名称-金额对：\n\n") for i, pair in enumerate(merged_result['name_money'], 1): f.write(f"{i}. {pair}\n") f.write("\n" + "=" * 80 + "\n") logger.info(f"TXT结果已保存到: {result_txt_path}") # 清理临时目录 try: shutil.rmtree(pdf_folder) logger.info("已清理临时PDF目录") except Exception as e: logger.warning(f"清理临时目录时出错: {str(e)}") return { 'pdf_file': pdf_filename, 'result': merged_result, 'stats': stats } except Exception as e: logger.error(f"处理PDF {pdf_path} 时出错: {str(e)}") logger.error(f"详细错误: {traceback.format_exc()}") return { 'pdf_file': os.path.basename(pdf_path), 'error': str(e) } def process_single_image_data(result, image_prefix): """ 处理单张图片的名称和金额数据，解决匹配问题 Args: result: 包含name和money的字典 image_prefix: 图片前缀标识 Returns: 处理后的结果字典，包含匹配好的name_money对 """ logger.info(f"开始处理图片 {image_prefix} 的数据匹配") # 复制原始数据，避免修改原始数据 processed_result = { 'name': result.get('name', []).copy(), 'money': result.get('money', []).copy(), 'name_money': [] } # 对金额应用后处理函数 - 合并被分割的金额（从pdf_to_table模块导入） try: # 对金额进行专门的后处理 processed_result['money'] = post_process_money_texts(processed_result['money']) logger.info(f"图片 {image_prefix} 金额后处理后数量: {len(processed_result['money'])}") except ImportError: logger.warning("无法导入post_process_money_texts函数，将使用内部金额处理逻辑") # 内部处理逻辑 - 合并分散的金额 processed_money = [] i = 0 while i < len(processed_result['money']): current = processed_result['money'][i].strip() # 检查当前金额是否需要与下一个合并 if i < len(processed_result['money']) - 1: next_money = processed_result['money'][i+1].strip() # 情况1: 当前金额以小数点结尾 if current.endswith('.'): merged_value = current + next_money processed_money.append(merged_value) logger.info(f"合并金额: {current} + {next_money} = {merged_value}") i += 2 # 跳过下一个已合并的金额 continue # 情况2: 当前金额包含小数点但后面不足两位数字 elif '.' in current: decimal_part = current.split('.')[-1] if len(decimal_part) < 2: merged_value = current + next_money processed_money.append(merged_value) logger.info(f"合并小数金额: {current} + {next_money} = {merged_value}") i += 2 continue # 情况3: 下一个是纯数字且长度<=2(可能是分离的小数部分) elif next_money.isdigit() and len(next_money) <= 2: # 检查当前值是否为整数或以小数点结尾的数 if current.isdigit() or current.endswith('.'): merged_value = current + next_money processed_money.append(merged_value) logger.info(f"合并可能的小数部分: {current} + {next_money} = {merged_value}") i += 2 continue # 情况4: 可能是分散的小数点(例如"6500"和".00"分开) elif not '.' in current and next_money.startswith('.') and next_money[1:].isdigit(): merged_value = current + next_money processed_money.append(merged_value) logger.info(f"合并分散小数点: {current} + {next_money} = {merged_value}") i += 2 continue # 如果没有合并，直接添加当前金额 processed_money.append(current) i += 1 # 更新处理后的金额列表 processed_result['money'] = processed_money # 金额格式化和验证 - 确保金额格式正确 validated_money = [] for money in processed_result['money']: # 尝试清理和标准化金额格式 cleaned_money = money.strip() # 移除可能的非数字字符(保留数字、小数点和逗号) cleaned_money = re.sub(r'[^\d.,]', '', cleaned_money) # 检查是否包含小数点，如果没有可能需要添加 if '.' not in cleaned_money and cleaned_money.isdigit(): cleaned_money = cleaned_money + '.00' logger.info(f"格式化金额: {money} -> {cleaned_money}") # 确保小数点后有两位数字 elif '.' in cleaned_money: parts = cleaned_money.split('.') if len(parts) == 2: integer_part, decimal_part = parts # 如果小数部分长度不足2，补零 if len(decimal_part) < 2: decimal_part = decimal_part.ljust(2, '0') # 如果小数部分超过2位，可能需要截断或四舍五入(根据需求) elif len(decimal_part) > 2: decimal_part = decimal_part[:2] cleaned_money = integer_part + '.' + decimal_part logger.info(f"标准化金额: {money} -> {cleaned_money}") # 检查金额格式是否有效 try: float_value = float(cleaned_money.replace(',', '')) if float_value < 0.01: logger.warning(f"忽略异常小的金额值: {cleaned_money}") continue validated_money.append(cleaned_money) except ValueError: logger.warning(f"忽略无效金额格式: {cleaned_money}") # 更新为验证后的金额 processed_result['money'] = validated_money # 处理name和money数量不匹配的情况 max_name_len = len(processed_result['name']) max_money_len = len(processed_result['money']) # 记录name和money配对前的情况 logger.info(f"图片 {image_prefix} 配对前: {max_name_len}个名称, {max_money_len}个金额") # 判断是否存在不匹配 if max_name_len != max_money_len: logger.warning(f"图片 {image_prefix} 名称和金额数量不匹配: {max_name_len}个名称 vs {max_money_len}个金额") # 处理姓名中可能存在的误识别情况 cleaned_names = [] for name in processed_result['name']: # 判断姓名中是否包含错误识别的数字和非中文字符 clean_name = re.sub(r'[0-9a-zA-Z.,:;!@#$%^&*()_+={}\[\]|\\/<>?~`-]', '', name) if clean_name != name: logger.info(f"清理姓名中的误识别: {name} -> {clean_name}") # 如果清理后不为空，则保留 if clean_name.strip(): cleaned_names.append(clean_name.strip()) else: logger.warning(f"姓名'{name}'在清理后为空，已移除") # 更新清理后的姓名列表 processed_result['name'] = cleaned_names # 再次检查数量 max_name_len = len(processed_result['name']) logger.info(f"清理后: {max_name_len}个名称, {max_money_len}个金额") # 创建最终的配对结果 # 采用截断方式匹配，确保数据不会错位 match_len = min(max_name_len, max_money_len) if match_len > 0: for i in range(match_len): name = processed_result['name'][i] money = processed_result['money'][i] pair = f"{name}-{money}" processed_result['name_money'].append(pair) logger.info(f"图片 {image_prefix} 匹配对 #{i+1}: {pair}") # 如果有多余的名称或金额，记录日志但不匹配 if max_name_len > match_len: extra_names = processed_result['name'][match_len:] logger.warning(f"图片 {image_prefix} 有 {len(extra_names)} 个未匹配的名称: {extra_names}") if max_money_len > match_len: extra_money = processed_result['money'][match_len:] logger.warning(f"图片 {image_prefix} 有 {len(extra_money)} 个未匹配的金额: {extra_money}") logger.info(f"图片 {image_prefix} 处理完成，共生成 {len(processed_result['name_money'])} 个名称-金额对") return processed_result def process_pdfs_in_directory(input_dir, output_dir, confidence_threshold=0.6): """ 处理目录中的所有PDF文件 Args: input_dir: 输入目录，包含PDF文件 output_dir: 输出目录 confidence_threshold: 置信度阈值 Returns: 处理结果统计 """ # 确保输出目录存在 os.makedirs(output_dir, exist_ok=True) # 查找所有PDF文件 pdf_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.pdf')] if not pdf_files: logger.warning(f"在 {input_dir} 目录中未找到PDF文件") return {'processed': 0, 'success': 0, 'failed': 0, 'files': []} logger.info(f"在 {input_dir} 目录中找到 {len(pdf_files)} 个PDF文件") # 按照文件名排序，确保处理顺序一致 pdf_files.sort() results = [] success_count = 0 failed_count = 0 # 顺序处理每个PDF for pdf_file in pdf_files: pdf_path = os.path.join(input_dir, pdf_file) pdf_name = os.path.splitext(pdf_file)[0] # 为每个PDF创建单独的输出目录 pdf_output_dir = os.path.join(output_dir, pdf_name) logger.info(f"开始处理 {pdf_file} ({len(results)+1}/{len(pdf_files)})") # 处理PDF result = process_pdf(pdf_path, pdf_output_dir, confidence_threshold) if 'error' in result: logger.error(f"处理 {pdf_file} 失败: {result['error']}") failed_count += 1 else: logger.info(f"成功处理 {pdf_file}") success_count += 1 results.append(result) # 创建汇总报告 summary = { 'processed': len(pdf_files), 'success': success_count, 'failed': failed_count, 'files': [r['pdf_file'] for r in results] } # 保存汇总报告 summary_path = os.path.join(output_dir, "processing_summary.json") with open(summary_path, 'w', encoding='utf-8') as f: json.dump(summary, f, ensure_ascii=False, indent=2) logger.info(f"处理完成，共处理 {len(pdf_files)} 个PDF，成功 {success_count} 个，失败 {failed_count} 个") logger.info(f"汇总报告已保存到: {summary_path}") return summary def main(): """主函数""" parser = argparse.ArgumentParser(description='OCR处理工具：处理PDF并提取姓名和金额') # 添加参数 parser.add_argument('--input', '-i', required=True, help='输入PDF文件路径或包含PDF文件的目录') parser.add_argument('--output', '-o', required=True, help='输出目录') parser.add_argument('--confidence', '-c', type=float, default=0.6, help='置信度阈值，默认0.6') # 解析参数 args = parser.parse_args() # 记录开始时间 start_time = time.time() # 检查输入路径 if not os.path.exists(args.input): logger.error(f"输入路径不存在: {args.input}") return 1 # 创建输出目录 os.makedirs(args.output, exist_ok=True) # 处理单个文件或目录 if os.path.isfile(args.input): if not args.input.lower().endswith('.pdf'): logger.error(f"输入文件不是PDF: {args.input}") return 1 logger.info(f"开始处理单个PDF文件: {args.input}") result = process_pdf(args.input, args.output, args.confidence) if 'error' in result: logger.error(f"处理失败: {result['error']}") return 1 else: logger.info("处理成功") else: logger.info(f"开始处理目录中的PDF文件: {args.input}") summary = process_pdfs_in_directory(args.input, args.output, args.confidence) if summary['success'] == 0: logger.warning("没有成功处理任何PDF文件") if summary['processed'] > 0: return 1 # 计算总处理时间 total_time = time.time() - start_time logger.info(f"总处理时间: {total_time:.2f} 秒") return 0 if __name__ == "__main__": try: sys.exit(main()) except Exception as e: logger.error(f"程序执行过程中发生错误: {str(e)}") logger.error(f"详细错误: {traceback.format_exc()}") sys.exit(1)