How to create a loss function in Keras that uses opencv?

Question

I am developing a machine Learning model using the Keras library and realize that the available loss functions are not giving the best results in my test suite.

I’m using an Unet architecture, where I enter an image (16,16,3) and the network also generates an image (16,16,3) (auto-Ncoder). I realized that perhaps a way to improve the model would be if I used a loss function that compares pixel to pixel in the image gradients (Laplaciano) between the network output and the Ground Truth set. However, I didn’t find any tutorial that could handle this type of application because it would need to use opencv’s Laplacian function on each network output image.

The loss function would be something like this:

def laplacian_loss(y_true, y_pred):

  # y_true already is the calculated gradients, only needs to compute on the y_pred
  # calculates the gradients for each predicted image
  y_pred_lap = []
  for img in y_pred:
    laplacian = cv2.Laplacian( np.float64(img), cv2.CV_64F )
    y_pred_lap.append( laplacian )

  y_pred_lap = np.array(y_pred_lap)

  # mean squared error, according to keras losses documentation
  return K.mean(K.square(y_pred_lap - y_true), axis=-1)

Someone has done something similar for the loss calculation?

Answer 1

I was able to come up with a very easy solution. The main issue was that the gradient filter is nothing more than a 2D filter applied to the image. For further information on the filter, refer to the documentation of the Opencv. Therefore, it is necessary that the result of my network is filtered by a filter as described in the documentation. To do so, just create an extra convolutional layer at the end of the network and have the values of the Laplaciano kernel. After that, as the network will have two outputs (one being the image you want to get, and the other the image of the gradients), it is necessary to define the losses for both outputs that, automatically, Keras will propagate to the whole network.

To be more clear what should be done, I will exemplify with code. So, at the end of your network you will do something like:

channels = 3 # number of channels of network output  
lap = Conv2D(channels , (3,3), padding='same', name='laplacian') (net_output)
model = Model(inputs=[net_input], outputs=[net_out, lap])

Compile your own template:

model.compile(optimizer=Adam(), loss=losses, loss_weights=lossWeights)

Define the Laplaciano kernel and apply its values to the convolutional layer weights and specify that the layer should not be trained:

bias = np.asarray([0]*3)
# laplacian kernel
l = np.asarray([
        [[[1,1,1],
        [1,-8,1],
        [1,1,1]
    ]]*channels
    ]*channels).astype(np.float32)
bias = np.asarray([0]*3).astype(np.float32)
wl = [l,bias] 
model.get_layer('laplacian').set_weights(wl)
model.get_layer('laplacian').trainable = False

Define how the losses will be calculated for each of the outputs and their weights (if necessary):

# losses for output, laplacian and gaussian
losses = {
    "enhanced": "mse",
    "laplacian": "mse"
}
lossWeights = {"enhanced": 1.0, "laplacian": 0.6}

When training, remember that you will need two values for y, ie:

model.fit(x=X, y = {"out": y_out, "laplacian": y_lap})

Observing: Do not use the Batchnormalization layer! If you use it, the weights of the Batchnormalization layer will be updated (which should not occur!).