tensorflow根据label图画出对应的热图 语义分割

在使用tensorflow对图像进行分割预测的时候，得到的值对应的是一个二维数组，如果直接将其保存为图片，就是一个黑色的图，区分不大，因此我要将其转成带颜色的图，每个类别对应一种颜色，对应的下图中的左道右。

   

def value_to_RGB(value):
	a=[None]*len(value)
	for i in range(len(value)):
		a[i]=[None]*len(value[i])
		for  j in range(len(value[i])):

			# 设置对应类别的颜色
			if value[i][j]==0:
				#lv zhibei
				a[i][j]=[0,255,127]
			elif value[i][j]==1:
				#hui  daolu 
				a[i][j]=[112,128,144]
			elif value[i][j]==2:
				# lan fangwu 
				a[i][j]=[70,130,180]
			elif value[i][j][0]==3:
				#heliu 
				a[i][j]=[255,215,0]
			else:
				#qita hong 
				a[i][j]=[220,20,60]
			
	a = np.array(a)

	return a

这里由于不能直接对数组中的元素进行修改，因此设计了一个临时列表，最后在转为数组，将其保存为图片就好了。
初学，方法可能比较笨重，欢迎指正。