High image segmentation metrics after training but poor results in prediction

Question

I'm trying to build a model with Keras that predicts four classes of features from microscopy noisy images which cover about 10 - 30 % of the image. I'm using U-net because my dataset is small (150 images for training and 30 images for validation). As metrics, I'm using accuracy, loss, intersection-Over-Union and dice coefficient with the following results after 100 epochs of training:

loss: 0.0518 - accuracy: 0.9555 - dice_coef: 0.9480 - iou_coef: 0.9038 - val_loss: 0.0922 - val_accuracy: 0.9125 - val_dice_coef: 0.9079 - val_iou_coef: 0.8503

Unfortunately, when I display the original and the predicted image don't match each other as much as I expected based on the metrics above while it seems that cannot recognize the difference between the classes.

Is that possible ? Below is the code from the metrics and U-net model that I built:

# IOU metric
def iou_coef(y_true, y_pred, smooth=1):
    intersection = K.sum(K.abs(y_true * y_pred), axis=[1,2,3])
    union = K.sum(y_true,[1,2,3])+K.sum(y_pred,[1,2,3])-intersection
    iou = K.mean((intersection + smooth) / (union + smooth), axis=0)
    return iou

def dice_coef(y_true, y_pred, smooth=1):
    intersection = K.sum(y_true * y_pred, axis=[1,2,3])
    union = K.sum(y_true, axis=[1,2,3]) + K.sum(y_pred, axis=[1,2,3])
    dice = K.mean((2. * intersection + smooth)/(union + smooth), axis=0)
    return dice


