Upload files to "/"
upload OCR code
This commit is contained in:
339
batchOCR.py
Normal file
339
batchOCR.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user