Display_D_HP/TeamAAS-VM/Core/YoloOpenVINOInference.cs

using System;
using System.Collections.Generic;
using System.Drawing;
using System.Drawing.Drawing2D;
using System.Drawing.Imaging;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using OpenVinoSharp;
using YamlDotNet.Serialization;
using YamlDotNet.Serialization.NamingConventions;

namespace WpfTest.core
{
    /// <summary>
    /// 检测结果类
    /// </summary>
    public class DetectionResult
    {
        public float[] Box { get; set; }  // [x1, y1, x2, y2]
        public float Confidence { get; set; }
        public int ClassId { get; set; }
        public string ClassName { get; set; }

        // 用于调试的字符串表示
        public override string ToString()
        {
            return $"{ClassName}: {Confidence:F3} [{Box[0]:F0}, {Box[1]:F0}, {Box[2]:F0}, {Box[3]:F0}]";
        }
    }

    /// <summary>
    /// 模型元数据配置类
    /// </summary>
    public class ModelMetadata
    {
        public string description { get; set; }
        public string author { get; set; }
        public string date { get; set; }
        public string version { get; set; }
        public string license { get; set; }
        public string docs { get; set; }
        public int stride { get; set; }
        public string task { get; set; }
        public int batch { get; set; }
        public List<int> imgsz { get; set; }
        public Dictionary<int, string> names { get; set; }
        public Dictionary<string, object> args { get; set; }
        public int channels { get; set; }
        public bool end2end { get; set; }
    }

    /// <summary>
    /// OpenVINO YOLO 推理引擎
    /// </summary>
    public class YoloOpenVINOInference : IDisposable
    {
        private Core _core;
        private CompiledModel _compiledModel;
        private InferRequest _inferRequest;
        private readonly string[] _classNames;
        private readonly int _inputWidth;
        private readonly int _inputHeight;
        private readonly float _confThreshold;
        private readonly float _iouThreshold;
        private readonly bool _isEnd2End;
        private bool _disposed = false;

        // 存储预处理信息
        private int _padLeft, _padTop;
        private float _scale;

        /// <summary>
        /// 构造函数
        /// </summary>
        /// <param name="modelPath">模型文件路径 (.xml)</param>
        /// <param name="metadataPath">元数据文件路径 (metadata.yaml)</param>
        /// <param name="confThreshold">置信度阈值</param>
        /// <param name="iouThreshold">IOU阈值</param>
        public YoloOpenVINOInference(string modelPath, string metadataPath,
            float confThreshold = 0.25f, float iouThreshold = 0.45f)
        {
            _confThreshold = confThreshold;
            _iouThreshold = iouThreshold;

            // 加载类别名称和模型配置
            var metadata = LoadMetadata(metadataPath);
            _classNames = metadata.names.OrderBy(x => x.Key).Select(x => x.Value).ToArray();
            _isEnd2End = metadata.end2end;

            // 初始化 OpenVINO
            _core = new Core();

            // 读取模型
            var model = _core.read_model(modelPath);

            // 编译模型到指定设备
            _compiledModel = _core.compile_model(model, "CPU");

            // 创建推理请求
            _inferRequest = _compiledModel.create_infer_request();

            // 获取输入尺寸
            var input = _compiledModel.input();
            var inputShape = input.get_shape();
            _inputHeight = (int)inputShape[2];
            _inputWidth = (int)inputShape[3];
        }

        /// <summary>
        /// 加载元数据
        /// </summary>
        private ModelMetadata LoadMetadata(string metadataPath)
        {
            try
            {
                var yaml = File.ReadAllText(metadataPath);
                var deserializer = new DeserializerBuilder()
                    .WithNamingConvention(UnderscoredNamingConvention.Instance)
                    .Build();
                return deserializer.Deserialize<ModelMetadata>(yaml);
            }
            catch (Exception ex)
            {
                System.Diagnostics.Debug.WriteLine($"加载元数据失败: {ex.Message}");
                // 返回默认配置
                return new ModelMetadata
                {
                    names = new Dictionary<int, string>
                    {
                        {0, "NG1"}, {1, "NG2"}, {2, "NG3"}, {3, "NG4"}
                    },
                    end2end = true,
                    imgsz = new List<int> { 640, 640 }
                };
            }
        }

