Reference

https://stackoverflow.com/a/46667294

the first dim is always the batch size, so the gt and pr have a dimension 3 of (B, H*W, C)

from keras.metrics import categorical_accuracy

import keras.backend as K
import tensorflow as tf

def custom_categorical_accuracy(gt, pr):
    mask = 1 - gt[:,:,0]
    gt = tf.boolean_mask(gt,mask)
    pr = tf.boolean_mask(pr,mask)

    return K.mean(K.equal(K.argmax(gt, axis=-1),
                          K.argmax(pr, axis=-1)))