# 图像分割
本示例主要通过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文件是通过使用具有普通背景的汽车图像来训练的。因此,“现实世界“图像的分割结果不完美是意料之中的。为了获得更好的结果,建议对现实世界示例数据集上的模型进行微调。