场景分割中 lable_value_img 和 label_co

raw image

1. lable_value_img 转成 label_color_img

lable_value_img 很像深度图，但是它的每个像素点的值表示 class idx

以 ADE20K 提供的数据为例，如下图：ADE_train_00000065.png

label_value_img

输出这个 img 的 value

label_img = sio.imread('img/label_value_img.png')
print(label_img)
print(label_img.shape)

512 x 624 的矩阵，每个元素表示 label_value

[[0 0 0 ... 1 1 0]
 [0 1 1 ... 1 1 0]
 [0 1 1 ... 1 1 0]
 ...
 [0 4 4 ... 4 4 4]
 [0 4 4 ... 4 4 4]
 [0 0 0 ... 0 0 0]]
(512, 624)

把它变成 label_color_img，只要将每个值对应 RGB 颜色就行了

import numpy as np
import pandas as pd
import cv2


def get_label_info(csv_path):
    ann = pd.read_csv(csv_path)
    label = {'null': [0, 0, 0]}
    for i, row in ann.iterrows():  # DataFrame
        label_name = row['name']
        r = row['r']
        g = row['g']
        b = row['b']
        label[label_name] = [int(r), int(g), int(b)]
    return label


def label_value2color(label_value_img, label_color_dict):
    """
    :param label_value_img: label values image, grayscale
    :param label_color_dict: label color dict, { 'name': [r,g,b], .. }
    :return: label_color_image
    """
    label_color_list = [label_color_dict[key] for key in list(label_color_dict.keys())]  # rgb 颜色 list
    color_codes = np.array(label_color_list)  # 转成 np list
    label_color_image = color_codes[label_value_img.astype(int)]
    return label_color_image


if __name__ == '__main__':
    label_img = cv2.imread('img/label_value_img.png', cv2.IMREAD_ANYDEPTH)
    color_img = label_value2color(label_img, get_label_info('csv/ade150.csv'))  # 读入 color dict 并转成 color label img
    cv2.imwrite('img/label_color_img.png', color_img)

转换之后的结果

label_color_img.png

2. lable_color_img 转成 label_value_img

同样的道理，需要找到对应关系，将 rgb 映射到 label value

def label_color2value(label_color_img, label_color_dict):
    """
    :param label_color_img: label values image, grayscale
    :param label_color_dict: label color dict, { 'name': [r,g,b], .. }
    :return: label_value_image
    """
    label_color_list = [label_color_dict[key] for key in list(label_color_dict.keys())]  # rgb 颜色 list
    rows, cols = label_color_img.shape[0], label_color_img.shape[1]
    label_value_img = np.zeros((rows, cols)).astype('uint8')
    for i in range(rows):
        for j in range(cols):
            label_value_img[i][j] = label_color_list.index(list(label_color_img[i][j]))
    return label_value_img

这个过程明显慢多了，因为每次都要先在 list 中 查找 idx
而 value -> color 只需要 索引出值即可

所以我们可以同样用索引的方法，先建立一个以 'rgb' 为 key 的，以 label_value 为 value 的字典，使用 rgb 索引出 label_value

def label_color2value2(label_color_img, label_color_dict):
    """
    :param label_color_img: label values image, grayscale
    :param label_color_dict: label color dict, { 'name': [r,g,b], .. }
    :return: label_value_image
    """
    label_color_list = [label_color_dict[key] for key in list(label_color_dict.keys())]  # rgb 颜色 list

    # 重建一个 dict，{ 'rgb': idx, .. }
    my_color_dict = {}
    for i, color in enumerate(label_color_list):
        my_color_dict[str(color[0]) + str(color[1]) + str(color[2])] = i

    rows, cols = label_color_img.shape[0], label_color_img.shape[1]
    label_value_img = np.zeros((rows, cols)).astype('uint8')
    for i in range(rows):
        for j in range(cols):
            color = list(label_color_img[i][j])
            label_value_img[i][j] = my_color_dict[str(color[0]) + str(color[1]) + str(color[2])]
    return label_value_img

时间对比：

color_img = sio.imread('img/label_color_img.png')
print(datetime.now())
label_value_img = label_color2value(color_img, get_label_info('csv/ade150.csv'))  # 读入 color dict 并转成 color label img
print(datetime.now())
label_value_img = label_color2value2(color_img, get_label_info('csv/ade150.csv'))  # 读入 color dict 并转成 color label img
print(datetime.now())

2019-01-10 23:00:04.828832
2019-01-10 23:00:11.944107
2019-01-10 23:00:13.055654

7.1 VS 1.1

生成的 label_value_img

my_label_value_img.png