using Check.Main.Common;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using YoloDotNet;
using YoloDotNet.Models;
namespace Check.Main.Infer
{
///
/// 静态全局YOLO模型管理器。
/// 负责在程序启动时加载所有模型,并在关闭时释放资源。
///
public static class YoloModelManager
{
// 使用 ConcurrentDictionary 保证线程安全
private static readonly ConcurrentDictionary _loadedModels = new ConcurrentDictionary();
///
/// 根据模型配置列表初始化所有YOLO模型实例。
/// 应在程序启动时调用一次。
///
/// 模型配置列表。
public static void Initialize(List modelSettings)
{
Shutdown(); // 先确保清理掉旧实例
if (modelSettings == null) return;
ThreadSafeLogger.Log("开始加载YOLO模型...");
// 筛选出启用的深度学习模型进行加载
foreach (var setting in modelSettings.Where(s => s.IsEnabled && s.M_AType == AlgorithmType.DeepLearning))
{
bool gpuUse = setting.CheckDevice == DetectDevice.GPU;
if (string.IsNullOrEmpty(setting.Path) || !File.Exists(setting.Path))
{
ThreadSafeLogger.Log($"[警告] YOLO模型 '{setting.Name}' (ID: {setting.Id}) 路径无效或文件不存在,已跳过加载。");
continue;
}
try
{
var yolo = new Yolo(new YoloOptions
{
OnnxModel = setting.Path,
ModelType = (YoloDotNet.Enums.ModelType)setting.MType,
Cuda = gpuUse,
//Confidence = setting.YoloConfidenceThreshold, // 从 ModelSettings 读取
//Nms = setting.YoloNmsThreshold, // 从 ModelSettings 读取
PrimeGpu = false // 根据需求设置
});
// 保存阈值配置
var conf = setting.YoloConfidenceThreshold;
var nms = setting.YoloNmsThreshold;
if (_loadedModels.TryAdd(setting.Id, yolo))
{
ThreadSafeLogger.Log($"成功加载YOLO模型 '{setting.Name}' (ID: {setting.Id})。");
}
}
catch (Exception ex)
{
ThreadSafeLogger.Log($"[错误] 加载YOLO模型 '{setting.Name}' (ID: {setting.Id}) 失败: {ex.Message}");
}
}
ThreadSafeLogger.Log($"YOLO模型加载完成,共成功加载 {_loadedModels.Count} 个模型。");
}
///
/// 获取指定ID的已加载YOLO模型。
///
/// 模型编号。
/// Yolo实例,如果未找到则返回null。
public static Yolo GetModel(int modelId)
{
_loadedModels.TryGetValue(modelId, out var yolo);
return yolo;
}
///
/// 释放所有已加载的YOLO模型资源。
/// 应在程序关闭时调用。
///
public static void Shutdown()
{
if (_loadedModels.Count > 0)
{
ThreadSafeLogger.Log("正在释放所有YOLO模型...");
foreach (var yolo in _loadedModels.Values)
{
yolo?.Dispose();
}
_loadedModels.Clear();
ThreadSafeLogger.Log("所有YOLO模型已释放。");
}
}
}
}