多颜色绘制语义分割/变化检测结果图

在论文绘图时，传统的二元语义分割结果图颜色单一（下图左），所以论文中常根据混淆矩阵类别使用多颜色进行绘制（下图右），可以看到，结果的可视化效果更好。

     

以下是绘制代码：

import os
import cv2
import argparse
import numpy as np
from tqdm import tqdm

def intersection_color(binary_result, ground_truth):
   
    if len(binary_result.shape) != 2 or len(ground_truth.shape) != 2:
        raise ValueError(
            f"The dim numbers of binary_result and ground_truth must be 2!"
        )
    # 将255的值转换为1
    binary_result = np.where(binary_result > 128, 1, 0)
    ground_truth = np.where(ground_truth > 128, 1, 0)
   
    # 创建RGB图像，根据TP、FP和FN的位置使用不同的颜色
    rgb_image = np.zeros(
        (binary_result.shape[0], binary_result.shape[1], 3), dtype=np.uint8
    )

    # True Positives (TP) - 白色
    rgb_image[(binary_result == 1) & (ground_truth == 1)] = [255, 255, 255]

    # False Positives (FP) - 红色
    rgb_image[(binary_result == 1) & (ground_truth == 0), 2] = 255

    # False Negatives (FN) - 绿色
    rgb_image[(binary_result == 0) & (ground_truth == 1), 1] = 255


    # rgb_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
    return rgb_image

def visual_label(args):
    pred_list = os.listdir(args.pred_root)
    gt_list = os.listdir(args.gt_root)
    if len(pred_list)!= len(gt_list):
        raise ValueError(
            f"The number of predicted result is not equal to that of ground truth!"
        )
        
    if not os.path.exists(args.out_root):
        os.makedirs(args.out_root)
        
    for file_name in tqdm(pred_list):
        pred_file = os.path.join(args.pred_root, file_name)
        gt_file = os.path.join(args.gt_root, file_name)
        pred = cv2.imread(pred_file, cv2.IMREAD_GRAYSCALE)
        gt = cv2.imread(gt_file, cv2.IMREAD_GRAYSCALE)
        rgb_label = intersection_color(pred, gt)
        out_path = os.path.join(args.out_root, file_name)
        cv2.imwrite(out_path, rgb_label)
        
def parse_args():
    parser = argparse.ArgumentParser(
        description='Open-CD test (and eval) a model')
    parser.add_argument('--pred_root', help='predict results path')
    parser.add_argument('--gt_root', help='gt path')
    parser.add_argument(
        '--out_root',
        help=('if specified, the evaluation metric results will be dumped'
              'into the directory as json'))
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    visual_label(args)
    
if __name__ == "__main__":
    main()

欢迎关注大地主的Github仓库（ABCnutter (PengChen) (github.com)），谢谢。