282 lines
12 KiB
Python
282 lines
12 KiB
Python
import json
|
||
import requests
|
||
import base64
|
||
import time
|
||
from typing import Dict, Any, Optional, List
|
||
|
||
|
||
class UmiOcrClient:
|
||
"""
|
||
封装了 Umi-OCR HTTP API 的客户端。
|
||
提供了文档识别的完整异步流程，以及图片和二维码识别接口。
|
||
"""
|
||
|
||
def __init__(self, base_url: str = "http://154.219.106.93:1224"):
|
||
"""
|
||
初始化客户端。
|
||
:param base_url: Umi-OCR 服务的基地址。
|
||
"""
|
||
self.base_url = base_url.rstrip('/')
|
||
self.headers_json = {"Content-Type": "application/json"}
|
||
print(f"UmiOcrClient 初始化。目标地址: {self.base_url}")
|
||
|
||
def _send_request(self, method: str, endpoint: str, data: Optional[Dict[str, Any]] = None,
|
||
files: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
发送 HTTP 请求的通用内部方法。
|
||
"""
|
||
url = self.base_url + endpoint
|
||
|
||
try:
|
||
if method.upper() == 'GET':
|
||
response = requests.get(url)
|
||
elif method.upper() == 'POST':
|
||
if files:
|
||
# 使用 multipart/form-data for file upload (Step 1)
|
||
response = requests.post(url, data=data, files=files, timeout=300)
|
||
else:
|
||
# 使用 application/json for other POST requests
|
||
data_str = json.dumps(data)
|
||
response = requests.post(url, data=data_str, headers=self.headers_json, timeout=300)
|
||
else:
|
||
raise ValueError(f"不支持的 HTTP 方法: {method}")
|
||
|
||
# 检查 HTTP 状态码
|
||
response.raise_for_status()
|
||
|
||
# 解析响应 JSON
|
||
return response.json()
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
print(f"请求 {url} 失败。")
|
||
raise Exception(f"请求错误: {e}")
|
||
except json.JSONDecodeError as e:
|
||
# 某些接口（如 /api/doc/clear）可能返回空响应体，这里需要容错
|
||
if response.status_code == 200 and not response.text:
|
||
return {"code": 100, "data": "Success (No content)"}
|
||
print(f"响应内容不是有效的 JSON。原始内容:\n{response.text}")
|
||
raise Exception(f"JSON 解析错误: {e}")
|
||
|
||
# ==========================================================================
|
||
# 📚 文档识别（PDF等文件）接口封装
|
||
# ==========================================================================
|
||
|
||
def doc_get_options(self) -> Dict[str, Any]:
|
||
""" 0. 准备工作：参数查询 (/api/doc/get_options) """
|
||
endpoint = "/api/doc/get_options"
|
||
return self._send_request("GET", endpoint)
|
||
|
||
def doc_upload(self, file_path: str, options: Optional[Dict[str, Any]] = None) -> str:
|
||
""" 1. 上传待识别文档 (/api/doc/upload)，返回任务ID """
|
||
print(f"🚀 上传文件: {file_path}")
|
||
endpoint = "/api/doc/upload"
|
||
|
||
# 准备 form-data 的 'json' 部分
|
||
json_data = json.dumps(options if options is not None else {})
|
||
|
||
# 准备 form-data 的 'file' 部分
|
||
files = {'file': open(file_path, 'rb')}
|
||
data = {'json': json_data}
|
||
|
||
try:
|
||
response = self._send_request("POST", endpoint, data=data, files=files)
|
||
|
||
if response.get('code') == 100:
|
||
return response.get('data') # 任务ID
|
||
else:
|
||
raise Exception(f"文件上传失败: {response.get('data')}")
|
||
finally:
|
||
# 确保文件流关闭
|
||
if 'file' in files and files['file'].closed == False:
|
||
files['file'].close()
|
||
|
||
def doc_query_result(self, task_id: str, is_data: bool = False, is_unread: bool = True,
|
||
data_format: str = "dict") -> Dict[str, Any]:
|
||
""" 2. 查询任务状态 (/api/doc/result) """
|
||
endpoint = "/api/doc/result"
|
||
request_data = {
|
||
"id": task_id,
|
||
"is_data": is_data,
|
||
"is_unread": is_unread,
|
||
"format": data_format,
|
||
}
|
||
return self._send_request("POST", endpoint, request_data)
|
||
|
||
def doc_get_download_link(self, task_id: str, file_types: List[str] = ["pdfLayered"], ignore_blank: bool = True) -> \
|
||
Dict[str, str]:
|
||
""" 3. 获取结果下载链接 (/api/doc/download) """
|
||
endpoint = "/api/doc/download"
|
||
request_data = {
|
||
"id": task_id,
|
||
"file_types": file_types,
|
||
"ignore_blank": ignore_blank,
|
||
}
|
||
response = self._send_request("POST", endpoint, request_data)
|
||
|
||
if response.get('code') == 100:
|
||
return {
|
||
"download_url": self.base_url + response.get('data'), # 完整的下载链接
|
||
"file_name": response.get('name')
|
||
}
|
||
else:
|
||
raise Exception(f"生成下载链接失败: {response.get('data')}")
|
||
|
||
def doc_clear_task(self, task_id: str) -> bool:
|
||
""" 5. 任务清理 (/api/doc/clear/<id>) """
|
||
endpoint = f"/api/doc/clear/{task_id}"
|
||
response = self._send_request("GET", endpoint)
|
||
|
||
if response.get('code') == 100:
|
||
print(f"🧹 任务ID {task_id} 清理成功。")
|
||
return True
|
||
else:
|
||
print(f"⚠️ 任务ID {task_id} 清理失败或不存在: {response.get('data')}")
|
||
return False
|
||
|
||
# --- 高级同步流程函数 ---
|
||
|
||
def run_doc_ocr_sync(self, file_path: str, options: Optional[Dict[str, Any]] = None,
|
||
download_file_types: List[str] = ["pdfLayered"], interval: int = 5) -> Dict[str, Any]:
|
||
"""
|
||
同步执行整个文档OCR流程 (上传 -> 轮询 -> 下载链接 -> 清理)。
|
||
"""
|
||
task_id = None
|
||
try:
|
||
# 1. 上传文件，获取任务ID
|
||
task_id = self.doc_upload(file_path, options)
|
||
print(f"📝 任务启动成功，任务ID: {task_id}")
|
||
|
||
# 2. 轮询任务状态
|
||
print("⏳ 正在轮询任务状态...")
|
||
while True:
|
||
status_res = self.doc_query_result(task_id, is_data=False)
|
||
|
||
is_done = status_res.get('is_done', False)
|
||
state = status_res.get('state', 'unknown')
|
||
processed = status_res.get('processed_count', 0)
|
||
total = status_res.get('pages_count', '?')
|
||
|
||
print(f" - 状态: {state.upper()} ({processed}/{total} 页已处理)")
|
||
|
||
if is_done:
|
||
if state == "success":
|
||
print("✅ 任务完成！")
|
||
break
|
||
else:
|
||
message = status_res.get('message', '未知错误')
|
||
raise Exception(f"任务执行失败 (State: {state}): {message}")
|
||
|
||
time.sleep(interval) # 休息指定间隔时间后继续查询
|
||
|
||
# 3. 获取下载链接
|
||
download_info = self.doc_get_download_link(task_id, download_file_types)
|
||
print(f"🔗 下载链接生成成功。文件: {download_info['file_name']}")
|
||
|
||
# 4. 获取识别文本（可选步骤，通过查询获取）
|
||
text_result = self.doc_query_result(task_id, is_data=True, data_format="text").get('data')
|
||
|
||
return {
|
||
"code": 100,
|
||
"task_id": task_id,
|
||
"download_info": download_info,
|
||
"recognized_text": text_result,
|
||
}
|
||
|
||
except Exception as e:
|
||
# 捕捉所有错误并报告
|
||
raise Exception(f"文档识别流程失败: {e}")
|
||
|
||
finally:
|
||
# 5. 任务清理 (无论成功失败都清理)
|
||
if task_id:
|
||
self.doc_clear_task(task_id)
|
||
|
||
# ==========================================================================
|
||
# 📸 其他接口（保留自上次封装）
|
||
# ==========================================================================
|
||
|
||
def ocr_base64_image(self, base64_image: str, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
""" 图片OCR：Base64 识别接口 (/api/ocr) """
|
||
print("\n--- 调用图片OCR Base64 识别接口 (/api/ocr) ---")
|
||
endpoint = "/api/ocr"
|
||
request_data = {
|
||
"base64": base64_image,
|
||
"options": options if options is not None else {}
|
||
}
|
||
return self._send_request("POST", endpoint, request_data)
|
||
|
||
def recognize_qrcode_base64(self, base64_image: str) -> Dict[str, Any]:
|
||
""" 【推测接口】二维码：Base64 识别 (/api/qrcode/recognize) """
|
||
print("\n--- 调用二维码 Base64 识别接口（推测: /api/qrcode/recognize）---")
|
||
endpoint = "/api/qrcode/recognize"
|
||
request_data = {"base64": base64_image}
|
||
return self._send_request("POST", endpoint, request_data)
|
||
|
||
def generate_qrcode(self, text: str, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
""" 【推测接口】二维码：从文本生成图片 (/api/qrcode/generate) """
|
||
print("\n--- 调用二维码生成接口（推测: /api/qrcode/generate）---")
|
||
endpoint = "/api/qrcode/generate"
|
||
request_data = {
|
||
"text": text,
|
||
"options": options if options is not None else {}
|
||
}
|
||
return self._send_request("POST", endpoint, request_data)
|
||
|
||
|
||
# ==============================================================================
|
||
# 示例用法
|
||
# ==============================================================================
|
||
|
||
if __name__ == '__main__':
|
||
# 请根据您的实际情况修改以下参数
|
||
TEST_PDF_PATH = r"D:\daima\java\gdyd_zhpb_zgf\第2章 方案建议书.pdf"
|
||
|
||
# 实例化客户端
|
||
client = UmiOcrClient()
|
||
|
||
# --- 演示文档识别流程 ---
|
||
print("\n" + "=" * 50)
|
||
print(" 演示：同步文档识别流程 (Step 1-5)")
|
||
print("=" * 50)
|
||
|
||
try:
|
||
# 尝试使用 run_doc_ocr_sync 接口
|
||
# 假设我们想要中文识别，并生成双层PDF和纯文本txt
|
||
doc_options = {
|
||
"ocr.language": "models/config_chinese.txt",
|
||
"pageRangeEnd": 1 # 限制只识别前5页
|
||
}
|
||
|
||
# ⚠️ 注意: 运行此段代码前，请确保您的目录下存在 TEST_PDF_PATH 所指的文件
|
||
final_result = client.run_doc_ocr_sync(
|
||
file_path=TEST_PDF_PATH,
|
||
options=doc_options,
|
||
download_file_types=["pdfLayered", "txtPlain"],
|
||
interval=3 # 轮询间隔3秒
|
||
)
|
||
|
||
print("\n🎉 整个文档识别任务执行完毕。")
|
||
print("--- 总结结果 ---")
|
||
print(f"任务ID: {final_result['task_id']}")
|
||
print(f"下载文件名: {final_result['download_info']['file_name']}")
|
||
print(f"下载链接: {final_result['download_info']['download_url']}")
|
||
print(f"部分识别文本 (前100字): \n{final_result['recognized_text']}...")
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ 文档识别流程中断/失败。错误信息: {e}")
|
||
print("提示：请确认 Umi-OCR 是否运行在目标 IP/端口，并且文件路径正确。")
|
||
|
||
print("\n" + "=" * 50)
|
||
print(" 演示：图片 Base64 识别")
|
||
print("=" * 50)
|
||
|
||
# 演示图片识别 (与上次封装相同，确认接口可用)
|
||
EXAMPLE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAC4AAAAXCAIAAAD7ruoFAAAACXBIWXMAABnWAAAZ1gEY0crtAAAAEXRFWHRTb2Z0d2FyZQBTbmlwYXN0ZV0Xzt0AAAHjSURBVEiJ7ZYrcsMwEEBXnR7FLuj0BPIJHJOi0DAZ2qSsMCxEgjYrDQqJdALrBJ2ASndRgeNI8ledutOCLrLl1e7T/mRkjIG/IXe/DWBldRTNEoQSpgNURe5puiiaJehrMuJSXSTgbaby0A1WzLrCCQCmyn0FwoN0V06QONWAt1nUxfnjHYA8p65GjhDKxcjedVH6JOejBPwYh21eE0Wzfe0tqIsEkGXcVcpoMH4CRZ+P0lsQp/pWJ4ripf1XFDFe8GHSHlYcSo9Es31t60RdFlN1RUmrma5oTzTVB8ZUaeeYEC9GmL6kNkDw9BANAQYo3xTNdqUkvHq+rYhDKW0Bj3RSEIpmyWyBaZaMTCrCK+tJ5Jsa07fs3E7esE66HzralRLgJKp0/BD6fJRSxvmDsb6joqkcFXGqMVVFFEHDL2gTxwCAaTabnkFUWhDCHTd9iYrGcAL1ZnqIp5Vpiqh7bCfua7FA4qN0INMcN1+cgCzj+UFxtbmvwdZvGIrI41JiqhZBWhhF8WxorkYPpQwJiWYJeA3rXE4hzcwJ+B96F9zCFHC0FcVegghvFul7oeEE8PvHeJqC0w0AUbbFIT8JnEwGbPKcS2OxU3HMTqD0r4wgEIuiKJ7i4MS16+og8/+bPZRPLa+6Ld2DSzcAAAAASUVORK5CYII="
|
||
try:
|
||
ocr_result = client.ocr_base64_image(
|
||
base64_image=EXAMPLE_BASE64,
|
||
options={"data.format": "text"}
|
||
)
|
||
print("✅ Base64 识别成功。识别结果:", ocr_result.get('data'))
|
||
except Exception as e:
|
||
print(f"❌ Base64 识别失败。") |