Loss function for a regression model in image segmentation task

I am training a model to segment an image to predict the degree of damage (ranging from 0: no damage, to 5: severe damage) for each pixel of an image. I have approached it this way:

def simple_loss(pred, mask):   # regression case
pred = torch.sigmoid(pred)

def structure_loss(pred, mask):   # binary case: damaged vs undamaged
wbce = (weit * wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))

pred = torch.sigmoid(pred)
inter = ((pred * mask) * weit).sum(dim=(2, 3))
union = ((pred + mask) * weit).sum(dim=(2, 3))
wiou = 1 - (inter + 1) / (union - inter + 1)

return (wbce + wiou).mean()

Binary case yields IoU > 0.6, but the regression model is inaccurate. My datset is imbalanced (100:1) with the majority of the pixels belonging to the undamaged class. Hence, the optimization is driven towards accurate prediction of undamaged pixels. The confusion matrix in the (1..5) region shows no correlation between the label and the predicted value.

I cannot balance the set because the undamaged region next to the damaged area is informative to humans, trained to examine the damage.

How can I modify the loss function to assign higher cost to regression errors regarding the degree of damage?

We can encode irrelevant pixels with -1. Then modify the loss function to ignore irrelevant classes this way:

from keras import backend as K

def f(y_true, y_pred):
return f

y_pred = K.constant([[ 1, 1, 1, 1],
[ 1, 1, 1, 3],
[ 1, 1, 1, 3],
[ 1, 1, 1, 3],
[ 1, 1, 1, 3],
[ 1, 1, 1, 3]])
y_true = K.constant([[ 1, 1, 1, 1],
[ 1, 1, 1, 1],
[-1, 1, 1, 1],
[-1,-1, 1, 1],
[-1,-1,-1, 1],
[-1,-1,-1,-1]])

true = K.eval(y_true)
pred = K.eval(y_pred)