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