二分类语义分割结果彩色可视化代码v2.0

# -*- encoding: utf-8 -*-
# @ModuleName: FP_FN_color
# @Function: 结果彩色可视化
# @Author: Yokon
# @Time: 2024/5/16 下午4:48
import os

from PIL import Image
import cv2
import numpy as np


def visualize_errors(image_path, prediction_path, vis):
    norm_path = os.path.normpath(image_path)
    # 按'/'分割路径
    new = norm_path.split('\\')[-1]
    new_path = os.path.join(vis, new)

    image = cv2.imread(image_path)
    mask = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    prediction = cv2.imread(prediction_path, cv2.IMREAD_GRAYSCALE)

    # 创建 FP 和 FN 的掩膜
    fp_mask = np.zeros_like(mask)
    fn_mask = np.zeros_like(mask)

    # 判断每个像素的分类是否正确，并标记 FP 和 FN 区域
    for i in range(mask.shape[0]):
        for j in range(mask.shape[1]):
            if mask[i, j] == 255 and prediction[i, j] != 255:  # FN
                fn_mask[i, j] = 255

    for i in range(mask.shape[0]):
        for j in range(mask.shape[1]):
            if mask[i, j] == 0 and prediction[i, j] == 255:  # FP
                fp_mask[i, j] = 255

    fn_visualized = np.stack((fn_mask, fn_mask, fn_mask), axis=2)
    fp_visualized = np.stack((fp_mask, fp_mask, fp_mask), axis=2)

    mask_3d_fp = np.all(fp_visualized == [255, 255, 255], axis=2)
    fp_visualized[mask_3d_fp] = [255, 0, 0]  # FP blue
    result = cv2.add(fp_visualized, image)

    for i in range(fn_visualized.shape[0]):
        for j in range(fn_visualized.shape[1]):
            if fn_visualized[i, j][0] == 255:
                result[i, j] = [0, 0, 255]  # FN red

    # cv2.imshow("result Image", result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    cv2.imwrite(new_path, result)
    print("visible img %s" % new_path)


if __name__ == '__main__':

    # 标签
    label_path = "./levirlabel"
    # 预测 8位24位均可
    pre_path = "./levir/hanet_256x256_40k_levircd/vis_data/vis_image"

    vis_result = "./levir_vis/hanet_256x256_40k_levircd"

    if not os.path.exists(vis_result):
        os.makedirs(vis_result)

    for file_name in os.listdir(label_path):
        label = os.path.join(label_path, file_name)
        pre = os.path.join(pre_path, file_name)
        visualize_errors(label, pre, vis_result)