180 lines
8.5 KiB
Python
180 lines
8.5 KiB
Python
"""Small async client for the Baidu OCR endpoints used by the application."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import time
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
|
|
import httpx
|
|
|
|
from app.core.config import get_settings
|
|
from app.core.exceptions import AppException
|
|
|
|
|
|
class BaiduOcrClient:
|
|
token_url = "https://aip.baidubce.com/oauth/2.0/token"
|
|
|
|
def __init__(self, transport: httpx.AsyncBaseTransport | None = None) -> None:
|
|
settings = get_settings()
|
|
self.api_key = settings.baidu_ocr_api_key.get_secret_value() if settings.baidu_ocr_api_key else ""
|
|
self.secret_key = settings.baidu_ocr_secret_key.get_secret_value() if settings.baidu_ocr_secret_key else ""
|
|
self.image_url = settings.baidu_ocr_image_url.strip()
|
|
self.document_url = settings.baidu_ocr_document_url.strip().rstrip("/")
|
|
self.timeout_seconds = settings.baidu_ocr_timeout_seconds
|
|
self.poll_interval = settings.baidu_ocr_poll_interval_seconds
|
|
self.transport = transport
|
|
self._access_token = ""
|
|
self._token_expires_at = 0.0
|
|
|
|
@property
|
|
def configured(self) -> bool:
|
|
return bool(self.api_key and self.secret_key)
|
|
|
|
async def recognize_image(self, content: bytes) -> tuple[str, float | None, int]:
|
|
token = await self._get_access_token()
|
|
payload = {
|
|
"image": base64.b64encode(content).decode("ascii"),
|
|
"language_type": "CHN_ENG",
|
|
"detect_direction": "true",
|
|
"probability": "true",
|
|
}
|
|
body = await self._post_form(self.image_url, token, payload)
|
|
words = body.get("words_result") or []
|
|
lines = [str(item.get("words") or "").strip() for item in words if isinstance(item, dict)]
|
|
confidences = [
|
|
float(item["probability"]["average"])
|
|
for item in words
|
|
if isinstance(item, dict)
|
|
and isinstance(item.get("probability"), dict)
|
|
and isinstance(item["probability"].get("average"), (int, float))
|
|
]
|
|
text = "\n".join(line for line in lines if line)
|
|
if not text:
|
|
raise AppException("图片中未识别到文字，请重新拍摄或手动填写", code=422, status_code=422)
|
|
confidence = sum(confidences) / len(confidences) if confidences else None
|
|
return text, confidence, 1
|
|
|
|
async def recognize_document(self, content: bytes, file_name: str) -> tuple[str, float | None, int]:
|
|
token = await self._get_access_token()
|
|
body = await self._post_form(
|
|
self.document_url,
|
|
token,
|
|
{
|
|
"file_data": base64.b64encode(content).decode("ascii"),
|
|
"file_name": file_name,
|
|
"merge_tables": "true",
|
|
"return_span_boxes": "true",
|
|
"angle_adjust": "true",
|
|
},
|
|
)
|
|
task_id = str((body.get("result") or {}).get("task_id") or "")
|
|
if not task_id:
|
|
raise AppException("百度文档识别未返回任务编号", code=502, status_code=502)
|
|
|
|
deadline = time.monotonic() + self.timeout_seconds
|
|
query_url = f"{self.document_url}/query"
|
|
result: dict[str, Any] = {}
|
|
while time.monotonic() < deadline:
|
|
await asyncio.sleep(self.poll_interval)
|
|
query_body = await self._post_form(query_url, token, {"task_id": task_id})
|
|
result = query_body.get("result") or {}
|
|
status = result.get("status")
|
|
if status == "success":
|
|
break
|
|
if status == "failed":
|
|
detail = str(result.get("task_error") or "文档解析失败")
|
|
raise AppException(f"百度文档识别失败：{detail}", code=502, status_code=502)
|
|
else:
|
|
raise AppException("文档识别超时，请稍后重试", code=504, status_code=504)
|
|
|
|
parse_url = str(result.get("parse_result_url") or "")
|
|
markdown_url = str(result.get("markdown_url") or "")
|
|
if parse_url:
|
|
parsed = await self._download_result(parse_url, expect_json=True)
|
|
text, pages = self._text_from_parse_result(parsed)
|
|
elif markdown_url:
|
|
text = str(await self._download_result(markdown_url, expect_json=False))
|
|
pages = max(1, text.count("\n---\n") + 1)
|
|
else:
|
|
raise AppException("百度文档识别未返回解析结果", code=502, status_code=502)
|
|
if not text.strip():
|
|
raise AppException("文档中未识别到文字，请重新上传或手动填写", code=422, status_code=422)
|
|
return text[:100_000], None, pages
|
|
|
|
async def _get_access_token(self) -> str:
|
|
if not self.configured:
|
|
raise AppException("未配置百度 OCR API Key 或 Secret Key", code=503, status_code=503)
|
|
if self._access_token and time.monotonic() < self._token_expires_at:
|
|
return self._access_token
|
|
timeout = httpx.Timeout(15, connect=10)
|
|
async with httpx.AsyncClient(timeout=timeout, transport=self.transport) as client:
|
|
try:
|
|
response = await client.post(
|
|
self.token_url,
|
|
params={
|
|
"grant_type": "client_credentials",
|
|
"client_id": self.api_key,
|
|
"client_secret": self.secret_key,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
except (httpx.HTTPError, ValueError) as exc:
|
|
raise AppException("无法连接百度 OCR 鉴权服务", code=502, status_code=502) from exc
|
|
token = str(body.get("access_token") or "")
|
|
if not token:
|
|
raise AppException("百度 OCR 密钥无效或服务尚未开通", code=502, status_code=502)
|
|
expires_in = int(body.get("expires_in") or 2_592_000)
|
|
self._access_token = token
|
|
self._token_expires_at = time.monotonic() + max(60, expires_in - 300)
|
|
return token
|
|
|
|
async def _post_form(self, url: str, token: str, data: dict[str, str]) -> dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout_seconds, connect=15)
|
|
async with httpx.AsyncClient(timeout=timeout, transport=self.transport) as client:
|
|
try:
|
|
response = await client.post(url, params={"access_token": token}, data=data)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
except httpx.TimeoutException as exc:
|
|
raise AppException("百度 OCR 请求超时，请重试", code=504, status_code=504) from exc
|
|
except (httpx.HTTPError, ValueError) as exc:
|
|
raise AppException("百度 OCR 服务暂时不可用", code=502, status_code=502) from exc
|
|
if body.get("error_code") not in (None, 0, "0"):
|
|
message = str(body.get("error_msg") or "识别失败")
|
|
raise AppException(f"百度 OCR 返回错误：{message}", code=502, status_code=502)
|
|
return body
|
|
|
|
async def _download_result(self, url: str, *, expect_json: bool) -> dict[str, Any] | str:
|
|
parsed = urlparse(url)
|
|
allowed = parsed.scheme == "https" and (
|
|
parsed.hostname == "baidubce.com"
|
|
or bool(parsed.hostname and parsed.hostname.endswith((".baidubce.com", ".bcebos.com")))
|
|
)
|
|
if not allowed:
|
|
raise AppException("百度 OCR 返回了不受信任的结果地址", code=502, status_code=502)
|
|
timeout = httpx.Timeout(self.timeout_seconds, connect=15)
|
|
async with httpx.AsyncClient(timeout=timeout, transport=self.transport, follow_redirects=True) as client:
|
|
try:
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
return response.json() if expect_json else response.text
|
|
except (httpx.HTTPError, ValueError) as exc:
|
|
raise AppException("无法下载百度 OCR 解析结果", code=502, status_code=502) from exc
|
|
|
|
@staticmethod
|
|
def _text_from_parse_result(body: dict[str, Any]) -> tuple[str, int]:
|
|
pages = body.get("pages") or []
|
|
chunks: list[str] = []
|
|
for page in pages:
|
|
if not isinstance(page, dict):
|
|
continue
|
|
page_text = str(page.get("text") or "").strip()
|
|
if page_text:
|
|
chunks.append(page_text)
|
|
for table in page.get("tables") or []:
|
|
if isinstance(table, dict) and table.get("markdown"):
|
|
chunks.append(str(table["markdown"]).strip())
|
|
return "\n\n".join(chunks), max(1, len(pages))
|