tensorflow计算psnr与ssim

本文主要是介绍tensorflow计算psnr与ssim，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

加载mnist进行测试

import mathimport tensorflow as tf
import numpy as npload_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 100return 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 + 1bath.append(images)if i>2:break# psnr = tf.image.psnr(images, images, 1)# print(psnr)# breakprint(tf.image.psnr(bath[0],bath[1],1))
print(tf.image.ssim(bath[0],bath[1],1))