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 слоя изображения 1 и 0, которые я хотел.

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

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

Спасибо,

Answer 1

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

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)

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