# 图像分割 本示例主要通过Unet模型说明如何使用MIGraphX C++ API进行图像分割模型的推理,包括模型初始化、预处理、模型推理。 ## 模型简介 本示例采用了经典的Unet模型进行图像分割,模型下载地址:https://www.dropbox.com/s/3ntkhyk30x05uuv/unet_13_256.onnx, 将unet_13_256.onnx文件保存在Resource/Models文件夹下。模型结构如下图所示,可以通过netron工具, 链接:https://netron.app/, 查看具体的模型结构,该模型的输入shape为[batch_size,3,256,256],输出shape为[batch_size,1,256,256],数据排布为NCHW。 ## 模型初始化 在模型初始化的过程中,首先采用parse_onnx()函数根据提供的模型地址加载图像分割Unet的onnx模型,保存在net中。其次,通过net.get_parameter_shapes()获取Unet模型的输入属性,包含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.计算sigmoid值,当计算值大于0.996时值为1,小于等于0.996时值为0,保存在数组value_mask中。 2.保存结果,创建一个cv::Mat将value_mask数组中的值按行依次赋值到对应位置,得到最终的分割图像。 ```c++ ErrorCode Unet::Segmentation(const cv::Mat &srcImage, cv::Mat &maskImage) { ... // 计算sigmoid值,并且当大于0.996时,值为1,当小于0.996时,值为0,存储在value_mask[]数组中 int value_mask[numberOfOutput]; for(int i=0; i 0.996) { value_mask[i] = 1; } else { value_mask[i] = 0; } } // 将对应的value_mask[]数组中的值依次赋值到outputImage对应位置处 cv::Mat outputImage = cv::Mat_(Size(outputShape.lens()[3], outputShape.lens()[2]), CV_32S); for(int i=0;i(i,j)=value_mask[256*i+j]; // 其中,256代表了outputShape.lens()[3]的值 } } ... } ``` 注:本次采用的模型权重onnx文件是通过使用具有普通背景的汽车图像来训练的。因此,“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果,建议对现实世界示例数据集上的模型进行微调。