372 lines
16 KiB
Python
372 lines
16 KiB
Python
# -*- encoding: utf-8 -*-
|
||
'''
|
||
@File : tencentcloud_ocr.py
|
||
@Time : @
|
||
@Author : eason
|
||
@Version : 1.0
|
||
@Contact :
|
||
@License : (C)Copyright 2021-2024
|
||
@Desc : 将图片转换为base64编码，并使用内部Youtu OCR服务进行文字识别（app_id自动生成）
|
||
'''
|
||
|
||
import json
|
||
import os
|
||
import base64
|
||
import random
|
||
import threading
|
||
import time
|
||
import requests
|
||
from config.config_utils import config # 如果不再需要配置可以移除，但保留其他配置可能使用
|
||
|
||
|
||
# ---- OCR 端点（多 IP 负载均衡）配置 ----
|
||
OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"
|
||
OCR_DEFAULT_PORT = 18081 # config 里只填 IP 时自动补的端口
|
||
DEFAULT_ENDPOINTS = ["http://asr.ocr.general.smart.l20.production.polaris:18081"]
|
||
WEIGHTS_FILE = "./config/ocr_weights.json"
|
||
|
||
|
||
def normalize_endpoint(e):
|
||
"""把 'IP' / 'IP:port' / 'http://IP:port' 统一成 'http://IP:port'。
|
||
config 里只填裸 IP 时，自动补 http:// 和默认端口。"""
|
||
if not e:
|
||
return None
|
||
e = str(e).strip().rstrip("/")
|
||
if not e:
|
||
return None
|
||
if not (e.startswith("http://") or e.startswith("https://")):
|
||
e = "http://" + e
|
||
host = e.split("://", 1)[1].split("/", 1)[0] # 取 host[:port] 部分
|
||
if ":" not in host: # 没写端口就补默认端口
|
||
e = e + f":{OCR_DEFAULT_PORT}"
|
||
return e
|
||
|
||
|
||
def img2base64(img_path: str) -> str:
|
||
"""将图片转换为base64编码"""
|
||
with open(img_path, 'rb') as f:
|
||
image_base64 = base64.b64encode(f.read()).decode()
|
||
return image_base64
|
||
|
||
|
||
def get_ocr_endpoints():
|
||
"""从配置读取 OCR 端点列表，自动归一化（只填 IP 也行）；缺省回退到内置默认。"""
|
||
ocr_cfg = config.get_tencent_ocr_config() or {}
|
||
eps = ocr_cfg.get("endpoints")
|
||
if isinstance(eps, str):
|
||
eps = [eps]
|
||
if not eps:
|
||
eps = DEFAULT_ENDPOINTS
|
||
return [u for u in (normalize_endpoint(e) for e in eps) if u]
|
||
|
||
|
||
def load_endpoint_weights():
|
||
"""读取压测产出的 QPS 权重文件 {base_url: qps}；不存在则返回 {}。"""
|
||
if not os.path.exists(WEIGHTS_FILE):
|
||
return {}
|
||
try:
|
||
with open(WEIGHTS_FILE, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
return {k.rstrip("/"): float(v) for k, v in data.items() if float(v) > 0}
|
||
except (json.JSONDecodeError, ValueError, TypeError, OSError):
|
||
return {}
|
||
|
||
|
||
class OcrDispatcher:
|
||
"""线程安全的 OCR 端点分发器：用 P-of-d-Choices(默认d=3)做【负载/延迟感知】派发——
|
||
随机取 sample_size 个端点, 把请求发给其中【在途更少 × 更快】的那个。慢/积压端点自动少分。
|
||
每个进程独立按后端真实快慢反应 → 无需任何协调/中心LB, 也能在 120 台 OCR 上涌现出全局均衡。
|
||
d 从 2 提到 3(甚至更多): 端点数越多, 随机抽2个都恰好较差的概率越高, 抽样数多一个能明显
|
||
收窄这个方差, 之后收益递减, 不用无限往上加。
|
||
失败端点冷却摘除(冷却时长带随机抖动, 防止多个进程对同一端点同时解冷却又一起打回去,
|
||
造成"整体一起失败-一起恢复-一起打爆"的同步震荡)。多个并发 OCR 线程共享一个实例。"""
|
||
|
||
def __init__(self, endpoints, weights=None, cooldown=30.0, sample_size=3):
|
||
self.endpoints = list(endpoints)
|
||
# PID 打乱(历史原因保留): 让各进程初始顺序不同; 现在 P-of-d 是随机取样, 影响不大但无害。
|
||
try:
|
||
random.Random(os.getpid()).shuffle(self.endpoints)
|
||
except Exception:
|
||
pass
|
||
weights = weights or {}
|
||
# 压测权重(QPS): 作为"处理能力", 强端点 score 更低 → 多分。无权重默认 1.0。
|
||
self.weights = {e: max(float(weights.get(e, 1.0)), 0.0001) for e in self.endpoints}
|
||
self.cooldown = cooldown
|
||
self.sample_size = max(2, int(sample_size)) # 每次候选抽样个数(P-of-d 的 d)
|
||
self._lock = threading.Lock()
|
||
self._inflight = {e: 0 for e in self.endpoints} # 当前在途请求数(本进程视角)
|
||
self._lat = {e: 0.05 for e in self.endpoints} # 延迟 EWMA(秒), 初始乐观→新端点也会被试
|
||
self._cooldown_until = {e: 0.0 for e in self.endpoints}
|
||
self._rng = random.Random(os.getpid() ^ 0x9E3779B1)
|
||
# —— 派发计数(验证分发是否均匀): 每端点累计派发数, 每 _dump_every 次落盘一次本进程文件 ——
|
||
self._count = {e: 0 for e in self.endpoints}
|
||
self._since_dump = 0
|
||
try:
|
||
self._dump_every = int((config.get_tencent_ocr_config() or {}).get('dispatch_log_every', 0) or 0)
|
||
except (ValueError, TypeError, AttributeError):
|
||
self._dump_every = 0
|
||
|
||
def _score(self, e):
|
||
"""负载分: (在途+1) × 延迟EWMA / 权重。越低越优先。"""
|
||
return (self._inflight[e] + 1) * self._lat[e] / self.weights[e]
|
||
|
||
def pick(self):
|
||
"""P-of-d: 随机不重复抽 sample_size(默认3) 个可用端点, 选负载分最低的;
|
||
选中即占一个在途名额(必须配对 done())。全部冷却时退化为全部可用。"""
|
||
now = time.time()
|
||
snapshot = None
|
||
with self._lock:
|
||
available = [e for e in self.endpoints if self._cooldown_until[e] <= now]
|
||
if not available:
|
||
available = list(self.endpoints) # 全挂了也得发, 给一次机会
|
||
if len(available) == 1:
|
||
chosen = available[0]
|
||
else:
|
||
k = min(self.sample_size, len(available))
|
||
cands = self._rng.sample(available, k) # 不重复抽样, 比各自 choice() 更有效(不会抽到同一个)
|
||
chosen = min(cands, key=self._score)
|
||
self._inflight[chosen] = self._inflight.get(chosen, 0) + 1
|
||
self._count[chosen] = self._count.get(chosen, 0) + 1
|
||
if self._dump_every > 0:
|
||
self._since_dump += 1
|
||
if self._since_dump >= self._dump_every:
|
||
self._since_dump = 0
|
||
snapshot = (dict(self._count),
|
||
{e: round(self._lat[e], 4) for e in self._lat})
|
||
if snapshot is not None: # 落盘在锁外, 不阻塞其它线程派发
|
||
self._write_dump(snapshot[0], snapshot[1])
|
||
return chosen
|
||
|
||
def _write_dump(self, counts, lats):
|
||
"""把本进程的累计派发计数写到 ./ocr_dispatch/<pid>.json(覆盖式, 单进程单文件无竞争)。"""
|
||
try:
|
||
d = "./ocr_dispatch"
|
||
os.makedirs(d, exist_ok=True)
|
||
path = os.path.join(d, f"{os.getpid()}.json")
|
||
data = {"pid": os.getpid(), "ts": time.time(),
|
||
"total": sum(counts.values()), "counts": counts, "lat_ewma": lats}
|
||
tmp = path + ".tmp"
|
||
with open(tmp, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, ensure_ascii=False)
|
||
os.replace(tmp, path)
|
||
except Exception:
|
||
pass
|
||
|
||
def done(self, endpoint, latency=None, ok=True):
|
||
"""请求结束回调: 释放在途名额; 成功则更新延迟EWMA, 失败则冷却摘除(冷却时长带±30%抖动,
|
||
避免同一时刻失败的多个进程在同一时刻解冷却、一起冲回去打同一个刚失败的端点)。"""
|
||
with self._lock:
|
||
if self._inflight.get(endpoint, 0) > 0:
|
||
self._inflight[endpoint] -= 1
|
||
if endpoint not in self._cooldown_until:
|
||
return
|
||
if ok:
|
||
if latency is not None:
|
||
self._lat[endpoint] = 0.7 * self._lat[endpoint] + 0.3 * max(latency, 0.001)
|
||
self._cooldown_until[endpoint] = 0.0
|
||
else:
|
||
jitter = self._rng.uniform(0.7, 1.3)
|
||
self._cooldown_until[endpoint] = time.time() + self.cooldown * jitter
|
||
|
||
# —— 兼容旧接口(内部都走 done) ——
|
||
def report_failure(self, endpoint):
|
||
self.done(endpoint, ok=False)
|
||
|
||
def report_success(self, endpoint):
|
||
self.done(endpoint, ok=True)
|
||
|
||
|
||
_dispatcher = None
|
||
_dispatcher_lock = threading.Lock()
|
||
|
||
|
||
def get_dispatcher():
|
||
"""惰性单例分发器（首次调用时读取端点+权重）。"""
|
||
global _dispatcher
|
||
if _dispatcher is None:
|
||
with _dispatcher_lock:
|
||
if _dispatcher is None:
|
||
eps = get_ocr_endpoints()
|
||
_dispatcher = OcrDispatcher(eps, load_endpoint_weights())
|
||
return _dispatcher
|
||
|
||
|
||
def _build_payload(image_base64: str) -> dict:
|
||
timestamp = int(time.time() * 1000)
|
||
rand_suffix = random.randint(100000, 999999)
|
||
app_id = f"{timestamp}{rand_suffix}"
|
||
return {
|
||
"app_id": app_id,
|
||
"image": image_base64,
|
||
"ocr_template": "OCR",
|
||
"options": {"pdf_page_index": 0, "detect_text": True},
|
||
}
|
||
|
||
|
||
# 跨进程【全局 OCR 并发上限】信号量。进程池下由 main 在父进程创建、经 initializer 注入每个 worker;
|
||
# 不管多少进程×多少线程,全局同时在飞的 OCR 请求数被这个信号量硬性卡住,防止打爆端点。
|
||
# 为 None 时不限流(线程池模式 / 未注入)。
|
||
_ocr_semaphore = None
|
||
|
||
|
||
def set_ocr_semaphore(sem):
|
||
"""由 worker 进程初始化时调用,注入全局并发信号量。"""
|
||
global _ocr_semaphore
|
||
_ocr_semaphore = sem
|
||
|
||
|
||
# 每个进程一个持久 Session(连接池/keep-alive 复用连接)。
|
||
# 原来用 requests.post() 每次新建+关闭连接 → 高并发下海量 TIME_WAIT → 本地临时端口耗尽
|
||
# → "Cannot assign requested address" / 连接超时(表现成"所有端点超时", 其实是本机端口用完)。
|
||
# 复用连接后, 到每个端点只保持少量长连接, 端口不再churn。惰性创建 → 进程池下各进程独立 Session。
|
||
_session = None
|
||
_session_lock = threading.Lock()
|
||
|
||
|
||
def _get_session():
|
||
global _session
|
||
if _session is None:
|
||
with _session_lock:
|
||
if _session is None:
|
||
s = requests.Session()
|
||
adapter = requests.adapters.HTTPAdapter(
|
||
pool_connections=256, # 连接池数(≥端点数, 每端点一个池)
|
||
pool_maxsize=64, # 每端点最大并发连接(复用)
|
||
max_retries=0,
|
||
pool_block=False)
|
||
s.mount("http://", adapter)
|
||
s.mount("https://", adapter)
|
||
_session = s
|
||
return _session
|
||
|
||
|
||
def ocr_request_once(base_url: str, image_base64: str, timeout: int = 30) -> str:
|
||
"""向指定端点发一次 OCR 请求，返回原始响应文本；失败抛异常。
|
||
受全局信号量限流;用持久 Session 复用连接(防端口耗尽)。"""
|
||
url = base_url.rstrip("/") + OCR_PATH
|
||
headers = {"Content-Type": "application/json"}
|
||
sem = _ocr_semaphore
|
||
if sem is not None:
|
||
sem.acquire()
|
||
try:
|
||
resp = _get_session().post(url, headers=headers, json=_build_payload(image_base64),
|
||
timeout=timeout)
|
||
resp.raise_for_status()
|
||
return resp.text
|
||
finally:
|
||
if sem is not None:
|
||
sem.release()
|
||
|
||
|
||
def _normalize_youtu_to_tencent(raw: str) -> str:
|
||
"""将 Youtu OCR 返回格式转为项目其他模块期望的腾讯云 OCR 格式。"""
|
||
try:
|
||
data = json.loads(raw)
|
||
except (json.JSONDecodeError, TypeError):
|
||
return raw
|
||
|
||
# 已经是腾讯云格式则直接返回
|
||
if "TextDetections" in data:
|
||
return raw
|
||
|
||
if data.get("error_code") != 0:
|
||
return json.dumps({"TextDetections": [], "Language": "zh",
|
||
"RequestId": data.get("session_id", "")}, ensure_ascii=False)
|
||
|
||
text_detections = []
|
||
for recognize_item in data.get("recognize_list", []):
|
||
item_list = recognize_item.get("item_content", {}).get("item_list", [])
|
||
for item in item_list:
|
||
content = item.get("content", "")
|
||
confidence = int(item.get("conf", 0) * 100) # 0-1 → 0-100
|
||
|
||
vertex = item.get("vertex", {})
|
||
if vertex:
|
||
polygon = [
|
||
{"X": vertex.get("x1", 0), "Y": vertex.get("y1", 0)},
|
||
{"X": vertex.get("x2", 0), "Y": vertex.get("y2", 0)},
|
||
{"X": vertex.get("x3", 0), "Y": vertex.get("y3", 0)},
|
||
{"X": vertex.get("x4", 0), "Y": vertex.get("y4", 0)},
|
||
]
|
||
else:
|
||
coord = item.get("coord", {})
|
||
x, y = coord.get("x", 0), coord.get("y", 0)
|
||
w, h = coord.get("width", 0), coord.get("height", 0)
|
||
polygon = [
|
||
{"X": x, "Y": y},
|
||
{"X": x + w, "Y": y},
|
||
{"X": x + w, "Y": y + h},
|
||
{"X": x, "Y": y + h},
|
||
]
|
||
|
||
text_detections.append({
|
||
"DetectedText": content,
|
||
"Confidence": confidence,
|
||
"Polygon": polygon,
|
||
})
|
||
|
||
return json.dumps({
|
||
"TextDetections": text_detections,
|
||
"Language": "zh",
|
||
"RequestId": data.get("session_id", ""),
|
||
}, ensure_ascii=False)
|
||
|
||
|
||
def _get_ocr_call_params():
|
||
"""从 config 读【单次OCR的重试上限/超时】, 带默认。
|
||
封顶重试是为了防"卡1小时": 端点都超时时, 一张图最坏耗时 = max_attempts × timeout,
|
||
而不是原来的 len(endpoints)(119) × 30s ≈ 3570s。"""
|
||
try:
|
||
cfg = config.get_tencent_ocr_config() or {}
|
||
max_attempts = int(cfg.get("max_attempts", 4) or 4)
|
||
timeout = int(cfg.get("timeout", 12) or 12)
|
||
except (ValueError, TypeError, AttributeError):
|
||
max_attempts, timeout = 4, 12
|
||
return max(1, max_attempts), max(1, timeout)
|
||
|
||
|
||
def youtu_ocr_api(image_base64: str) -> str:
|
||
"""
|
||
调用内部 Youtu OCR 服务进行文字识别。
|
||
按 QPS 权重自动在多个 OCR IP 间分发；某端点失败时摘除冷却并换下一个端点重试。
|
||
重试次数封顶(config: ocr.max_attempts), 单次超时可配(ocr.timeout), 避免端点全超时时拖死 worker。
|
||
"""
|
||
dispatcher = get_dispatcher()
|
||
max_attempts, timeout = _get_ocr_call_params()
|
||
# 最多试 max_attempts 个端点(且不超过端点总数), 而不是把 119 个全试一遍
|
||
attempts = min(max(1, len(dispatcher.endpoints)), max_attempts)
|
||
last_err = None
|
||
for _ in range(attempts):
|
||
url = dispatcher.pick()
|
||
t0 = time.time()
|
||
try:
|
||
raw = ocr_request_once(url, image_base64, timeout=timeout)
|
||
# 把真实延迟反馈给分发器 → 慢端点延迟EWMA升高 → 后续自动少分(负载感知核心)
|
||
dispatcher.done(url, latency=time.time() - t0, ok=True)
|
||
return _normalize_youtu_to_tencent(raw)
|
||
except requests.exceptions.RequestException as e:
|
||
last_err = e
|
||
dispatcher.done(url, latency=time.time() - t0, ok=False)
|
||
continue
|
||
print(f"Youtu OCR 请求失败（已尝试 {attempts} 个端点, 单次超时 {timeout}s）: {last_err}")
|
||
return ""
|
||
|
||
|
||
# 为了保持原函数名兼容，将新接口赋值给原名称
|
||
tencentcloud_ocr_api = youtu_ocr_api
|
||
|
||
|
||
def tencent_ocr_video_all_frames(pic_dir_path: str,
|
||
jsonl_save_path: str = 'frames_ocr.jsonl') -> list:
|
||
"""
|
||
对一个视频抽帧后的全部图片进行 ocr处理并保存起来
|
||
"""
|
||
with open(jsonl_save_path, "w", encoding='utf-8') as f:
|
||
for root, _, files in os.walk(pic_dir_path):
|
||
for pic_name in sorted(files):
|
||
if pic_name.endswith(('jpg')):
|
||
pic_path = os.path.join(root, pic_name)
|
||
img_base64 = img2base64(img_path=pic_path)
|
||
json_string = tencentcloud_ocr_api(image_base64=img_base64)
|
||
f.write("pic_path: {}\n".format(pic_path))
|
||
f.write(json_string + "\n") |