tensorflow2.0 语义分割 评价指标代码（tensor版）

np编写的语义分割评价指标不能通过GPU 加速，故改编了一份tensor的代码，在tf2.0测试通过。第一次写博客，将就着看吧

def get_hist_t(predictions, labels, num_class):
# labels:[b.h,w] tensor
# predictions:[b.h,w,c] tensor
num_class = predictions.shape[3]
batch_size = predictions.shape[0]
hist = tf.zeros((num_class, num_class), dtype=‘int32’)

for i in range(batch_size):
    hist += fast_hist_t(tf.squeeze(tf.reshape(labels[i], (-1, 1)), 1),
                        tf.squeeze(tf.reshape(tf.argmax(predictions[i], axis=2), (-1, 1)), 1), num_class)
return hist

def per_class_acc_t(predictions, label_tensor):
labels = label_tensor
size = predictions.shape[0]
num_class = predictions.shape[3]
hist = tf.zeros((num_class, num_class), dtype=‘int32’)

#     num_class = predictions.shape[3]
#     batch_size = predictions.shape[0]
#     hist =tf.zeros((num_class, num_class),dtype='int32')
for i in range(size):
    hist += fast_hist_t(tf.squeeze(tf.reshape(labels[i], (-1, 1)), 1),
                        tf.squeeze(tf.reshape(tf.argmax(predictions[i], axis=2), (-1, 1)), 1), num_class)
#           hist += fast_hist(labels[i].flatten(), predictions[i].argmax(2).flatten(), num_class)

acc_total = tf.reduce_sum(tf.linalg.tensor_diag_part(hist)) / tf.reduce_sum(hist)
print('accuracy = %f' % np.nanmean(acc_total))
iu = tf.linalg.tensor_diag_part(hist) / (
            tf.reduce_sum(hist, 1) + tf.reduce_sum(hist, 0) - tf.linalg.tensor_diag_part(hist))
print('mean IU  = %f' % np.nanmean(iu))
for ii in range(num_class):
    if float(tf.reduce_sum(hist, 1)[ii]) == 0:
        acc = 0.0
    else:
        acc = tf.cast(tf.linalg.tensor_diag_part(hist)[ii], dtype='float32') / tf.cast(tf.reduce_sum(hist, 1)[ii],
                                                                                       dtype='float32')
    print("    class # %d accuracy = %f " % (ii, acc))

def fast_hist_t(a, b, n):#tensor版
k = (a >= 0) & (a < n)
d=tf.math.bincount(n * tf.cast(a[k],dtype=‘int32’) + tf.cast(b[k],dtype=‘int32’), minlength=n2)
# return np.bincount(n * a[k].astype(int) + b[k], minlength=n2).reshape(n, n)
return tf.reshape(d,(n, n))