IoU.result raises an axis error when axis is not the last dimension
Description
IoU.result() reduces the per class IoU vector with ops.sum(iou, axis=self.axis):
iou is one dimensional, one entry per target class, so any axis other than 0 or -1 raises an out of bounds error. self.axis is the class dimension of the inputs and is only meaningful for the argmax in update_state. The bug affects IoU, MeanIoU, OneHotIoU and OneHotMeanIoU whenever the inputs are channels first and axis=1 is passed, which is exactly the case the axis argument exists for. update_state succeeds and the error only shows up in result(), so it also breaks model.fit and model.evaluate with these metrics at the end of the first epoch.
Reproduction
import numpy as np
import keras
# Class probabilities on axis 1, i.e. channels first predictions of
# shape (batch_size, num_classes, width).
y_true = np.array([[2, 0], [1, 0]])
y_pred = np.array(
[
[[0.2, 0.1], [0.3, 0.2], [0.5, 0.7]],
[[0.5, 0.1], [0.3, 0.4], [0.1, 0.5]],
]
)
m = keras.metrics.MeanIoU(num_classes=3, sparse_y_pred=False, axis=1)
m.update_state(y_true, y_pred) # fine
m.result() # ValueError: axis 1 is out of bounds for array of dimension 1
The JAX, TensorFlow and PyTorch backends all fail the same way. The expected result is 1 / 9: the argmax over axis 1 gives the predictions [[2, 2], [0, 2]], the confusion matrix is [[0, 0, 2], [1, 0, 0], [0, 0, 1]], and only class 2 has a non zero IoU of 1 / 3.
Environment
Keras at master (5edcf00a, version 3.16.0), Python 3.11, CPU, JAX / TensorFlow / PyTorch backends.
Fix
Reducing over the only dimension of iou, i.e. ops.sum(iou), fixes it. I have the one line fix plus regression tests for MeanIoU, OneHotIoU and OneHotMeanIoU with axis=1 ready on a branch, passing on all three backends. I would like to be assigned so I can open the PR.
Source: keras-team/keras