拥有prediction和label之后对TP TN FP FN进行标色（ 直接可用）

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上。

需要特别注意的是颜色的注入顺序很有讲究，为了让红绿不被覆盖，我们需要先涂红绿再涂白色。