tensorflow计算psnr与ssim

加载mnist进行测试

import math

import tensorflow as tf
import numpy as np

load_engine = tf.keras.datasets.mnist
# load_engine = tf.keras.datasets.fashion_mnist
(x_train, y_train), (x_test, y_test) = load_engine.load_data()
print(x_train.shape)

#将标签转one_hot编码
y_train = tf.one_hot(y_train,10)
y_test = tf.one_hot(y_test,10)

#扩展一维，将灰度图变成单通道图
x_train = np.expand_dims(x_train.astype(np.float32) / 255.0, axis=-1)  # [60000, 28, 28, 1]
# x_train = tf.expand_dims(x_train,axis=-1)
# x_train = tf.reshape(x_train,(-1,28,28,1))

方便的预处理操作

洗牌、批量、转换等操作都可以。

mnist_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
mnist_dataset = mnist_dataset.batch(4)

计算指标

#自定义计算psnr
def get_psnr2(img1, img2):
   mse = np.mean((img1/1.0 - img2/1.0) ** 2 )
   if mse < 1.0e-10:
      return 100
   return 10 * math.log10(1.0/mse)

测试1

for images, labels in mnist_dataset:    # image: [4, 28, 28, 1], labels: [4]
    print(labels.shape)
    print(images.shape)

    print(images[0].shape)
    print(images[1].shape)
    print(get_psnr2(images[0], images[1]))
    psnr = tf.image.psnr(images[0], images[1], 1)
    print(psnr)
    break

结果正常！

测试2

for images, labels in mnist_dataset:    # image: [4, 28, 28, 1], labels: [4]
    print(labels.shape)
    print(images.shape)

    print(images[0].shape)
    print(images[1].shape)
    print(get_psnr2(images[0], images[0]))
    psnr = tf.image.psnr(images[0], images[0], 1)
    print(psnr)
    break

不会吧不会吧，不会tensorflow连除数为0都不考虑吧。。。
计算psnr需要计算mse，当两张一样的图片时，mse几乎为0，在计算psnr那就无穷大了。
tensorflow万万没想到我会计算两张一样的图片？？？。。。

测试3，批量计算

bath=[]
i = 0
for images, labels in mnist_dataset:    # image: [4, 28, 28, 1], labels: [4]
    i = i + 1
    bath.append(images)
    if i>2:
        break

    # psnr = tf.image.psnr(images, images, 1)
    # print(psnr)
    # break

print(tf.image.psnr(bath[0],bath[1],1))
print(tf.image.ssim(bath[0],bath[1],1))