python - 如何有效地将图像标签转换为 RGB 值？

Question

我有一本将类映射到 RGB 值的字典：

label_to_color = {
    0: [128,  64, 128],
    1: [244,  35, 232],
    2: [ 70,  70,  70],
    3: [102, 102, 156],
    4: [190, 153, 153],
    5: [153, 153, 153]
}

我将图像中的类替换为 RGB 值，如下所示：

def mask2pixel(image):
    h,w = image.shape
    im = np.zeros((h,w,3))
    for i in range(h):
        for j in range(w):
            im[i,j,:] = label_to_color[image[i,j]]
    return im.astype(int)

image = cv2.imread(im_path,-1)

print(image.shape) # 1200x720
print(image[0,0]) # 0

colormap = mask2pixel(image)

print(colormap.shape) # 1200x720x3
print(colormap[0,0]) # array([128, 64,128])

有没有更有效的方法来做到这一点？

标签： pythonnumpyopencvimage-processingpython-imaging-library

Answer 1

不要迭代像素！迭代颜色，并使用 NumPy 的布尔索引数组。

那将是一个纯粹的 NumPy 解决方案：

import numpy as np

label_to_color = {
    0: [128,  64, 128],
    1: [244,  35, 232],
    2: [ 70,  70,  70],
    3: [102, 102, 156],
    4: [190, 153, 153],
    5: [153, 153, 153]
}

img_classes = np.random.randint(0, 6, (20, 30), dtype=np.uint8)
h, w = img_classes.shape
img_rgb = np.zeros((h, w, 3), dtype=np.uint8)

for gray, rgb in label_to_color.items():
    img_rgb[img_classes == gray, :] = rgb

import matplotlib.pyplot as plt
plt.figure(figsize=(16, 8))
plt.subplot(1, 2, 1), plt.imshow(img_classes)
plt.subplot(1, 2, 2), plt.imshow(img_rgb)
plt.tight_layout(), plt.show()

对应的输出：

由于您的标签中有python-imaging-library，因此还可以选择使用 Pillow 的图像模式P来创建索引图像，并将颜色作为调色板提供：

import numpy as np
from PIL import Image

label_to_color = {
    0: [128,  64, 128],
    1: [244,  35, 232],
    2: [ 70,  70,  70],
    3: [102, 102, 156],
    4: [190, 153, 153],
    5: [153, 153, 153]
}

img_classes = np.random.randint(0, 6, (20, 30), dtype=np.uint8)
img_classes = Image.fromarray(img_classes, 'L')

img_classes_p = img_classes.convert('P')
img_classes_p.putpalette(
    [rgb for pixel in label_to_color.values() for rgb in pixel])

img_rgb = img_classes_p.convert('RGB')

import matplotlib.pyplot as plt
plt.figure(figsize=(24, 8))
plt.subplot(1, 3, 1), plt.imshow(img_classes)
plt.subplot(1, 3, 2), plt.imshow(img_classes_p)
plt.subplot(1, 3, 3), plt.imshow(img_rgb)
plt.tight_layout(), plt.show()

输出：

----------------------------------------
System information
----------------------------------------
Platform:      Windows-10-10.0.19041-SP0
Python:        3.9.1
PyCharm:       2021.1.2
Matplotlib:    3.4.2
NumPy:         1.20.3
Pillow:        8.2.0
----------------------------------------