pytorch 转换为ONNX后的YOLOv8输出格式

wlwcrazw  于 2023-11-19  发布在  其他
关注(0)|答案(1)|浏览(96)

我转换YOLOv8检测(特别是best.pt检查点)模型到onnx formate,但我不知道如何获得边界框和信心。我也不知道模型是否正确转换。
我期待的形状是这样的|1、25200、7|最后一个数字是x,y,w,h,confidence,class0,class1。|1、5、5040|我不知道它是否正确,只是转置和7而不是5,因为我只有一个类,或者如果它是完全错误的,我搞砸了,而转换我尝试了2种类型的转换

import torch
    import torchvision
    import torch.onnx
    from ultralytics import YOLO
    from torchsummary import summary
    
    m = YOLO("best.pt")
    m.model.eval()
    
    dummy_input = torch.randn(1, 3, 384, 640)
    
    #first way of convertion
    m.export(format="onnx",imgsz=[384,640], opset=12)
    
    #second way of coversion
    torch.onnx.export(m.model, dummy_input, "yolo5.onnx", verbose=True, opset_version=12)

字符串
这两种情况下我都得到了输出形状:|1、5、5040|这是型号modelsummary的netron图像这是正确的吗,只是移调它感觉不对什么是错的?

nvbavucw

nvbavucw1#

您的描述遗漏了问题。我假设您想了解Yolov 8 ONNX转换模型的输出格式以及如何在代码中使用它。我将选择C++实现作为答案,因为您也没有提到您正在使用的编程语言。
有两种替代方案:
1.让我们看一下https://github.com/mallumoSK/yolov8/blob/master/yolo/YoloDet.cppdetect函数实现了将原始Yolov 8输出Tensor转换为检测对象的逻辑:

std::vector<YoloDet::Detection> YoloDet::detect(cv::Mat &mat) {
    static cv::Mat blob;
    static std::vector<cv::Mat> outputs;

    cv::dnn::blobFromImage(mat, blob, 1.0 / 255.0, modelShape, cv::Scalar(), true, false);
    net.setInput(blob);
    net.forward(outputs, net.getUnconnectedOutLayersNames());

    int rows = outputs[0].size[2];
    int dimensions = outputs[0].size[1];

    outputs[0] = outputs[0].reshape(1, dimensions);
    cv::transpose(outputs[0], outputs[0]);

    auto data = (float *) outputs[0].data;

    std::vector<int> class_ids{};
    std::vector<float> confidences{};
    std::vector<cv::Rect> boxes{};

    for (int i = 0; i < rows; ++i) {
        float *classes_scores = data + 4;

        cv::Mat scores(1, classesCount, CV_32FC1, classes_scores);
        cv::Point class_id;
        double maxClassScore;

        minMaxLoc(scores, nullptr, &maxClassScore, nullptr, &class_id);

        if (maxClassScore > modelScoreThreshold) {
            confidences.push_back(float(maxClassScore));
            class_ids.push_back(class_id.x);

            float x = data[0];
            float y = data[1];
            float w = data[2];
            float h = data[3];

            int left = int((x - 0.5 * w) * x_factor);
            int top = int((y - 0.5 * h) * y_factor);

            int width = int(w * x_factor);
            int height = int(h * y_factor);

            boxes.emplace_back(left, top, width, height);
        }
        data += dimensions;
    }

    std::vector<int> nms_result;
    cv::dnn::NMSBoxes(boxes, confidences, modelScoreThreshold, modelNMSThreshold, nms_result);

    std::vector<Detection> detections{};

    for (int idx: nms_result) {
        Detection result;
        result.class_id = class_ids[idx];
        result.confidence = confidences[idx];
        result.box = boxes[idx];
        detections.push_back(result);
    }

    return detections;
}

字符串
它在引擎盖下使用cv::dnn::NMSBoxes。请注意,OpenCV DNN模块在默认情况下并不总是可用的。
1.作为上述自定义脚本的替代方案,ONNXRuntime团队已经准备了一个使用NMS的优化Yolov 8检测模型(非最大抑制)* 直接内置于模型中 *(通过ONNXRuntime-Extensions包)。您可以找到模型here和对应的Android example。使用这种方法,您甚至不需要深入到兔子洞里去尝试理解Yolov 8的输出格式,因为该模型输出具有来自输入图像的分数的边界框。
我还没有测试哪一个更快,但我认为ONNXRuntime +内置的NMS应该会产生更好的性能。如果有人做基准测试,请告诉我。
请注意,与前几代产品相比,Yolov 8已经更改了输出格式(请参见以下行):

// yolov5 has an output of shape (batchSize, 25200, 85) (Num classes + box[x,y,w,h] + confidence[c])
    // yolov8 has an output of shape (batchSize, 84,  8400) (Num classes + box[x,y,w,h])

相关问题