103 lines
3.6 KiB
C#
103 lines
3.6 KiB
C#
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
|
||
{
|
||
/// <summary>
|
||
/// 静态全局YOLO模型管理器。
|
||
/// 负责在程序启动时加载所有模型，并在关闭时释放资源。
|
||
/// </summary>
|
||
public static class YoloModelManager
|
||
{
|
||
// 使用 ConcurrentDictionary 保证线程安全
|
||
private static readonly ConcurrentDictionary<int, Yolo> _loadedModels = new ConcurrentDictionary<int, Yolo>();
|
||
|
||
/// <summary>
|
||
/// 根据模型配置列表初始化所有YOLO模型实例。
|
||
/// 应在程序启动时调用一次。
|
||
/// </summary>
|
||
/// <param name="modelSettings">模型配置列表。</param>
|
||
public static void Initialize(List<ModelSettings> modelSettings)
|
||
{
|
||
Shutdown(); // 先确保清理掉旧实例
|
||
|
||
if (modelSettings == null) return;
|
||
|
||
ThreadSafeLogger.Log("开始加载YOLO模型...");
|
||
foreach (var setting in modelSettings)
|
||
{
|
||
bool gpuUse = false;
|
||
if (setting.CheckDevice == DetectDevice.GPU)
|
||
{
|
||
gpuUse = true;
|
||
}
|
||
|
||
if (string.IsNullOrEmpty(setting.Path) || !File.Exists(setting.Path))
|
||
{
|
||
ThreadSafeLogger.Log($"[警告] 模型 '{setting.Name}' (ID: {setting.Id}) 路径无效或文件不存在，已跳过加载。");
|
||
continue;
|
||
}
|
||
|
||
try
|
||
{
|
||
// 创建YOLO实例
|
||
var yolo = new Yolo(new YoloOptions
|
||
{
|
||
OnnxModel = setting.Path,
|
||
// 您可以根据需要从配置中读取这些值
|
||
ModelType = (YoloDotNet.Enums.ModelType)setting.MType,
|
||
Cuda = gpuUse, // 推荐使用GPU
|
||
PrimeGpu = false
|
||
});
|
||
|
||
if (_loadedModels.TryAdd(setting.Id, yolo))
|
||
{
|
||
ThreadSafeLogger.Log($"成功加载模型 '{setting.Name}' (ID: {setting.Id})。");
|
||
}
|
||
}
|
||
catch (Exception ex)
|
||
{
|
||
ThreadSafeLogger.Log($"[错误] 加载模型 '{setting.Name}' (ID: {setting.Id}) 失败: {ex.Message}");
|
||
}
|
||
}
|
||
ThreadSafeLogger.Log($"YOLO模型加载完成，共成功加载 {_loadedModels.Count} 个模型。");
|
||
}
|
||
|
||
/// <summary>
|
||
/// 获取指定ID的已加载YOLO模型。
|
||
/// </summary>
|
||
/// <param name="modelId">模型编号。</param>
|
||
/// <returns>Yolo实例，如果未找到则返回null。</returns>
|
||
public static Yolo GetModel(int modelId)
|
||
{
|
||
_loadedModels.TryGetValue(modelId, out var yolo);
|
||
return yolo;
|
||
}
|
||
|
||
/// <summary>
|
||
/// 释放所有已加载的YOLO模型资源。
|
||
/// 应在程序关闭时调用。
|
||
/// </summary>
|
||
public static void Shutdown()
|
||
{
|
||
if (_loadedModels.Count > 0)
|
||
{
|
||
ThreadSafeLogger.Log("正在释放所有YOLO模型...");
|
||
foreach (var yolo in _loadedModels.Values)
|
||
{
|
||
yolo?.Dispose();
|
||
}
|
||
_loadedModels.Clear();
|
||
ThreadSafeLogger.Log("所有YOLO模型已释放。");
|
||
}
|
||
}
|
||
}
|
||
}
|