博客

简介

在做图像分类或者语义分割时候，大家都会接触到标签图像。标签图像也称gt或者mask图，其可能为单通道图像，也可能为3通道图像（RGB通道值相同），由于标签通常是0,1,2,3,…所以打开看着就是全黑，有时候做实验自己想看看标签图或者想看看自己生成的标签，就需要将标签进行可视化表达。

网上也有相关方法，但是看着有点复杂，随即写了几行代码，很简单，第一次入手的同学可以看的懂。因为本人也非科班，代码有瑕疵欢迎随时指出。

流程

1、确定颜色映射关系

这里主要根据需求来设定不同label值可视化后呈现颜色。可以通过查询RGB色板，或者其他方法，根据个人喜好设定。我们可以预设多个数据集的对应关系，需要使用的时候用哪个选哪个即可。

def color_map(dataname):
    “”“
    这里可以预设多个对应关系，疯狂if... elif...即可
    这里以LoveDa数据集为例：
    该方法返回的是对应关系，类型为字典dict
    注意图像通道 （B G R）
    ”“”
    if dataname == "LoveDa":
        color_map_ = {
            0: [0, 0, 0],  # ignore
            1: [255, 248, 220],  # background
            2: [100, 149, 237],  # building
            3: [102, 205, 170],  # road
            4: [205, 133, 63],  # water
            5: [160, 32, 240],  # barren
            6: [255, 64, 64],  # forest
            7: [139, 69, 19],  # agriculture
            }
        return color_map_
    else:
        print("DATA name error!")

2、读取label图像并生成带色彩的图像

这里首先获取label图像A，在生成一个新的图像B，A与B尺寸相同。然后读取A图像（x,y）处像素的label值，然后再查询映射字典在B图像相同位置替换成预设RGB值。本质就是操作矩阵（or Tensor）。

def gt2color(path):
    # read image （np.array）
    img = cv2.imread(path)
    # show image
    cv2.imshow("GT", img)
    # 生成一个尺寸相同的图像，注意这里label图像是3通道的，如果你label是单通道，稍微修改下
    img_out = np.zeros(img.shape, np.uint8)
    # 获取图像宽*高
    img_x, img_y = img.shape[0], img.shape[1]
    # 得到该数据色彩映射关系
    label_color = color_map(dataname=dataname)
    # 每行每列像素值依次替换，这里写的简单，你也可以用for ... enumerate...获取索引和值 
    for x in range(img_x):
        for y in range(img_y):
            label = img[x][y][0]    # get img label
            img_out[x][y] = label_color[label]
    cv2.imshow("RGB", img_out)
    return img_out
    # 另存图片
    # cv2.imwrite(r"D:\...\1367_rgb.png", img_out)

3 剩余代码

这三个代码copy到一起就是完整代码了，直接在你电脑上运行即可。萌新需要注意的是cv库路径中不要包含中文，不然会报错。

if __name__ == "__main__":
    # 定义数据集，图像路径
    dataname = r"LoveDa"
    img_path = r"...\1367.png"
    img_fold = r"...\datasets" # 批处理根目录
    get_color_img = gt2color(img_path)
    cv2.waitKey()

EX 批处理

上面都是针对一个label图像的可视化，批处理主要提取出每个图像的路径，再针对每个路径执行gt2color()方法即可。这里给出路径列表方法。

def get_filename_list(path):
    file_list = []
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            file_ = {}
            file = filename.split(".png")[0]
            file_['file'] = file
            file_list.append(file_)
     
     return file_list

Result