python - 将图像分割类输出转换为 tensorflow 中的图像

Question

我有图像分割问题，我的类列表在 (num_class,3) 数组中，其中包含每个类的颜色。在 u-net 之后，我将得到一个 (width,height,num_class) 形状的概率张量，我想将其转换为图像 (width,height,3)。我该怎么做？

class_colors=[[128,0,0],[0,128,0],...] #(num_class,3)
logit=unet(img) # (W,H,num_class)
probs=tf.nn.softmax(logit)
predictions=tf.argmax(probs)
prediction_image= ? # (W,H,3)

Answer 1

您可以使用函数tf.gather_nd，但首先需要将class_colors声明为 tensorflow 变量。检查以下示例(图像大小 50x50，2 个类别):

import tensorflow as tf

predictions = tf.argmax(tf.nn.softmax(tf.random_normal([50,50,2])),axis=-1) #(50,50)
class_colors = tf.Variable([[255,0,0],[0,255,0]]) #(2,3)
prediction_image = tf.gather_nd(class_colors, tf.expand_dims(predictions,axis=-1)) #(50,50,3)

sess = tf.Session()
sess.run(tf.global_variables_initializer())
print(sess.run(prediction_image).shape) #(50, 50, 3)

或者，您可以评估预测张量并使用 numpy 运算。