基本思路
做语义分割的同学们,经常需要在原图上可视化预测结果,我们可以借助matplot这个库来进行可视化。
基本原理是首先在画图上展示原图,然后不要清除画布,接着展示标签,为标签增加一个透明度通道,将标签不为零的部分透明度设为0.5,为零的部分透明度设为0
。
代码
import matplotlib.pyplot as plt
# 定义类别颜色
color_mask = [[0, 255, 0], [255, 255, 0], [255, 0, 0]]
for name in os.listdir(image_path):
image_name = os.path.join(image_path, name)
label_name = os.path.join(label_path, name)
image = cv2.imread(image_name)
label = cv2.imread(label_name, cv2.IMREAD_GRAYSCALE)
# label = np.expand_dims(label,2).repeat(3, axis=2)
image = image[:, :, ::-1]
plt.axis('off')
plt.imshow(image)
for i in range(3):
m = np.array(label)
# 这里需要改成符合自己的数值,一定要大于分类数,比如有3类,这里可以写4,5...
m[m != i] = 4
m[m == i] = 1
m[m == 4] = 0
img = np.ones( (m.shape[0], m.shape[1], 3) )
# 随机生成颜色
# color = np.random.random((1, 3)).tolist()[0]
color = color_mask[i]
for j in range(3):
img[:,:,j] = color[j]
# 增加透明度通道,叠加到原图上
plt.imshow(np.dstack( (img, m*0.5) ))
# 显示图像
# plt.show()
# 保存图像,去除白边
plt.savefig(os.path.join('E:/baidu/myproject1/figure', name), bbox_inches='tight',pad_inches = -0.1)
# 清除画布
plt.clf()