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: """在子线程中执行检测逻辑(返回4元素tuple:检测是否成功、结果数据、检测器类型、任务ID)""" thread_name = threading.current_thread().name print(f"任务[{task_id}] 开始执行、线程: {thread_name}") try: # 按照配置顺序执行检测 for detector in DETECTION_ORDER: if detector == "yolo": success, result = yoloDetect(frame) elif detector == "face": success, result = faceDetect(frame) elif detector == "ocr": success, result = ocrDetect(frame) else: success, result = False, None print(f"任务[{task_id}] {detector}检测状态: {'成功' if success else '未检测到内容'}") if success: print(f"任务[{task_id}] 完成检测、使用检测器: {detector}") return (success, result, detector, task_id) # 4元素tuple # 所有检测器均未检测到结果 print(f"任务[{task_id}] 所有检测器均未检测到有效内容") return (False, "未检测到任何有效内容", "none", task_id) # 4元素tuple except Exception as e: print(f"任务[{task_id}] 检测过程发生错误: {str(e)}") return (False, f"检测错误: {str(e)}", "error", task_id) # 4元素tuple # -------------------------- 外部调用接口 -------------------------- def detect(frame: np.ndarray) -> Future: """ 提交检测任务到线程池(返回Future对象,需调用result()获取4元素结果) 参数: frame: 待检测图像(ndarray格式、cv2.imdecode生成) 返回: Future对象、result()返回tuple: (success, data, detector_type, task_id) success: 布尔值,表示是否检测到有效内容 data: 检测结果数据(成功时为具体结果,失败时为错误信息) detector_type: 使用的检测器类型("yolo"/"face"/"ocr"/"none"/"error") task_id: 任务唯一标识 """ # 确保模型已加载 if not _model_loaded: print("警告: 模型尚未加载、将自动加载") load_model() # 生成任务ID task_id = _get_next_task_id() # 提交任务到线程池(返回Future) future = _executor.submit(_detect_in_thread, frame, task_id) print(f"任务[{task_id}]: 已提交到线程池") return future