Migrate `IoU` segmentaiton metrics from `keras.metrics` to `keras_cv.metrics`

[innat]

Intersection-Over-Union is a common evaluation metric for semantic image segmentation.

• tf.keras.metrics.MeanIoU
• tf.keras.metrics.IoU
• tf.keras.metrics.BinaryIoU
• tf.keras.metrics.OneHotMeanIoU
• tf.keras.metrics.OneHotIoU

tf.keras.metrics.OneHotIoU(
    num_classes: int,
    target_class_ids: Union[List[int], Tuple[int, ...]],
    name=None,
    dtype=None,
    ignore_class: Optional[int] = None,
    sparse_y_pred: bool = False,
    axis: int = -1
)

[DavidLandup0]

Small note - MeanIoU has a known bug that most people fix by overriding it with the fix. We might want to do that while porting it over

[DavidLandup0]

It's reported here: https://github.com/tensorflow/tensorflow/issues/32875 Sparse categorical crossentropy doesn't play well with it, apparently.

Most fixes boil down to subclassing it modifying one line that throws the exception:

class MeanIoU(tf.keras.metrics.MeanIoU):
  def __init__(self,
               y_true=None,
               y_pred=None,
               num_classes=None,
               name=None,
               dtype=None):
    super(UpdatedMeanIoU, self).__init__(num_classes = num_classes,name=name, dtype=dtype)

  def update_state(self, y_true, y_pred, sample_weight=None):
    y_pred = tf.math.argmax(y_pred, axis=-1)
    return super().update_state(y_true, y_pred, sample_weight)

[LukeWood]

An easy solution might be to upstream the fixes then just import these metrics from keras in keras_cv.metrics.

[tanzhenyu]

It's reported here: tensorflow/tensorflow#32875 Sparse categorical crossentropy doesn't play well with it, apparently.

Most fixes boil down to subclassing it modifying one line that throws the exception:

class MeanIoU(tf.keras.metrics.MeanIoU):
  def __init__(self,
               y_true=None,
               y_pred=None,
               num_classes=None,
               name=None,
               dtype=None):
    super(UpdatedMeanIoU, self).__init__(num_classes = num_classes,name=name, dtype=dtype)

  def update_state(self, y_true, y_pred, sample_weight=None):
    y_pred = tf.math.argmax(y_pred, axis=-1)
    return super().update_state(y_true, y_pred, sample_weight)

Depending on what we want. So if this metrics is particularly for segmentation, we should have it at keras_cv. The Keras version MeanIOU is intended for a wider usage which doesn't really think about the segmentation map tensor shape. So I'd say in this case, we should do the migration.

[LukeWood]

It's reported here: tensorflow/tensorflow#32875 Sparse categorical crossentropy doesn't play well with it, apparently. Most fixes boil down to subclassing it modifying one line that throws the exception:

class MeanIoU(tf.keras.metrics.MeanIoU):
  def __init__(self,
               y_true=None,
               y_pred=None,
               num_classes=None,
               name=None,
               dtype=None):
    super(UpdatedMeanIoU, self).__init__(num_classes = num_classes,name=name, dtype=dtype)

  def update_state(self, y_true, y_pred, sample_weight=None):
    y_pred = tf.math.argmax(y_pred, axis=-1)
    return super().update_state(y_true, y_pred, sample_weight)

Depending on what we want. So if this metrics is particularly for segmentation, we should have it at keras_cv. The Keras version MeanIOU is intended for a wider usage which doesn't really think about the segmentation map tensor shape. So I'd say in this case, we should do the migration.

Yeah gotcha, that makes sense to me. We also could use a bounding box variant of IOU too.

[tanzhenyu]

As of tf 2.10, this is fixed by accepting sparse_y_true and sparse_y_pred -- @innat can you verify it before I close this?

[innat]

@tanzhenyu I'm on leave. I will check and get back to you soon.

[innat]

@tanzhenyu Could you please elaborate more? Is it not needed to migrate these metrices from keras to keras-cv for now?

[tanzhenyu]

@tanzhenyu Could you please elaborate more? Is it not needed to migrate these metrices from keras to keras-cv for now?

I have been using the MeanIOU and haven't run into any issues yet, see this PR. Are there any other functionality required, or this is a pure migration?

[innat]

Yes, this is for migration only. @DavidLandup0 pointed a issue though, but that is not directly related to this ticket.

[tanzhenyu]

Yes, this is for migration only. @DavidLandup0 pointed a issue though, but that is not directly related to this ticket.

I'd like to understand a bit more about the issue, if this is worth fixing in KerasCV then we can do it. @DavidLandup0

[DavidLandup0]

The issue with the class was that it was expecting exclusively categorical vectors, so when using sparse categorical crossentropy and an output fit for it - there would be a shape mismatch. The common fix was to simply get the argmax and thus convert it into a sparse representation.

However, as @tanzhenyu pointed out in the previous comment:

As of tf 2.10, this is fixed by accepting sparse_y_true and sparse_y_pred -- @innat can you verify it before I close this?

Not sure if it's a better API choice to have MeanIoU(arg) or SparseMeanIoU and CategoricalMeanIoU. KerasCV losses have gone with the latter format so far, which also allows for more specialization down the line.

[ianstenbit]

These metrics are going to continue to be widely available in Keras Core moving forward, so I don't think there's a strong argument for having a second copy in KerasCV