From 500c74c759480099673dd6283c234da93b84f2aa Mon Sep 17 00:00:00 2001 From: tony-claw Date: Sat, 9 May 2026 14:51:13 +0800 Subject: [PATCH] Upload files to "/" upload OCR code --- batchOCR.py | 339 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 339 insertions(+) create mode 100644 batchOCR.py diff --git a/batchOCR.py b/batchOCR.py new file mode 100644 index 0000000..efb026c --- /dev/null +++ b/batchOCR.py @@ -0,0 +1,339 @@ +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() \ No newline at end of file