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

Question

I have a dictionary that maps classes to RGB values:

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]
}

I'm replacing the classes in the image with RGB values like below:

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])

Is there any more efficient way of doing this?

Answer 1

Don't iterate the pixels! Iterate the colors, and use NumPy's boolean index arrays .

That'd be a pure NumPy solution:

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()

The corresponding output:

Since you have python-imaging-library in your tags, there'd be also the option to use Pillow's image mode P to create an indexed image, and provide the colors as a palette :

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()

The output:

----------------------------------------
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
----------------------------------------