OpenVINO + UNet模型部署，实现道路裂纹检测

点击上方蓝字关注我们

微信公众号：OpenCV学堂 关注获取更多计算机视觉与深度学习知识

UNet模型

这个模型来自我之前的文章，基于CrackForest数据集训练生成的模型，如何训练道理裂纹数据集，生成UNet模型并导出为ONNX看这里，这个系列文章就可以知道：

轻松学Pytorch – 构建UNet实现道路裂纹检测

代码语言：javascript复制

https://mp.weixin.qq.com/s/xeUdW2l71RsHe1Zdzr5a7Q

然后我把模型转换ONNX格式了，然后我用OpenVINO ONNX做个部署演示。之前一直有例子是Python版本，谁知道Pytorch视频课程群里有位小伙伴说想要C 版本的，怎么说，安排吧！

代码语言：javascript复制

模型输入图像：1x1x320x480输出格式：1x2x320x480

输入图像预处理是要转换为float，0~1之间的数据。

C 版本关键注意点

首先输入格式设置FP32，然后在设置输入数据的时候获取的是float类型指针，不是之前的uchar类型指针了。注意输入的是灰度图像，单通道的，这部分不出错，基本对一半了，相关的代码如下：

代码语言：javascript复制

// 创建IE插件, 查询支持硬件设备
Core ie;
vector<string> availableDevices = ie.GetAvailableDevices();
for (int i = 0; i < availableDevices.size(); i  ) {
    printf("supported device name : %s n", availableDevices[i].c_str());
}

//  加载检测模型
auto network = ie.ReadNetwork(model_onnx);

// 请求网络输入与输出信息
InferenceEngine::InputsDataMap input_info(network.getInputsInfo());
InferenceEngine::OutputsDataMap output_info(network.getOutputsInfo());
// 设置输入格式
for (auto &item : input_info) {
    auto input_data = item.second;
    input_data->setPrecision(Precision::FP32);
    input_data->setLayout(Layout::NCHW);
}
printf("get it n");

// 设置输出格式
std::string output_name = "";
for (auto &item : output_info) {
    output_name = item.first;
    auto output_data = item.second;
    output_data->setPrecision(Precision::FP32);
}

然后模型加载跟推理是跟其它模型完全相似，这里就不在说了，直接贴代码吧

代码语言：javascript复制

// 创建可执行网络对象
auto executable_network = ie.LoadNetwork(network, "CPU");

// 请求推断图
auto infer_request = executable_network.CreateInferRequest();

/** Iterating over all input blobs **/
for (auto & item : input_info) {
    auto input_name = item.first;

    /** Getting input blob **/
    auto input = infer_request.GetBlob(input_name);
    size_t num_channels = input->getTensorDesc().getDims()[1];
    size_t h = input->getTensorDesc().getDims()[2];
    size_t w = input->getTensorDesc().getDims()[3];
    Mat blob_image;
    resize(src, blob_image, Size(w, h));
    blob_image.convertTo(blob_image, CV_32F);
    blob_image = blob_image / 255.0;

    // NCHW
    float* data = static_cast<float*>(input->buffer());
    for (size_t row = 0; row < h; row  ) {
        for (size_t col = 0; col < w; col  ) {
            data[row * w   col] = blob_image.at<float>(row, col);
        }
    }
}

// 执行预测
infer_request.Infer();

然后就是解析部分，C 中没python那么好用的函数，怎么办/? 直接指针索引访问就得了，后处理部分得代码如下：

代码语言：javascript复制

// 获取输出数据
auto output = infer_request.GetBlob(output_name);
const float* detection = static_cast<PrecisionTrait<Precision::FP32>::value_type*>(output->buffer());
const SizeVector outputDims = output->getTensorDesc().getDims();
const int out_c = outputDims[1];
const int out_h = outputDims[2];
const int out_w = outputDims[3];
cv::Mat result = cv::Mat::zeros(cv::Size(out_w, out_h), CV_32FC1);
// 解析输出结果
for (int row = 0; row < out_h; row  ) {
    for (int col = 0; col < out_w; col  ) {
        float c1 = detection[row * out_w  col];
        float c2 = detection[out_h*out_w   row * out_w   col];
        if (c1 > c2) {
            result.at<float>(row, col) = 0;
        }
        else {
            result.at<float>(row, col) = 1;
        }
    }
}
result = result * 255;
result.convertTo(result, CV_8U);
imshow("input", src);
std::vector<std::vector<cv::Point>> contours;
std::vector<cv::Vec4i> hierarchy;
cv::findContours(result, contours, hierarchy, cv::RETR_TREE, cv::CHAIN_APPROX_SIMPLE, cv::Point());
cv::cvtColor(src, src, cv::COLOR_GRAY2BGR);
cv::drawContours(src, contours, -1, cv::Scalar(0, 0, 255), -1, 8);
imshow("OpenVINO UNet", src);

最终运行结果如下：