99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
"""SDS 面单 OCR enrichment:对已发货的 SDS 订单,用面单 PDF URL 调 OCR 服务回填标签字段。
|
|
|
|
对应 sds-label-ocr。独立于同步 flow,每 4 小时跑一批。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
from prefect import flow, task
|
|
|
|
from api.ocr import OcrClient, country_to_ocr_code
|
|
from flows.shared import SyncDB, build_stats_rows, emit_sync_summary
|
|
|
|
OCR_CONCURRENCY = 5
|
|
SHIPPED_STATUSES = ["已完成", "部分发货", "备货中", "待确认"]
|
|
|
|
_OCR_FIELD_MAP = [
|
|
("layout", "ocr_layout"), ("courier", "ocr_courier"), ("tracking", "ocr_tracking"),
|
|
("marker", "ocr_marker"), ("s_name", "ocr_s_name"), ("s_street", "ocr_s_street"),
|
|
("s_city", "ocr_s_city"), ("s_state", "ocr_s_state"), ("s_zip", "ocr_s_zip"),
|
|
("r_name", "ocr_r_name"), ("r_street", "ocr_r_street"), ("r_city", "ocr_r_city"),
|
|
("r_state", "ocr_r_state"), ("r_zip", "ocr_r_zip"),
|
|
]
|
|
|
|
|
|
def _process_one(db: SyncDB, ocr: OcrClient, order_no: str, created_at, pdf_url: str, dest_country: str | None) -> str:
|
|
country_code = country_to_ocr_code(dest_country)
|
|
if not country_code:
|
|
db.update_ocr_fields(order_no, created_at, {"ocr_layout": "UNSUPPORTED"})
|
|
return "unsupported"
|
|
try:
|
|
r = ocr.recognize_url(pdf_url, country_code)
|
|
except Exception as err: # noqa: BLE001
|
|
print(f"[OCR] {order_no} 请求失败: {str(err)[:100]}")
|
|
db.update_ocr_fields(order_no, created_at, {"ocr_layout": "ERROR"})
|
|
return "error"
|
|
if not r.ok:
|
|
print(f"[OCR] {order_no} 识别失败: {(r.error or '')[:80]}")
|
|
db.update_ocr_fields(order_no, created_at, {"ocr_layout": "ERROR"})
|
|
return "error"
|
|
fields = {col: getattr(r, attr) or "" for attr, col in _OCR_FIELD_MAP}
|
|
db.update_ocr_fields(order_no, created_at, fields)
|
|
return "success"
|
|
|
|
|
|
@task
|
|
def process_batch(pending: list) -> dict[str, int]:
|
|
db = SyncDB()
|
|
ocr = OcrClient()
|
|
stats = {"success": 0, "error": 0, "unsupported": 0}
|
|
try:
|
|
done = 0
|
|
total = len(pending)
|
|
with ThreadPoolExecutor(max_workers=min(OCR_CONCURRENCY, total)) as pool:
|
|
futures = {
|
|
pool.submit(_process_one, db, ocr, no, ca, url, dc): no
|
|
for no, ca, url, dc in pending
|
|
}
|
|
for fut in as_completed(futures):
|
|
result = fut.result()
|
|
stats[result] += 1
|
|
done += 1
|
|
if done % 50 == 0 or done == total:
|
|
print(f"[OCR] 进度: {done}/{total} 成功 {stats['success']} / 失败 {stats['error']} / 不支持 {stats['unsupported']}")
|
|
stats["processed"] = total
|
|
return stats
|
|
finally:
|
|
ocr.close()
|
|
|
|
|
|
@flow(name="sds-label-ocr", log_prints=True)
|
|
def sds_label_ocr(batch_size: int = 200) -> dict:
|
|
db = SyncDB()
|
|
before = db.count_ocr_status()
|
|
print(f"OCR 状态(开始): {before}")
|
|
|
|
pending = db.get_ocr_pending(batch_size)
|
|
print(f"待 OCR: {len(pending)} 单 (batch_size={batch_size})")
|
|
if not pending:
|
|
emit_sync_summary("SDS 面单 OCR", build_stats_rows("SDS", {}, [
|
|
("待处理", 0), ("累计已完成", before.get("done", 0)), ("累计失败", before.get("error", 0)),
|
|
]))
|
|
return {"processed": 0, **before}
|
|
|
|
stats = process_batch(pending)
|
|
after = db.count_ocr_status()
|
|
print(f"完成: {stats} 累计: {after}")
|
|
emit_sync_summary("SDS 面单 OCR", build_stats_rows("SDS", {}, [
|
|
("本批处理", stats["processed"]), ("本批成功", stats["success"]),
|
|
("本批失败", stats["error"]), ("本批不支持", stats["unsupported"]),
|
|
("累计已完成", after.get("done", 0)), ("累计待处理", after.get("pending", 0)),
|
|
]))
|
|
return {**stats, **after}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sds_label_ocr.serve(name="sds-label-ocr", cron="0 */4 * * *")
|