重构OCR引擎管理器:双引擎并发+1分钟超时释放
1. 新增OCREngineManager类: - 预初始化EasyOCR(GPU=False) + Windows OCR两套引擎 - 双引擎并发识别,取得分最高者结果 - 1分钟超时自动释放内存(消息监控运行中不释放) - 支持监控激活时重置计时器 2. 修复截屏API: - GetSystemMetrics改用ctypes.windll.user32,兼容性更好 3. GUI集成: - 启动时后台预初始化OCR引擎 - 监控开关时通知OCR管理器激活/停止状态 - 移除旧_win_ocr方法(已由管理器统一处理)
This commit is contained in:
@@ -20,6 +20,252 @@ from auto_reply import AutoReply
|
||||
logger = logging.getLogger("msgbot.gui")
|
||||
|
||||
|
||||
class OCREngineManager:
|
||||
"""双引擎OCR管理器:EasyOCR + Windows OCR,得分最高者胜出。
|
||||
|
||||
功能:
|
||||
1. 预初始化两套引擎
|
||||
2. 双引擎并发识别,取得分最高者的结果
|
||||
3. 1分钟超时释放机制:无操作则杀进程释放内存
|
||||
4. 支持消息监控调度时直接使用预热的引擎
|
||||
"""
|
||||
|
||||
def __init__(self, log_callback=None):
|
||||
self._log = log_callback or (lambda msg: None)
|
||||
self._easyocr_reader = None
|
||||
self._winocr_engine = None
|
||||
self._initialized = False
|
||||
self._lock = threading.Lock()
|
||||
self._last_use_time = 0
|
||||
self._timeout_seconds = 60 # 1分钟超时
|
||||
self._consecutive_success = 0
|
||||
self._consecutive_fail = 0
|
||||
self._monitor_active = False # 消息监控是否激活
|
||||
self._timer_thread = None
|
||||
self._timer_stop = threading.Event()
|
||||
|
||||
def initialize(self):
|
||||
"""预初始化两套OCR引擎"""
|
||||
if self._initialized:
|
||||
return True
|
||||
|
||||
with self._lock:
|
||||
if self._initialized:
|
||||
return True
|
||||
|
||||
self._log("🔧 初始化OCR引擎...")
|
||||
|
||||
# 初始化EasyOCR
|
||||
try:
|
||||
import easyocr
|
||||
self._log(" ⏳ 初始化EasyOCR (CPU模式)...")
|
||||
self._easyocr_reader = easyocr.Reader(['ch_sim', 'en'], gpu=False, verbose=False)
|
||||
self._log(" ✅ EasyOCR就绪")
|
||||
except Exception as e:
|
||||
self._log(f" ⚠ EasyOCR初始化失败: {e}")
|
||||
self._easyocr_reader = None
|
||||
|
||||
# 初始化Windows OCR
|
||||
try:
|
||||
from winrt.windows.media.ocr import OcrEngine
|
||||
from winrt.windows.globalization import Language
|
||||
self._winocr_engine = OcrEngine.try_create_from_user_profile_languages()
|
||||
if self._winocr_engine:
|
||||
self._log(" ✅ Windows OCR就绪")
|
||||
else:
|
||||
self._log(" ⚠ Windows OCR不可用")
|
||||
except Exception as e:
|
||||
self._log(f" ⚠ Windows OCR初始化失败: {e}")
|
||||
self._winocr_engine = None
|
||||
|
||||
if not self._easyocr_reader and not self._winocr_engine:
|
||||
self._log("❌ 无可用OCR引擎")
|
||||
return False
|
||||
|
||||
self._initialized = True
|
||||
self._last_use_time = time.time()
|
||||
self._start_timeout_timer()
|
||||
self._log("✅ OCR引擎初始化完成")
|
||||
return True
|
||||
|
||||
def _start_timeout_timer(self):
|
||||
"""启动超时计时器线程"""
|
||||
if self._timer_thread and self._timer_thread.is_alive():
|
||||
return
|
||||
|
||||
self._timer_stop.clear()
|
||||
|
||||
def _timer_loop():
|
||||
while not self._timer_stop.is_set():
|
||||
self._timer_stop.wait(10) # 每10秒检查一次
|
||||
if self._timer_stop.is_set():
|
||||
break
|
||||
|
||||
# 消息监控激活中,不释放
|
||||
if self._monitor_active:
|
||||
self._last_use_time = time.time() # 重置计时
|
||||
continue
|
||||
|
||||
# 有新截图动作,不释放
|
||||
if time.time() - self._last_use_time < self._timeout_seconds:
|
||||
continue
|
||||
|
||||
# 超时,释放资源
|
||||
self._log("⏰ OCR引擎1分钟无操作,释放内存...")
|
||||
self._release()
|
||||
break
|
||||
|
||||
self._timer_thread = threading.Thread(target=_timer_loop, daemon=True)
|
||||
self._timer_thread.start()
|
||||
|
||||
def _release(self):
|
||||
"""释放OCR引擎资源"""
|
||||
with self._lock:
|
||||
self._easyocr_reader = None
|
||||
self._winocr_engine = None
|
||||
self._initialized = False
|
||||
self._log("🗑 OCR引擎已释放")
|
||||
|
||||
def set_monitor_active(self, active: bool):
|
||||
"""设置消息监控激活状态"""
|
||||
self._monitor_active = active
|
||||
if active and self._initialized:
|
||||
self._last_use_time = time.time() # 激活时重置计时
|
||||
|
||||
def recognize(self, image) -> list:
|
||||
"""双引擎并发识别,返回得分最高的结果
|
||||
|
||||
Args:
|
||||
image: PIL Image对象
|
||||
|
||||
Returns:
|
||||
识别到的文本列表
|
||||
"""
|
||||
if not self._initialized:
|
||||
self.initialize()
|
||||
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
self._last_use_time = time.time()
|
||||
|
||||
results_easyocr = []
|
||||
results_winocr = []
|
||||
errors = []
|
||||
|
||||
# 并发识别
|
||||
def _run_easyocr():
|
||||
nonlocal results_easyocr
|
||||
try:
|
||||
if self._easyocr_reader:
|
||||
results_easyocr = self._easyocr_reader.readtext(image, detail=1)
|
||||
except Exception as e:
|
||||
errors.append(f"EasyOCR: {e}")
|
||||
|
||||
def _run_winocr():
|
||||
nonlocal results_winocr
|
||||
try:
|
||||
if self._winocr_engine:
|
||||
results_winocr = self._recognize_winocr(image)
|
||||
except Exception as e:
|
||||
errors.append(f"WinOCR: {e}")
|
||||
|
||||
threads = []
|
||||
if self._easyocr_reader:
|
||||
threads.append(threading.Thread(target=_run_easyocr, daemon=True))
|
||||
if self._winocr_engine:
|
||||
threads.append(threading.Thread(target=_run_winocr, daemon=True))
|
||||
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join(timeout=30) # 30秒超时
|
||||
|
||||
# 选择得分最高的结果
|
||||
best_texts = self._select_best_result(results_easyocr, results_winocr)
|
||||
|
||||
if best_texts:
|
||||
self._consecutive_success += 1
|
||||
self._consecutive_fail = 0
|
||||
else:
|
||||
self._consecutive_success = 0
|
||||
self._consecutive_fail += 1
|
||||
if errors:
|
||||
self._log(f"⚠ OCR识别失败: {'; '.join(errors)}")
|
||||
|
||||
return best_texts
|
||||
|
||||
def _recognize_winocr(self, image) -> list:
|
||||
"""Windows OCR识别,返回带得分的结果"""
|
||||
import asyncio
|
||||
import io
|
||||
import tempfile
|
||||
|
||||
try:
|
||||
# 保存为临时文件
|
||||
tmp_path = os.path.join(tempfile.gettempdir(), '_ocr_nick_temp.png')
|
||||
image.save(tmp_path)
|
||||
|
||||
async def _do_ocr():
|
||||
from winrt.windows.media.ocr import OcrEngine
|
||||
from winrt.windows.graphics.imaging import BitmapDecoder
|
||||
from winrt.windows.storage import StorageFile
|
||||
|
||||
file = await StorageFile.get_file_from_path_async(tmp_path)
|
||||
stream = await file.open_read_async()
|
||||
decoder = await BitmapDecoder.create_async(stream)
|
||||
bitmap = await decoder.get_software_bitmap_async()
|
||||
result = await self._winocr_engine.recognize_async(bitmap)
|
||||
return result.text
|
||||
|
||||
text = asyncio.run(_do_ocr())
|
||||
# 按行分割,每行视为一个结果(得分0.8作为默认)
|
||||
lines = [line.strip() for line in text.split('\n') if line.strip()]
|
||||
return [(line, 0.8) for line in lines]
|
||||
except Exception as e:
|
||||
logger.error(f"Windows OCR失败: {e}")
|
||||
return []
|
||||
|
||||
def _select_best_result(self, easyocr_results, winocr_results) -> list:
|
||||
"""选择得分最高的结果
|
||||
|
||||
EasyOCR返回: [(bbox, text, confidence), ...]
|
||||
WinOCR返回: [(text, confidence), ...]
|
||||
"""
|
||||
# 统一格式为 [(text, confidence), ...]
|
||||
eo = [(r[1], r[2]) for r in easyocr_results if len(r) >= 3]
|
||||
wo = winocr_results # 已经是 [(text, confidence)]
|
||||
|
||||
all_results = eo + wo
|
||||
if not all_results:
|
||||
return []
|
||||
|
||||
# 按置信度排序,取最高的
|
||||
all_results.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# 返回所有结果的文本(去重)
|
||||
seen = set()
|
||||
texts = []
|
||||
for text, conf in all_results:
|
||||
text = text.strip()
|
||||
if text and text not in seen:
|
||||
seen.add(text)
|
||||
texts.append(text)
|
||||
|
||||
return texts
|
||||
|
||||
def get_status(self) -> dict:
|
||||
"""获取引擎状态"""
|
||||
return {
|
||||
'initialized': self._initialized,
|
||||
'easyocr': self._easyocr_reader is not None,
|
||||
'winocr': self._winocr_engine is not None,
|
||||
'consecutive_success': self._consecutive_success,
|
||||
'consecutive_fail': self._consecutive_fail,
|
||||
'monitor_active': self._monitor_active,
|
||||
}
|
||||
|
||||
|
||||
def _clean_query_for_kb(msg: str) -> str:
|
||||
"""清理查询文本中的用户名/头衔/时间戳等噪音,提升KB搜索精度。"""
|
||||
if not msg:
|
||||
@@ -77,6 +323,11 @@ class MsgBotApp(ctk.CTk):
|
||||
# 自动回复
|
||||
self.ar = AutoReply(self.monitor)
|
||||
|
||||
# ⭐ OCR引擎管理器(双引擎:EasyOCR + Windows OCR)
|
||||
self._ocr_manager = OCREngineManager(log_callback=self._log)
|
||||
# 预初始化OCR引擎(后台线程,不阻塞UI)
|
||||
threading.Thread(target=self._ocr_manager.initialize, daemon=True).start()
|
||||
|
||||
# ⭐ 跨平台消息处理锁:防止同一消息被QQ和微信同时处理
|
||||
self._processing_lock = threading.Lock()
|
||||
self._processing_messages = set() # 正在处理的消息hash集合
|
||||
@@ -304,10 +555,14 @@ class MsgBotApp(ctk.CTk):
|
||||
self.monitor.start(full_restart=True)
|
||||
self._status_label.configure(text="▶ 运行中", text_color="#2ecc71")
|
||||
self._log("✅ 消息监控已启动")
|
||||
# 通知OCR管理器:监控已激活,不释放引擎
|
||||
self._ocr_manager.set_monitor_active(True)
|
||||
else:
|
||||
self.monitor.stop()
|
||||
self._status_label.configure(text="⏸ 已暂停")
|
||||
self._log("⏸ 消息监控已暂停")
|
||||
# 通知OCR管理器:监控已停止,开始计时
|
||||
self._ocr_manager.set_monitor_active(False)
|
||||
|
||||
def _browse_kb(self):
|
||||
path = filedialog.askdirectory(title="选择知识库文件夹")
|
||||
@@ -430,7 +685,7 @@ class MsgBotApp(ctk.CTk):
|
||||
self._log(f"⚠ 删除失败: {e}")
|
||||
|
||||
def _ocr_nicknames_from_screenshot(self):
|
||||
"""截图OCR识别昵称:最小化窗口→截屏→OCR→提取昵称→入库"""
|
||||
"""截图OCR识别昵称:截屏→双引擎OCR→提取昵称→入库"""
|
||||
self._log("📸 正在截屏,请确保QQ聊天窗口可见...")
|
||||
self._ocr_nick_btn.configure(state="disabled", text="截屏中...")
|
||||
|
||||
@@ -445,16 +700,10 @@ class MsgBotApp(ctk.CTk):
|
||||
self._log_queue.put("⚠ 截屏失败")
|
||||
return
|
||||
|
||||
self._log_queue.put("🔍 OCR识别中...")
|
||||
self._log_queue.put("🔍 OCR识别中(双引擎并发)...")
|
||||
|
||||
# OCR识别
|
||||
try:
|
||||
import easyocr
|
||||
reader = easyocr.Reader(['ch_sim', 'en'], gpu=False, verbose=False)
|
||||
results = reader.readtext(screenshot, detail=0)
|
||||
except Exception as e:
|
||||
self._log_queue.put(f"⚠ EasyOCR失败: {e},尝试Windows OCR...")
|
||||
results = self._win_ocr(screenshot)
|
||||
# 使用OCR管理器识别(双引擎并发,得分最高者胜出)
|
||||
results = self._ocr_manager.recognize(screenshot)
|
||||
|
||||
if not results:
|
||||
self._log_queue.put("⚠ OCR未识别到任何文字")
|
||||
@@ -499,11 +748,13 @@ class MsgBotApp(ctk.CTk):
|
||||
import win32gui
|
||||
import win32ui
|
||||
import win32con
|
||||
import ctypes
|
||||
from PIL import Image
|
||||
|
||||
# 获取屏幕尺寸
|
||||
screen_w = win32gui.GetSystemMetrics(win32con.SM_CXSCREEN)
|
||||
screen_h = win32gui.GetSystemMetrics(win32con.SM_CYSCREEN)
|
||||
# 获取屏幕尺寸(ctypes方式,兼容性更好)
|
||||
user32 = ctypes.windll.user32
|
||||
screen_w = user32.GetSystemMetrics(0) # SM_CXSCREEN
|
||||
screen_h = user32.GetSystemMetrics(1) # SM_CYSCREEN
|
||||
|
||||
# 创建设备上下文
|
||||
hdesktop = win32gui.GetDesktopWindow()
|
||||
@@ -536,37 +787,6 @@ class MsgBotApp(ctk.CTk):
|
||||
logger.error(f"截屏失败: {e}")
|
||||
return None
|
||||
|
||||
def _win_ocr(self, image):
|
||||
"""Windows OCR后备方案(无需GPU)"""
|
||||
try:
|
||||
import io
|
||||
from PIL import Image
|
||||
# 保存为临时文件
|
||||
tmp_path = os.path.join(os.environ.get('TEMP', '.'), '_ocr_temp.png')
|
||||
image.save(tmp_path)
|
||||
|
||||
# 使用Windows OCR(需要winrt)
|
||||
from winrt.windows.media.ocr import OcrEngine
|
||||
from winrt.windows.graphics.imaging import BitmapDecoder
|
||||
from winrt.windows.storage import StorageFile
|
||||
import asyncio
|
||||
|
||||
async def _do_ocr():
|
||||
file = await StorageFile.get_file_from_path_async(tmp_path)
|
||||
stream = await file.open_read_async()
|
||||
decoder = await BitmapDecoder.create_async(stream)
|
||||
bitmap = await decoder.get_software_bitmap_async()
|
||||
engine = OcrEngine.try_create_from_user_profile_languages()
|
||||
result = await engine.recognize_async(bitmap)
|
||||
return result.text
|
||||
|
||||
text = asyncio.run(_do_ocr())
|
||||
# 按行分割
|
||||
return [line.strip() for line in text.split('\n') if line.strip()]
|
||||
except Exception as e:
|
||||
logger.error(f"Windows OCR失败: {e}")
|
||||
return []
|
||||
|
||||
def _extract_nicknames_from_ocr(self, ocr_results):
|
||||
"""从OCR识别结果中提取QQ昵称模式"""
|
||||
nicks = set()
|
||||
|
||||
Reference in New Issue
Block a user