video/core/all.py

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