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