├── .gitignore ├── LICENSE.txt ├── README.md ├── einstein.png ├── max_ssim.gif ├── max_ssim.py ├── pytorch_ssim └── __init__.py ├── setup.cfg └── setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.gif 3 | *.png 4 | *.jpg 5 | test* 6 | !einstein.png 7 | !max_ssim.gif 8 | MANIFEST 9 | dist/* 10 | .sync-config.cson 11 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pytorch-ssim (This repo is not maintained) 2 | 3 | The code doesn't work because it is on super old pytorch. 4 | 5 | ### Differentiable structural similarity (SSIM) index. 6 | ![einstein](https://raw.githubusercontent.com/Po-Hsun-Su/pytorch-ssim/master/einstein.png) ![Max_ssim](https://raw.githubusercontent.com/Po-Hsun-Su/pytorch-ssim/master/max_ssim.gif) 7 | 8 | ## Installation 9 | 1. Clone this repo. 10 | 2. Copy "pytorch_ssim" folder in your project. 11 | 12 | ## Example 13 | ### basic usage 14 | ```python 15 | import pytorch_ssim 16 | import torch 17 | from torch.autograd import Variable 18 | 19 | img1 = Variable(torch.rand(1, 1, 256, 256)) 20 | img2 = Variable(torch.rand(1, 1, 256, 256)) 21 | 22 | if torch.cuda.is_available(): 23 | img1 = img1.cuda() 24 | img2 = img2.cuda() 25 | 26 | print(pytorch_ssim.ssim(img1, img2)) 27 | 28 | ssim_loss = pytorch_ssim.SSIM(window_size = 11) 29 | 30 | print(ssim_loss(img1, img2)) 31 | 32 | ``` 33 | ### maximize ssim 34 | ```python 35 | import pytorch_ssim 36 | import torch 37 | from torch.autograd import Variable 38 | from torch import optim 39 | import cv2 40 | import numpy as np 41 | 42 | npImg1 = cv2.imread("einstein.png") 43 | 44 | img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0)/255.0 45 | img2 = torch.rand(img1.size()) 46 | 47 | if torch.cuda.is_available(): 48 | img1 = img1.cuda() 49 | img2 = img2.cuda() 50 | 51 | 52 | img1 = Variable( img1, requires_grad=False) 53 | img2 = Variable( img2, requires_grad = True) 54 | 55 | 56 | # Functional: pytorch_ssim.ssim(img1, img2, window_size = 11, size_average = True) 57 | ssim_value = pytorch_ssim.ssim(img1, img2).data[0] 58 | print("Initial ssim:", ssim_value) 59 | 60 | # Module: pytorch_ssim.SSIM(window_size = 11, size_average = True) 61 | ssim_loss = pytorch_ssim.SSIM() 62 | 63 | optimizer = optim.Adam([img2], lr=0.01) 64 | 65 | while ssim_value < 0.95: 66 | optimizer.zero_grad() 67 | ssim_out = -ssim_loss(img1, img2) 68 | ssim_value = - ssim_out.data[0] 69 | print(ssim_value) 70 | ssim_out.backward() 71 | optimizer.step() 72 | 73 | ``` 74 | 75 | ## Reference 76 | https://ece.uwaterloo.ca/~z70wang/research/ssim/ 77 | -------------------------------------------------------------------------------- /einstein.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Po-Hsun-Su/pytorch-ssim/3add4532d3f633316cba235da1c69e90f0dfb952/einstein.png -------------------------------------------------------------------------------- /max_ssim.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Po-Hsun-Su/pytorch-ssim/3add4532d3f633316cba235da1c69e90f0dfb952/max_ssim.gif -------------------------------------------------------------------------------- /max_ssim.py: -------------------------------------------------------------------------------- 1 | import pytorch_ssim 2 | import torch 3 | from torch.autograd import Variable 4 | from torch import optim 5 | import cv2 6 | import numpy as np 7 | 8 | npImg1 = cv2.imread("einstein.png") 9 | 10 | img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0)/255.0 11 | img2 = torch.rand(img1.size()) 12 | 13 | if torch.cuda.is_available(): 14 | img1 = img1.cuda() 15 | img2 = img2.cuda() 16 | 17 | 18 | img1 = Variable( img1, requires_grad=False) 19 | img2 = Variable( img2, requires_grad = True) 20 | 21 | 22 | # Functional: pytorch_ssim.ssim(img1, img2, window_size = 11, size_average = True) 23 | ssim_value = pytorch_ssim.ssim(img1, img2).data[0] 24 | print("Initial ssim:", ssim_value) 25 | 26 | # Module: pytorch_ssim.SSIM(window_size = 11, size_average = True) 27 | ssim_loss = pytorch_ssim.SSIM() 28 | 29 | optimizer = optim.Adam([img2], lr=0.01) 30 | 31 | while ssim_value < 0.95: 32 | optimizer.zero_grad() 33 | ssim_out = -ssim_loss(img1, img2) 34 | ssim_value = - ssim_out.data[0] 35 | print(ssim_value) 36 | ssim_out.backward() 37 | optimizer.step() 38 | -------------------------------------------------------------------------------- /pytorch_ssim/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from torch.autograd import Variable 4 | import numpy as np 5 | from math import exp 6 | 7 | def gaussian(window_size, sigma): 8 | gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) 9 | return gauss/gauss.sum() 10 | 11 | def create_window(window_size, channel): 12 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 13 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 14 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 15 | return window 16 | 17 | def _ssim(img1, img2, window, window_size, channel, size_average = True): 18 | mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel) 19 | mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel) 20 | 21 | mu1_sq = mu1.pow(2) 22 | mu2_sq = mu2.pow(2) 23 | mu1_mu2 = mu1*mu2 24 | 25 | sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq 26 | sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq 27 | sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2 28 | 29 | C1 = 0.01**2 30 | C2 = 0.03**2 31 | 32 | ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2)) 33 | 34 | if size_average: 35 | return ssim_map.mean() 36 | else: 37 | return ssim_map.mean(1).mean(1).mean(1) 38 | 39 | class SSIM(torch.nn.Module): 40 | def __init__(self, window_size = 11, size_average = True): 41 | super(SSIM, self).__init__() 42 | self.window_size = window_size 43 | self.size_average = size_average 44 | self.channel = 1 45 | self.window = create_window(window_size, self.channel) 46 | 47 | def forward(self, img1, img2): 48 | (_, channel, _, _) = img1.size() 49 | 50 | if channel == self.channel and self.window.data.type() == img1.data.type(): 51 | window = self.window 52 | else: 53 | window = create_window(self.window_size, channel) 54 | 55 | if img1.is_cuda: 56 | window = window.cuda(img1.get_device()) 57 | window = window.type_as(img1) 58 | 59 | self.window = window 60 | self.channel = channel 61 | 62 | 63 | return _ssim(img1, img2, window, self.window_size, channel, self.size_average) 64 | 65 | def ssim(img1, img2, window_size = 11, size_average = True): 66 | (_, channel, _, _) = img1.size() 67 | window = create_window(window_size, channel) 68 | 69 | if img1.is_cuda: 70 | window = window.cuda(img1.get_device()) 71 | window = window.type_as(img1) 72 | 73 | return _ssim(img1, img2, window, window_size, channel, size_average) 74 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.md 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | setup( 3 | name = 'pytorch_ssim', 4 | packages = ['pytorch_ssim'], # this must be the same as the name above 5 | version = '0.1', 6 | description = 'Differentiable structural similarity (SSIM) index', 7 | author = 'Po-Hsun (Evan) Su', 8 | author_email = 'evan.pohsun.su@gmail.com', 9 | url = 'https://github.com/Po-Hsun-Su/pytorch-ssim', # use the URL to the github repo 10 | download_url = 'https://github.com/Po-Hsun-Su/pytorch-ssim/archive/0.1.tar.gz', # I'll explain this in a second 11 | keywords = ['pytorch', 'image-processing', 'deep-learning'], # arbitrary keywords 12 | classifiers = [], 13 | ) 14 | --------------------------------------------------------------------------------