Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchvision import models | |
| from torchvision.models import VGG19_Weights | |
| class VGGPerceptualLoss(nn.Module): | |
| """ | |
| Perceptual loss using VGG19 pretrained on ImageNet. | |
| We extract features at: | |
| - relu1_2 (index: 3) | |
| - relu2_2 (index: 8) | |
| - relu3_2 (index: 17) | |
| - relu4_2 (index: 26) | |
| Then compute L1 distance between those feature maps. | |
| Input images are in [-1,1]. We convert to [0,1], then normalize with ImageNet stats. | |
| """ | |
| def __init__(self, device): | |
| super(VGGPerceptualLoss, self).__init__() | |
| vgg19_features = models.vgg19(weights=VGG19_Weights.DEFAULT).features.to(device).eval() | |
| # We only need layers up to 26 (relu4_2) | |
| self.slices = nn.ModuleDict({ | |
| "relu1_2": nn.Sequential(*list(vgg19_features.children())[:4]), # conv1_1, relu1_1, conv1_2, relu1_2 | |
| "relu2_2": nn.Sequential(*list(vgg19_features.children())[4:9]), # pool1, conv2_1, relu2_1, conv2_2, relu2_2 | |
| "relu3_2": nn.Sequential(*list(vgg19_features.children())[9:18]), # pool2, conv3_1, relu3_1, conv3_2, relu3_2, ... | |
| "relu4_2": nn.Sequential(*list(vgg19_features.children())[18:27]) # pool3, conv4_1, relu4_1, conv4_2, relu4_2 | |
| }) | |
| for param in self.parameters(): | |
| param.requires_grad = False | |
| self.l1 = nn.L1Loss() | |
| self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)) | |
| self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)) | |
| def forward(self, img_gen, img_ref): | |
| """ | |
| img_gen, img_ref: [B,3,H,W] in range [-1,1]. | |
| Return: sum of L1 distances between VGG feature maps at chosen layers. | |
| """ | |
| # Convert to [0,1] | |
| gen = (img_gen + 1.0) / 2.0 | |
| ref = (img_ref + 1.0) / 2.0 | |
| # Normalize | |
| gen_norm = (gen - self.mean) / self.std | |
| ref_norm = (ref - self.mean) / self.std | |
| loss = 0.0 | |
| x_gen = gen_norm | |
| x_ref = ref_norm | |
| for slice_mod in self.slices.values(): | |
| x_gen = slice_mod(x_gen) | |
| x_ref = slice_mod(x_ref) | |
| loss += self.l1(x_gen, x_ref) | |
| return loss | |
| class SobelLoss(nn.Module): | |
| """ | |
| Computes the Sobel loss between two images, which encourages edge similarity. | |
| This loss operates on the grayscale versions of the input images. | |
| """ | |
| def __init__(self): | |
| super(SobelLoss, self).__init__() | |
| # Sobel kernels for edge detection | |
| self.kernel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3) | |
| self.kernel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3) | |
| self.l1 = nn.L1Loss() | |
| # Grayscale conversion weights (ITU-R BT.601) | |
| self.rgb_to_gray_weights = torch.tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1) | |
| def _get_edges(self, img): | |
| """ | |
| Converts an RGB image to grayscale and applies Sobel filters. | |
| Args: | |
| img: [B, 3, H, W] image tensor in range [-1, 1]. | |
| Returns: | |
| Gradient magnitude map [B, 1, H, W]. | |
| """ | |
| # Convert from [-1, 1] to [0, 1] | |
| img = (img + 1.0) / 2.0 | |
| # Convert to grayscale | |
| grayscale_img = F.conv2d(img, self.rgb_to_gray_weights.to(img.device)) | |
| # Apply Sobel filters | |
| grad_x = F.conv2d(grayscale_img, self.kernel_x.to(img.device), padding=1) | |
| grad_y = F.conv2d(grayscale_img, self.kernel_y.to(img.device), padding=1) | |
| # Compute gradient magnitude | |
| edges = torch.sqrt(grad_x**2 + grad_y**2 + 1e-6) # add epsilon for stability | |
| return edges | |
| def forward(self, img_gen, img_ref): | |
| """ | |
| img_gen, img_ref: [B, 3, H, W] in range [-1, 1]. | |
| Returns: L1 loss between the edge maps of the two images. | |
| """ | |
| edges_gen = self._get_edges(img_gen) | |
| edges_ref = self._get_edges(img_ref) | |
| return self.l1(edges_gen, edges_ref) | |