        /// <summary>
        /// Letterbox 预处理（保持宽高比，添加填充）
        /// </summary>
        private (Bitmap resizedImage, int padLeft, int padTop, float scale) LetterboxImage(Bitmap image, int targetWidth, int targetHeight)
        {
            // 计算缩放比例
            float scale = Math.Min((float)targetWidth / image.Width, (float)targetHeight / image.Height);
            int newWidth = (int)(image.Width * scale);
            int newHeight = (int)(image.Height * scale);

            // 调整图像尺寸
            var resizedImage = new Bitmap(newWidth, newHeight);
            using (var g = Graphics.FromImage(resizedImage))
            {
                g.InterpolationMode = InterpolationMode.HighQualityBicubic;
                g.DrawImage(image, 0, 0, newWidth, newHeight);
            }

            // 计算填充位置
            int padLeft = (targetWidth - newWidth) / 2;
            int padTop = (targetHeight - newHeight) / 2;

            // 创建带填充的目标图像
            var result = new Bitmap(targetWidth, targetHeight);
            using (var g = Graphics.FromImage(result))
            {
                // 使用与训练时相同的填充颜色 (114, 114, 114)
                g.Clear(Color.FromArgb(114, 114, 114));
                g.DrawImage(resizedImage, padLeft, padTop, newWidth, newHeight);
            }

            resizedImage.Dispose();
            return (result, padLeft, padTop, scale);
        }

        /// <summary>
        /// 预处理图像
        /// </summary>
        private float[] PreprocessImage(Bitmap image)
        {
            // Letterbox 处理
            var (resizedImage, padLeft, padTop, scale) = LetterboxImage(image, _inputWidth, _inputHeight);

            // 存储填充信息用于后续坐标转换
            _padLeft = padLeft;
            _padTop = padTop;
            _scale = scale;

            var inputData = new float[3 * _inputHeight * _inputWidth];
            var bitmapData = resizedImage.LockBits(
                new Rectangle(0, 0, _inputWidth, _inputHeight),
                ImageLockMode.ReadOnly,
                PixelFormat.Format24bppRgb);

            try
            {
                int stride = bitmapData.Stride;
                IntPtr scan0 = bitmapData.Scan0;
                byte[] pixelData = new byte[stride * _inputHeight];
                Marshal.Copy(scan0, pixelData, 0, pixelData.Length);

                // NCHW 格式，BGR 顺序
                for (int y = 0; y < _inputHeight; y++)
                {
                    for (int x = 0; x < _inputWidth; x++)
                    {
                        int pixelIndex = y * stride + x * 3;
                        float b = pixelData[pixelIndex] / 255.0f;
                        float g = pixelData[pixelIndex + 1] / 255.0f;
                        float r = pixelData[pixelIndex + 2] / 255.0f;

                        int index = (y * _inputWidth + x);
                        inputData[index] = b;                                    // B
                        inputData[_inputHeight * _inputWidth + index] = g;       // G
                        inputData[2 * _inputHeight * _inputWidth + index] = r;   // R
                    }
                }
            }
            finally
            {
                resizedImage.UnlockBits(bitmapData);
                resizedImage.Dispose();
            }

            return inputData;
        }

