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
|
|
|