下面的代码是将单张图片喂入JSCC,计算PSNR。实验目的是:比较不同图像预处理方式对PSNR的影响。
from torch.nn import MSELoss
from PIL import Image
img_PIL = Image.open("/home/xxx/xxx/dataset/test_cifar/test_cifar10/cat/cat_2054.jpg")
transform_train = transforms.Compose(
[transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomCrop(32, padding=5, pad_if_needed=True, fill=0, padding_mode='reflect'),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
img_tensor = transform_train(img_PIL)
img_tensor = img_tensor.view(1,3,32,32).to(args.device)
checkpoints = torch.load(workdir+"/models/"+"deepjscc_cifar10_awgn_200epochs_10db.pth")
jscc = Jscc(32, args.P, 10).to(device)
jscc.load_state_dict(checkpoints)
output_imgs = jscc(img_tensor)
loss = MSELoss()
mse = loss(img_tensor*255., output_imgs*255.)
psnr = 10 * (torch.log(255. * 255. / mse) / np.log(10))
print(f"psnr = {psnr}")
print(f"mse = {mse}")
结果如下:
psnr = 17.981430053710938
mse = 1034.992431640625
如果将transforms.Normalize((0.5, 0.5, 0.5)
删去,结果如下:
psnr = 33.61720657348633
mse = 28.272192001342773
结论:
transforms.Normalize((0.5, 0.5, 0.5)将输入图像的每个信道都作了归一化(均值=0.5,方差=0.5),(input[channel]-0.5)/0.5=2*input[channel]-1,原来图像∈[0,1],现在∈[-1,1],mse增加,psnr减少,验证如下
from PIL import Image
img_PIL = Image.open("/home/caobin/dongb/dataset/test_cifar/test_cifar10/cat/cat_1454.jpg")
print(img_PIL)
transform_totensor = transforms.ToTensor()
img_totensor= transform_totensor(img_PIL)
print(img_totensor[0][0][0])
transform_norm = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
img_norm= transform_norm(img_totensor)
print(img_norm[0][0][0])
结果如下:
tensor(0.1412)
tensor(-0.7176)
备注:
mse的计算方式对结果影响不大,只要统一就好,比如按照如下的方式计算:
mse = torch.mean((output_imgs-img_tensor)**2)
psnr = 20*torch.log10(255. / torch.sqrt(mse))
含有归一化:
psnr = 17.55906867980957
mse = 1140.70556640625
不含归一化:
psnr = 33.05939483642578
mse = 32.14704132080078
与上述结果几乎没有差别,只要统一就好。