语义分割结果可视化(原图+语义掩码+图例)
由于实习工作需要把语义分割结果可视化出来,要使用自定义颜色来区分不同的label,并绘制出图例并插入在图像右端。本文将介绍如何实现这样的语义分割结果图。
1 确定不同Label的RGB颜色
为了使视觉效果尽可能的好,我们需要选择一组相互之间差异最大的颜色,用于填涂不同Label对应的mask区域。笔者已在这篇博文中详细介绍:生成一组差异较大的RGB颜色
其中,我们利用了CIEDE2000来度量颜色的差异,它考虑了人眼的视觉感知特性,能够更好地反映实际的颜色差异情况。具体来说,它将三个颜色属性分别转化到LCH(明度、色度、色调)颜色空间,然后计算两个LCH颜色向量之间的距离。具体代码可见如上博文。
2 生成语义图和mask边缘图
读取标注图后,利用对应的colors生成相应的语义图,并检测每个label的mask的边缘,将labels、语义图和mask边缘图列表返回。
# 获取labels、语义图和mask边缘图列表
def get_semantic_edges(ann_path, colors):
mask = cv2.imread(ann_path, cv2.IMREAD_ANYDEPTH)
chosen_labels = np.unique(mask)
chosen_colors = [colors[label] for label in chosen_labels]
contours_list = []
# 新建一个图
color_masks = np.zeros((mask.shape[0],mask.shape[1],3))
for i,label in enumerate(chosen_labels):
if label!=0: # label为0表示background
idx = np.nonzero(mask == label)
color_masks[idx[0],idx[1]] = chosen_colors[i] # 填图
# get contours
new_mask = np.zeros((mask.shape[0],mask.shape[1]))
new_mask[idx[0],idx[1]] = 255
contours,_ = cv2.findContours(new_mask.astype('uint8'), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
contours_list.append(contours)
color_masks = color_masks.astype('uint8')
color_masks = cv2.cvtColor(color_masks, cv2.COLOR_RGBA2BGR)
return chosen_labels, color_masks, contours_list
3 合并语义图+原图+边缘
通过cv2.addWeighted合并语义图和原图,注意背景不要融合进去了,否则会出现不佳的合并效果。再直接用cv2.drawContours把边缘绘制上去。
# 合并mask和原图,并绘制边缘
def merge_mask_img(masks, img, contours_list,original_img_ratio):
idx = np.nonzero(masks != 0)
masks[idx[0],idx[1]] = cv2.addWeighted(img[idx[0],idx[1]], original_img_ratio, masks[idx[0],idx[1]], 1, 0)
idx = np.nonzero(masks == 0)
masks[idx[0],idx[1]] = img[idx[0],idx[1]]
for contours in contours_list:
cv2.drawContours(masks, contours, -1, (255, 255, 255), 5)
return masks
4 绘制图例
输入原图、labels、colors和原图中包含的Labels,输出与原图高度相同的图例,以便后续将第3节的输出结果与图例直接横向拼接。
# 绘制图例
def get_resized(image, labels, colors, chosen_labels):
# 选择颜色和label
new_labels = [label for i,label in enumerate(labels) if i in chosen_labels]
new_colors = [color for i,color in enumerate(colors) if i in chosen_labels]
# 将RGB颜色转换为归一化值,并创建自定义颜色映射
cmap = ListedColormap([np.array(new_colors[i])/255 for i in range(len(new_colors))])
# 绘制语义分割图和图例
legend_elements = []
for i in range(len(new_labels)):
rectangle = plt.Rectangle((0, 0), 1, 1, fc=cmap(i / len(new_labels)), edgecolor='black')
legend_elements.append(rectangle)
# 仅保存图例
fig_leg = plt.figure(figsize=(1.5, 5), dpi=500)
legend = fig_leg.legend(legend_elements, new_labels, loc='upper left',fontsize=8)
legend.get_frame().set_linewidth(0.0)
# plt转
canvas = FigureCanvas(fig_leg)
canvas.draw()
buf = canvas.buffer_rgba()
img = np.asarray(buf)
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
height = image.shape[0]
print(img.shape,height)
if img.shape[0] >= height:
resized = cv2.resize(img, (height*img.shape[1]//img.shape[0], height))
else:
resized = np.pad(img, ((0, height-img.shape[0]), (0,0), (0, 0)), constant_values=255)
return resized
在拼接了图例之后,效果图如下所示:
如果想获取完整代码,请见:语义分割结果可视化(原图+语义掩码+图例)