# Tversky Loss function for RGB masks

I have a very imbalanced dataset for my semantic segmentation problem (monitoring deforestation using setellite images) and I found Tversky Loss to be much better than categorical crossentropy (due to dataset imbalance). My problem is that the masks for my images are in RGB. Here's an example of an image and a matching mask: If I understand correctly, our masks should have the following form: H x W x Num_classes. In my case, I have 3 classes(forest - green, deforestation - red, other - blue) and the masks are 512 x 512 x 3, but instead of pixel values like this [1; 0; 0], [0; 1; 0], [0; 0; 1] I have [226,19,15] (red), [11,195,72] (green), [51, 15, 200] (blue). And categorical cross-entropy works with it, but I want to try Tversky Loss and with such masks it is impossible. It looks like a one-hot encoding issue to me, but I'm not sure.

Here is an example of code for my Tversky Loss (https://github.com/keras-team/keras/issues/9395 taken from here):

ALPHA = 0.5
BETA = 0.5
# First example
def TverskyLoss(targets, inputs, alpha=ALPHA, beta=BETA, smooth=1e-6):
#flatten label and prediction tensors
inputs = tf.squeeze(inputs)
targets = tf.squeeze(targets)

#True Positives, False Positives & False Negatives
TP = tf.reduce_sum((inputs * targets))
FP = tf.reduce_sum(((1-targets) * inputs))
FN = tf.reduce_sum((targets * (1-inputs)))

Tversky = (TP + smooth) / (TP + alpha*FP + beta*FN + smooth)

return 1 - Tversky
# Second example
def tversky_loss(y_true, y_pred):
alpha = 0.5
beta  = 0.5

ones = tf.ones(tf.shape(y_true))
p0 = y_pred      # proba that voxels are class i
p1 = ones - y_pred # proba that voxels are not class i
g0 = y_true
g1 = ones - y_true

num = tf.math.reduce_sum(p0 * g0, (0,1,2))
den = num + alpha*tf.math.reduce_sum(p0*g1,(0,1,2)) + beta*tf.math.reduce_sum(p1*g0,(0,1,2))

T = tf.math.reduce_sum(num/den) # when summing over classes, T has dynamic range [0 Ncl]

Ncl = tf.cast(tf.shape(y_true)[-1], 'float32')
return Ncl-T


Here is an example of predicting a model while training using Tversky Loss. Obviously the loss function doesn't work: I am using standard U-net architecture, with this output layer (3 filters for 3 output classes). Looks like I should change activation to "sigmoid" if I'll convert my masks:

conv = Conv2D(3, 1, padding='same', activation='softmax')(conv_prev)


What should I do to use my Tversky Loss, convert each mask to values like [1; 0; 0] and to see my results do the inverse conversion to RGB colors like this [226,19,15]? Is it possible to do such transformations using tf.keras.preprocessing.image.ImageDataGenerator? And what if I had 5 or 10 classes, how to view results of output tensor with shape H x W x 10 ? I tried to figure out this problem myself and in this tutorial (https://www.tensorflow.org/tutorials/images/segmentation) the masks also had 3 classes. Thanks for any help in advance :)

def preprocess_one_hot_encode(image_rgb):