How to add engineered features to an image segmentation model

Question

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.

Answer 1

Adding distance information before the final 1x1 convolutional layer implies that that distance information can be used to linearly separate the two similar classes, given the similar visual features.

If that is the case, then your implementation should work: from my understanding of your task, given an image with class A and class B, it could be that the visual features (d4) for the class A image and class B are similar, but the distance values are different enough to classify them.

You can check to see if this is the case by finding failure modes when the model isn't given the distance information and is struggling to distinguish between the two classes, then seeing if the distance values are different enough in those cases to distinguish the images linearly.

If the relationship between distance and the classes is more complex than that, then you can try adding more layers after your concatenation or adding the distance information earlier in the model