video/core/yolo.py

import os
import numpy as np
from ultralytics import YOLO
from service.model_service import get_current_yolo_model  # 从模型管理模块获取模型

# 全局模型变量
_yolo_model = None


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模型加载失败（指定路径）：{str(e)}")
        return False


def detect(frame, conf_threshold=0.2):
    """执行目标检测，返回(是否成功, 结果字符串)"""
    global _yolo_model

    # 确保模型已加载
    if not _yolo_model:
        if not load_model():
            return (False, "模型未初始化")

    if frame is None:
        return (False, "无效输入帧")

    try:
        # 执行检测（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:
            return (False, "未检测到目标")

        # 构建结果字符串
        result_parts = []
        for box in results[0].boxes:
            cls = int(box.cls[0])
            conf = float(box.conf[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}）")

        return (True, "; ".join(result_parts))

    except Exception as e:
        print(f"检测过程出错：{str(e)}")
        return (False, f"检测错误：{str(e)}")