125 lines
5.3 KiB
Python
125 lines
5.3 KiB
Python
"""面单 OCR enrichment:对所有系统(RIIN+SDS)有待处理面单的订单,用 v2 batch API 批量识别。
|
|
|
|
对应 label-ocr。每小时跑一次,跟上增量同步节奏。
|
|
按国家分组,每组 200 单/批调 /api/v2/batch(服务端并发下载+重试)。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
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
|
|
|
|
BATCH_SIZE = 200
|
|
|
|
_OCR_FIELD_MAP = [
|
|
("layout", "ocrLayout"), ("courier", "ocrCourier"), ("tracking", "ocrTracking"),
|
|
("marker", "ocrMarker"), ("s_name", "ocrSName"), ("s_street", "ocrSStreet"),
|
|
("s_city", "ocrSCity"), ("s_state", "ocrSState"), ("s_zip", "ocrSZip"),
|
|
("r_name", "ocrRName"), ("r_street", "ocrRStreet"), ("r_city", "ocrRCity"),
|
|
("r_state", "ocrRState"), ("r_zip", "ocrRZip"),
|
|
]
|
|
|
|
|
|
def _result_to_fields(r) -> dict[str, str | None]:
|
|
if not r.ok:
|
|
return {"ocrLayout": "ERROR"}
|
|
return {col: (getattr(r, attr) or "") for attr, col in _OCR_FIELD_MAP}
|
|
|
|
|
|
@task
|
|
def process_all(pending: list) -> dict[str, int]:
|
|
db = SyncDB()
|
|
ocr = OcrClient()
|
|
stats = {"success": 0, "error": 0, "unsupported": 0, "batches": 0}
|
|
try:
|
|
# 1. 按国家分组;不支持的国家直接标 UNSUPPORTED
|
|
by_country: dict[str, list[tuple[str, object, str]]] = {}
|
|
unsupp: list[tuple[str, object]] = []
|
|
for order_no, created_at, pdf_url, dest_country in pending:
|
|
code = country_to_ocr_code(dest_country)
|
|
if code:
|
|
by_country.setdefault(code, []).append((order_no, created_at, pdf_url))
|
|
else:
|
|
unsupp.append((order_no, created_at))
|
|
|
|
if unsupp:
|
|
db.update_ocr_batch([(no, ca, {"ocrLayout": "UNSUPPORTED"}) for no, ca in unsupp])
|
|
stats["unsupported"] = len(unsupp)
|
|
print(f"[OCR] 不支持国家: {len(unsupp)} 单 → UNSUPPORTED")
|
|
|
|
# 2. 每个国家分批调 v2 batch API
|
|
for country_code, orders in by_country.items():
|
|
total = len(orders)
|
|
for i in range(0, total, BATCH_SIZE):
|
|
batch = orders[i : i + BATCH_SIZE]
|
|
created_map = {no: ca for no, ca, _ in batch}
|
|
urls = [{"url": url, "filename": no} for no, _, url in batch]
|
|
batch_no = i // BATCH_SIZE + 1
|
|
total_batches = -(-total // BATCH_SIZE)
|
|
try:
|
|
results = ocr.recognize_batch(urls, country_code)
|
|
except Exception as err: # noqa: BLE001
|
|
print(f"[OCR] {country_code} batch {batch_no}/{total_batches} 请求失败: {str(err)[:100]}")
|
|
db.update_ocr_batch([(no, created_map[no], {"ocrLayout": "ERROR"}) for no in created_map])
|
|
stats["error"] += len(batch)
|
|
continue
|
|
|
|
updates = []
|
|
for r in results:
|
|
order_no = r.file
|
|
created_at = created_map.get(order_no)
|
|
if created_at is None:
|
|
continue
|
|
fields = _result_to_fields(r)
|
|
updates.append((order_no, created_at, fields))
|
|
stats["success" if r.ok else "error"] += 1
|
|
db.update_ocr_batch(updates)
|
|
stats["batches"] += 1
|
|
ok = sum(1 for r in results if r.ok)
|
|
print(f"[OCR] {country_code} batch {batch_no}/{total_batches}: "
|
|
f"{ok}/{len(results)} 成功")
|
|
return stats
|
|
finally:
|
|
ocr.close()
|
|
|
|
|
|
@flow(name="label-ocr", log_prints=True)
|
|
def label_ocr(destination_country: str | None = 'GB', courier_source: str | None = None) -> dict:
|
|
"""面单 OCR enrichment。默认只处理英国(GB),可选 courierSource 进一步筛选。
|
|
|
|
手动跑其他国家: label_ocr(destination_country='US')
|
|
"""
|
|
db = SyncDB()
|
|
before = db.count_ocr_status()
|
|
print(f"OCR 状态(开始): {before}")
|
|
|
|
pending = db.get_ocr_pending(destination_country=destination_country, courier_source=courier_source)
|
|
scope = ""
|
|
if destination_country or courier_source:
|
|
scope = f" [{destination_country or ''} {courier_source or ''}]".strip()
|
|
print(f"待 OCR: {len(pending)} 单{scope} (RIIN+SDS)")
|
|
if not pending:
|
|
emit_sync_summary(f"面单 OCR{scope}", build_stats_rows("ALL", {}, [
|
|
("待处理", 0), ("累计已完成", before.get("done", 0)),
|
|
]))
|
|
return {"processed": 0, **before}
|
|
|
|
stats = process_all(pending)
|
|
after = db.count_ocr_status()
|
|
total_processed = stats["success"] + stats["error"] + stats["unsupported"]
|
|
print(f"完成{scope}: 处理 {total_processed} 单 / {stats['batches']} 批 | "
|
|
f"成功 {stats['success']} 失败 {stats['error']} 不支持 {stats['unsupported']}")
|
|
print(f"累计: {after}")
|
|
emit_sync_summary(f"面单 OCR{scope}", build_stats_rows("ALL", {}, [
|
|
("本批处理", total_processed), ("成功", stats["success"]),
|
|
("失败", stats["error"]), ("不支持", stats["unsupported"]),
|
|
("累计已完成", after.get("done", 0)), ("累计待处理", after.get("pending", 0)),
|
|
]))
|
|
return {"processed": total_processed, **stats, **after}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
label_ocr.serve(name="label-ocr", cron="0 * * * *")
|