Python小技巧 - 颜色和类之间的映射

在做语义分割的时候，我们往往会遇到彩色的 ground truth 和 label map 之间的相互映射。这个时候最常用的方法就是用循环来做，这样是最简单的实现方法，但是效率低，循环过大会直接导致电脑卡死，而且代码还不美观；知道用矩阵来操作的话效率会高很多，但苦于不会实现。经过一番参考学习之后提供给大家一个简单高效的实现方法。

import numpy as np

def label_mapping(gts):

    colorize = np.zeros([4,3],'uint8')
    colorize[0,:] = [ 30, 100,  30]
    colorize[1,:] = [200, 200, 200]
    colorize[2,:] = [  0, 255,   0]
    colorize[3,:] = [  0,   0, 255]

    ims = colorize[gts,:].reshape([gts.shape[0],gts.shape[1],3])
    return ims

对于一个有4个类的图像来说，输入 gts 为一张只含有 label 的图像， 输出 ims 为映射后的彩色图像，非常简洁美观，而且效率也高。

做了一个测试，500x500 的 label map，对于用 for 循环来实现映射的方法，10张图像的每张平均预测时间为 0.412 秒，相同条件下用上述方法时间为 0.003 秒。