在论文绘图时,传统的二元语义分割结果图颜色单一(下图左),所以论文中常根据混淆矩阵类别使用多颜色进行绘制(下图右),可以看到,结果的可视化效果更好。
以下是绘制代码:
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)),谢谢。