将图像分割 class 输出转换为张量流中的图像

Question

我有图像分割问题，我在 (num_class,3) 数组中有我的 class 列表，其中包含每个 class 的颜色。在 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)

或者，您可以计算 predictions 张量并使用 numpy 运算。

将图像分割 class 输出转换为张量流中的图像

convert image-segmentation class output to image in tensorflow

python

machine-learning

image-segmentation

tensorflow