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模型已释放。"); } } } }