|
|
|
|
|
|
|
import math |
|
from math import exp |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from torch.autograd import Variable |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def gaussian(window_size, sigma): |
|
gauss = torch.Tensor( |
|
[ |
|
exp(-((x - window_size // 2) ** 2) / float(2 * sigma**2)) |
|
for x in range(window_size) |
|
] |
|
) |
|
return gauss / gauss.sum() |
|
|
|
|
|
def create_window(window_size, channel): |
|
_1D_window = gaussian(window_size, 1.5).unsqueeze(1) |
|
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) |
|
window = Variable( |
|
_2D_window.expand(channel, 1, window_size, window_size).contiguous() |
|
) |
|
return window |
|
|
|
|
|
def _ssim(img1, img2, window, window_size, channel, size_average=True, full=False): |
|
padd = 0 |
|
|
|
mu1 = F.conv2d(img1, window, padding=padd, groups=channel) |
|
mu2 = F.conv2d(img2, window, padding=padd, groups=channel) |
|
|
|
mu1_sq = mu1.pow(2) |
|
mu2_sq = mu2.pow(2) |
|
mu1_mu2 = mu1 * mu2 |
|
|
|
sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq |
|
sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq |
|
sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2 |
|
|
|
C1 = 0.01**2 |
|
C2 = 0.03**2 |
|
|
|
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ( |
|
(mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2) |
|
) |
|
|
|
v1 = 2.0 * sigma12 + C2 |
|
v2 = sigma1_sq + sigma2_sq + C2 |
|
cs = torch.mean(v1 / v2) |
|
|
|
if size_average: |
|
ret = ssim_map.mean() |
|
else: |
|
ret = ssim_map.mean(1).mean(1).mean(1) |
|
|
|
if full: |
|
return ret, cs |
|
return ret |
|
|
|
|
|
class SSIM(torch.nn.Module): |
|
def __init__(self, window_size=11, size_average=True): |
|
super(SSIM, self).__init__() |
|
self.window_size = window_size |
|
self.size_average = size_average |
|
self.channel = 1 |
|
self.window = create_window(window_size, self.channel) |
|
|
|
def forward(self, img1, img2): |
|
(_, channel, _, _) = img1.size() |
|
|
|
if channel == self.channel and self.window.data.type() == img1.data.type(): |
|
window = self.window |
|
else: |
|
window = create_window(self.window_size, channel) |
|
|
|
if img1.is_cuda: |
|
window = window.cuda(img1.get_device()) |
|
window = window.type_as(img1) |
|
|
|
self.window = window |
|
self.channel = channel |
|
|
|
return _ssim(img1, img2, window, self.window_size, channel, self.size_average) |
|
|
|
|
|
def ssim(img1, img2, window_size=11, size_average=True, full=False): |
|
(_, channel, height, width) = img1.size() |
|
|
|
real_size = min(window_size, height, width) |
|
window = create_window(real_size, channel) |
|
|
|
if img1.is_cuda: |
|
window = window.cuda(img1.get_device()) |
|
window = window.type_as(img1) |
|
|
|
return _ssim(img1, img2, window, real_size, channel, size_average, full=full) |
|
|
|
|
|
def msssim(img1, img2, window_size=11, size_average=True): |
|
|
|
if img1.size() != img2.size(): |
|
raise RuntimeError( |
|
"Input images must have the same shape (%s vs. %s)." |
|
% (img1.size(), img2.size()) |
|
) |
|
if len(img1.size()) != 4: |
|
raise RuntimeError( |
|
"Input images must have four dimensions, not %d" % len(img1.size()) |
|
) |
|
|
|
weights = torch.tensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333], dtype=img1.dtype) |
|
if img1.is_cuda: |
|
weights = weights.cuda(img1.get_device()) |
|
|
|
levels = weights.size()[0] |
|
mssim = [] |
|
mcs = [] |
|
for _ in range(levels): |
|
sim, cs = ssim( |
|
img1, img2, window_size=window_size, size_average=size_average, full=True |
|
) |
|
mssim.append(sim) |
|
mcs.append(cs) |
|
|
|
img1 = F.avg_pool2d(img1, (2, 2)) |
|
img2 = F.avg_pool2d(img2, (2, 2)) |
|
|
|
mssim = torch.stack(mssim) |
|
mcs = torch.stack(mcs) |
|
return torch.prod(mcs[0 : levels - 1] ** weights[0 : levels - 1]) * ( |
|
mssim[levels - 1] ** weights[levels - 1] |
|
) |
|
|
|
|
|
class MSSSIM(torch.nn.Module): |
|
def __init__(self, window_size=11, size_average=True, channel=3): |
|
super(MSSSIM, self).__init__() |
|
self.window_size = window_size |
|
self.size_average = size_average |
|
self.channel = channel |
|
|
|
def forward(self, img1, img2): |
|
|
|
return msssim( |
|
img1, img2, window_size=self.window_size, size_average=self.size_average |
|
) |
|
|
|
|
|
def calc_psnr(sr, hr, scale=0, benchmark=False): |
|
|
|
diff = (sr - hr).data |
|
if benchmark: |
|
shave = scale |
|
if diff.size(1) > 1: |
|
convert = diff.new(1, 3, 1, 1) |
|
convert[0, 0, 0, 0] = 65.738 |
|
convert[0, 1, 0, 0] = 129.057 |
|
convert[0, 2, 0, 0] = 25.064 |
|
diff.mul_(convert).div_(256) |
|
diff = diff.sum(dim=1, keepdim=True) |
|
else: |
|
shave = scale + 6 |
|
|
|
valid = diff[:, :, shave:-shave, shave:-shave] |
|
mse = valid.pow(2).mean() |
|
|
|
return -10 * math.log10(mse) |
|
|
|
|
|
|
|
|
|
|
|
from torch import nn |
|
|
|
|
|
def psnr(predict, target): |
|
with torch.no_grad(): |
|
criteria = nn.MSELoss() |
|
mse = criteria(predict, target) |
|
return -10 * torch.log10(mse) |
|
|
