语义分割后处理

对于语义分割网络，其输出为（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内置调色板方法，

语义分割后处理

原文：https://www.cnblogs.com/learningcaiji/p/14197841.html