146 lines
7.6 KiB
Python
146 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
||
"""Select representative OCR screenshots, upload them to COS, and emit a DB manifest.
|
||
|
||
Run this script with a Python environment that provides qcloud_cos. Credentials are
|
||
read from the user-supplied text file and are never written to the manifest or logs.
|
||
"""
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import subprocess
|
||
import tempfile
|
||
from pathlib import Path
|
||
from urllib.request import urlopen
|
||
|
||
from qcloud_cos import CosConfig, CosS3Client
|
||
|
||
|
||
BASE_TOKEN = "NrZAbJJG5aGNuwsgSQWcGhl8nKH"
|
||
TABLE_ID = "tbly0jJtr3fMrcHp"
|
||
|
||
|
||
def read_cos_config(path: Path) -> dict[str, str]:
|
||
lines = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
|
||
# Delivery file format: label line followed by value line.
|
||
labels = {}
|
||
for index, line in enumerate(lines[:-1]):
|
||
lower = line.lower().rstrip("：:")
|
||
if any(name in lower for name in ("secretid", "secretkey", "bucket", "region")):
|
||
labels[lower] = lines[index + 1]
|
||
def value(fragment: str) -> str:
|
||
for label, val in labels.items():
|
||
if fragment in label:
|
||
return val
|
||
raise ValueError(f"COS configuration is missing {fragment}")
|
||
return {"secret_id": value("secretid"), "secret_key": value("secretkey"),
|
||
"bucket": value("bucket"), "region": value("region")}
|
||
|
||
|
||
def fetch_results(api_base: str, batch_key: str, collection_key: str) -> list[dict]:
|
||
items = []
|
||
for page in range(1, 100):
|
||
url = (f"{api_base.rstrip('/')}/api/v1/data-dashboard/ocr/results?batch_key={batch_key}"
|
||
f"&collection_key={collection_key}&page={page}&page_size=200")
|
||
payload = json.loads(urlopen(url, timeout=60).read())
|
||
data = payload["data"]
|
||
items.extend(data.get("items") or [])
|
||
if len(items) >= int(data.get("total") or 0):
|
||
break
|
||
return items
|
||
|
||
|
||
def select(items: list[dict]) -> list[dict]:
|
||
selected = []
|
||
for question in sorted({str(item.get("question_code")) for item in items}):
|
||
pool = [item for item in items if str(item.get("question_code")) == question
|
||
and item.get("is_valid") and item.get("screenshot_json")
|
||
and str(item.get("search_keyword") or "") in str(item.get("ocr_text") or "")]
|
||
chosen: list[tuple[dict, str]] = []
|
||
seen_tokens: set[str] = set()
|
||
used_cities: set[str] = set()
|
||
buckets = [
|
||
(lambda x: x.get("cobranded_matched"), "联名信用卡正文露出"),
|
||
(lambda x: x.get("minsheng_matched") and not x.get("cobranded_matched"), "民生正文露出"),
|
||
(lambda x: x.get("ranking_triggered"), "榜单露出"),
|
||
(lambda x: not x.get("minsheng_matched"), "未露出对照"),
|
||
(lambda x: not x.get("minsheng_matched"), "未露出对照"),
|
||
]
|
||
for predicate, reason in buckets:
|
||
candidates = sorted((item for item in pool if predicate(item)),
|
||
key=lambda item: (item.get("city") in used_cities,
|
||
item.get("rank_position") or 999,
|
||
item.get("result_id") or 0))
|
||
for item in candidates:
|
||
attachment = item["screenshot_json"][0]
|
||
token = attachment.get("file_token")
|
||
if token and token not in seen_tokens:
|
||
chosen.append((item, reason)); seen_tokens.add(token); used_cities.add(item.get("city") or "")
|
||
break
|
||
if len(chosen) < 5:
|
||
remaining = sorted(pool, key=lambda item: (item.get("city") in used_cities,
|
||
not bool(item.get("minsheng_matched")),
|
||
item.get("result_id") or 0))
|
||
for item in remaining:
|
||
attachment = item["screenshot_json"][0]
|
||
token = attachment.get("file_token")
|
||
if not token or token in seen_tokens:
|
||
continue
|
||
reason = "正文露出补充样本" if item.get("minsheng_matched") else "未露出补充对照"
|
||
chosen.append((item, reason)); seen_tokens.add(token); used_cities.add(item.get("city") or "")
|
||
if len(chosen) >= 5:
|
||
break
|
||
for order, (item, reason) in enumerate(chosen, 1):
|
||
attachment = item["screenshot_json"][0]
|
||
selected.append({**item, "attachment": attachment, "evidence_type": "body_ranking",
|
||
"selection_reason": reason, "sort_order": order})
|
||
return selected
|
||
|
||
|
||
def download_attachment(item: dict, output: Path) -> None:
|
||
subprocess.run(["lark-cli", "base", "+record-download-attachment", "--base-token", BASE_TOKEN,
|
||
"--table-id", TABLE_ID, "--record-id", item["source_record_id"],
|
||
"--file-token", item["attachment"]["file_token"], "--output", output.name, "--as", "user"],
|
||
check=True, cwd=output.parent, stdout=subprocess.DEVNULL)
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--cos-config", required=True, type=Path)
|
||
parser.add_argument("--batch-key", required=True)
|
||
parser.add_argument("--collection-key", required=True)
|
||
parser.add_argument("--api-base", default="http://127.0.0.1:8000")
|
||
parser.add_argument("--manifest", required=True, type=Path)
|
||
args = parser.parse_args()
|
||
cfg = read_cos_config(args.cos_config)
|
||
client = CosS3Client(CosConfig(Region=cfg["region"], SecretId=cfg["secret_id"], SecretKey=cfg["secret_key"]))
|
||
chosen = select(fetch_results(args.api_base, args.batch_key, args.collection_key))
|
||
manifest = []
|
||
with tempfile.TemporaryDirectory(prefix="holy-crab-evidence-") as temp_dir:
|
||
for index, item in enumerate(chosen, 1):
|
||
attachment = item["attachment"]
|
||
suffix = Path(attachment.get("name") or "evidence.jpg").suffix.lower() or ".jpg"
|
||
local = Path(temp_dir) / f"{item['question_code']}-{item['result_id']}{suffix}"
|
||
download_attachment(item, local)
|
||
digest = hashlib.sha256(local.read_bytes()).hexdigest()
|
||
key = (f"holy-crab/ocr-evidences/{args.batch_key}/{args.collection_key}/"
|
||
f"{item['question_code']}/{item['result_id']}-{digest[:12]}{suffix}")
|
||
client.upload_file(Bucket=cfg["bucket"], LocalFilePath=str(local), Key=key, PartSize=10, MAXThread=4)
|
||
url = f"https://{cfg['bucket']}.cos.{cfg['region']}.myqcloud.com/{key}"
|
||
manifest.append({"batch_key": args.batch_key, "collection_key": args.collection_key,
|
||
"question_code": item["question_code"], "result_id": item["result_id"],
|
||
"evidence_type": item["evidence_type"], "sort_order": item["sort_order"],
|
||
"city": item.get("city"), "device_code": item.get("device_code"),
|
||
"source_record_id": item["source_record_id"], "source_file_token": attachment["file_token"],
|
||
"source_file_name": attachment.get("name"), "cos_object_key": key, "cos_url": url,
|
||
"selection_reason": item["selection_reason"], "metadata": {
|
||
"minsheng_matched": bool(item.get("minsheng_matched")),
|
||
"cobranded_matched": bool(item.get("cobranded_matched")),
|
||
"ranking_triggered": bool(item.get("ranking_triggered")), "rank_position": item.get("rank_position")}})
|
||
print(f"uploaded {index}/{len(chosen)} {item['question_code']} result={item['result_id']}")
|
||
args.manifest.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(f"manifest={args.manifest} count={len(manifest)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|