在论文绘图时,传统的二元语义分割结果图颜色单一(下图左),所以论文中常根据混淆矩阵类别使用多颜色进行绘制(下图右),可以看到,结果的可视化效果更好。
以下是绘制代码:
python">import os
import cv2
import argparse
import numpy as np
from tqdm import tqdmdef 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的值转换为1binary_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_imagedef 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 argsdef main():args = parse_args()visual_label(args)if __name__ == "__main__":main()
欢迎关注大地主的Github仓库(ABCnutter (PengChen) (github.com)),谢谢。