преобразовать вывод класса сегментации изображения в изображение в тензорном потоке

Question

У меня проблема с сегментацией изображения, у меня есть список моих классов в массиве (num_class,3), который содержит цвет каждого класса.после сети у меня будет тензор вероятности по форме (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 как переменную тензорного потока. Проверьте следующий пример (размер изображения 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 и использовать пустые операции.