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

时间:2019-05-15 13:47:07

标签: python tensorflow machine-learning image-segmentation

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

1 个答案:

答案 0 :(得分:1)

您可以使用函数tf.gather_nd,但首先需要将class_colors声明为张量流变量。检查以下示例(图像尺寸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操作。