使用numpy、颜色映射表 将灰度图变彩图

目录

问题：

代码：

实例信息：

---

问题：

在制作语义分割图像数据集过程中，输出gt可能是表示类别的单通道灰度图，为了可视化，需要使用颜色映射转化为3通道彩色图（灰度转单通道彩色可参考，python将灰度图保存为8bit彩色图）

使用遍历像素的方法将灰度图转彩图较慢，本文使用np.vectorize，可以极大提升转换速度

同样数据，遍历方法平均0.45秒，vectorize方法平均0.04秒，提升10倍以上的处理速度

代码：

# -*- coding: utf-8 -*-
# @Time : 2020/4/7 16:44 
# @Author : Zhao HL
# @File : gt_visualizetion.py 
import os,sys,time,cv2
import numpy as np
from collections import namedtuple

# 类别信息
Cls = namedtuple('cls', ['name', 'id', 'one_hot', 'color'])
Clss = [
    Cls('road', 1, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], (128, 0, 0)),
    Cls('build_resident', 2, [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], (75, 0, 130)),
    Cls('build_industrial', 3, [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], (255, 215, 0)),
    Cls('build_service', 4, [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0], (34, 139, 34)),
    Cls('build_village', 5, [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], (255, 0, 0)),
    Cls('greenland', 6, [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], (0, 0, 128)),
    Cls('forest', 7, [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], (0, 128, 128)),
    Cls('bareland', 8, [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0], (128, 128, 0)),
    Cls('farmland', 9, [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], (160, 82, 45)),
    Cls('waterbody', 10, [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0], (72, 209, 204)),
    Cls('other', 11, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], (255, 0, 255)),
]
gts_gray_path = r'E:\_Python\01_DeepLearning\dataset\NB_S11\test\gt_test'
gts_color_path = r'E:\_Python\01_DeepLearning\dataset\NB_S11\test\gt_test_color'

def gray_color(color_dict,gray_path=gts_gray_path,color_path=gts_color_path):
    '''
    swift gray image to color, by color mapping relationship
    :param color_dict:color mapping relationship, dict format
    :param gray_path:gray imgs path
    :param color_path:color imgs path
    :return:
    '''
    pass
    t1 = time.time()
    gt_list = os.listdir(gray_path)
    for index,gt_name in enumerate(gt_list):
        gt_gray_path = os.path.join(gray_path,gt_name)
        gt_color_path = os.path.join(color_path, gt_name)
        # gt_gray = tif.imread(gt_gray_path)
        gt_gray = cv2.imread(gt_gray_path,cv2.IMREAD_GRAYSCALE)
        assert len(gt_gray.shape) == 2                          # make sure gt_gray is 1band

        # # region method 1: swift by pix, slow
        # gt_color = np.zeros((gt_gray.shape[0],gt_gray.shape[1],3),np.uint8)
        # for i in range(gt_gray.shape[0]):
        #     for j in range(gt_gray.shape[1]):
        #         gt_color[i][j] = color_dict[gt_gray[i][j]]      # gray to color
        # # endregion

        # region method 2: swift by array, fast
        gt_color = np.array(np.vectorize(color_dict.get)(gt_gray),np.uint8).transpose(1,2,0)
        # endregion

        gt_color = cv2.cvtColor(gt_color,cv2.COLOR_RGB2BGR)
        cv2.imwrite(gt_color_path,gt_color,)
        process_show(index+1,len(gt_list))
    print(time.time()-t1)



def nt_dic(nt=Clss):
    '''
    swift nametuple to color dict
    :param nt: nametuple
    :return:
    '''
    pass
    color_dict= {}
    for cls in nt:
        color_dict[cls.id] = cls.color
    return color_dict

def process_show(num, nums, pre_fix='', suf_fix=''):
    '''
    auxiliary function, print work progress
    :param num:
    :param nums:
    :param pre_fix:
    :param suf_fix:
    :return:
    '''
    rate = num / nums
    ratenum = round(rate, 3) * 100
    bar = '\r%s %g/%g [%s%s]%.1f%% %s' % \
          (pre_fix, num, nums, '#' * (int(ratenum) // 5), '_' * (20 - (int(ratenum) // 5)), ratenum, suf_fix)
    sys.stdout.write(bar)
    sys.stdout.flush()

if __name__ == '__main__':
    pass
    color_dict = nt_dic()
    gray_color(color_dict)

实例信息：

类别信息

图像

gt gray

gt color