不要迭代像素!迭代颜色,并使用 NumPy 的boolean index arrays。
那将是一个纯粹的 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 的图像mode P 创建索引图像,并以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()
输出:
----------------------------------------
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
----------------------------------------