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