| 329 | |
| 330 | |
| 331 | def _ssim(img1, img2, window, window_size, channel, size_average=True): |
| 332 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) |
| 333 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) |
| 334 | |
| 335 | mu1_sq = mu1.pow(2) |
| 336 | mu2_sq = mu2.pow(2) |
| 337 | mu1_mu2 = mu1 * mu2 |
| 338 | |
| 339 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq |
| 340 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq |
| 341 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 |
| 342 | |
| 343 | C1 = 0.01 ** 2 |
| 344 | C2 = 0.03 ** 2 |
| 345 | |
| 346 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) |
| 347 | |
| 348 | if size_average: |
| 349 | return ssim_map.mean() |
| 350 | else: |
| 351 | return ssim_map.mean(1) |
| 352 | |
| 353 | |
| 354 | class SSIM(torch.nn.Module): |