        /// <summary>
        /// 解析 End2End 模型输出 [1, 300, 6]
        /// 注意：模型输出的是归一化坐标 (0-1之间)
        /// </summary>
        /// <summary>
        /// 解析 End2End 模型输出 [1, 300, 6]
        /// 模型输出格式: [cx, cy, w, h, confidence, classId]
        /// </summary>
        private List<DetectionResult> ParseEnd2EndOutput(float[] outputData, int[] shape, int imageWidth, int imageHeight)
        {
            int numDetections = shape[1];
            int numAttributes = shape[2];  // 应该是 6

            var detections = new List<DetectionResult>();

            for (int i = 0; i < numDetections; i++)
            {
                int offset = i * numAttributes;

                if (offset + 5 >= outputData.Length)
                    break;

                // 模型输出的是预处理后图像（640x640）的坐标
                float x1_pre = outputData[offset];
                float y1_pre = outputData[offset + 1];
                float x2_pre = outputData[offset + 2];
                float y2_pre = outputData[offset + 3];
                float confidence = outputData[offset + 4];
                int classId = (int)outputData[offset + 5];

                // 过滤低置信度检测
                if (confidence < _confThreshold)
                    continue;

                // 关键：将预处理图像的坐标还原到原始图像坐标
                // 使用之前存储的 _padLeft, _padTop, _scale
                float x1 = (x1_pre - _padLeft) / _scale;
                float y1 = (y1_pre - _padTop) / _scale;
                float x2 = (x2_pre - _padLeft) / _scale;
                float y2 = (y2_pre - _padTop) / _scale;

                // 确保坐标在原始图像范围内
                x1 = Math.Max(0, Math.Min(imageWidth, x1));
                y1 = Math.Max(0, Math.Min(imageHeight, y1));
                x2 = Math.Max(0, Math.Min(imageWidth, x2));
                y2 = Math.Max(0, Math.Min(imageHeight, y2));

                // 检查有效性
                if (x2 <= x1 || y2 <= y1)
                    continue;

                detections.Add(new DetectionResult
                {
                    Box = new float[] { x1, y1, x2, y2 },
                    Confidence = confidence,
                    ClassId = classId,
                    ClassName = classId < _classNames.Length ? _classNames[classId] : $"Class{classId}"
                });
            }

            return ApplyNMS(detections);
        }

        /// <summary>
        /// 解析标准 YOLO 输出 [1, 84, 8400] 或类似格式
        /// </summary>
        private List<DetectionResult> ParseStandardOutput(float[] outputData, int[] shape, int imageWidth, int imageHeight)
        {
            int numClasses = _classNames.Length;
            // shape[0] = batch size (1)
            // shape[1] = number of attributes (4 + numClasses)
            // shape[2] = number of predictions (8400)
            int numPredictions = shape[2];  // 8400
            int numAttributes = shape[1];   // 4 + numClasses = 8

            var detections = new List<DetectionResult>();

            for (int i = 0; i < numPredictions; i++)
            {
                // 获取边界框坐标 (cx, cy, w, h) - 归一化坐标
                float cx = outputData[0 * numPredictions + i];
                float cy = outputData[1 * numPredictions + i];
                float w = outputData[2 * numPredictions + i];
                float h = outputData[3 * numPredictions + i];

                // 转换为 x1, y1, x2, y2
                float x1 = cx - w / 2;
                float y1 = cy - h / 2;
                float x2 = cx + w / 2;
                float y2 = cy + h / 2;

                // 获取类别置信度
                float maxConf = 0;
                int bestClass = 0;
                for (int c = 0; c < numClasses; c++)
                {
                    int confIndex = (4 + c) * numPredictions + i;
                    if (confIndex >= outputData.Length)
                        break;

                    float conf = outputData[confIndex];
                    if (conf > maxConf)
                    {
                        maxConf = conf;
                        bestClass = c;
                    }
                }

                if (maxConf < _confThreshold)
                    continue;

                // 坐标转换为像素坐标
                float finalX1 = x1 * imageWidth;
                float finalY1 = y1 * imageHeight;
                float finalX2 = x2 * imageWidth;
                float finalY2 = y2 * imageHeight;

                finalX1 = Math.Max(0, Math.Min(imageWidth, finalX1));
                finalY1 = Math.Max(0, Math.Min(imageHeight, finalY1));
                finalX2 = Math.Max(0, Math.Min(imageWidth, finalX2));
                finalY2 = Math.Max(0, Math.Min(imageHeight, finalY2));

                detections.Add(new DetectionResult
                {
                    Box = new float[] { finalX1, finalY1, finalX2, finalY2 },
                    Confidence = maxConf,
                    ClassId = bestClass,
                    ClassName = _classNames[bestClass]
                });
            }

            return ApplyNMS(detections);
        }

