wac.py

def wac_score(y_true, y_pred):
    """
    Parameters
    ----------
    y_true: numpy.ndarray
        True labels
    y_pred: numpy.ndarray
        Estimated labels
    Returns
    -------
    wac_score: float
    """
    cm = confusion_matrix(y_true, y_pred)
    if cm.shape != (2, 2):
        return accuracy_score(y_true, y_pred)

    tp, fn, fp, tn = cm[1, 1], \
                     cm[1, 0], \
                     cm[0, 1], \
                     cm[0, 0]

    if tp == 0 and fn == 0:
        return 0.5*tn/float(tn+fp)

    elif tn == 0 and fp == 0:
        return 0.5*tp/float(tp+fn)

    return 0.5*tp/float(tp+fn) + 0.5*tn/float(tn+fp)