IOU metric capped at 0.49 in a two class segmentation problem

Hi, I'm working on a two class problem (background & region of interest). I need to set the loss function's weights to [0, 1] to ensure the model is learning from the region of interest, otherwise it overfits to the background and does not generalize well (I'm using synthetic images and the background does not have enough variety).

However, setting the weights like this results in incorrect IOU score (I'm using sm.metrics.IOUScore), meaning the IOU will be capped at 0.49 during training, while the performance of the trained model on actual images is almost perfect (> 0.90 IOU).

Feeding the class weights to the metric doesn't change anything; if I set the class_indexes to [1] the IOU metric will remain close to zero, if I set the class_indexes to [0] the IOU metric will look like as if it was subtracted from 1 (meaning it goes down from 0.065 to 0.028 as the model is training). If I set the loss function's weights to [0.01, 0.99] then the IOU metric works as expected and it goes above 0.92 during training.

This is a cosmetic issue as the models do train and are working, I just would like to understand why this is happening. Any advice would be appreciated.

The background is encoded as 0 and the RoI as 1 in the segmentation masks.

versions:

segmentation_models 1.0.1
tensorflow 2.3.4

simplified code:

SM_FRAMEWORK=tf.keras NCLASSES = 2

params={ 'backbone':'efficientnetb5', 'loss': [sm.losses.JaccardLoss(class_weights=np.array([0,1]))], 'weights':[1], 'augmentation': True, 'weights_pretrained':'imagenet', 'batch_size': 256, 'steps_per_epoch': 237, 'n_epochs': 26, 'encoder_freeze': False, 'learning_rate': 0.000073, 'activation' : 'sigmoid', 'metrics' : ['accuracy',sm.metrics.IOUScore()], }

self.backbone=params['backbone'] self.loss =params['loss'] self.weights =params['weights'] self.augmentation =params['augmentation'] self.weights_pretrained =params['weights_pretrained'] self.batch_size =params['batch_size'] self.steps_per_epoch =params['steps_per_epoch'] self.n_epochs =params['n_epochs'] self.encoder_freeze =params['encoder_freeze'] self.learning_rate = params['learning_rate']

callbacks = [ keras.callbacks.ModelCheckpoint(save_to, save_weights_only=False,save_freq='epoch'), csv_logger ]

self.model = sm.Linknet(self.backbone, encoder_weights=self.weights_pretrained,classes=NCLASSES, activation=self.activation, encoder_freeze=self.encoder_freeze)

self.model.compile('adam', self.loss, self.metrics, loss_weights=self.weights)

self.model_history=self.model.fit(self.train_generator, epochs = self.n_epochs, steps_per_epoch = self.steps_per_epoch, validation_data = self.val_generator, validation_steps = 1, callbacks=callbacks)

qubvel / segmentation_models

IOU metric capped at 0.49 in a two class segmentation problem #508