|
- import torch
- import torch.nn.functional as F
- from torch.nn import Module
- import torch.nn as nn
-
-
- class SSIMLoss(Module):
- def __init__(self, kernel_size=11, sigma=1.5, as_loss=True):
- super().__init__()
- self.kernel_size = kernel_size
- self.sigma = sigma
- self.as_loss = as_loss
- self.gaussian_kernel = self._create_gaussian_kernel(self.kernel_size, self.sigma)
-
- def forward(self, x, y):
-
- if not self.gaussian_kernel.is_cuda:
- self.gaussian_kernel = self.gaussian_kernel.to(x.device)
-
- ssim_map = self._ssim(x, y)
-
- if self.as_loss:
- return 1 - ssim_map.mean()
- else:
- return 1 - ssim_map
-
- def _ssim(self, x, y):
-
- # Compute means
- y = y.unsqueeze(1)
- #print('line30:',x.shape,self.gaussian_kernel.shape,' Y.shape:',y.shape)
- ux = F.conv2d(x, self.gaussian_kernel, padding=self.kernel_size // 2, groups=1)
- uy = F.conv2d(y, self.gaussian_kernel, padding=self.kernel_size // 2, groups=1)
-
- # Compute variances
- uxx = F.conv2d(x * x, self.gaussian_kernel, padding=self.kernel_size // 2, groups=1)
- uyy = F.conv2d(y * y, self.gaussian_kernel, padding=self.kernel_size // 2, groups=1)
- uxy = F.conv2d(x * y, self.gaussian_kernel, padding=self.kernel_size // 2, groups=1)
- vx = uxx - ux * ux
- vy = uyy - uy * uy
- vxy = uxy - ux * uy
-
- c1 = 0.01**2
- c2 = 0.03**2
- numerator = (2 * ux * uy + c1) * (2 * vxy + c2)
- denominator = (ux**2 + uy**2 + c1) * (vx + vy + c2)
- return numerator / (denominator + 1e-12)
-
- def _create_gaussian_kernel(self, kernel_size, sigma):
-
- start = (1 - kernel_size) / 2
- end = (1 + kernel_size) / 2
- kernel_1d = torch.arange(start, end, step=1, dtype=torch.float)
- kernel_1d = torch.exp(-torch.pow(kernel_1d / sigma, 2) / 2)
- kernel_1d = (kernel_1d / kernel_1d.sum()).unsqueeze(dim=0)
-
- kernel_2d = torch.matmul(kernel_1d.t(), kernel_1d)
- kernel_2d = kernel_2d.expand(3, 1, kernel_size, kernel_size).contiguous()
- return kernel_2d
-
-
- class Multi_SSIM_loss(Module):
- def __init__(self, window_sizes=[3, 7, 11], sigma=1.5, as_loss=True):
- super(Multi_SSIM_loss, self).__init__()
- self.losses = list()
- for size in window_sizes:
- self.losses.append(SSIMLoss(size, sigma, as_loss))
- self.losses = nn.ModuleList(self.losses)
-
- def forward(self, img1, img2):
- loss_multi = list()
- for i in range(len(self.losses)):
- loss_multi.append(self.losses[i](img1, img2))
- total_loss = torch.stack(loss_multi, 0).sum()
-
- return total_loss
|