        /// <summary>
        /// 非极大值抑制
        /// </summary>
        private List<DetectionResult> ApplyNMS(List<DetectionResult> detections)
        {
            if (detections.Count == 0)
                return new List<DetectionResult>();

            // 按置信度降序排序
            detections = detections.OrderByDescending(d => d.Confidence).ToList();
            var results = new List<DetectionResult>();

            while (detections.Count > 0)
            {
                var best = detections[0];
                results.Add(best);
                detections.RemoveAt(0);

                for (int i = detections.Count - 1; i >= 0; i--)
                {
                    float iou = CalculateIOU(best.Box, detections[i].Box);
                    if (iou > _iouThreshold)
                    {
                        detections.RemoveAt(i);
                    }
                }
            }

            return results;
        }

        /// <summary>
        /// 计算两个边界框的 IOU
        /// </summary>
        private float CalculateIOU(float[] box1, float[] box2)
        {
            float x1 = Math.Max(box1[0], box2[0]);
            float y1 = Math.Max(box1[1], box2[1]);
            float x2 = Math.Min(box1[2], box2[2]);
            float y2 = Math.Min(box1[3], box2[3]);

            float intersection = Math.Max(0, x2 - x1) * Math.Max(0, y2 - y1);
            float area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]);
            float area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]);
            float union = area1 + area2 - intersection;

            return union > 0 ? intersection / union : 0;
        }

        /// <summary>
        /// 执行推理
        /// </summary>
        /// <param name="image">输入图像</param>
        /// <returns>检测结果列表</returns>
        public List<DetectionResult> Infer(Bitmap image)
        {
            if (image == null)
                throw new ArgumentNullException(nameof(image));

            // 预处理
            var inputData = PreprocessImage(image);

            // 获取输入张量并设置数据
            var inputTensor = _inferRequest.get_input_tensor();
            inputTensor.set_data(inputData);

            // 执行推理
            _inferRequest.infer();

            // 获取输出
            var outputTensor = _inferRequest.get_output_tensor();
            var outputShape = outputTensor.get_shape();
            var outputSize = outputTensor.get_size();

            // 获取输出数据
            var outputData = outputTensor.get_data<float>((int)outputSize);

            // 获取形状信息
            int[] shapeArray = new int[outputShape.Count];
            for (int i = 0; i < outputShape.Count; i++)
            {
                shapeArray[i] = (int)outputShape[i];
            }

            // 解析输出
            List<DetectionResult> detections;

            // 判断是否为 End2End 模型输出
            // End2End 模型输出形状为 [1, 300, 6]
            if (_isEnd2End || (shapeArray.Length == 3 && shapeArray[2] == 6))
            {
                detections = ParseEnd2EndOutput(outputData, shapeArray, image.Width, image.Height);
            }
            else
            {
                detections = ParseStandardOutput(outputData, shapeArray, image.Width, image.Height);
            }

            return detections;
        }

        /// <summary>
        /// 从文件路径执行推理
        /// </summary>
        public List<DetectionResult> InferFromFile(string imagePath)
        {
            using (var image = new Bitmap(imagePath))
            {
                return Infer(image);
            }
        }

        /// <summary>
        /// 释放资源
        /// </summary>
        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!_disposed)
            {
                if (disposing)
                {
                    _inferRequest?.Dispose();
                    _compiledModel?.Dispose();
                    _core?.Dispose();
                }
                _disposed = true;
            }
        }
    }
}