YOLOV5 6.0版本 ML.NET cpu onnx 推理的使用 c#

yolov5 6.0 版本引入了nano版本,CPU 的onnx推理快了很多,i5 4570 在68ms左右。github下载的代码,套用自己生成的onnx。

训练自己的数据集

python train.py --img 640 --batch 16 --epoch 1500 --data  data/myvoc.yaml --cfg  models/yolov5n.yaml --weights yolov5n.pt --workers 0

然后导出

python export.py --weights runs\train\exp10\weights\best.pt --include onnx --opset 12 --dynamic

训练代码下载:

yolov5 6.0版本yolov5-nano训练自己的数据集-互联网文档类资源-CSDN下载yolov56.0版本yolov5-nano训练自己的数据集,已经包含自己的数据集,直接训练更多下载资源、学习资料请访问CSDN下载频道.https://download.csdn.net/download/vokxchh/46266366

 注意 output=分类数+5

using Microsoft.ML;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.IO;
using static Microsoft.ML.Transforms.Image.ImageResizingEstimator;

namespace OnnxYoloV5
{
    class Program
        {

        const string modelPath = @"Assets\Models\yolov5n.onnx";

        const string imageFolder = @"Assets\Images";

        const string imageOutputFolder = @"Assets\Output";
        static readonly string[] classesNames = new string[] { "ok", "ng" };

