from core.ocr import load_model as ocrLoadModel, detect as ocrDetect from core.face import load_model as faceLoadModel, detect as faceDetect from core.yolo import load_model as yoloLoadModel, detect as yoloDetect from concurrent.futures import ThreadPoolExecutor, Future import threading import cv2 import numpy as np # -------------------------- 核心配置参数 -------------------------- MAX_WORKERS = 6 # 线程池最大线程数 DETECTION_ORDER = ["yolo", "face", "ocr"] # 检测优先级顺序 TIMEOUT = 30 # 检测超时时间(秒) # -------------------------- 全局状态管理 -------------------------- _executor = None # 线程池实例 _model_loaded = False # 模型加载状态标记 _model_lock = threading.Lock() # 模型加载线程锁 _executor_lock = threading.Lock() # 线程池初始化锁 _task_counter = 0 # 任务计数器 _task_counter_lock = threading.Lock() # 任务计数锁 # -------------------------- 工具函数 -------------------------- def _get_next_task_id(): """获取唯一任务ID、用于日志追踪""" global _task_counter with _task_counter_lock: _task_counter += 1 return _task_counter # -------------------------- 模型加载 -------------------------- def load_model(): """加载所有检测模型并初始化线程池(仅执行一次)""" global _model_loaded if not _model_loaded: with _model_lock: if not _model_loaded: print("=== 开始加载检测模型 ===") # 按顺序加载模型 print("加载YOLO模型...") yoloLoadModel() print("加载人脸检测模型...") faceLoadModel() print("加载OCR模型...") ocrLoadModel() _model_loaded = True print("=== 所有模型加载完成 ===") # 初始化线程池 _init_thread_pool() # -------------------------- 线程池管理 -------------------------- def _init_thread_pool(): """初始化线程池(仅内部调用)""" global _executor with _executor_lock: if _executor is None: _executor = ThreadPoolExecutor( max_workers=MAX_WORKERS, thread_name_prefix="DetectionThread" ) print(f"=== 线程池初始化完成、最大线程数: {MAX_WORKERS} ===") def shutdown(): """关闭线程池、释放资源""" global _executor with _executor_lock: if _executor is not None: _executor.shutdown(wait=True) _executor = None print("=== 线程池已安全关闭 ===") # -------------------------- 检测逻辑实现 -------------------------- def _detect_in_thread(frame: np.ndarray, task_id: int) -> tuple: """在子线程中执行检测逻辑""" thread_name = threading.current_thread().name print(f"任务[{task_id}] 开始执行、线程: {thread_name}") try: # 按照优先级执行检测 for detector in DETECTION_ORDER: if detector == "yolo": flag, result = yoloDetect(frame) elif detector == "face": flag, result = faceDetect(frame) elif detector == "ocr": flag, result = ocrDetect(frame) else: flag, result = False, None print(f"任务[{task_id}] {detector}检测结果: {'成功' if flag else '失败'}") if flag: print(f"任务[{task_id}] 完成检测、使用检测器: {detector}") return (True, result, detector, task_id) # 所有检测器均未检测到结果 print(f"任务[{task_id}] 所有检测器均未检测到内容") return (False, "未检测到任何内容", "none", task_id) except Exception as e: print(f"任务[{task_id}] 检测过程发生错误: {str(e)}") return (False, f"检测错误: {str(e)}", "error", task_id) # -------------------------- 外部调用接口 -------------------------- def detect(frame: np.ndarray) -> Future: """ 提交检测任务到线程池 参数: frame: 待检测图像(ndarray格式、cv2.imdecode生成) 返回: Future对象、通过result()方法获取检测结果 """ # 确保模型已加载 if not _model_loaded: print("警告: 模型尚未加载、将自动加载") load_model() # 生成任务ID task_id = _get_next_task_id() # 提交任务到线程池 future = _executor.submit(_detect_in_thread, frame, task_id) print(f"任务[{task_id}]: 已提交到线程池") return future