常常会遇到需要将分割的结果可视化的问题，这里采用修改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)])]
Original: https://www.cnblogs.com/camilia/p/16875343.html
Author: CAMILIA
Title: [python]-分割结果可视化
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/683355/
转载文章受原作者版权保护。转载请注明原作者出处！