最新可用

core/all.py

@@ -1,45 +1,136 @@
 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
 
-# 添加一个标记变量，用于监控load_model是否已被调用
-_model_loaded = False
+# -------------------------- 核心配置参数 --------------------------
+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("=== 开始加载检测模型 ===")
 
-    # 如果已经调用过，直接忽略
-    if _model_loaded:
-        return
+                # 按顺序加载模型
+                print("加载YOLO模型...")
+                yoloLoadModel()
 
-    # 首次调用时加载模型
-    ocrLoadModel()
-    faceLoadModel()
-    yoloLoadModel()
+                print("加载人脸检测模型...")
+                faceLoadModel()
 
-    # 标记为已调用
-    _model_loaded = True
+                print("加载OCR模型...")
+                ocrLoadModel()
+
+                _model_loaded = True
+                print("=== 所有模型加载完成 ===")
+
+                # 初始化线程池
+                _init_thread_pool()
 
 
-def detect(frame):
-    # 先进行YOLO检测
-    yolo_flag, yolo_result = yoloDetect(frame)
-    print("YOLO检测结果：", yolo_result)
-    if yolo_flag:
-        return (True, yolo_result, "yolo")
+# -------------------------- 线程池管理 --------------------------
+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} ===")
 
-    # YOLO未检测到，进行人脸检测
-    face_flag, face_result = faceDetect(frame)
-    print("人脸检测结果：", face_result)
-    if face_flag:
-        return (True, face_result, "face")
 
-    # 人脸未检测到，进行OCR检测
-    ocr_flag, ocr_result = ocrDetect(frame)
-    print("OCR检测结果：", ocr_result)
-    if ocr_flag:
-        return (True, ocr_result, "ocr")
+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
 
-    # 所有检测都未检测到
-    return (False, "未检测到任何内容", "none")