I have built a U-net model for image segmentation of 3-channel remote sensing images. I have a total of four classes; two of these classes look very similar and are hard to distinguish in the images without extra context. In this case, extra context would be the distance to something (i.e., distinguishing very similar looking buildings based on their proximity to a main road).
To add more context I've created an extra channel of values representing the distance to the nearest road, to try to improve segmentation results, but haven't got the results I expected. What I'm currently doing is using the RGB channels as normal in the U-net, then concatenating the 4th 'distance' channel before the final convolution layer, like so:
def conv_block(input, num_filters):
x = Conv2D(num_filters, 3, dilation_rate = 3, padding="same")(input)
x = BatchNormalization()(x)
x = Activation("relu")(x)
return x
def decoder(inputs, skip_features, num_filters):
up = Conv2DTranspose(num_filters, (2,2), strides=2, padding='same')(inputs)
x = conv_block(up, num_filters)
return x
inputs = Input((size,size,4))
x = Rescaling(1/255.0)(inputs)
sliced_input_1 = x[:, :, :, :3]
'''pretrained encoder'''
xception = Xception(include_top = False, weights = 'imagenet', input_tensor = sliced_input_1)
xception.trainable = False
s1 = xception.get_layer("input_1").output ##512
s2 = xception.get_layer("block1_conv1_act").output
s2 = ZeroPadding2D(( (1, 0), (1, 0) ))(s2) ##256
s3 = xception.get_layer("block3_sepconv2_bn").output
s3 = ZeroPadding2D(( (1, 0), (1, 0) ))(s3) ##128
s4 = xception.get_layer("block4_sepconv2_bn").output ##64
'''bridge'''
b1 = xception.get_layer("block13_sepconv2_bn").output ##32
'''decoder'''
d1 = decoder(b1, s4, int(size/1))
d2 = decoder(d1, s3, int(size/2))
d3 = decoder(d2, s2, int(size/4))
d4 = decoder(d3, s1, int(size/8))
'''outputs'''
sliced_input2 = tf.expand_dims(x[:, :, :, 3], axis=-1)
outputs = Concatenate()([sliced_input2, d4])
outputs = Conv2D(4, (1,1), padding='same', activation='softmax')(outputs)
model = Model(inputs, outputs)
model.summary()
The shape of the tensor before the final convolution is (512,512,65), which is the previous 64 filters from the RGB U-net and a concatenated 'distance' layer, but should the number of filters from the U-net be reduced even more prior to concatenation? Since it seems like the 'distance' layer is still being ignored.