EasyOCR GPU crash隔离:子进程隔离GPU推理,崩溃只杀子进程不冻电脑;新增easyocr_worker.py常驻worker;主进程降级到Windows OCR
This commit is contained in:
@@ -0,0 +1,92 @@
|
|||||||
|
#!/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
|
||||||
|
|
||||||
|
# ── 线程/进程限制:避免 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
|
||||||
|
|
||||||
|
|
||||||
|
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 识别 → 返回空格拼接文本。"""
|
||||||
|
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 " ".join([r[1] for r in results]).strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _main():
|
||||||
|
reader = _init_reader()
|
||||||
|
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 字段")
|
||||||
|
text = _process_image(img_b64, reader)
|
||||||
|
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__":
|
||||||
|
_main()
|
||||||
+103
-31
@@ -37,6 +37,8 @@ import win32process
|
|||||||
import re
|
import re
|
||||||
import ctypes
|
import ctypes
|
||||||
import difflib
|
import difflib
|
||||||
|
import subprocess
|
||||||
|
import json
|
||||||
import psutil
|
import psutil
|
||||||
import os
|
import os
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
@@ -182,6 +184,9 @@ class MessageMonitor:
|
|||||||
self._ocr_engine = None
|
self._ocr_engine = None
|
||||||
# EasyOCR 备用引擎(lazy init)
|
# EasyOCR 备用引擎(lazy init)
|
||||||
self._easyocr_reader = None
|
self._easyocr_reader = None
|
||||||
|
# EasyOCR 子进程(GPU crash 隔离,崩溃只杀子进程不冻电脑)
|
||||||
|
self._easyocr_proc: Optional['subprocess.Popen'] = None
|
||||||
|
self._easyocr_broken: bool = False
|
||||||
# 引擎选择计数器
|
# 引擎选择计数器
|
||||||
self._ocr_engine_wins: dict = {"windows": 0, "easyocr": 0}
|
self._ocr_engine_wins: dict = {"windows": 0, "easyocr": 0}
|
||||||
# 上次 OCR 原始文本(用于增量检测)
|
# 上次 OCR 原始文本(用于增量检测)
|
||||||
@@ -208,11 +213,7 @@ class MessageMonitor:
|
|||||||
if self._running:
|
if self._running:
|
||||||
return
|
return
|
||||||
self._running = True
|
self._running = True
|
||||||
# 在主线程预加载 EasyOCR,避免 daemon 线程 torch 段错误
|
# EasyOCR 已在子进程隔离(easyocr_worker.py),不需在主线程预加载
|
||||||
if _HAS_EASYOCR:
|
|
||||||
import torch
|
|
||||||
torch.set_num_threads(1)
|
|
||||||
self._init_easyocr()
|
|
||||||
self._thread = threading.Thread(target=self._run_loop, daemon=True, name="monitor")
|
self._thread = threading.Thread(target=self._run_loop, daemon=True, name="monitor")
|
||||||
self._thread.start()
|
self._thread.start()
|
||||||
logger.info("消息监控已启动")
|
logger.info("消息监控已启动")
|
||||||
@@ -766,6 +767,82 @@ class MessageMonitor:
|
|||||||
logger.error(f"[ocr] EasyOCR 初始化失败: {e}")
|
logger.error(f"[ocr] EasyOCR 初始化失败: {e}")
|
||||||
self._easyocr_reader = None
|
self._easyocr_reader = None
|
||||||
|
|
||||||
|
# ── EasyOCR 子进程隔离 ────────────────────────────────────
|
||||||
|
# 直接 GPU 推理 crash(access violation)会杀整个进程 + 冻电脑。
|
||||||
|
# 子进程方案:GPU crash 只杀子进程,主进程降级到 Windows OCR。
|
||||||
|
|
||||||
|
def _ensure_easyocr_subprocess(self) -> bool:
|
||||||
|
"""确保 EasyOCR 子进程在运行。返回 True = 可用。"""
|
||||||
|
if self._easyocr_broken:
|
||||||
|
return False
|
||||||
|
if self._easyocr_proc is not None:
|
||||||
|
if self._easyocr_proc.poll() is None:
|
||||||
|
return True # 还在跑
|
||||||
|
logger.warning("[ocr] EasyOCR 子进程已退出,重启...")
|
||||||
|
self._easyocr_proc = None
|
||||||
|
|
||||||
|
worker_path = os.path.join(
|
||||||
|
os.path.dirname(os.path.abspath(__file__)),
|
||||||
|
'easyocr_worker.py'
|
||||||
|
)
|
||||||
|
if not os.path.exists(worker_path):
|
||||||
|
logger.error(f"[ocr] EasyOCR worker 不存在: {worker_path}")
|
||||||
|
self._easyocr_broken = True
|
||||||
|
return False
|
||||||
|
|
||||||
|
try:
|
||||||
|
self._easyocr_proc = subprocess.Popen(
|
||||||
|
[sys.executable, '-u', worker_path],
|
||||||
|
stdin=subprocess.PIPE,
|
||||||
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.DEVNULL, # 防 stderr 管道满阻塞
|
||||||
|
text=True,
|
||||||
|
bufsize=1, # 行缓冲
|
||||||
|
)
|
||||||
|
logger.debug("[ocr] EasyOCR 子进程已启动")
|
||||||
|
return True
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[ocr] 启动 EasyOCR 子进程失败: {e}")
|
||||||
|
self._easyocr_broken = True
|
||||||
|
self._easyocr_proc = None
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _call_easyocr_safe(self, img_np) -> Optional[str]:
|
||||||
|
"""通过子进程安全调用 EasyOCR。返回识别文本,失败/崩溃时返回 None。"""
|
||||||
|
import io
|
||||||
|
import base64
|
||||||
|
|
||||||
|
if not self._ensure_easyocr_subprocess():
|
||||||
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# numpy → PIL → PNG bytes → base64
|
||||||
|
img_pil = Image.fromarray(img_np)
|
||||||
|
buf = io.BytesIO()
|
||||||
|
img_pil.save(buf, format='PNG')
|
||||||
|
img_b64 = base64.b64encode(buf.getvalue()).decode()
|
||||||
|
|
||||||
|
self._easyocr_proc.stdin.write(json.dumps({"img": img_b64}) + "\n")
|
||||||
|
self._easyocr_proc.stdin.flush()
|
||||||
|
|
||||||
|
line = self._easyocr_proc.stdout.readline()
|
||||||
|
if not line:
|
||||||
|
raise ConnectionError("EasyOCR 子进程 stdout 关闭")
|
||||||
|
result = json.loads(line)
|
||||||
|
if result.get("ok"):
|
||||||
|
return result["text"]
|
||||||
|
else:
|
||||||
|
logger.debug(f"[ocr] EasyOCR worker 错误: {result.get('error', '?')[:100]}")
|
||||||
|
return None
|
||||||
|
except (BrokenPipeError, ConnectionError, OSError) as e:
|
||||||
|
logger.debug(f"[ocr] EasyOCR 子进程通信断开: {e}")
|
||||||
|
self._easyocr_broken = True
|
||||||
|
self._easyocr_proc = None
|
||||||
|
return None
|
||||||
|
except Exception as e:
|
||||||
|
logger.debug(f"[ocr] EasyOCR 子进程异常: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _score_ocr_text(text: str) -> float:
|
def _score_ocr_text(text: str) -> float:
|
||||||
"""为 OCR 结果质量打分(越高越好)。对比引擎用。"""
|
"""为 OCR 结果质量打分(越高越好)。对比引擎用。"""
|
||||||
@@ -924,21 +1001,18 @@ class MessageMonitor:
|
|||||||
win_text = win_result.text.strip()
|
win_text = win_result.text.strip()
|
||||||
logger.debug(f"[ocr] Windows OCR 完成: {len(win_text)} chars")
|
logger.debug(f"[ocr] Windows OCR 完成: {len(win_text)} chars")
|
||||||
|
|
||||||
# EasyOCR 备用
|
# EasyOCR 备用(GPU crash 隔离:通过子进程调用)
|
||||||
easy_text = ""
|
easy_text = ""
|
||||||
if _HAS_EASYOCR:
|
if _HAS_EASYOCR and not self._easyocr_broken:
|
||||||
logger.debug("[ocr] EasyOCR 检查/初始化...")
|
logger.debug("[ocr] EasyOCR 子进程...")
|
||||||
self._init_easyocr()
|
img_np = _np.array(im)
|
||||||
logger.debug("[ocr] EasyOCR reader 就绪状态")
|
easy_text = self._call_easyocr_safe(img_np)
|
||||||
if self._easyocr_reader:
|
if easy_text is None:
|
||||||
try:
|
self._easyocr_broken = True
|
||||||
logger.debug("[ocr] EasyOCR 识别开始...")
|
easy_text = ""
|
||||||
img_np = _np.array(im)
|
logger.warning("[ocr] EasyOCR 子进程崩溃,已禁用(本次会话不再尝试)")
|
||||||
easy_results = self._easyocr_reader.readtext(img_np)
|
elif easy_text:
|
||||||
easy_text = ' '.join([r[1] for r in easy_results]).strip()
|
logger.debug(f"[ocr] EasyOCR 识别完成: {len(easy_text)} chars")
|
||||||
logger.debug(f"[ocr] EasyOCR 识别完成: {len(easy_text)} chars")
|
|
||||||
except Exception as e:
|
|
||||||
logger.debug(f"[ocr] EasyOCR 识别失败: {e}")
|
|
||||||
|
|
||||||
# 引擎原始输出逐行标记日志(用于长期对比两引擎能力)
|
# 引擎原始输出逐行标记日志(用于长期对比两引擎能力)
|
||||||
if win_text:
|
if win_text:
|
||||||
@@ -1041,18 +1115,16 @@ class MessageMonitor:
|
|||||||
logger.debug(f"[wechat] Windows OCR 完成: {len(win_text)} chars")
|
logger.debug(f"[wechat] Windows OCR 完成: {len(win_text)} chars")
|
||||||
|
|
||||||
easy_text = ""
|
easy_text = ""
|
||||||
if _HAS_EASYOCR:
|
if _HAS_EASYOCR and not self._easyocr_broken:
|
||||||
logger.debug("[wechat] EasyOCR 检查/初始化...")
|
logger.debug("[wechat] EasyOCR 子进程...")
|
||||||
self._init_easyocr()
|
img_np = _np.array(im)
|
||||||
logger.debug("[wechat] EasyOCR reader 就绪状态")
|
easy_text = self._call_easyocr_safe(img_np)
|
||||||
if self._easyocr_reader:
|
if easy_text is None:
|
||||||
try:
|
self._easyocr_broken = True
|
||||||
logger.debug("[wechat] EasyOCR 识别开始...")
|
easy_text = ""
|
||||||
img_np = _np.array(im)
|
logger.warning("[wechat] EasyOCR 子进程崩溃,已禁用(本次会话不再尝试)")
|
||||||
easy_results = self._easyocr_reader.readtext(img_np)
|
elif easy_text:
|
||||||
easy_text = ' '.join([r[1] for r in easy_results]).strip()
|
logger.debug(f"[wechat] EasyOCR 识别完成: {len(easy_text)} chars")
|
||||||
except Exception as e:
|
|
||||||
logger.debug(f"[wechat] EasyOCR 识别失败: {e}")
|
|
||||||
|
|
||||||
# 引擎原始输出逐行标记日志(用于长期对比两引擎能力)
|
# 引擎原始输出逐行标记日志(用于长期对比两引擎能力)
|
||||||
if win_text:
|
if win_text:
|
||||||
|
|||||||
Reference in New Issue
Block a user