def visualize_prediction(output, prediction,path=None):
# 创建一个空的 RGB 图像,大小为 [3, 256, 256]
result_image = torch.zeros(3, 256, 256, dtype=torch.uint8).cuda()
# 计算 TP、FN、TN 和 FP 区域
TP = (output == 1) & (prediction == 1)
FN = (output == 1) & (prediction == 0)
TN = (output == 0) & (prediction == 0)
FP = (output == 0) & (prediction == 1)
# 将对应的颜色填充到 RGB 图像中
result_image[0] = torch.where(FP, torch.tensor(255, dtype=torch.uint8).cuda(), result_image[0]) # 红色: [255, 0, 0]
result_image[1] = torch.where(FN, torch.tensor(255, dtype=torch.uint8).cuda(), result_image[1]) # 绿色: [0, 255, 0]
result_image[0] = torch.where(TP, torch.tensor(255, dtype=torch.uint8).cuda(),result_image[0]) # 白色: [255, 255, 255]
result_image[1] = torch.where(TP, torch.tensor(255, dtype=torch.uint8).cuda(),result_image[1]) # 白色: [255, 255, 255]
result_image[2] = torch.where(TP, torch.tensor(255, dtype=torch.uint8).cuda(),result_image[2]) # 白色: [255, 255, 255]
#黑色就是最后没标的那些,因为本身像素点是黑色的
# 从 GPU 转移到 CPU 并转换为 NumPy 数组
result_image_np = result_image.cpu().numpy()
# 将 NumPy 数组转换为 PIL 图像
result_image_pil = Image.fromarray(result_image_np.transpose(1, 2, 0)) # 转置为 [H, W, C]
# 保存图像
result_image_pil.save(path)#path_name 中输入自己需要的路径
在该函数示例中,output和prediction都是尺寸为[1,256,256] 、数据类型为tensor且都在cuda上。
需要特别注意的是颜色的注入顺序很有讲究,为了让红绿不被覆盖,我们需要先涂红绿再涂白色。