Tensorflow查找具有匹配值的像素

Question

我正在尝试创建用于训练2类语义分段网络的数据批。目标分割图像具有2层，第一层的1类所有像素均具有1层，否则为0层。第二层具有倒置的像素。

在数据集中，我的输出图像是带有[255,255,255]和[0,0,0]的3通道rgb图像。输入和输出图像存储在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图像。

我正在尝试在tensorflow中重复它，我是新手。我尝试了c1_pix = tf.where(tf.equal(op_img, [[[255,255,255]]]))。它似乎可以正常工作，但是它返回1和0的3通道int64张量，而我却无法求反。

有人可以帮我解决这个问题吗？

谢谢

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)

这对我有用。