120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
||
"""EasyOCR 子进程 worker。
|
||
|
||
通过 stdin 读 base64 图片、stdout 写 JSON 结果。GPU crash 只杀此进程，不冻主进程。
|
||
协议：
|
||
stdin: {"img":"<base64_png_bytes>"}
|
||
stdout: {"ok":true,"text":"识别结果"}
|
||
stdout: {"ok":false,"error":"错误描述"}
|
||
|
||
初始化慢（首次 3-10s），初始化完成后常驻处理。
|
||
"""
|
||
|
||
import json
|
||
import sys
|
||
import base64
|
||
import io
|
||
import os
|
||
import traceback
|
||
|
||
# ── 顶层异常捕获：确保崩溃原因写入 stderr（否则父进程看到空 stderr） ──
|
||
def _flush_exit(msg, code=1):
|
||
sys.stderr.write(f"[easyocr_worker] FATAL: {msg}\n")
|
||
sys.stderr.flush()
|
||
sys.exit(code)
|
||
|
||
try:
|
||
# ── 线程/进程限制：避免 MKL/OpenMP 在子进程多线程冲突 ──
|
||
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
||
os.environ.setdefault("MKL_NUM_THREADS", "1")
|
||
|
||
from PIL import Image
|
||
import numpy as np
|
||
import easyocr
|
||
except Exception as e:
|
||
_flush_exit(f"import 失败: {e}\n{traceback.format_exc()}")
|
||
|
||
|
||
def _init_reader():
|
||
"""初始化 EasyOCR reader（GPU 优先）。"""
|
||
try:
|
||
import torch
|
||
use_gpu = torch.cuda.is_available()
|
||
except Exception:
|
||
use_gpu = False
|
||
|
||
try:
|
||
reader = easyocr.Reader(["ch_sim", "en"], gpu=use_gpu)
|
||
mode = "GPU" if use_gpu else "CPU"
|
||
sys.stderr.write(f"[easyocr_worker] Reader 就绪（{mode}模式）\n")
|
||
sys.stderr.flush()
|
||
return reader
|
||
except Exception as e:
|
||
# GPU 不可用 → CPU 兜底
|
||
sys.stderr.write(f"[easyocr_worker] GPU 初始化失败({e})，尝试 CPU\n")
|
||
sys.stderr.flush()
|
||
reader = easyocr.Reader(["ch_sim", "en"], gpu=False)
|
||
sys.stderr.write("[easyocr_worker] Reader 就绪（CPU 模式）\n")
|
||
sys.stderr.flush()
|
||
return reader
|
||
|
||
|
||
def _process_image(img_b64: str, reader) -> str:
|
||
"""解码 base64 图片 → EasyOCR 识别 → 返回空格拼接文本。"""
|
||
if reader is None:
|
||
return ""
|
||
raw = base64.b64decode(img_b64)
|
||
img = Image.open(io.BytesIO(raw))
|
||
# 保持与 monitor.py 一致：RGBA → numpy
|
||
if img.mode != "RGBA":
|
||
img = img.convert("RGBA")
|
||
img_np = np.array(img)
|
||
results = reader.readtext(img_np)
|
||
return "\n".join([r[1] for r in results]).strip()
|
||
|
||
|
||
def _main():
|
||
reader = None
|
||
try:
|
||
reader = _init_reader()
|
||
except Exception as e:
|
||
sys.stderr.write(f"[easyocr_worker] Reader 初始化失败: {e}\n")
|
||
sys.stderr.flush()
|
||
# ⭐ 就绪握手：即使 Reader 失败也要发 ready，否则父进程永远等握手
|
||
sys.stdout.write(json.dumps({"ready": True}, ensure_ascii=False) + "\n")
|
||
sys.stdout.flush()
|
||
sys.stderr.write("[easyocr_worker] 就绪，等待 stdin 输入\n")
|
||
sys.stderr.flush()
|
||
|
||
for line in sys.stdin:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
try:
|
||
payload = json.loads(line)
|
||
img_b64 = payload.get("img", "")
|
||
if not img_b64:
|
||
raise ValueError("缺少 img 字段")
|
||
sys.stderr.write(f"[easyocr_worker] 收到请求，处理中... (img长度={len(img_b64)})\n")
|
||
sys.stderr.flush()
|
||
text = _process_image(img_b64, reader)
|
||
sys.stderr.write(f"[easyocr_worker] 处理完成，结果长度={len(text)}\n")
|
||
sys.stderr.flush()
|
||
out = json.dumps({"ok": True, "text": text}, ensure_ascii=False)
|
||
sys.stdout.write(out + "\n")
|
||
sys.stdout.flush()
|
||
except Exception as e:
|
||
out = json.dumps({"ok": False, "error": str(e)}, ensure_ascii=False)
|
||
sys.stdout.write(out + "\n")
|
||
sys.stdout.flush()
|
||
sys.stderr.write(f"[easyocr_worker] 处理失败: {e}\n")
|
||
sys.stderr.flush()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
_main()
|
||
except Exception as e:
|
||
_flush_exit(f"_main 崩溃: {e}\n{traceback.format_exc()}")
|