OpenCV4学习笔记（68）——dnn模块之调用基于Cityscapes数据集的图像分割模型
本次笔记要整理记录的是基于Cityscapes数据集训练而成的图像分割模型model-best.net，该模型能够对行人、汽车、红绿灯、马路、地形、天空等等二十种类别进行分类，可用于城市景观分割。
下面通过代码整理如何在OpenCV中利用dnn模块对该模型进行加载调用，并实现城市景观分割。
首先加载模型，并且设置计算后台和目标设备
string model_path = "D:\\opencv_c++\\opencv_tutorial\\data\\models\\enet\\model-best.net";
Net net = readNetFromTorch(model_path);
net.setPreferableBackend(DNN_BACKEND_OPENCV);
net.setPreferableTarget(DNN_TARGET_CPU);
然后加载测试图像，将测试图像转换成blob后传入网络的输入层，并进行前向传播。这里输出前向传播结果矩阵的四个维度数。
Mat test_image = imread("D:\\opencv_c++\\opencv_tutorial\\data\\images\\cityscapes_test.jpg");
imshow("test_image", test_image);
Mat inputBlob = blobFromImage(test_image, 0.00392, Size(1024, 512), Scalar(0, 0, 0), true, false);
net.setInput(inputBlob);
Mat prob = net.forward();
cout << prob.size[0] << endl; //输入图像数目
cout << prob.size[1] << endl; //输出的类别数
cout << prob.size[2] << endl; //输出图像行数
cout << prob.size[3] << endl; //输出图像列数
输出结果如下：
其中，“1”是输入图像的数量，“20”表示总共有二十个类别，“512”表示输出图像的高度，“1024”表示输出图像的宽度。
接着使用生成随机数，来为每个分类随机生成代表颜色，这里将前后两个分类的颜色像素值设置一定关联，能使分割后的图像效果更好些。
//为输出的二十个分类生成代表颜色
vector<Vec3b>class_colors;
RNG rng(0);
for (int i = 1;i < prob.size[1];i++)
{
class_colors.push_back(Vec3b());
int b = (rng.uniform(0, 256) + class_colors[i - 1][0]) / 2;
int g = (rng.uniform(0, 256) + class_colors[i - 1][1]) / 2;
int r = (rng.uniform(0, 256) + class_colors[i - 1][2]) / 2;
Vec3b color = Vec3b(b, g, r);
class_colors.push_back(color);
}
然后通过每个类别的输出矩阵来寻找最大置信度的分类ID和置信度。
Mat classID = Mat::zeros(prob.size[2], prob.size[3], CV_8UC1);
Mat max_score = Mat::zeros(prob.size[2], prob.size[3], CV_16F);
for (int ch = 0; ch < prob.size[1];ch++)
{
for (int row = 0;row < prob.size[2];row++)
{
float *prob_ptr = prob.ptr<float>(0, ch, row); //定位到第ch个通道的第row行
float *max_ptr = max_score.ptr<float>(row);
for (int col = 0; col < prob.size[3];col++)
{
if (prob_ptr[col] > max_ptr[col])
{
max_ptr[col] = prob_ptr[col];
classID.at<uchar>(row, col) = uchar(ch);
}
}
}
}
到这里得到的classID这个Mat矩阵，每个像素点对应原图像中相同位置的像素点，其像素值表示该像素点所属的类别ID。
接下来通过classID这个矩阵来判断原图像中每个像素点的所属类别是什么，并为每个像素点赋值为其所属类别的代表颜色，实现不同类别的分割。
//将每个像素点的最大可能分类所代表颜色赋给该像素点
Mat result = Mat::zeros(prob.size[2], prob.size[3], CV_8UC3);
for (int row = 0;row < result.rows;row++)
{
for (int col = 0;col < result.cols;col++)
{
int index = classID.at<uchar>(row, col);
Vec3b class_color = class_colors[index];
result.at<Vec3b>(row, col) = class_color;
}
}
最后将原图像和分割后的图像进行加权融合，并进行显示。
resize(result, result, test_image.size(),0,0, INTER_NEAREST);
addWeighted(test_image, 0.1, result, 0.9, 0.0, result);
imshow("result", result);
显示结果如下：
原图
分割图像
可以看到分割结果中，道路两边的汽车、骑行者、天空、建筑等都被分割出来，但是也能很明显地看到一些噪声或分割错误的地方。通过这一张测试图像的结果来说，我感觉这个网络模型的效果并没有特别的好，而且在其他测试图像上的尝试也比较难得到很好的效果，不知是该模型的鲁棒性较差的原因还是我解码过程存在问题的原因。如果有朋友需要进行城市景观分割的话，也可以用这个模型去做些尝试。
那本次笔记到此结束。
PS：本人的注释比较杂，既有自己的心得体会也有网上查阅资料时摘抄下的知识内容，所以如有雷同，纯属我向前辈学习的致敬，如果有前辈觉得我的笔记内容侵犯了您的知识产权，请和我联系，我会将涉及到的博文内容删除，谢谢！