| 191 | |
| 192 | |
| 193 | class SSIMLoss(torch.nn.Module): |
| 194 | def __init__(self, window_size=11, size_average=True): |
| 195 | super(SSIMLoss, self).__init__() |
| 196 | self.window_size = window_size |
| 197 | self.size_average = size_average |
| 198 | self.channel = 1 |
| 199 | self.window = create_window(window_size, self.channel) |
| 200 | |
| 201 | def forward(self, img1, img2): |
| 202 | (_, channel, _, _) = img1.size() |
| 203 | if channel == self.channel and self.window.data.type() == img1.data.type(): |
| 204 | window = self.window |
| 205 | else: |
| 206 | window = create_window(self.window_size, channel) |
| 207 | if img1.is_cuda: |
| 208 | window = window.cuda(img1.get_device()) |
| 209 | window = window.type_as(img1) |
| 210 | self.window = window |
| 211 | self.channel = channel |
| 212 | return 1 - _ssim(img1, img2, window, self.window_size, channel, self.size_average) |
| 213 | |
| 214 | |
| 215 | def gaussian(window_size, sigma): |