137 lines
4.6 KiB
Python
137 lines
4.6 KiB
Python
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
|
||
|