# 图像分割 本示例主要通过DeepLabv3模型说明如何使用MIGraphX C++ API进行图像分割模型的推理,包括模型初始化、预处理、模型推理。 ## 模型简介 本示例采用了经典的DeepLabv3模型进行图像分割, 模型deeplabv3_resnet101.onnx文件保存在Resource/Models文件夹下。模型结构如下图所示,可以通过netron工具, 链接:https://netron.app/, 查看具体的模型结构,该模型的输入shape为[batch_size,3,513,513],输出shape为[batch_size,513,513],数据排布为NCHW。 ## 模型初始化 在模型初始化的过程中,首先采用parse_onnx()函数根据提供的模型地址加载图像分割deeplabv3的onnx模型,保存在net中。其次,通过net.get_parameter_shapes()获取deeplabv3模型的输入属性,包含inputName和inputShape。最后,完成模型加载后使用migraphx::gpu::target{}设置编译模式为GPU模式,并使用compile()函数编译模型,完成模型的初始化过程。 其中,模型地址设置在/Resource/Configuration.xml文件中的Unet节点中。 ```C++ ErrorCode Unet::Initialize(InitializationParameterOfSegmentation initParamOfSegmentationUnet) { ... // 加载模型 net = migraphx::parse_onnx(modelPath); LOG_INFO(logFile,"succeed to load model: %s\n",GetFileName(modelPath).c_str()); // 获取模型输入/输出节点信息 std::cout<<"inputs:"< inputs=net.get_inputs(); for(auto i:inputs) { std::cout< outputs=net.get_outputs(); for(auto i:outputs) { std::cout< results = net.eval(inputData); // 如果想要指定输出节点,可以给eval()函数中提供outputNames参数来实现 //std::vector outputNames = {"outputs"}; //std::vector inferenceResults = net.eval(inputData, outputNames); // 获取输出节点的属性 migraphx::argument result = results[0]; // 获取第一个输出节点的数据 migraphx::shape outputShape=result.get_shape(); // 输出节点的shape std::vector outputSize=outputShape.lens(); // 每一维大小,维度顺序为(N,C,H,W) int numberOfOutput=outputShape.elements(); // 输出节点元素的个数 float *data = (float *)result.data(); // 输出节点数据指针 ... } ``` 模型得到的推理结果并不能直接作为图像的分割结果,还需要做如下处理: 1.计算softmax值,计算不同通道同一[H,W]位置的softmax值,找出概率最高的通道。 2.保存结果,创建一个cv::Mat,根据不同的通道索引在颜色映射表取值并按行依次赋值到Mat对应位置,得到最终的分割图像。 ```c++ ErrorCode Unet::Segmentation(const cv::Mat &srcImage, cv::Mat &maskImage) { ... cv::Mat outputImage(cv::Size(W, H), CV_8UC3); // 创建颜色映射表 std::vector color_map = create_color_map(); for(int i = 0;i < H; i++){ for(int j = 0;j < W;j++){ std::vector channel_value; for(int k = 0;k < C;k++){ channel_value.push_back(data[k*(H*W)+i*W+j]); } std::vector probs = softmax(channel_value); // 找到概率最高的类别索引 int max_index = std::max_element(probs.begin(),probs.end())-probs.begin(); cv::Scalar sc = color_map[max_index]; outputImage.at(i, j)[0]= sc.val[0]; outputImage.at(i, j)[1]= sc.val[1]; outputImage.at(i, j)[2]= sc.val[2]; } } maskImage = outputImage.clone(); ... } ``` 注:本次采用的模型权重onnx文件是通过使用PASCAL VOC 2012数据集来训练的。因此,“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果,建议对现实世界示例数据集上的模型进行微调。