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