Intersection over union (IOU) metric for multi-class semantic segmentation task

Question

I have a semantic segmentation task to predict 5 channel mask using UNET for example mask shape is (224,244,5).

I'm using this function for IOU :

    def mean_iou(y_true, y_pred):
        y_pred = tf.round(tf.cast(y_pred, tf.int32))
        intersect = tf.reduce_sum(tf.cast(y_true, tf.float32) * tf.cast(y_pred, tf.float32), axis=[1])
        union = tf.reduce_sum(tf.cast(y_true, tf.float32),axis=[1]) + tf.reduce_sum(tf.cast(y_pred, tf.float32),axis=[1])
        smooth = tf.ones(tf.shape(intersect))
        return tf.reduce_mean((intersect + smooth) / (union - intersect + smooth))

 def iou_loss(y_true, y_pred):
      y_true = tf.reshape(y_true, [-1])
      y_pred = tf.reshape(y_pred, [-1])
      intersection = tf.reduce_sum(tf.cast(y_true, tf.float32) * tf.cast(y_pred, tf.float32))
      score = (intersection + 1.) / (tf.reduce_sum(tf.cast(y_true, tf.float32)) + 
      tf.reduce_sum(tf.cast(y_pred, tf.float32)) - intersection + 1.)
      return 1 - score`

and the output layer of UNET model :

  outputs = tf.keras.layers.Conv2D(5, (1, 1), activation='softmax')(c9)
   model = tf.keras.Model(inputs=[input_img], outputs=[outputs])
   opti =  tf.keras.optimizers.Adam(lr=0.003, clipvalue=0.7)
   model.compile(optimizer=opti, loss=iou_loss, metrics=['accuracy',mean_iou])

But I'm not sure is the IOU function the correct implementation or not,

could you clarify this.

score 4 · Accepted Answer · edited Nov 16 '21 at 11:43

Let us break it into smaller parts to understand what is happening:

tf.reshape(y_true, [-1]),y_pred = tf.reshape(y_pred, [-1]); The predictions and ground truths are converted into a 1-dimensional array. This can take place because, in essence, although you have multiple ground truth masks, they are all composed of only 1s and 0s.
intersection = tf.reduce_sum(tf.cast(y_true, tf.float32) * tf.cast(y_pred, tf.float32)). We multiply here since multiplication will yield 1 only if both the predictions and the ground truth have 1s on a certain position... 0x1 or 1x0 or 0x0 of course do not belong to intersection.
tf.reduce_sum(). We just sum up the 1s from the intersection.
score = (intersection + 1.) / (tf.reduce_sum(tf.cast(y_true, tf.float32)) + tf.reduce_sum(tf.cast(y_pred, tf.float32)) - intersection + 1.). This is the definition of IoU. Note that 1. is added both to the numerator and denominator to avoid division by 0. At the denominator level, since the Union operation in itself already contains the intersection, in order to correctly compute the IoU, we need to remember to subtract the intersection, thus yielding the correct IoU value.
1 - score. We return 1-score, since if the IoU is 0.75, for example, the loss is 0.25 (perfect IoU == 1)

thanks very much, @timbus Calin for clarification, if I wanted to track the losses for the IOU and accuracy, what loss should I monitor, CategoricalCrossentropy() or IOU loss what happens if I choose either one of them? — N.zay, Jan 31 '21 at 11:25
You should track both, but pay a greater attention to IoU loss if your dataset is imbalanced; the accuracy metric may be misleading in case your dataset is very imbalanced.\ — Timbus Calin, Jan 31 '21 at 11:40

Intersection over union (IOU) metric for multi-class semantic segmentation task

1 Answers1