        static void Main(string[] args)
        {
            


            Directory.CreateDirectory(imageOutputFolder);
            MLContext mlContext = new MLContext();
            var pipeline = mlContext.Transforms.ResizeImages(inputColumnName: "bitmap", outputColumnName: "images", imageWidth: 640, imageHeight: 640, resizing: ResizingKind.Fill)
                .Append(mlContext.Transforms.ExtractPixels(outputColumnName: "images", scaleImage: 1f / 255f, interleavePixelColors: false))
                .Append(mlContext.Transforms.ApplyOnnxModel(
                    shapeDictionary: new Dictionary()
                    {
                        { "images", new[] { 1, 3, 640, 640 } },
                        { "output", new[] { 1, 25200, 7 } },
                    },
                    inputColumnNames: new[]
                    {
                        "images"
                    },
                    outputColumnNames: new[]
                    {
                        "output"
                    },
                    modelFile: modelPath, recursionLimit : 100));
         

            // Fit on empty list to obtain input data schema
            var model = pipeline.Fit(mlContext.Data.LoadFromEnumerable(new List()));

            // Create prediction engine
            var predictionEngine = mlContext.Model.CreatePredictionEngine(model);
          

            // save model
            //mlContext.Model.Save(model, predictionEngine.OutputSchema, Path.ChangeExtension(modelPath, "zip"));

            foreach (string imageName in Directory.GetFiles(imageFolder))
            {
                using (var bitmap = new Bitmap(imageName))
                {
                    // predict
                    Stopwatch _stopWath2 = new Stopwatch();
                    _stopWath2.Start();
                    var predict = predictionEngine.Predict(new YoloV5BitmapData() { Image = bitmap });
                    var results = predict.GetResults(classesNames, 0.3f, 0.7f);
                
                    using (var g = Graphics.FromImage(bitmap))
                    {
                        foreach (var res in results)
                        {
                            // draw predictions
                            var x1 = res.BBox[0];
                            var y1 = res.BBox[1];
                            var x2 = res.BBox[2];
                            var y2 = res.BBox[3];
                            g.DrawRectangle(Pens.Red, x1, y1, x2 - x1, y2 - y1);
                            using (var brushes = new SolidBrush(Color.FromArgb(50, Color.Red)))
                            {
                                g.FillRectangle(brushes, x1, y1, x2 - x1, y2 - y1);
                            }

                            g.DrawString(res.Label + " " + res.Confidence.ToString("0.00"),
                                         new Font("Arial", 12), Brushes.Blue, new PointF(x1, y1));
                        }
                        var ss =Path.Combine(imageOutputFolder, Path.GetFileNameWithoutExtension(imageName)+"_Processed"+Path.GetExtension(imageName));
                        bitmap.Save(ss);
                        _stopWath2.Stop();

                        Console.WriteLine(_stopWath2.ElapsedMilliseconds.ToString()+"ms");
                        Console.WriteLine(Path.GetFileNameWithoutExtension(imageName) +" Processed ");
                    }
                }
            }
            Console.ReadKey();
           
        }
    }
}
using Microsoft.ML.Data;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace OnnxYoloV5
{
     class YoloV5Prediction
    {/// 
     /// Identity
     /// 
        [VectorType(1, 25200, 7)]
        [ColumnName("output")]
        public float[] Output { get; set; }

        [ColumnName("width")]
        public float ImageWidth { get; set; }

        [ColumnName("height")]
        public float ImageHeight { get; set; }

        public IReadOnlyList GetResults(string[] categories, float scoreThres = 0.5f, float iouThres = 0.5f)
        {

            // Probabilities + Characteristics
            int characteristics = categories.Length + 5;

            // Needed info
            float modelWidth = 640.0F;
            float modelHeight = 640.0F;
            float xGain = modelWidth / ImageWidth;
            float yGain = modelHeight / ImageHeight;
            float[] results = Output;

            List postProcessedResults = new List();

            // For every cell of the image, format for NMS
            for (int i = 0; i < 25200; i++)
            {
                // Get offset in float array
                int offset = characteristics * i;

                // Get a prediction cell
                var predCell = results.Skip(offset).Take(characteristics).ToList();

                // Filter some boxes
                var objConf = predCell[4];
                if (objConf <= scoreThres) continue;

                // Get corners in original shape
                var x1 = (predCell[0] - predCell[2] / 2) / xGain; //top left x
                var y1 = (predCell[1] - predCell[3] / 2) / yGain; //top left y
                var x2 = (predCell[0] + predCell[2] / 2) / xGain; //bottom right x
                var y2 = (predCell[1] + predCell[3] / 2) / yGain; //bottom right y

                // Get real class scores
                var classProbs = predCell.Skip(5).Take(categories.Length).ToList();
                var scores = classProbs.Select(p => p * objConf).ToList();

                // Get best class and index
                float maxConf = scores.Max();
                float maxClass = scores.ToList().IndexOf(maxConf);

                postProcessedResults.Add(new[] { x1, y1, x2, y2, maxConf, maxClass });
            }

            var resultsNMS = ApplyNMS(postProcessedResults, categories, iouThres);

            return resultsNMS;
        }

        private List ApplyNMS(List postProcessedResults, string[] categories, float iouThres = 0.5f)
        {
            postProcessedResults = postProcessedResults.OrderByDescending(x => x[4]).ToList(); // sort by confidence
            List resultsNms = new List();

            int f = 0;
            while (f < postProcessedResults.Count)
            {
                var res = postProcessedResults[f];
                if (res == null)
                {
                    f++;
                    continue;
                }

                var conf = res[4];
                string label = categories[(int)res[5]];

                resultsNms.Add(new YoloV5Result(res.Take(4).ToArray(), label, conf));
                postProcessedResults[f] = null;

                var iou = postProcessedResults.Select(bbox => bbox == null ? float.NaN : BoxIoU(res, bbox)).ToList();
                for (int i = 0; i < iou.Count; i++)
                {
                    if (float.IsNaN(iou[i])) continue;
                    if (iou[i] > iouThres)
                    {
                        postProcessedResults[i] = null;
                    }
                }
                f++;
            }

            return resultsNms;
        }

        /// 
        /// Return intersection-over-union (Jaccard index) of boxes.
        /// Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
        /// 
        private static float BoxIoU(float[] boxes1, float[] boxes2)
        {
            static float box_area(float[] box)
            {
                return (box[2] - box[0]) * (box[3] - box[1]);
            }

            var area1 = box_area(boxes1);
            var area2 = box_area(boxes2);

            Debug.Assert(area1 >= 0);
            Debug.Assert(area2 >= 0);

            var dx = Math.Max(0, Math.Min(boxes1[2], boxes2[2]) - Math.Max(boxes1[0], boxes2[0]));
            var dy = Math.Max(0, Math.Min(boxes1[3], boxes2[3]) - Math.Max(boxes1[1], boxes2[1]));
            var inter = dx * dy;

            return inter / (area1 + area2 - inter);
        }
    }
}

工程所有代码下载:

YOLOv5-ML.NET-yolov5-nano 6.0版本-互联网文档类资源-CSDN文库https://download.csdn.net/download/vokxchh/46273789

你可能感兴趣的:(c#,图像识别)