Question

我想用Python将3通道RGB图像转换为索引图像。它用于处理深度网络训练标签以进行语义分割。索引图像我的意思是它有一个通道，每个像素是索引，它应该从零开始。当然它们应该具有相同的尺寸。转换基于Python dict中的以下映射：

color2index = {
        (255, 255, 255) : 0,
        (0,     0, 255) : 1,
        (0,   255, 255) : 2,
        (0,   255,   0) : 3,
        (255, 255,   0) : 4,
        (255,   0,   0) : 5
    }

我实现了一个天真的功能：

def im2index(im):
    """
    turn a 3 channel RGB image to 1 channel index image
    """
    assert len(im.shape) == 3
    height, width, ch = im.shape
    assert ch == 3
    m_lable = np.zeros((height, width, 1), dtype=np.uint8)
    for w in range(width):
        for h in range(height):
            b, g, r = im[h, w, :]
            m_lable[h, w, :] = color2index[(r, g, b)]
    return m_lable

输入im是由cv2.imread()创建的 numpy 数组。但是，这段代码非常慢。由于im处于numpy数组中，我首先尝试使用 numpy 的ufunc：

RGB2index = np.frompyfunc(lambda x: color2index(tuple(x)))
indices = RGB2index(im)

但事实证明ufunc每次只需要一个元素。我被允许一次给函数三个参数（RGB值）。

还有其他方法可以进行优化吗？如果存在更有效的数据结构，则映射不必那样。我注意到Python dict的访问并不需要花费太多时间，但是从 numpy array 到 tuple （可以清除）的转换就可以了。

PS：我得到的一个想法是在CUDA中实现一个内核。但它会更复杂。

UPDATA1： Dan Mašek's Answer运行正常。但首先我们必须将RGB图像转换为灰度。当两种颜色具有相同的灰度值时，可能会出现问题。

我在这里粘贴工作代码。希望它可以帮助别人。

lut = np.ones(256, dtype=np.uint8) * 255
lut[[255,29,179,150,226,76]] = np.arange(6, dtype=np.uint8)
im_out = cv2.LUT(cv2.cvtColor(im, cv2.COLOR_BGR2GRAY), lut)

Answer 1

那呢？

color2index = {
    (255, 255, 255) : 0,
    (0,     0, 255) : 1,
    (0,   255, 255) : 2,
    (0,   255,   0) : 3,
    (255, 255,   0) : 4,
    (255,   0,   0) : 5
}

def rgb2mask(img):

    assert len(img.shape) == 3
    height, width, ch = img.shape
    assert ch == 3

    W = np.power(256, [[0],[1],[2]])

    img_id = img.dot(W).squeeze(-1) 
    values = np.unique(img_id)

    mask = np.zeros(img_id.shape)

    for i, c in enumerate(values):
        try:
            mask[img_id==c] = color2index[tuple(img[img_id==c][0])] 
        except:
            pass
    return mask

然后致电：

mask = rgb2mask(ing)

Answer 2

类似于Armali和Mendrika的建议，我不得不以某种方式对其进行一些调整才能使其正常工作（也许完全是我的错）。所以我只想分享一个有效的代码段。

COLORS = np.array([
    [0, 0, 0],
    [0, 0, 255],
    [255, 0, 0]
])
W = np.power(255, [0, 1, 2])

HASHES = np.sum(W * COLORS, axis=-1)
HASH2COLOR = {h : c for h, c in zip(HASHES, COLORS)}
HASH2IDX = {h: i for i, h in enumerate(HASHES)}


def rgb2index(segmentation_rgb):
    """
    turn a 3 channel RGB color to 1 channel index color
    """
    s_shape = segmentation_rgb.shape
    s_hashes = np.sum(W * segmentation_rgb, axis=-1)
    func = lambda x: HASH2IDX[int(x)]
    segmentation_idx = np.apply_along_axis(func, 0, s_hashes.reshape((1, -1)))
    segmentation_idx = segmentation_idx.reshape(s_shape[:2])
    return segmentation_idx

segmentation = np.array([[0, 0, 0], [0, 0, 255], [255, 0, 0]] * 3).reshape((3, 3, 3))
rgb2index(segmentation)

Example plot

代码也可以在这里找到： https://github.com/theRealSuperMario/supermariopy/blob/dev/scripts/rgb2labels.py

Answer 3

您是否检查了枕头库https://python-pillow.org/？我记得，它有一些类和方法来处理颜色转换。请参阅：https://pillow.readthedocs.io/en/4.0.x/reference/Image.html#PIL.Image.Image.convert

Answer 4

这是一个小的实用程序函数，用于将图像（np.array）转换为每像素标签（索引），这也可以是一个热门的编码：

def rgb2label(img, color_codes = None, one_hot_encode=False):
    if color_codes is None:
        color_codes = {val:i for i,val in enumerate(set( tuple(v) for m2d in img for v in m2d ))}
    n_labels = len(color_codes)
    result = np.ndarray(shape=img.shape[:2], dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        result[(img==rgb).all(2)] = idx

    if one_hot_encode:
        one_hot_labels = np.zeros((img.shape[0],img.shape[1],n_labels))
        # one-hot encoding
        for c in range(n_labels):
            one_hot_labels[: , : , c ] = (result == c ).astype(int)
        result = one_hot_labels

    return result, color_codes


img = cv2.imread("input_rgb_for_labels.png")
img_labels, color_codes = rgb2label(img)
print(color_codes) # e.g. to see what the codebook is

img1 = cv2.imread("another_rgb_for_labels.png")
img1_labels, _ = rgb2label(img1, color_codes) # use the same codebook

如果提供None，它会计算（并返回）颜色代码簿。

Answer 5

实际上，循环需要很多时间。

binary_mask = (im_array[:,:,0] == 255) & (im_array[:,:,1] == 255) & (im_array[:,:,2] == 0)

也许上面的代码可以帮助您

Answer 6

我实现了一个简单的功能：… 我首先尝试使用 numpy 的ufunc，例如：…

我建议使用更幼稚的功能，该功能仅转换一个像素：

def rgb2index(rgb):
    """
    turn a 3 channel RGB color to 1 channel index color
    """
    return color2index[tuple(rgb)]

然后使用 numpy 例程是一个好主意，但是我们不需要ufunc：

np.apply_along_axis(rgb2index, 2, im)

此处numpy.apply_along_axis()用于将rgb2index()函数应用于整个图像im沿三个轴（0、1、2）的最后一个RGB切片。

我们甚至可以不用该函数而只写：

np.apply_along_axis(lambda rgb: color2index[tuple(rgb)], 2, im)

Answer 7

如果您喜欢使用 MATLAB - 也许将结果保存为 *.mat 并使用 scipy.io.loadmat 加载 - MATLAB 中有 rgb2ind 函数，它完全符合您的要求。如果没有，它可以作为 Python 中类似实现的灵感。

将RGB图像转换为索引图像

7 个答案: