Tensorflow находит пиксели с соответствующим значением

Question

Я пытаюсь создать пакеты данных для обучения сети семантической сегментации 2 класса.Целевое сегментированное изображение имеет 2 слоя, первый слой с 1 для всех пикселей класса 1 и 0 в противном случае.Второй слой имеет инвертированные пиксели.

В наборе данных, который у меня есть, выходные изображения представляют собой 3-канальные изображения RGB с [255,255,255] и [0,0,0].Входные и выходные изображения хранятся в файлах tf-записи.

Когда я экспериментировал с numpy, я создал 2-канальное двоичное изображение с кодом ниже:

c1_pix = np.all(op_img == np.array([255,255,255]), axis=2)
c1_pix = c1_pix.reshape(*(h,w), 1)
op_arr = np.concatenate((c1_pix, np.invert(c1_pix)), axis=2)

Это дало мне 2изображения слоев 1s и 0s, которые я хотел.

Я пытаюсь повторить его в тензорном потоке, и я новичок в этом.Я пытался c1_pix = tf.where(tf.equal(op_img, [[[255,255,255]]])).Вроде работает, но возвращает 3-х канальный int64 тензор 1 с и 0, и я не могу его инвертировать.

Может кто-нибудь помочь мне разобраться?

Спасибо,

Answer 1

Пройдя некоторое время через документацию по tf, я нашел следующее решение

c1_pix, _, _ = tf.split(tf.equal(op_img,[[[255,255,255]]]), 3, axis=-1)
c2_pix = tf.logical_not(c1_pix)
new_op_img = tf.concat([c1_pix, c2_pix], -1)

Это сработало для меня.