Оценка Tensorflow изображения один на канале

Question

Я использую Keras с бэкэндом Tensorflow и хочу оценить PSNR между двумя изображениями. Я могу оценить все три канала RGB, например:

def psnr(hr, sr):
    return tf.image.psnr(hr, sr, max_val=255)

с функцией psnr от tenorflow ( tf.image.psnr

Но что бы я сделал, чтобы оценить только на первом канале? Я предполагаю, что мне нужно извлечь эти значения из тензора каким-то образом. В python обычно можно сделать что-то вроде hr [:, 0, 0], но это, очевидно, здесь не работает.

Answer 1

im1 = tf.image.convert_image_dtype(np.random.randn(64,64,3), tf.float32)
im2 = tf.image.convert_image_dtype(np.random.randn(64,64,3), tf.float32)

psnr = tf.image.psnr(tf.expand_dims(im1[:,:,0], 2), 
                     tf.expand_dims(im2[:,:,0], 2), max_val=255)

with tf.Session() as sess:
    print (sess.run(psnr))

Получите первый канал, используя im1[:,:,0], и измените его до h x w x 1, добавив один канал, используя expand_dims