单通道二维数组指定颜色编码转RGB显示

公众号：清纯世纪 

代码：

        法一：

import numpy as np
color = {0 : [0, 0, 0],
         1 : [128, 0, 0],
         2 : [0, 128, 128],
         3 : [128, 128, 0]} # 对应标签的颜色编码

predict = np.zeros((5,5,3)) # 转RGB用的，接收相应的预测RGB值
pre = np.array([[0,1,1,3,2],
                [0,3,2,1,2],
                [1,1,1,3,3],
                [0,2,2,2,2],
                [1,2,1,3,0]])  # 真正的网络输出二维数组值

for i in range(pre.shape[0]):
    for j in range(pre.shape[1]):
        predict[i,j] = color[pre[i,j]]

print(np.array(predict).transpose((2,0,1)))

        法二：

import numpy as np

color = {0: [0, 0, 0],
         1: [128, 0, 0],
         2: [0, 128, 128],
         3: [128, 128, 0]}  # 对应标签的颜色编码

predict = np.zeros((5, 5, 3))  # 转RGB用的，接收相应的预测RGB值
pre = np.array([[0, 1, 1, 3, 2],
                [0, 3, 2, 1, 2],
                [1, 1, 1, 3, 3],
                [0, 2, 2, 2, 2],
                [1, 2, 1, 3, 0]])  # 真正的网络输出二维数组值

# predict 是用来赋色的：color[i] 代表一类的颜色值；colors[i][0] 代表该类的 0 通道的值
# 下面的三行分别来对每个通道上的各类的颜色值进行赋值，pre只有一个通道(二维），所以才一直是 pre[:,:,0]
for i in range(len(color)):
    predict[:, :, 0] += ((pre[:, :] == i) * (color[i][0])).astype('uint8')
    predict[:, :, 1] += ((pre[:, :] == i) * (color[i][1])).astype('uint8')
    predict[:, :, 2] += ((pre[:, :] == i) * (color[i][2])).astype('uint8')
print(np.array(predict).transpose((2, 0, 1)))