I am trying to create data batches for training 2 class semantic segmentation network. The target segmented image has 2 layers, first layer with 1 for all pixels of class-1 and 0 otherwise. The second layer has pixels inverted.
In the dataset I have the output images are 3 channel rgb images with [255,255,255] and [0,0,0]. The input and output images are stored in tf-record files.
When I was experimenting in numpy I created a 2 channel binary image with code below:
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)
This gave me a 2 layer 1s and 0s images I wanted.
I am trying to repeat it in tensorflow and I am new to it. I tried c1_pix = tf.where(tf.equal(op_img, [[[255,255,255]]])). It seems to work but it returns 3 channel int64 tensor of 1s and 0s and I am not able to inver it.
Could someone please help me figure this out ?
Thanks,
After going through tf documentation for some time, I came up with the following solution
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)
This worked for me.