import torch
import torch.nn.functional as F
import torch.nn as nn
from models.vgg import Vgg19
from utils.image_processing import gram
def to_gray_scale(image):
# https://github.com/pytorch/vision/blob/main/torchvision/transforms/v2/functional/_color.py#L33
# Image are assum in range 1, -1
image = (image + 1.0) / 2.0 # To [0, 1]
r, g, b = image.unbind(dim=-3)
l_img = r.mul(0.2989).add_(g, alpha=0.587).add_(b, alpha=0.114)
l_img = l_img.unsqueeze(dim=-3)
l_img = l_img.to(image.dtype)
l_img = l_img.expand(image.shape)
l_img = l_img / 0.5 - 1.0 # To [-1, 1]
return l_img
class ColorLoss(nn.Module):
def __init__(self):
super(ColorLoss, self).__init__()
self.l1 = nn.L1Loss()
self.huber = nn.SmoothL1Loss()
# self._rgb_to_yuv_kernel = torch.tensor([
# [0.299, -0.14714119, 0.61497538],
# [0.587, -0.28886916, -0.51496512],
# [0.114, 0.43601035, -0.10001026]
# ]).float()
self._rgb_to_yuv_kernel = torch.tensor([
[0.299, 0.587, 0.114],
[-0.14714119, -0.28886916, 0.43601035],
[0.61497538, -0.51496512, -0.10001026],
]).float()
def to(self, device):
new_self = super(ColorLoss, self).to(device)
new_self._rgb_to_yuv_kernel = new_self._rgb_to_yuv_kernel.to(device)
return new_self
def rgb_to_yuv(self, image):
'''
https://en.wikipedia.org/wiki/YUV
output: Image of shape (H, W, C) (channel last)
'''
# -1 1 -> 0 1
image = (image + 1.0) / 2.0
image = image.permute(0, 2, 3, 1) # To channel last
yuv_img = image @ self._rgb_to_yuv_kernel.T
return yuv_img
def forward(self, image, image_g):
image = self.rgb_to_yuv(image)
image_g = self.rgb_to_yuv(image_g)
# After convert to yuv, both images have channel last
return (
self.l1(image[:, :, :, 0], image_g[:, :, :, 0])
+ self.huber(image[:, :, :, 1], image_g[:, :, :, 1])
+ self.huber(image[:, :, :, 2], image_g[:, :, :, 2])
)
class AnimeGanLoss:
def __init__(self, args, device, gray_adv=False):
if isinstance(device, str):
device = torch.device(device)
self.content_loss = nn.L1Loss().to(device)
self.gram_loss = nn.L1Loss().to(device)
self.color_loss = ColorLoss().to(device)
self.wadvg = args.wadvg
self.wadvd = args.wadvd
self.wcon = args.wcon
self.wgra = args.wgra
self.wcol = args.wcol
self.wtvar = args.wtvar
# If true, use gray scale image to calculate adversarial loss
self.gray_adv = gray_adv
self.vgg19 = Vgg19().to(device).eval()
self.adv_type = args.gan_loss
self.bce_loss = nn.BCEWithLogitsLoss()
def compute_loss_G(self, fake_img, img, fake_logit, anime_gray):
'''
Compute loss for Generator
@Args:
- fake_img: generated image
- img: real image
- fake_logit: output of Discriminator given fake image
- anime_gray: grayscale of anime image
@Returns:
- Adversarial Loss of fake logits
- Content loss between real and fake features (vgg19)
- Gram loss between anime and fake features (Vgg19)
- Color loss between image and fake image
- Total variation loss of fake image
'''
fake_feat = self.vgg19(fake_img)
gray_feat = self.vgg19(anime_gray)
img_feat = self.vgg19(img)
# fake_gray_feat = self.vgg19(to_gray_scale(fake_img))
return [
# Want to be real image.
self.wadvg * self.adv_loss_g(fake_logit),
self.wcon * self.content_loss(img_feat, fake_feat),
self.wgra * self.gram_loss(gram(gray_feat), gram(fake_feat)),
self.wcol * self.color_loss(img, fake_img),
self.wtvar * self.total_variation_loss(fake_img)
]
def compute_loss_D(
self,
fake_img_d,
real_anime_d,
real_anime_gray_d,
real_anime_smooth_gray_d=None
):
if self.gray_adv:
# Treat gray scale image as real
return (
self.adv_loss_d_real(real_anime_gray_d)
+ self.adv_loss_d_fake(fake_img_d)
+ 0.3 * self.adv_loss_d_fake(real_anime_smooth_gray_d)
)
else:
return (
# Classify real anime as real
self.adv_loss_d_real(real_anime_d)
# Classify generated as fake
+ self.adv_loss_d_fake(fake_img_d)
# Classify real anime gray as fake
# + self.adv_loss_d_fake(real_anime_gray_d)
# Classify real anime as fake
# + 0.1 * self.adv_loss_d_fake(real_anime_smooth_gray_d)
)
def total_variation_loss(self, fake_img):
"""
A smooth loss in fact. Like the smooth prior in MRF.
V(y) = || y_{n+1} - y_n ||_2
"""
# Channel first -> channel last
fake_img = fake_img.permute(0, 2, 3, 1)
def _l2(x):
# sum(t ** 2) / 2
return torch.sum(x ** 2) / 2
dh = fake_img[:, :-1, ...] - fake_img[:, 1:, ...]
dw = fake_img[:, :, :-1, ...] - fake_img[:, :, 1:, ...]
return _l2(dh) / dh.numel() + _l2(dw) / dw.numel()
def content_loss_vgg(self, image, recontruction):
feat = self.vgg19(image)
re_feat = self.vgg19(recontruction)
feature_loss = self.content_loss(feat, re_feat)
content_loss = self.content_loss(image, recontruction)
return feature_loss# + 0.5 * content_loss
def adv_loss_d_real(self, pred):
"""Push pred to class 1 (real)"""
if self.adv_type == 'hinge':
return torch.mean(F.relu(1.0 - pred))
elif self.adv_type == 'lsgan':
# pred = torch.sigmoid(pred)
return torch.mean(torch.square(pred - 1.0))
elif self.adv_type == 'bce':
return self.bce_loss(pred, torch.ones_like(pred))
raise ValueError(f'Do not support loss type {self.adv_type}')
def adv_loss_d_fake(self, pred):
"""Push pred to class 0 (fake)"""
if self.adv_type == 'hinge':
return torch.mean(F.relu(1.0 + pred))
elif self.adv_type == 'lsgan':
# pred = torch.sigmoid(pred)
return torch.mean(torch.square(pred))
elif self.adv_type == 'bce':
return self.bce_loss(pred, torch.zeros_like(pred))
raise ValueError(f'Do not support loss type {self.adv_type}')
def adv_loss_g(self, pred):
"""Push pred to class 1 (real)"""
if self.adv_type == 'hinge':
return -torch.mean(pred)
elif self.adv_type == 'lsgan':
# pred = torch.sigmoid(pred)
return torch.mean(torch.square(pred - 1.0))
elif self.adv_type == 'bce':
return self.bce_loss(pred, torch.ones_like(pred))
raise ValueError(f'Do not support loss type {self.adv_type}')
class LossSummary:
def __init__(self):
self.reset()
def reset(self):
self.loss_g_adv = []
self.loss_content = []
self.loss_gram = []
self.loss_color = []
self.loss_d_adv = []
def update_loss_G(self, adv, gram, color, content):
self.loss_g_adv.append(adv.cpu().detach().numpy())
self.loss_gram.append(gram.cpu().detach().numpy())
self.loss_color.append(color.cpu().detach().numpy())
self.loss_content.append(content.cpu().detach().numpy())
def update_loss_D(self, loss):
self.loss_d_adv.append(loss.cpu().detach().numpy())
def avg_loss_G(self):
return (
self._avg(self.loss_g_adv),
self._avg(self.loss_gram),
self._avg(self.loss_color),
self._avg(self.loss_content),
)
def avg_loss_D(self):
return self._avg(self.loss_d_adv)
def get_loss_description(self):
avg_adv, avg_gram, avg_color, avg_content = self.avg_loss_G()
avg_adv_d = self.avg_loss_D()
return f'loss G: adv {avg_adv:2f} con {avg_content:2f} gram {avg_gram:2f} color {avg_color:2f} / loss D: {avg_adv_d:2f}'
@staticmethod
def _avg(losses):
return sum(losses) / len(losses)