YOLO-NAS C++部署 2023.5.17

这不最近新出了网络，YOLO-NAS，听过性能和速度都不错，而且int8量化后效果也不错。

一、吐槽

但是我打开该项目阅读readme.txt时候，发现这些示例网站一个都打不开！

而且readme.txt很不详细，你想训练自己的模型都无从下手。

二、ONXX格式导出

先安装super-gradients

pip install super-gradients  

然后写个python文件，运行

from super_gradients.training import models
from super_gradients.common.object_names import Models

# Load model with pretrained weights
net = models.get(Models.YOLO_NAS_S, pretrained_weights='coco')
# net.predict("/home/jason/work/01-img/dog.png").show()



models.convert_to_onnx(model=net, input_shape=(3, 640, 640), out_path='yolo-nas-s.onnx')

用netron 查看下输出：

需要注意以下两点，通过官方工具导出的"bboxes"已经是"xyxy"格式了，所以不需要再执行make_grid, distance2bbox等操作了，直接进行"NMS"即可。但是，OpenCV的NMS要求输入的BBOXES格式为xywh，所以还需要再改一下

可以看到有两个输出

1135：float32[1,8400,4]。8400 是框子的数量，4是位置信息

1126：float32[1,8400,80]。 80是每个类别的概率。

三、onnx runtime C++部署

最大类别概率是178,索引是78,人就在摄像头面前。。。。位置信息也不对。。。

 这是关于后处理的那个函数，有大佬试过这个网络的部署吗，望指点一二！

bool Yolov8Onnx::OnnxBatchDetect(std::vector<cv::Mat> &srcImgs, std::vector<std::vector<OutputDet> > &output)
{
    vector<Vec4d> params;
    vector<Mat> input_images;
    cv::Size input_size(_netWidth, _netHeight);

    //preprocessing (信封处理)
    Preprocessing(srcImgs, input_images, params);
    // [0~255] --> [0~1]; BGR2RGB
    Mat blob = cv::dnn::blobFromImages(input_images, 1 / 255.0, input_size, Scalar(0,0,0), true, false);

    // 前向传播得到推理结果
    int64_t input_tensor_length = VectorProduct(_inputTensorShape);// ?
    std::vector<Ort::Value> input_tensors;
    std::vector<Ort::Value> output_tensors;
    input_tensors.push_back(Ort::Value::CreateTensor<float>(_OrtMemoryInfo, (float*)blob.data,
                                                            input_tensor_length, _inputTensorShape.data(),
                                                            _inputTensorShape.size()));

    output_tensors = _OrtSession->Run(Ort::RunOptions{ nullptr },
        _inputNodeNames.data(),
        input_tensors.data(),
        _inputNodeNames.size(),
        _outputNodeNames.data(),
        _outputNodeNames.size()
    );

    //post-process
//    int net_width = _className.size() + 4; // 不用加4
    int net_width = _className.size();
    float* all_data0 = output_tensors[0].GetTensorMutableData<float>(); // 位置信息
    float* all_data1 = output_tensors[1].GetTensorMutableData<float>(); // 80类别概率

    _output0TensorShape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape(); //  [1, 8400, 4]
    _output1TensorShape = output_tensors[1].GetTensorTypeAndShapeInfo().GetShape(); // [1,8400, 80]


    int64_t one_output1_length = VectorProduct(_output1TensorShape) / _output1TensorShape[0]; // 一张图片输出所占内存长度 8400*80
    int64_t one_output0_length = VectorProduct(_output0TensorShape) / _output0TensorShape[0]; // 一张图片输出所占内存长度 8400*4

    int net_height = _output0TensorShape[1]; // 8400

    for (int img_index = 0; img_index < srcImgs.size(); ++img_index){

        // 一张图片的预测框
        vector<int> class_ids;
        vector<float> confidences;
        vector<Rect> boxes;
        for (int r=0; r<net_height; ++r){//如果是yolov5则需要做修改
            Mat scores(1, _className.size(), CV_32F, all_data1); // 80个类别的概率

            // 得到最大类别概率、类别索引
            Point classIdPoint;
            double max_class_soces;
            minMaxLoc(scores, 0, &max_class_soces, 0, &classIdPoint);
            max_class_soces = (float)max_class_soces;
            printf("max_clasee:%d id:%d\n",(int)max_class_soces,(int)classIdPoint.x);

            // 预测框坐标映射到原图上
            if (max_class_soces >= _classThreshold){
                // rect [x,y,w,h]

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



                x = (x - params[img_index][2]) / params[img_index][0]; //x
                y = (y - params[img_index][3]) / params[img_index][1]; //y
                w = w / params[img_index][0]; //w
                h = h / params[img_index][1]; //h

                int left = MAX(int(x - 0.5 *w +0.5), 0);
                int top = MAX(int(y - 0.5*h + 0.5), 0);

                class_ids.push_back(classIdPoint.x);
                confidences.push_back(max_class_soces);
                boxes.push_back(Rect(left, top, int(w + 0.5), int(h + 0.5)));
            }

        }
        all_data0 += one_output0_length; //指针指向下一个图片的地址
        all_data1 += one_output1_length;

        // 对一张图的预测框执行NMS处理
        vector<int> nms_result;
        cv::dnn::NMSBoxes(boxes, confidences, _classThreshold, _nmsThrehold, nms_result); // 还需要classThreshold？

        // 对一张图片：依据NMS处理得到的索引，得到类别id、confidence、box，并置于结构体OutputDet的容器中
        vector<OutputDet> temp_output;
        for (size_t i=0; i<nms_result.size(); ++i){
            int idx = nms_result[i];
            OutputDet result;
            result.id = class_ids[idx];
            result.confidence = confidences[idx];
            result.box = boxes[idx];
            temp_output.push_back(result);
        }
        output.push_back(temp_output); // 多张图片的输出；添加一张图片的输出置于此容器中

    }
    if (output.size())
        return true;
    else
        return false;

}

参考：

YOLO-NAS_whaosoft143的博客-CSDN博客