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

我有图像分割问题，我在 (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)

最佳答案

您可以使用 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)

或者，您可以评估 predictions 张量并使用 numpy 操作。

关于python - 将图像分割类输出转换为 tensorflow 中的图像，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/56150958/