如何在 ML.Net 中实现 yolo v3 或 v4 onnx 模型的后处理

Question

我按照微软的教程做的，没有问题。但我想将模型更改为 yolo v3 或 v4。我从onnx/models获取 YOLOv4 onnx 模型，并且能够获取 yolov4 onnx 模型的所有三个浮点输出数组，但问题在于后处理，我无法从这些输出中获得正确的边界框。

我更改了微软教程 src 代码中的所有内容，例如锚点、步幅、输出网格大小、一些功能和...以与 yolov4 兼容。但我无法得到正确的结果。我用python 实现检查了所有代码，但我不知道问题出在哪里。有谁有链接或知道如何使用 ML.Net 在 c# 中实现 yolo v3 或 v4 onnx 模型

任何帮助将不胜感激

Answer 1

我认为不可能直接将微软的教程从 YOLO v2 移植到 v3，因为它依赖于每个模型的输入和输出。

附带说明一下，我在此 GitHub 存储库中将另一个 YOLO v3 模型移植到了 ML.Net：“YOLOv3MLNet ”。它包含一个功能齐全的 ML.Net 管道。

我还在这里提供了这个答案的代码：

• YOLO v3 与 ML.Net
• YOLO v4 与 ML.Net

回到模型，我将以 YOLO v3（可在 onnx/models 存储库中找到）为例。可以在此处找到对该模型的详细解释。

第一个建议是使用Netron查看模型。这样做，您将看到输入层和输出层。他们还在 onnx/models 文档中描述了这些层。

Netron的yolov3-10截图

（我在 Netron 中看到这个特定的 YOLO v3 模型还通过非极大值抑制步骤进行了一些后处​​理。）

• 输入层名称：input_1,image_shape
• 输出层名称：yolonms_layer_1/ExpandDims_1:0, yolonms_layer_1/ExpandDims_3:0,yolonms_layer_1/concat_2:0

根据模型文档，输入形状为：

调整大小的图像 (1x3x416x416) 原始图像大小 (1x2)，即 [image.size['1], image.size[0]]

我们首先需要定义 ML.Net 输入和输出类，如下所示：

public class YoloV3BitmapData
{
    [ColumnName("bitmap")]
    [ImageType(416, 416)]
    public Bitmap Image { get; set; }

    [ColumnName("width")]
    public float ImageWidth => Image.Width;

    [ColumnName("height")]
    public float ImageHeight => Image.Height;
}

public class YoloV3Prediction
{
    /// <summary>
    /// ((52 x 52) + (26 x 26) + 13 x 13)) x 3 = 10,647.
    /// </summary>
    public const int YoloV3BboxPredictionCount = 10_647;

    /// <summary>
    /// Boxes
    /// </summary>
    [ColumnName("yolonms_layer_1/ExpandDims_1:0")]
    public float[] Boxes { get; set; }

    /// <summary>
    /// Scores
    /// </summary>
    [ColumnName("yolonms_layer_1/ExpandDims_3:0")]
    public float[] Scores { get; set; }

    /// <summary>
    /// Concat
    /// </summary>
    [ColumnName("yolonms_layer_1/concat_2:0")]
    public int[] Concat { get; set; }
}

然后，我们创建 ML.Net 管道并加载预测引擎：

// Define scoring pipeline
var pipeline = mlContext.Transforms.ResizeImages(inputColumnName: "bitmap", outputColumnName: "input_1", imageWidth: 416, imageHeight: 416, resizing: ResizingKind.IsoPad)
    .Append(mlContext.Transforms.ExtractPixels(outputColumnName: "input_1", outputAsFloatArray: true, scaleImage: 1f / 255f))
    .Append(mlContext.Transforms.Concatenate("image_shape", "height", "width"))
    .Append(mlContext.Transforms.ApplyOnnxModel(shapeDictionary: new Dictionary<string, int[]>() { { "input_1", new[] { 1, 3, 416, 416 } } },
                    inputColumnNames: new[]
                    {
                        "input_1",
                        "image_shape"
                    },
                    outputColumnNames: new[]
                    {
                        "yolonms_layer_1/ExpandDims_1:0",
                        "yolonms_layer_1/ExpandDims_3:0",
                        "yolonms_layer_1/concat_2:0"
                    },
                    modelFile: @"D:\yolov3-10.onnx"));

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

// Create prediction engine
var predictionEngine = mlContext.Model.CreatePredictionEngine<YoloV3BitmapData, YoloV3Prediction>(model);

注意：我们需要定义shapeDictionary参数，因为它们在模型中没有完整定义。

根据模型文档，输出形状为：

该模型有 3 个输出。框：（1x'n_candidates'x4），所有锚框的坐标，分数：（1x80x'n_candidates'），每个类的所有锚框的分数，索引：（'nbox'x3），从框张量中选择的索引。选择的索引格式为(batch_index, class_index, box_index)。

下面的函数将帮助您处理结果，我将其留给您进行微调。

public IReadOnlyList<YoloV3Result> GetResults(YoloV3Prediction prediction, string[] categories)
{
    if (prediction.Concat == null || prediction.Concat.Length == 0)
    {
        return new List<YoloV3Result>();
    }

    if (prediction.Boxes.Length != YoloV3Prediction.YoloV3BboxPredictionCount * 4)
    {
        throw new ArgumentException();
    }

    if (prediction.Scores.Length != YoloV3Prediction.YoloV3BboxPredictionCount * categories.Length)
    {
        throw new ArgumentException();
    }

    List<YoloV3Result> results = new List<YoloV3Result>();

    // Concat size is 'nbox'x3 (batch_index, class_index, box_index)
    int resulstCount = prediction.Concat.Length / 3;
    for (int c = 0; c < resulstCount; c++)
    {
        var res = prediction.Concat.Skip(c * 3).Take(3).ToArray();

        var batch_index = res[0];
        var class_index = res[1];
        var box_index = res[2];

        var label = categories[class_index];
        var bbox = new float[]
        {
            prediction.Boxes[box_index * 4],
            prediction.Boxes[box_index * 4 + 1],
            prediction.Boxes[box_index * 4 + 2],
            prediction.Boxes[box_index * 4 + 3],
        };
        var score = prediction.Scores[box_index + class_index * YoloV3Prediction.YoloV3BboxPredictionCount];

        results.Add(new YoloV3Result(bbox, label, score));
    }

    return results;
}

在此版本的模型中，它们有 80 个类（有关链接，请参阅模型的 GitHub 文档）。

您可以像这样使用上面的内容：

// load image
string imageName = "dog_cat.jpg";
using (var bitmap = new Bitmap(Image.FromFile(Path.Combine(imageFolder, imageName))))
{
    // predict
    var predict = predictionEngine.Predict(new YoloV3BitmapData() { Image = bitmap });

    var results = GetResults(predict, classesNames);

    // draw predictions
    using (var g = Graphics.FromImage(bitmap))
    {
        foreach (var result in results)
        {
            var y1 = result.BBox[0];
            var x1 = result.BBox[1];
            var y2 = result.BBox[2];
            var x2 = result.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(result.Label + " " + result.Confidence.ToString("0.00"),
                         new Font("Arial", 12), Brushes.Blue, new PointF(x1, y1));
        }

        bitmap.Save(Path.Combine(imageOutputFolder, Path.ChangeExtension(imageName, "_processed" + Path.GetExtension(imageName))));
    }
}

您可以在此处找到结果示例。