[python]-分割结果可视化

常常会遇到需要将分割的结果可视化的问题，这里采用修改dataset的方式，传入原图路径，最终保存的图片是在原图上叠加分割的结果，可以根据需要修改颜色。
注意：这里的颜色三元组顺序是BGR。

示例1：
这是每张图包含两类分割标签的可视化，每个类别需要分别保存，而且加入gt和预测的overlap。

import matplotlib.pyplot as plt
import numpy as np

def masks_to_colorimg(masks,color):
    colors = np.asarray([color]) 
    colorimg = np.zeros((masks.shape[1], masks.shape[2], 3), dtype=np.float32) * 255
    channels, height, width = masks.shape
    for y in range(height):
        for x in range(width):
            selected_colors = colors[masks[:,y,x] > 0.5]
            if len(selected_colors) > 0:
                colorimg[y,x,:] = np.mean(selected_colors, axis=0)

    return colorimg.astype(np.uint8)



c = 0
with torch.no_grad():
    for inputs, labels, path in test_loader:
        c += 1
        inputs = inputs.to(device)
        labels = labels.to(device)
        pred = model(inputs)
        pred = torch.sigmoid(pred)
        pred = pred.data.cpu().numpy()
        labels = labels.data.cpu().numpy()

        pred_mask = pred[0,:,:,:,]
        label_mask = labels[0,:,:,:,]
        iou = np.zeros((2,512,512), dtype=np.float32)
	# overlap大小为512*512，两个标签类别
        for i in range(0,512):
            for j in range(0,512):
                if pred_mask[0][i][j] > 0.5 and label_mask[0][i][j] > 0.5:
                    iou[0][i][j] = 1
                if pred_mask[1][i][j] > 0.5 and label_mask[1][i][j] > 0.5:
                    iou[1][i][j] = 1

        target_0 = [helper.masks_to_colorimg(np.expand_dims(label_mask[0,:,:],0), (0,255,255))]
        target_1 = [helper.masks_to_colorimg(np.expand_dims(label_mask[1,:,:],0), (0,255,153))]
        pred_0 = [helper.masks_to_colorimg(np.expand_dims(pred_mask[0,:,:],0), (0,0,255))]
        pred_1 = [helper.masks_to_colorimg(np.expand_dims(pred_mask[1,:,:],0), (255,0,255))]

        image_origin_path = path[0]
        img = cv2.imread(image_origin_path)
        s = 0.5 # 阈值

        iou_0 = [helper.masks_to_colorimg(np.expand_dims(iou[0,:,:],0), (0,102,255))]
        iou_1 = [helper.masks_to_colorimg(np.expand_dims(iou[1,:,:],0), (255,0,0))]

        target_0 = np.array(target_0) * s + img
        target_1 = np.array(target_1) * s + img
        pred_0 = np.array(pred_0) * s + img
        pred_1 = np.array(pred_1) * s + img
        iou_0 = np.array(iou_0) * s + img
        iou_1 = np.array(iou_1) * s + img

        cv2.imwrite('save_path' + '{}_gt_1'.format(c)+ '.jpg', np.squeeze(target_1))
        cv2.imwrite('save_path/' + '{}_gt_0'.format(c)+ '.jpg', np.squeeze(target_0))
        cv2.imwrite('save_path' + '{}_pred_1'.format(c)+ '.jpg', np.squeeze(pred_1))
        cv2.imwrite('save_path' + '{}_pred_0'.format(c)+ '.jpg', np.squeeze(pred_0))
        cv2.imwrite('save_path' + '{}_iou_1'.format(c)+ '.jpg', np.squeeze(iou_1))
        cv2.imwrite('save_path' + '{}_iou_0'.format(c)+ '.jpg', np.squeeze(iou_0))

示例2：
如果不需要将不同标签分开保存，那么可以相应地修改传入维度和颜色列表

def masks_to_colorimg(masks,color):
    colors = np.asarray(color) # 注意这里的color是颜色元组的列表
    colorimg = np.zeros((masks.shape[1], masks.shape[2], 3), dtype=np.float32) * 255
    channels, height, width = masks.shape
    for y in range(height):
        for x in range(width):
            selected_colors = colors[masks[:,y,x] > 0.5]
            if len(selected_colors) > 0:
                colorimg[y,x,:] = np.mean(selected_colors, axis=0)

    return colorimg.astype(np.uint8)

# 注意这里的label_mask是三维数组
target = [helper.masks_to_colorimg(label_mask, [(0,255,255), (255,0,0)])]