对于语义分割网络，其输出为（b, h, w, classes）,对索引求最大值，得到维度为（b, h, w, 1）
相对于得到一个灰度图，其亮度值为类别index。因为类别值为[1, num_classes]， 如果对输出
结果直接显示，会的到一副纯黑的图。
所以需要进行预测结果可视化
将预测结果转化为RGB图像
首先建立预测类别和相应rgb颜色的映射
Label = namedtuple('Label', [
'name',
'trainId',
'category',
'categoryId',
'hasInstances',
'ignoreInEval',
'color',
])
labels = [
# name id trainId category catId hasInstances ignoreInEval color
Label('unlabeled', 0, 255, 'void', 0, False, True, (0, 0, 0)),
Label('ego vehicle', 1, 255, 'void', 0, False, True, (0, 0, 0)),
Label('rectification border', 2, 255, 'void', 0, False, True, (0, 0, 0)),
Label('out of roi', 3, 255, 'void', 0, False, True, (0, 0, 0)),
Label('static', 4, 255, 'void', 0, False, True, (0, 0, 0)),
Label('dynamic', 5, 255, 'void', 0, False, True, (111, 74, 0)),
Label('ground', 6, 255, 'void', 0, False, True, (81, 0, 81)),
Label('road', 7, 0, 'flat', 1, False, False, (128, 64, 128)),
Label('sidewalk', 8, 1, 'flat', 1, False, False, (244, 35, 232)),
Label('parking', 9, 255, 'flat', 1, False, True, (250, 170, 160)),
Label('rail track', 10, 255, 'flat', 1, False, True, (230, 150, 140)),
Label('building', 11, 2, 'construction', 2, False, False, (70, 70, 70)),
Label('wall', 12, 3, 'construction', 2, False, False, (102, 102, 156)),
Label('fence', 13, 4, 'construction', 2, False, False, (190, 153, 153)),
Label('guard rail', 14, 255, 'construction', 2, False, True, (180, 165, 180)),
Label('bridge', 15, 255, 'construction', 2, False, True, (150, 100, 100)),
Label('tunnel', 16, 255, 'construction', 2, False, True, (150, 120, 90)),
Label('pole', 17, 5, 'object', 3, False, False, (153, 153, 153)),
Label('polegroup', 18, 255, 'object', 3, False, True, (153, 153, 153)),
Label('traffic light', 19, 6, 'object', 3, False, False, (250, 170, 30)),
Label('traffic sign', 20, 7, 'object', 3, False, False, (220, 220, 0)),
Label('vegetation', 21, 8, 'nature', 4, False, False, (107, 142, 35)),
Label('terrain', 22, 9, 'nature', 4, False, False, (152, 251, 152)),
Label('sky', 23, 10, 'sky', 5, False, False, (70, 130, 180)),
Label('person', 24, 11, 'human', 6, True, False, (220, 20, 60)),
Label('rider', 25, 12, 'human', 6, True, False, (255, 0, 0)),
Label('car', 26, 13, 'vehicle', 7, True, False, (0, 0, 142)),
Label('truck', 27, 14, 'vehicle', 7, True, False, (0, 0, 70)),
Label('bus', 28, 15, 'vehicle', 7, True, False, (0, 60, 100)),
Label('caravan', 29, 255, 'vehicle', 7, True, True, (0, 0, 90)),
Label('trailer', 30, 255, 'vehicle', 7, True, True, (0, 0, 110)),
Label('train', 31, 16, 'vehicle', 7, True, False, (0, 80, 100)),
Label('motorcycle', 32, 17, 'vehicle', 7, True, False, (0, 0, 230)),
Label('bicycle', 33, 18, 'vehicle', 7, True, False, (119, 11, 32)),
Label('license plate', -1, -1, 'vehicle', 7, False, True, (0, 0, 142)),
]
trainId2label = {label.trainId: label for label in reversed(labels)}
// {-1： Label（）， 18：Label（），。。。。}
生成一个与原图大小一样的三维矩阵
colored_image = np.zeros(
(class_id_image.shape[0], class_id_image.shape[1], 3), np.uint8)
将对应位置填补为类别对应的RGB
for row in range(class_id_image.shape[0]):
for col in range(class_id_image.shape[1]):
try:
colored_image[row, col, :] = class_id_to_rgb_map[
int(class_id_image[row, col])].color
所以全过程为
probs = pspnet.predict(img)
cm = np.argmax(probs, axis=2)
colored_class_image = color_class_image(cm)
alpha_blended = 0.5 * colored_class_image + 0.5 * img
与原图混合
补充：还可以用PIL内置调色板方法，
new_mask = PIL.Image.fromarray(mask.astype(np.uint8)).convert('P')
new_mask.putpalette(palette)
from PIL import Image
Image.open('PennFudanPed/PNGImages/FudanPed00001.png')
mask = Image.open('PennFudanPed/PedMasks/FudanPed00001_mask.png')
mask.putpalette([
0, 0, 0, # black background
255, 0, 0, # index 1 is red
255, 255, 0, # index 2 is yellow
255, 153, 0, # index 3 is orange
])