从二维整数数组中生成具有固定颜色的彩色图像

Question

Segmentation output gives me a 2D array with each pixel having a unique integer value corresponding to a class. I want to create a coloured image from this array with fixed colours for every class. Please help. If I just stack up the 2D array to create 3 channel image, the image is only different shades of grey for different classes.

Answer 1

You can check the following code and modify according to your need:

import numpy as np
import matplotlib.pyplot as plt

arr = np.array([[2,6,8,9], [1,2,7,8], [4,5,1,7]]) # this may be your array
unique_values = set(np.unique(arr).tolist()) # getting unique classes
colors = [plt.cm.Spectral(each) for each in np.linspace(0, 1, len(unique_values))] # generating random colors for each unique classes

Rarr = np.zeros_like(arr, dtype = 'float64') # Red
Garr = np.zeros_like(arr, dtype = 'float64') # Green
Barr = np.zeros_like(arr, dtype = 'float64') # Blue
for val, col in zip(unique_values, colors):
    Rarr[arr == val ] = col[0]
    Garr[arr == val ] = col[1]
    Barr[arr == val ] = col[2]

rgb = np.dstack((Rarr,Garr,Barr)) # Combining three channels


plt.subplot(1, 4, 1)
plt.imshow(Rarr, 'Reds')
plt.subplot(1, 4, 2)
plt.imshow(Garr, 'Greens')
plt.subplot(1, 4, 3)
plt.imshow(Garr, 'Blues')
plt.subplot(1, 4, 4)
plt.imshow(rgb)
plt.show()