语义分割结果可视化（原图+语义掩码+图例）

语义分割结果可视化（原图+语义掩码+图例）

由于实习工作需要把语义分割结果可视化出来，要使用自定义颜色来区分不同的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

在拼接了图例之后，效果图如下所示：

如果想获取完整代码，请见：语义分割结果可视化（原图+语义掩码+图例）