forked from as3ert/myPytorchSSIM
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
57 lines (44 loc) · 2.03 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import torch
import ssim
import numpy as np
from torch.utils.data import DataLoader
from torchvision.io import read_image
from ssim import SSIMLoss
from skimage.metrics import structural_similarity as skSSIM
def tensor2numpy(img):
return np.array(img.cpu().detach())
mySSIM = SSIMLoss()
# Test Bench #1
print("Test #1")
print("Calculate Skimage Benchmark")
benchmark = [1.000, 0.988, 0.913, 0.840, 0.694, 0.662]
originalImg = read_image("data\image4.jpg").to(torch.float32)
inputImg = originalImg.unsqueeze(0)
for i in range(6):
testImg = read_image("data\image{}.jpg".format(4+2*i)).to(torch.float32)
outputImg = testImg.unsqueeze(0)
calcSSIM = mySSIM(inputImg, outputImg).item()
testSSIM = skSSIM(tensor2numpy(originalImg), tensor2numpy(testImg), channel_axis=0)
print("{:.3f} {:.3f} {:.3f}".format(1-calcSSIM, testSSIM, benchmark[i]))
print()
# Test Bench #2
mySSIM = SSIMLoss(dtype=torch.float64)
print("Test #2")
print("Calculate Skimage Benchmark")
originalImg = read_image("data\simga_0_ssim_1.0000.png").to(torch.float64)
inputImg = originalImg.unsqueeze(0)
testImg = read_image("data\simga_0_ssim_1.0000.png").to(torch.float64)
outputImg = testImg.unsqueeze(0)
calcSSIM = mySSIM(inputImg, outputImg).item()
testSSIM = skSSIM(tensor2numpy(originalImg), tensor2numpy(testImg), channel_axis=0)
print("{:.4f} {:.4f} 1.0000".format(1-calcSSIM, testSSIM))
testImg = read_image("data\simga_50_ssim_0.4225.png").to(torch.float64)
outputImg = testImg.unsqueeze(0)
calcSSIM = mySSIM(inputImg, outputImg).item()
testSSIM = skSSIM(tensor2numpy(originalImg), tensor2numpy(testImg), channel_axis=0)
print("{:.4f} {:.4f} 0.4225".format(1-calcSSIM, testSSIM))
testImg = read_image("data\simga_100_ssim_0.1924.png").to(torch.float64)
outputImg = testImg.unsqueeze(0)
calcSSIM = mySSIM(inputImg, outputImg).item()
testSSIM = skSSIM(tensor2numpy(originalImg), tensor2numpy(testImg), channel_axis=0)
print("{:.4f} {:.4f} 0.1924".format(1-calcSSIM, testSSIM))