目前可以成功动态更换模型运行的

This commit is contained in:
2025-09-12 14:05:09 +08:00
parent 435b2a0e6c
commit 4be7f7bf14
13 changed files with 1518 additions and 325 deletions

View File

@ -1,37 +1,43 @@
import os
import numpy as np
from ultralytics import YOLO
from service.model_service import get_current_yolo_model # 从模型管理模块获取模型
# 全局变量
# 全局模型变量
_yolo_model = None
model_path = os.path.join(os.path.dirname(__file__), "models", "best.pt")
def load_model():
"""加载YOLO目标检测模型"""
def load_model(model_path=None):
"""加载YOLO模型优先使用模型管理模块的默认模型"""
global _yolo_model
if model_path is None:
_yolo_model = get_current_yolo_model()
return _yolo_model is not None
try:
_yolo_model = YOLO(model_path)
return True
except Exception as e:
print(f"YOLO model load failed: {e}")
print(f"YOLO模型加载失败(指定路径):{str(e)}")
return False
return True if _yolo_model else False
def detect(frame, conf_threshold=0.2):
"""YOLO目标检测返回(是否识别到, 结果字符串)"""
"""执行目标检测返回(是否成功, 结果字符串)"""
global _yolo_model
if not _yolo_model or frame is None:
return (False, "未初始化或无效帧")
# 确保模型已加载
if not _yolo_model:
if not load_model():
return (False, "模型未初始化")
if frame is None:
return (False, "无效输入帧")
try:
results = _yolo_model(frame, conf=conf_threshold)
# 检查是否有检测结果
# 执行检测frame应为numpy数组
results = _yolo_model(frame, conf=conf_threshold, verbose=False)
has_results = len(results[0].boxes) > 0 if results else False
if not has_results:
@ -42,13 +48,12 @@ def detect(frame, conf_threshold=0.2):
for box in results[0].boxes:
cls = int(box.cls[0])
conf = float(box.conf[0])
bbox = [float(x) for x in box.xyxy[0]]
bbox = [round(x, 2) for x in box.xyxy[0].tolist()] # 保留两位小数
class_name = _yolo_model.names[cls] if hasattr(_yolo_model, 'names') else f"类别{cls}"
result_parts.append(f"{class_name} (置信度: {conf:.2f}, 边界框: {bbox})")
result_parts.append(f"{class_name}置信度{conf:.2f},位置:{bbox}")
result_str = "; ".join(result_parts)
return (has_results, result_str)
return (True, "; ".join(result_parts))
except Exception as e:
print(f"YOLO detect error: {e}")
return (False, f"检测错误: {str(e)}")
print(f"检测过程出错:{str(e)}")
return (False, f"检测错误{str(e)}")