def unet(pretrained_weights=None, input_size=(IMG_SIZE, IMG_SIZE, 3),num_class=2):
    inputs = Input(input_size)
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
    conv1 = BatchNormalization()(conv1)
    conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
    conv1 = BatchNormalization()(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
    conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
    conv2 = BatchNormalization()(conv2)
    conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
    conv2 = BatchNormalization()(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
    conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
    conv3 = BatchNormalization()(conv3)
    conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
    conv3 = BatchNormalization()(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
    conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
    conv4 = BatchNormalization()(conv4)
    conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
    conv4 = BatchNormalization()(conv4)
    drop4 = Dropout(0.5)(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)

    conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool4)
    conv5 = BatchNormalization()(conv5)
    conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)
    conv5 = BatchNormalization()(conv5)
    drop5 = Dropout(0.5)(conv5)

    up6 = Conv2D(512, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
    UpSampling2D(size=(2, 2))(drop5))
    up6 = BatchNormalization()(up6)
    merge6 = concatenate([drop4, up6], axis=3)
    conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge6)
    conv6 = BatchNormalization()(conv6)
    conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)
    conv6 = BatchNormalization()(conv6)

    up7 = Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
    UpSampling2D(size=(2, 2))(conv6))
    up7 = BatchNormalization()(up7)
    merge7 = concatenate([conv3, up7], axis=3)
    conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge7)
    conv7 = BatchNormalization()(conv7)
    conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)
    conv7 = BatchNormalization()(conv7)

    up8 = Conv2D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
    UpSampling2D(size=(2, 2))(conv7))
    up8 = BatchNormalization()(up8)
    merge8 = concatenate([conv2, up8], axis=3)
    conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge8)
    conv8 = BatchNormalization()(conv8)
    conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)
    conv8 = BatchNormalization()(conv8)

    up9 = Conv2D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal')(
    UpSampling2D(size=(2, 2))(conv8))
    up9 = BatchNormalization()(up9)
    merge9 = concatenate([conv1, up9], axis=3)
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9)
    conv9 = BatchNormalization()(conv9)
    conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
    conv9 = BatchNormalization()(conv9)
    conv9 = Conv2D(num_class, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
    conv9 = BatchNormalization()(conv9)
    conv10 = Conv2D(num_class, 1, activation='softmax')(conv9)
    model = Model(inputs, conv10)

   # Compile the model
    model.compile(loss=dice_loss, optimizer=Adam(lr=1e-4), metrics=['accuracy', dice_coef, iou_coef])

  if (pretrained_weights):
      model.load_weights(pretrained_weights)
  return model

UPDATE

Thanks to the helpful answers and even more research, I think that the issue is in validation data (46 images) because the learning curves for 100 epochs, and learning rate 0.0001 are too noisy:

I tried to shuffle the data and decrease the learning rate to encounter the issue. Thus, I re-run the model with learning rate 0.00001 and 0.000001 but in smaller learning rates while the validation loss and accuracy were less noisy the validation IOU and dice coefficient stucked in 30% in all epochs.

It's worth to mention that I'm using early stopping callback with 'Accuracy' as a monitored parameter but 'Accuracy' is increased in all epochs, thus without early stopping. If I set any metric from validation the model stops in about 7 epochs without a useful result.

Any suggestion ?

Update

After more research and experimentation, I think that one of the main issues that contaminate the results was the batch size. I increased the batch size in a level of 9 and decrease the input size from 512, 512 --> 256, 256 while I decreased the learning rate in a level of 0.001 also. As a result the validation curve fluctuation limited in the first epochs. Specifically, I re-run the model with the following parameters:

• input size: 256, 256, 3
• batch size: 9
• train dataset size: 189
• validation dataset size: 46
• Learning rate: 0.001
• epochs: 64 (stopped with early stopping)

And the final results was:

loss: 0.0813 - accuracy: 0.9320 - dice_coef: 0.9182 - iou_coef: 0.8532 - val_loss: 0.0719 - val_accuracy: 0.9365 - val_dice_coef: 0.9273 - val_iou_coef: 0.8696

But again when I display the original and the predicted image don't match each other as much as I expected based on the metrics from the final experiment.

Is still the model over-fitting ? Why is this happens ?

UPDATE

After the recommendation of @fswings, I re-splitted the data maintaining an equal number of classes in training and validation. My final dataset size was 125 images for training and 34 for validation. The other parameters are exactly the same with the last update above:

• input size: 256, 256, 3
• batch size: 9
• train dataset size: 125
• validation dataset size: 34
• Learning rate: 0.001
• epochs: 44 (stopped with early stopping)

The metrics of the last epoch was: loss: 0.0969 - accuracy: 0.9350 - dice_coef: 0.9022 - iou_coef: 0.8279 - val_loss: 0.0818 - val_accuracy: 0.9299 - val_dice_coef: 0.9175 - val_iou_coef: 0.8537

which looks quite good but again the validation loss and accuracy curves are very noisy:

While the prediction results are quite worse from the last update above.

Any suggestion ?

Answer 1

The coefficients are reported on your 150 training examples? This looks like a textbook example of overfitting: good performance on training data, bad on test data. The U-net model has a large number of parameters that allow to fit your training data perfectly but it doesn't generalize to the examples the model hasn't seen yet.

You need to freeze all parameters of the model except the last couple of layers and see if that helps.

Answer 2

I agree the metrics between your test set and validation set are quite close, but looking at your code it seems you may have run for the full 100 epochs.

keras supports early stopping, i.e. when scores fail to improve meaningfully you can have the model revert to the best scores it has seen to date:

https://stackoverflow.com/q/48285129/1928322

You should also combine this with regular checkpoints should it crash when training, the API details are here:

https://keras.io/api/callbacks/model_checkpoint/

This should mitigate some of the overfitting. After, which it's the hyper parameter tweaking. Don't forget to you need to split the data 3 ways so you can test the model once you've found the 'optimal' one

Update The plots show significant differences between the training and validation data. Having tried different sets of samples the result is the same.

Given that you have 4 classes, I would look closely at which of the 4 classes has the biggest difference. Given the relatively small sample size I suspect their is a class imbalance between the training and validation set. Ie the validation set is not representative of the training data. Formally this can be established when significant changes in training accuracy are not reflected in the validation accuracy.

One possible corrective action is to stratify the splitting of data. The second as a sanity check is to test with a single class or fewer classes.