MAE Loss Function that penalizes specific colors

Hi everyone,

I’m struggling with an interesting problem. I want to create a loss function based on the principles of MAE but penalizes certain colors more than others. I’m creating a computer vision neural network which should produce images with certain locations marked with red, green or blue pixels. Since these colored pixels are the most important parts of the image I want to increase the loss of the network if it does not create a colored pixel on y_pred in a location which has a circle in y_true.

I created a function capable of doing this on a 2D image (my research uses 3D images) to test the concept. Basically what the function does it compares each pixel in all three channels and then determines whether the pixel is red, green or blue. If so it increases the value in a temporary array from 1 to 100. This temporary array is later multiplied with the original loss. I’m currently stuck at implementing this function in a neural network. The major problem being that Tensors don’t support item assignments and I can’t seem to find a workaround. Any and all advice or tips will be greatly appreciated!

James

def colour_specific_MAE(y, y_hat, pixel_threshold=0.8, penalty=100): #Pixel threshold is used to determine whether the pixel is red, green, blue or just plain white
shape = y.shape
x_scale = shape[1]
y_scale = shape[2]

MAE = np.absolute(y - y_hat)

RED = np.ones_like(y[:, :, 0].shape)
RED[(y[:, :, 0] >= pixel_threshold) & (y[:, :, 1] < pixel_threshold) & (y[:, :, 2] < pixel_threshold)] = penalty  

BLUE = np.ones_like(y[:, :, 1].shape)
BLUE[(y[:, :, 1] >= pixel_threshold) & (y[:, :, 0] < pixel_threshold) & (y[:, :, 2] < pixel_threshold)] = penalty

GREEN = np.ones_like(y[:, :, 2].shape)
GREEN[(y[:, :, 2] >= pixel_threshold) & (y[:, :, 0] < pixel_threshold) & (y[:, :, 1] < pixel_threshold)] = penalty

MAE[:, :, 0] = MAE[:, :, 0] * RED
MAE[:, :, 1] = MAE[:, :, 1] * BLUE
MAE[:, :, 2] = MAE[:, :, 2] * GREEN

return np.sum(MAE) / ( x_scale * y_scale * 3)