forked from Po-Hsun-Su/pytorch-ssim
-
Notifications
You must be signed in to change notification settings - Fork 0
/
max_ssim.py
37 lines (27 loc) · 941 Bytes
/
max_ssim.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
import pytorch_ssim
import torch
from torch.autograd import Variable
from torch import optim
import cv2
import numpy as np
npImg1 = cv2.imread("einstein.png")
img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0)/255.0
img2 = torch.rand(img1.size())
if torch.cuda.is_available():
img1 = img1.cuda()
img2 = img2.cuda()
img1 = Variable( img1, requires_grad=False)
img2 = Variable( img2, requires_grad = True)
# Functional: pytorch_ssim.ssim(img1, img2, window_size = 11, size_average = True)
ssim_value = pytorch_ssim.ssim(img1, img2).data[0]
print("Initial ssim:", ssim_value)
# Module: pytorch_ssim.SSIM(window_size = 11, size_average = True)
ssim_loss = pytorch_ssim.SSIM()
optimizer = optim.Adam([img2], lr=0.01)
while ssim_value < 0.95:
optimizer.zero_grad()
ssim_out = -ssim_loss(img1, img2)
ssim_value = - ssim_out.data[0]
print(ssim_value)
ssim_out.backward()
optimizer.step()