IOU error

[DanieleParravicini]

def IoU(y_true, y_pred):

    # from probability to predicted class {0, 1}
    y_pred = tf.cast(y_pred > 0.5, tf.float32) # when using sigmoid. Use argmax for softmax

    # A and B

    intersection = tf.reduce_sum(tf.math.multiply(y_true ,y_pred), axis=[1,2])

    # A or B
    union = tf.reduce_sum(y_true, axis=[1,2]) + tf.reduce_sum(y_pred, axis=[1,2]) #- intersection
    # IoU

    return tf.reduce_mean(intersection / union)