import os import torch import torch.nn as nn from .utils import init_net import torch.nn.functional as F from . import networks from . import losses from . import utils from torchvision.models import vgg19 vgg_activation = dict() def get_activation(name): def hook(model, input, output): vgg_activation[name] = output return hook def gram_matrix(x): a, b, c, d = x.size() features = x.view(a * b, c * d) G = torch.mm(features, features.t()) return G.div(a * b * c * d) class CSDMT(nn.Module): def __init__(self, opts): super(CSDMT, self).__init__() # parameters self.opts = opts self.lr = opts.lr self.batch_size = opts.batch_size self.gpu = torch.device('cuda:{}'.format(opts.gpu)) if opts.gpu >= 0 else torch.device('cpu') self.weight_semantic = opts.weight_semantic self.weight_corr = opts.weight_corr self.weight_identity = opts.weight_identity self.weight_back = opts.weight_back self.weight_adv = opts.weight_adv self.weight_contrastive=opts.weight_contrastive self.weight_self_recL1 = opts.weight_self_recL1 self.weight_cycleL1 = opts.weight_cycleL1 self.target_layer=['relu_3', 'relu_8'] if self.opts.phase == 'train': self.vgg = vgg19(pretrained=True) for layer in self.target_layer: self.vgg.features[int(layer.split('_')[-1])].register_forward_hook(get_activation(layer)) self.vgg.to(self.gpu) self.FloatTensor = torch.cuda.FloatTensor if opts.gpu >= 0 else torch.FloatTensor self.criterionL1 = nn.L1Loss() self.criterionL2 = nn.MSELoss() self.criterionIdentity = losses.GPLoss() self.criterionGAN = losses.GANLoss(gan_mode=opts.gan_mode, tensor=self.FloatTensor) self._build_model() def _build_model(self): print('start build model') # discriminators if self.opts.dis_scale > 1: self.dis = init_net(networks.MultiScaleDis(3, self.opts.dis_scale, norm=self.opts.dis_norm, sn=self.opts.dis_sn), self.gpu, init_type='normal', gain=0.02) else: self.dis = init_net(networks.Dis(3, norm=self.opts.dis_norm, sn=self.opts.dis_sn), self.gpu, init_type='normal', gain=0.02) self.gen = init_net(networks.Generator(input_dim=3,parse_dim=self.opts.semantic_dim,ngf=16), self.gpu, init_type='normal', gain=0.02) self.dis_opt = torch.optim.Adam(self.dis.parameters(), lr=self.lr, betas=(0.5, 0.999), weight_decay=0.0001) self.gen_opt = torch.optim.Adam(self.gen.parameters(), lr=self.lr, betas=(0.5, 0.999), weight_decay=0.0001) print('finish build model') def load_data(self, data): self.non_makeup_color_img = data['non_makeup_color_img'].to(self.gpu).detach() self.non_makeup_split_parse = data['non_makeup_split_parse'].to(self.gpu).detach() self.non_makeup_all_mask = data['non_makeup_all_mask'].to(self.gpu).detach() self.non_makeup_face_mask = data['non_makeup_face_mask'].to(self.gpu).detach() self.non_makeup_brow_mask = data['non_makeup_brow_mask'].to(self.gpu).detach() self.non_makeup_eye_mask = data['non_makeup_eye_mask'].to(self.gpu).detach() self.non_makeup_lip_mask = data['non_makeup_lip_mask'].to(self.gpu).detach() self.non_makeup_face_mask_no_neck_no_ear = self.non_makeup_face_mask+self.non_makeup_brow_mask+self.non_makeup_eye_mask+self.non_makeup_lip_mask self.makeup_color_img = data['makeup_color_img'].to(self.gpu).detach() self.makeup_split_parse = data['makeup_split_parse'].to(self.gpu).detach() self.makeup_all_mask = data['makeup_all_mask'].to(self.gpu).detach() self.makeup_face_mask = data['makeup_face_mask'].to(self.gpu).detach() self.makeup_brow_mask = data['makeup_brow_mask'].to(self.gpu).detach() self.makeup_eye_mask = data['makeup_eye_mask'].to(self.gpu).detach() self.makeup_lip_mask = data['makeup_lip_mask'].to(self.gpu).detach() self.makeup_face_mask_no_neck_no_ear = self.makeup_face_mask + self.makeup_brow_mask + self.makeup_eye_mask + self.makeup_lip_mask self.makeup_color_change_img = data['makeup_color_change_img'].to(self.gpu).detach() self.makeup_color_change_warp_img = data['makeup_color_change_warp_img'].to(self.gpu).detach() self.makeup_split_warp_parse = data['makeup_split_parse_warp'].to(self.gpu).detach() self.makeup_all_warp_mask = data['makeup_all_warp_mask'].to(self.gpu).detach() self.makeup_face_warp_mask = data['makeup_face_warp_mask'].to(self.gpu).detach() self.makeup_brow_warp_mask = data['makeup_brow_warp_mask'].to(self.gpu).detach() self.makeup_eye_warp_mask = data['makeup_eye_warp_mask'].to(self.gpu).detach() self.makeup_lip_warp_mask = data['makeup_lip_warp_mask'].to(self.gpu).detach() self.makeup_face_warp_mask_no_neck_no_ear = self.makeup_face_warp_mask + self.makeup_brow_warp_mask + self.makeup_eye_warp_mask + self.makeup_lip_warp_mask self.makeup_color_change2_img = data['makeup_color_change_img2'].to(self.gpu).detach() self.makeup_color_change3_img = data['makeup_color_change_img3'].to(self.gpu).detach() self.makeup_color_change4_img = data['makeup_color_change_img4'].to(self.gpu).detach() # def warp(self,x,map): # n, c, h, w = x.shape # x_warp = torch.bmm(x.view(n, c, h * w), map) # n*HW*1 # x_warp = x_warp.view(n, c, h, w) # return x_warp def forward(self): # makeup transfer self.transfer_output_data=self.gen(source_img=self.non_makeup_color_img, source_parse=self.non_makeup_split_parse, source_all_mask=self.non_makeup_all_mask, ref_img=self.makeup_color_img, ref_parse=self.makeup_split_parse, ref_all_mask=self.makeup_all_mask) # cycle consistent self.cycle_output_data=self.gen(source_img=self.makeup_color_img, source_parse=self.makeup_split_parse, source_all_mask=self.makeup_all_mask, ref_img=self.transfer_output_data['transfer_img'], ref_parse=self.non_makeup_split_parse, ref_all_mask=self.non_makeup_all_mask) # self supervision self.rec_output_data=self.gen(source_img=self.makeup_color_change_img, source_parse=self.makeup_split_parse, source_all_mask=self.makeup_all_mask, ref_img=self.makeup_color_change_warp_img, ref_parse=self.non_makeup_split_parse, ref_all_mask=self.makeup_all_warp_mask, replace_content_img=self.makeup_color_img) def update_D(self): self.forward() self.dis_opt.zero_grad() loss_dis = self.backward_D(self.dis, self.makeup_color_img, self.transfer_output_data['transfer_img']) self.loss_dis = loss_dis.item() self.dis_opt.step() def backward_D(self, netD, real, fake): pred_fake = netD.forward(fake.detach()) pred_real = netD.forward(real) ad_true_loss = self.criterionGAN(pred_real, target_is_real=True, for_discriminator=True) ad_fake_loss = self.criterionGAN(pred_fake, target_is_real=False, for_discriminator=True) loss_D = ad_true_loss + ad_fake_loss loss_D.backward() return loss_D def update_G(self): self.gen_opt.zero_grad() self.backward_G() self.gen_opt.step() def warp_basedon_corr(self, x, corr): n, c, h, w = x.shape x_warp = torch.bmm(x.view(n, c, h * w), corr) # n*HW*1 x_warp = x_warp.view(n, c, h, w) return x_warp def gram_style(self, reference, fake_images): fake_activation = dict() real_activation = dict() # percep_style g_loss_style = 0 self.vgg(reference) for layer in self.target_layer: real_activation[layer] = vgg_activation[layer] vgg_activation.clear() self.vgg(fake_images) for layer in self.target_layer: fake_activation[layer] = vgg_activation[layer] vgg_activation.clear() for layer in self.target_layer: g_loss_style += self.criterionL1(gram_matrix(fake_activation[layer]), gram_matrix(real_activation[layer])) return g_loss_style def backward_G(self): #corr non_makeup_split_parse_down4 = F.interpolate(self.non_makeup_split_parse, scale_factor=1 / 4, mode='nearest') makeup_split_parse_down4 = F.interpolate(self.makeup_split_parse, scale_factor=1 / 4, mode='nearest') makeup_split_parse_down4_warp=self.warp_basedon_corr(makeup_split_parse_down4,self.transfer_output_data['corr_ref2source']) makeup_split_warp_parse_down4 = F.interpolate(self.makeup_split_warp_parse, scale_factor=1 / 4, mode='nearest') makeup_split_warp_parse_down4_warp=self.warp_basedon_corr(makeup_split_warp_parse_down4,self.rec_output_data['corr_ref2source']) loss_G_sem=self.criterionL1(makeup_split_parse_down4_warp,non_makeup_split_parse_down4)+\ self.criterionL1(makeup_split_warp_parse_down4_warp,makeup_split_parse_down4) loss_G_sem=loss_G_sem*self.weight_semantic transfer_color_face_down4 = F.interpolate(self.transfer_output_data['transfer_img'] * self.non_makeup_face_mask_no_neck_no_ear, scale_factor=1 / 4, mode='bilinear') self.transfer_color_face_down4_warp = self.warp_basedon_corr(transfer_color_face_down4, self.cycle_output_data['corr_ref2source']) makeup_color_face_down4 = F.interpolate(self.makeup_color_img * self.makeup_face_mask_no_neck_no_ear, scale_factor=1 / 4, mode='bilinear') self_rec_color_face_down4 = F.interpolate(self.rec_output_data['transfer_img'] * self.makeup_face_mask_no_neck_no_ear, scale_factor=1 / 4, mode='bilinear') self.self_rec_color_face_down4_warp = self.gen.forward_cross_attention(source_parse=self.makeup_split_warp_parse, ref_parse=self.makeup_split_parse, ref_img=self_rec_color_face_down4) makeup_warp_color_face_down4 = F.interpolate(self.makeup_color_change_warp_img * self.makeup_face_warp_mask_no_neck_no_ear, scale_factor=1 / 4, mode='bilinear') loss_G_corr=self.criterionL1(self.transfer_color_face_down4_warp,makeup_color_face_down4)+\ self.criterionL1(self.self_rec_color_face_down4_warp,makeup_warp_color_face_down4) loss_G_corr=loss_G_corr*self.weight_corr # transfer # identity loss_G_identity = self.criterionIdentity(self.transfer_output_data['transfer_img'], self.non_makeup_color_img) loss_G_identity = loss_G_identity * self.weight_identity # back loss_G_back = self.criterionL1(self.transfer_output_data['transfer_img'] * (1 - self.non_makeup_all_mask[:,0:1,:,:]), self.non_makeup_color_img * (1 - self.non_makeup_all_mask[:,0:1,:,:])) loss_G_back = loss_G_back * self.weight_back # Contrastive loss positive_distance=self.gram_style(self.transfer_output_data['transfer_img']*self.non_makeup_face_mask_no_neck_no_ear[:,0:1,:,:], self.makeup_color_img*self.makeup_face_mask_no_neck_no_ear[:,0:1,:,:]) #print(self.makeup_color_change3_img.shape) negative_distance = self.gram_style(self.transfer_output_data['transfer_img'] * self.non_makeup_face_mask_no_neck_no_ear[:, 0:1, :, :], self.makeup_color_change_img * self.makeup_face_mask_no_neck_no_ear[:, 0:1, :, :])+\ self.gram_style(self.transfer_output_data['transfer_img'] * self.non_makeup_face_mask_no_neck_no_ear[:, 0:1, :, :], self.makeup_color_change2_img * self.makeup_face_mask_no_neck_no_ear[:, 0:1, :, :])+\ self.gram_style(self.transfer_output_data['transfer_img'] * self.non_makeup_face_mask_no_neck_no_ear[:, 0:1, :, :], self.makeup_color_change3_img * self.makeup_face_mask_no_neck_no_ear[:, 0:1, :, :])+\ self.gram_style(self.transfer_output_data['transfer_img'] * self.non_makeup_face_mask_no_neck_no_ear[:, 0:1, :, :], self.makeup_color_change4_img * self.makeup_face_mask_no_neck_no_ear[:, 0:1, :, :]) loss_G_contrastive=-torch.log(1.-positive_distance/negative_distance) loss_G_contrastive=loss_G_contrastive*self.weight_contrastive # adv loss_G_GAN = self.backward_G_GAN(self.transfer_output_data['transfer_img'], self.dis) loss_G_GAN = loss_G_GAN * self.weight_adv # Self supervision loss_G_selfL1 = self.criterionL1(self.rec_output_data['transfer_img'], self.makeup_color_change_img) loss_G_selfL1 = loss_G_selfL1 * self.weight_self_recL1 # Cycle loss_G_cycleL1 = self.criterionL1(self.cycle_output_data['transfer_img'], self.makeup_color_img) loss_G_cycleL1 = loss_G_cycleL1 * self.weight_cycleL1 loss_G = loss_G_identity + loss_G_GAN + loss_G_selfL1 + loss_G_back + loss_G_sem + loss_G_corr+loss_G_cycleL1+loss_G_contrastive loss_G.backward() self.G_loss = loss_G.item() self.loss_G_identity = loss_G_identity.item() self.loss_G_back = loss_G_back.item() self.loss_G_GAN = loss_G_GAN.item() self.loss_G_selfL1 = loss_G_selfL1.item() self.loss_G_cycleL1 = loss_G_cycleL1.item() self.loss_G_sem = loss_G_sem.item() self.loss_G_contrastive = loss_G_contrastive.item() def backward_G_GAN(self, fake, netD): outs_fake = netD.forward(fake) loss = self.criterionGAN(outs_fake, target_is_real=True, for_discriminator=False) return loss def set_scheduler(self, opts, last_ep=0): self.dis_sch = utils.get_scheduler(self.dis_opt, opts, last_ep) self.gen_sch = utils.get_scheduler(self.gen_opt, opts, last_ep) def update_lr(self): self.dis_sch.step() self.gen_sch.step() def resume(self, model_dir): checkpoint = torch.load(model_dir) self.gen.load_state_dict(checkpoint['gen']) # optimizer return checkpoint['ep'], checkpoint['total_it'] def save(self, filename, ep, total_it): state = { 'gen': self.gen.state_dict(), 'ep': ep, 'total_it': total_it } torch.save(state, filename) return def normalize_image(self, x): return x[:, 0:3, :, :] def assemble_outputs(self): # row1 non_makeup_color_img = self.normalize_image(self.non_makeup_color_img).detach() makeup_color_img = self.normalize_image(self.makeup_color_img).detach() transfer_content_img=self.normalize_image(self.transfer_output_data['source_face_content']).detach() transfer_makeup_img = self.normalize_image(self.transfer_output_data['ref_face_style']).detach() transfer_corr1 = F.interpolate(self.makeup_color_img, scale_factor=1 / 4, mode='nearest') transfer_corr1=self.gen.forward_cross_attention(source_parse=self.non_makeup_split_parse, ref_parse=self.makeup_split_parse, ref_img=transfer_corr1) transfer_corr1 = F.interpolate(transfer_corr1, scale_factor=4, mode='nearest') transfer_corr1 = self.normalize_image(transfer_corr1).detach() transfer_img = self.normalize_image(self.transfer_output_data['transfer_img']).detach() makeup_color_change2_img = self.normalize_image(self.makeup_color_change2_img).detach() cycle_content_img = self.normalize_image(self.cycle_output_data['source_face_content']).detach() cycle_makeup_img = self.normalize_image(self.cycle_output_data['ref_face_style']).detach() cycle_corr1 = F.interpolate(self.transfer_color_face_down4_warp, scale_factor=4, mode='nearest') cycle_img=self.normalize_image(self.cycle_output_data['transfer_img']) makeup_color_change3_img = self.normalize_image(self.makeup_color_change3_img).detach() makeup_color_change_img = self.normalize_image(self.makeup_color_change_img).detach() makeup_color_change_warp_img = self.normalize_image(self.makeup_color_change_warp_img).detach() rec_content_img = self.normalize_image(self.rec_output_data['source_face_content']).detach() rec_makeup_img = self.normalize_image(self.rec_output_data['ref_face_style']).detach() rec_corr1 = F.interpolate(self.makeup_color_change_warp_img, scale_factor=1 / 4, mode='nearest') rec_corr1 = self.gen.forward_cross_attention(source_parse=self.makeup_split_parse, ref_parse=self.makeup_split_warp_parse, ref_img=rec_corr1) rec_corr1 = F.interpolate(rec_corr1, scale_factor=4, mode='nearest') rec_corr1 = self.normalize_image(rec_corr1).detach() rec_img = self.normalize_image(self.rec_output_data['transfer_img']).detach() makeup_color_change4_img = self.normalize_image(self.makeup_color_change4_img).detach() # print(non_makeup_color_img[0:1, ::].shape) # print(makeup_color_img[0:1, ::].shape) # print(transfer_content_img[0:1, ::].shape) # print(transfer_makeup_img[0:1, ::].shape) row1 = torch.cat( (non_makeup_color_img[0:1, ::], makeup_color_img[0:1, ::], transfer_content_img[0:1, ::].repeat(1,3,1,1),transfer_makeup_img[0:1, ::], transfer_corr1[0:1, ::], transfer_img[0:1, ::],makeup_color_change2_img[0:1, ::]), 3) row2 = torch.cat( (makeup_color_img[0:1, ::], transfer_img[0:1, ::], cycle_content_img[0:1, ::].repeat(1,3,1,1), cycle_makeup_img[0:1, ::], cycle_corr1[0:1, ::],cycle_img[0:1, ::],makeup_color_change3_img[0:1, ::]), 3) row3 = torch.cat( (makeup_color_change_img[0:1, ::], makeup_color_change_warp_img[0:1, ::], rec_content_img[0:1, ::].repeat(1,3,1,1), rec_makeup_img[0:1, ::], rec_corr1[0:1, ::],rec_img[0:1, ::],makeup_color_change4_img[0:1, ::]), 3) return torch.cat((row1, row2,row3), 2) def test_pair(self, data): self.non_makeup_color_img = data['non_makeup_color_img'].to(self.gpu).detach() self.non_makeup_split_parse = data['non_makeup_split_parse'].to(self.gpu).detach() self.non_makeup_all_mask = data['non_makeup_all_mask'].to(self.gpu).detach() self.non_makeup_face_mask = data['non_makeup_face_mask'].to(self.gpu).detach() self.non_makeup_brow_mask = data['non_makeup_brow_mask'].to(self.gpu).detach() self.non_makeup_eye_mask = data['non_makeup_eye_mask'].to(self.gpu).detach() self.non_makeup_lip_mask = data['non_makeup_lip_mask'].to(self.gpu).detach() self.non_makeup_face_mask_no_neck_no_ear = self.non_makeup_face_mask + self.non_makeup_brow_mask + self.non_makeup_eye_mask + self.non_makeup_lip_mask self.makeup_color_img = data['makeup_color_img'].to(self.gpu).detach() self.makeup_split_parse = data['makeup_split_parse'].to(self.gpu).detach() self.makeup_all_mask = data['makeup_all_mask'].to(self.gpu).detach() self.makeup_face_mask = data['makeup_face_mask'].to(self.gpu).detach() self.makeup_brow_mask = data['makeup_brow_mask'].to(self.gpu).detach() self.makeup_eye_mask = data['makeup_eye_mask'].to(self.gpu).detach() self.makeup_lip_mask = data['makeup_lip_mask'].to(self.gpu).detach() self.makeup_face_mask_no_neck_no_ear = self.makeup_face_mask + self.makeup_brow_mask + self.makeup_eye_mask + self.makeup_lip_mask with torch.no_grad(): self.transfer_output_data = self.gen(source_img=self.non_makeup_color_img, source_parse=self.non_makeup_split_parse, source_all_mask=self.non_makeup_all_mask, ref_img=self.makeup_color_img, ref_parse=self.makeup_split_parse, ref_all_mask=self.makeup_all_mask) non_makeup_color_img = self.normalize_image(self.non_makeup_color_img).detach() makeup_color_img = self.normalize_image(self.makeup_color_img).detach() transfer_content_img = self.normalize_image(self.transfer_output_data['source_face_content']).detach() transfer_makeup_img = self.normalize_image(self.transfer_output_data['ref_face_style']).detach() transfer_corr1 = F.interpolate(self.makeup_color_img, scale_factor=1 / 4, mode='nearest') transfer_corr1 = self.gen.forward_cross_attention(source_parse=self.non_makeup_split_parse, ref_parse=self.makeup_split_parse, ref_img=transfer_corr1) non_makeup_all_mask_down4 = F.interpolate(self.non_makeup_all_mask, scale_factor=1 / 4, mode='nearest') transfer_corr1 = F.interpolate(transfer_corr1 * non_makeup_all_mask_down4, scale_factor=4, mode='nearest') transfer_corr1 = self.normalize_image(transfer_corr1).detach() transfer_img = self.normalize_image(self.transfer_output_data['transfer_img']).detach() row1 = torch.cat( (non_makeup_color_img[0:1, ::], makeup_color_img[0:1, ::], transfer_corr1[0:1, ::], transfer_img[0:1, ::]), 3) return row1