请问预测得到的mask怎么上色？

[wcl6]

预测出来的图像是灰度图，怎么上不同的颜色

[moguiyan]

同问

[LiNinghui-AI]

import cv2
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap

#在自建的colorabar ，得到自己想要的颜色
def get_spectral():

ALL = 0
VEGETATION = 1
ROAD = 4
BUILDING = 2
WATER = 3
colormap_float = np.zeros( (5, 3), np.float ) #新建一个数组用于存储颜色的数值

colormap_float[ALL, :] = [255, 255, 255]
colormap_float[WATER, :] = [34, 180, 238]
colormap_float[VEGETATION, :] = [159, 255, 84]
colormap_float[BUILDING, :] = [255, 191, 0]
# colormap_float[BUILDING, :] = [34, 180, 238]
colormap_float[ROAD, :] = [38, 71, 139]

return colormap_float

matplotlib.use("Qt5Agg")
path = Satellite-Segmentation-master/segnet/pre1.png"
p_mask = cv2.imread(path)

colormap_float = get_spectral()
rgb_table = LinearSegmentedColormap.from_list('sst cmap', colormap_float/255)

import cmaps
plt.imshow((p_mask[..., 0]), cmap=rgb_table)