import os import base64 import time import json import requests from datetime import datetime # ============================================================ # ⚙️ 配置变量 - 在这里设置 # ============================================================ INPUT_DIR = "D:\\新風\\國標GB\\OCR" # 存放 PDF 的目录 OUTPUT_DIR = "D:\\新風\\國標GB\\OCR\\電梯知識" # 输出 Markdown 的目录 # API 配置 API_KEY = "cea4984c93aee4c3" BASE_URL = "https://agent.imqimacau.com" # ============================================================ # 确保输出目录存在 os.makedirs(OUTPUT_DIR, exist_ok=True) def print_full_data(full_data, max_length=500): """打印完整数据结构的辅助函数""" if not full_data: print("⚠️ 没有返回完整数据 (full_data = None)") return print("\n" + "=" * 60) print("完整数据 (Full Data):") print("=" * 60) # 打印元数据 if 'metadata' in full_data: print("\n📊 元数据:") print(f" 总文本块数: {full_data['metadata'].get('total_blocks', 0)}") print(f" 总字符数: {full_data['metadata'].get('total_chars', 0)}") # 打印 OCR 文本 ocr_text = full_data.get('ocr_text', '') if ocr_text: print(f"\n📝 OCR 文本 (共 {len(ocr_text)} 字符):") print("-" * 60) print(ocr_text[:max_length]) if len(ocr_text) > max_length: print(f"... (还有 {len(ocr_text) - max_length} 字符)") # 打印文本块详情 blocks = full_data.get('blocks', []) if blocks: print(f"\n📦 文本块详情 (共 {len(blocks)} 个块):") print("-" * 60) for i, block in enumerate(blocks[:10]): # 只显示前10个块 print(f"\n块 {i+1}:") print(f" 内容: {block.get('content', '')[:100]}...") print(f" 类型: {block.get('block_type', 'unknown')}") print(f" 置信度: {block.get('confidence', 'N/A')}") if block.get('bbox'): bbox = block['bbox'] if len(bbox) == 4: print(f" 位置: [{bbox[0]:.1f}, {bbox[1]:.1f}, {bbox[2]:.1f}, {bbox[3]:.1f}]") else: print(f" 位置: {bbox}") if len(blocks) > 10: print(f"\n... (还有 {len(blocks) - 10} 个块未显示)") # 打印原始输出摘要 if 'raw_output' in full_data: print(f"\n🔧 原始输出: {'存在但可能包含不可序列化对象' if full_data['raw_output'] else 'None'}") def process_pdf_file(input_file, output_dir, api_url, api_key): """处理单个 PDF 文件""" print("\n" + "=" * 80) print(f"📄 开始处理: {input_file}") print("=" * 80) # 生成输出文件名 file_name = os.path.basename(input_file) base_name = os.path.splitext(file_name)[0] # 输出文件路径 output_md = os.path.join(output_dir, f"{base_name}.md") raw_response_file = os.path.join(output_dir, f"{base_name}_raw_response.json") full_data_file = os.path.join(output_dir, f"{base_name}_full_data.json") blocks_detail_file = os.path.join(output_dir, f"{base_name}_blocks_detail.json") spatial_json_file = os.path.join(output_dir, f"{base_name}_spatial_structure.json") pure_text_file = os.path.join(output_dir, f"{base_name}_ocr_text_only.txt") # 读取并编码 PDF print("读取图片...") try: with open(input_file, "rb") as f: file_base64 = base64.b64encode(f.read()).decode() except Exception as e: print(f"❌ 读取文件失败: {e}") return False print(f"文件编码完成，大小: {len(file_base64)} 字符") # 提交任务 print("提交任务...") try: response = requests.post( f"{api_url}/task/submit", headers={ "X-API-Key": api_key, "Content-Type": "application/json" }, json={ "file_base64": file_base64, "file_type": "pdf", "enable_ai_description": False, "output_type": "ocr_only", }, timeout=30 ) except Exception as e: print(f"❌ 提交任务失败: {e}") return False print(f"状态码: {response.status_code}") # 解析 JSON try: result = response.json() except Exception as e: print(f"JSON 解析失败: {e}") print(f"原始响应: {response.text[:200]}") return False if response.status_code != 200: print(f"请求失败: {result}") return False task_id = result.get('task_id') print(f"Task ID: {task_id}") # 轮询获取结果 print("等待处理...", end="") ocr_text = None ai_description = None spatial_structure = None full_data = None error_msg = None raw_result = None while True: try: status_response = requests.get( f"{api_url}/task/result/{task_id}", headers={"X-API-Key": api_key}, timeout=30 ) except Exception as e: print(f"\n❌ 状态查询失败: {e}") time.sleep(4) continue if status_response.status_code != 200: print(f"\n状态查询失败: {status_response.status_code}") time.sleep(4) continue try: data = status_response.json() except Exception as e: print(f"\nJSON 解析失败: {e}") time.sleep(4) continue if data.get('status') == 'completed': print("\n✅ 任务完成！") # 获取结果 if data.get('result'): ocr_text = data['result'].get('ocr_text', '') ai_description = data['result'].get('ai_description', '') spatial_structure = data['result'].get('spatial_structure', None) full_data = data['result'].get('full_data', None) raw_result = data.get('result') # 保存原始返回数据 # with open(raw_response_file, "w", encoding="utf-8") as f: # json.dump(raw_result, f, ensure_ascii=False, indent=2) # print(f"✓ 原始返回数据已保存到: {raw_response_file}") # 打印完整数据结构（可选，注释掉以减少输出） # print_full_data(full_data) # 保存完整数据 if full_data: with open(full_data_file, "w", encoding="utf-8") as f: json.dump(full_data, f, ensure_ascii=False, indent=2) print(f"✓ 完整数据已保存到: {full_data_file}") if full_data.get('blocks'): with open(blocks_detail_file, "w", encoding="utf-8") as f: json.dump(full_data['blocks'], f, ensure_ascii=False, indent=2) print(f"✓ 文本块详情已保存到: {blocks_detail_file}") break elif data.get('status') == 'failed': error_msg = data.get('error', 'Unknown error') print(f"\n❌ 任务失败: {error_msg}") return False else: status = data.get('status', 'unknown') print(f". ({status})", end="", flush=True) time.sleep(2) # 保存结果到 Markdown 文件 print("💾 正在保存结果到文件...") try: with open(output_md, 'w', encoding='utf-8') as f: f.write(f"# OCR 识别结果\n\n") f.write(f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n") f.write(f"**Task ID**: `{task_id}`\n\n") f.write(f"**源文件**: `{input_file}`\n\n") f.write("---\n\n") if error_msg: f.write(f"## ❌ 任务失败\n\n") f.write(f"错误信息: {error_msg}\n") else: # OCR 结果 f.write(f"## 📝 OCR 识别结果\n\n") if ocr_text: f.write(ocr_text) else: f.write("(无 OCR 结果)\n") # 空间结构信息 if spatial_structure: f.write(f"\n\n## 🗺️ 空间结构信息\n\n") if 'block_count' in spatial_structure: f.write(f"- **图片尺寸**: {spatial_structure.get('width', 'N/A')} x {spatial_structure.get('height', 'N/A')}\n") f.write(f"- **识别块数量**: {spatial_structure.get('block_count', 0)}\n") # AI 分析 if ai_description: f.write(f"\n\n## 🤖 AI 分析结果\n\n") f.write(ai_description) print(f"✓ Markdown 结果已保存到: {output_md}") # 保存空间结构 JSON if spatial_structure: with open(spatial_json_file, 'w', encoding='utf-8') as f: json.dump(spatial_structure, f, ensure_ascii=False, indent=2) print(f"✓ 空间结构已保存到: {spatial_json_file}") # 保存纯文本 # if ocr_text and not error_msg: # with open(pure_text_file, 'w', encoding='utf-8') as f: # f.write(ocr_text) # print(f"✓ OCR 纯文本已保存到: {pure_text_file}") return True except Exception as e: print(f"✗ 保存结果失败: {e}") return False def main(): """主函数：遍历目录处理所有 PDF 文件""" print("=" * 80) print("🚀 PDF OCR 批量处理工具") print("=" * 80) print(f"输入目录: {INPUT_DIR}") print(f"输出目录: {OUTPUT_DIR}") print("=" * 80) # 检查输入目录是否存在 if not os.path.exists(INPUT_DIR): print(f"❌ 错误：输入目录不存在 - {INPUT_DIR}") return # 获取所有 PDF 文件 pdf_files = [] for file in os.listdir(INPUT_DIR): if file.lower().endswith('.pdf'): pdf_files.append(os.path.join(INPUT_DIR, file)) if not pdf_files: print(f"⚠️ 在目录 {INPUT_DIR} 中没有找到 PDF 文件") return print(f"\n📁 找到 {len(pdf_files)} 个 PDF 文件:") for f in pdf_files: print(f" - {os.path.basename(f)}") print("-" * 80) # 统计处理结果 success_count = 0 fail_count = 0 # 逐个处理 PDF 文件 for i, pdf_file in enumerate(pdf_files, 1): print(f"\n[{i}/{len(pdf_files)}] 处理中...") result = process_pdf_file( input_file=pdf_file, output_dir=OUTPUT_DIR, api_url=BASE_URL, api_key=API_KEY ) if result: success_count += 1 print(f"✅ 完成: {os.path.basename(pdf_file)}") else: fail_count += 1 print(f"❌ 失败: {os.path.basename(pdf_file)}") # 在文件之间添加短暂延迟，避免请求过快 if i < len(pdf_files): print("\n⏳ 等待 2 秒后处理下一个文件...") time.sleep(2) # 打印最终统计 print("\n" + "=" * 80) print("📊 批量处理完成！") print("=" * 80) print(f"总计: {len(pdf_files)} 个文件") print(f"成功: {success_count} 个") print(f"失败: {fail_count} 个") print(f"输出目录: {OUTPUT_DIR}") print("=" * 80) if __name__ == "__main__": main()