APDrawingGAN2/models/apdrawingpp_style_model.py

import torch
from util.image_pool import ImagePool
from .base_model import BaseModel
from . import networks
import os
import math

W = 11
aa = int(math.floor(512. / W))
res = 512 - W * aa


def padpart(A, part, centers, opt, device):
    IMAGE_SIZE = opt.fineSize
    bs, nc, _, _ = A.shape
    ratio = IMAGE_SIZE / 256
    NOSE_W = opt.NOSE_W * ratio
    NOSE_H = opt.NOSE_H * ratio
    EYE_W = opt.EYE_W * ratio
    EYE_H = opt.EYE_H * ratio
    MOUTH_W = opt.MOUTH_W * ratio
    MOUTH_H = opt.MOUTH_H * ratio
    A_p = torch.ones((bs, nc, IMAGE_SIZE, IMAGE_SIZE)).to(device)
    padvalue = -1  # black
    for i in range(bs):
        center = centers[i]
        if part == 'nose':
            A_p[i] = torch.nn.ConstantPad2d((
                                            int(center[2, 0] - NOSE_W / 2), IMAGE_SIZE - int(center[2, 0] + NOSE_W / 2),
                                            int(center[2, 1] - NOSE_H / 2),
                                            IMAGE_SIZE - int(center[2, 1] + NOSE_H / 2)), padvalue)(A[i])
        elif part == 'eyel':
            A_p[i] = torch.nn.ConstantPad2d((int(center[0, 0] - EYE_W / 2), IMAGE_SIZE - int(center[0, 0] + EYE_W / 2),
                                             int(center[0, 1] - EYE_H / 2), IMAGE_SIZE - int(center[0, 1] + EYE_H / 2)),
                                            padvalue)(A[i])
        elif part == 'eyer':
            A_p[i] = torch.nn.ConstantPad2d((int(center[1, 0] - EYE_W / 2), IMAGE_SIZE - int(center[1, 0] + EYE_W / 2),
                                             int(center[1, 1] - EYE_H / 2), IMAGE_SIZE - int(center[1, 1] + EYE_H / 2)),
                                            padvalue)(A[i])
        elif part == 'mouth':
            A_p[i] = torch.nn.ConstantPad2d((int(center[3, 0] - MOUTH_W / 2),
                                             IMAGE_SIZE - int(center[3, 0] + MOUTH_W / 2),
                                             int(center[3, 1] - MOUTH_H / 2),
                                             IMAGE_SIZE - int(center[3, 1] + MOUTH_H / 2)), padvalue)(A[i])
    return A_p


import numpy as np


def nonlinearDt(dt, type='atan',
                xmax=torch.Tensor([10.0])):  # dt in [0,1], first multiply xmax(>1), then remap to [0,1]
    if type == 'atan':
        nldt = torch.atan(dt * xmax) / torch.atan(xmax)
    elif type == 'sigmoid':
        nldt = (torch.sigmoid(dt * xmax) - 0.5) / (torch.sigmoid(xmax) - 0.5)
    elif type == 'tanh':
        nldt = torch.tanh(dt * xmax) / torch.tanh(xmax)
    elif type == 'pow':
        nldt = torch.pow(dt * xmax, 2) / torch.pow(xmax, 2)
    elif type == 'exp':
        if xmax.item() > 1:
            xmax = xmax / 3
        nldt = (torch.exp(dt * xmax) - 1) / (torch.exp(xmax) - 1)
    # print("remap dt:", type, xmax.item())
    return nldt


class APDrawingPPStyleModel(BaseModel):
    def name(self):
        return 'APDrawingPPStyleModel'

    @staticmethod
    def modify_commandline_options(parser, is_train=True):

        # changing the default values to match the pix2pix paper
        # (https://phillipi.github.io/pix2pix/)
        parser.set_defaults(pool_size=0, no_lsgan=True, norm='batch')  # no_lsgan=True, use_lsgan=False
        parser.set_defaults(dataset_mode='aligned')
        parser.set_defaults(auxiliary_root='auxiliaryeye2o')
        parser.set_defaults(use_local=True, hair_local=True, bg_local=True)
        parser.set_defaults(discriminator_local=True, gan_loss_strategy=2)
        parser.set_defaults(chamfer_loss=True, dt_nonlinear='exp', lambda_chamfer=0.35, lambda_chamfer2=0.35)
        parser.set_defaults(nose_ae=True, others_ae=True, compactmask=True, MOUTH_H=56)
        parser.set_defaults(soft_border=1, batch_size=1, save_epoch_freq=25)
        parser.add_argument('--nnG_hairc', type=int, default=6, help='nnG for hair classifier')
        parser.add_argument('--use_resnet', action='store_true', help='use resnet for generator')
        parser.add_argument('--regarch', type=int, default=4, help='architecture for netRegressor')
        if is_train:
            parser.add_argument('--lambda_L1', type=float, default=100.0, help='weight for L1 loss')
            parser.add_argument('--lambda_local', type=float, default=25.0, help='weight for Local loss')
            parser.set_defaults(netG_dt='unet_512')
            parser.set_defaults(netG_line='unet_512')

        return parser

    def initialize(self, opt):
        BaseModel.initialize(self, opt)
        self.isTrain = opt.isTrain
        # specify the training losses you want to print out. The program will call base_model.get_current_losses
        self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
        if self.isTrain and self.opt.no_l1_loss:
            self.loss_names = ['G_GAN', 'D_real', 'D_fake']
        if self.isTrain and self.opt.use_local and not self.opt.no_G_local_loss:
            self.loss_names.append('G_local')
            self.loss_names.append('G_hair_local')
            self.loss_names.append('G_bg_local')
        if self.isTrain and self.opt.discriminator_local:
            self.loss_names.append('D_real_local')
            self.loss_names.append('D_fake_local')
            self.loss_names.append('G_GAN_local')
        if self.isTrain and self.opt.chamfer_loss:
            self.loss_names.append('G_chamfer')
            self.loss_names.append('G_chamfer2')
        if self.isTrain and self.opt.continuity_loss:
            self.loss_names.append('G_continuity')
        self.loss_names.append('G')
        print('loss_names', self.loss_names)
        # specify the images you want to save/display. The program will call base_model.get_current_visuals
        self.visual_names = ['real_A', 'fake_B', 'real_B']
        if self.opt.use_local:
            self.visual_names += ['fake_B0', 'fake_B1']
            self.visual_names += ['fake_B_hair', 'real_B_hair', 'real_A_hair']
            self.visual_names += ['fake_B_bg', 'real_B_bg', 'real_A_bg']
            if self.opt.region_enm in [0, 1]:
                if self.opt.nose_ae:
                    self.visual_names += ['fake_B_nose_v', 'fake_B_nose_v1', 'fake_B_nose_v2', 'cmask1no']
                if self.opt.others_ae:
                    self.visual_names += ['fake_B_eyel_v', 'fake_B_eyel_v1', 'fake_B_eyel_v2', 'cmask1el']
                    self.visual_names += ['fake_B_eyer_v', 'fake_B_eyer_v1', 'fake_B_eyer_v2', 'cmask1er']
                    self.visual_names += ['fake_B_mouth_v', 'fake_B_mouth_v1', 'fake_B_mouth_v2', 'cmask1mo']
            elif self.opt.region_enm in [2]:
                self.visual_names += ['fake_B_nose', 'fake_B_eyel', 'fake_B_eyer', 'fake_B_mouth']
        if self.isTrain and self.opt.chamfer_loss:
            self.visual_names += ['dt1', 'dt2']
            self.visual_names += ['dt1gt', 'dt2gt']
        if self.isTrain and self.opt.soft_border:
            self.visual_names += ['mask']
        if not self.isTrain and self.opt.save2:
            self.visual_names = ['real_A', 'fake_B']
        print('visuals', self.visual_names)
        # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
        self.auxiliary_model_names = []
        if self.isTrain:
            self.model_names = ['G', 'D']
            if self.opt.discriminator_local:
                self.model_names += ['DLEyel', 'DLEyer', 'DLNose', 'DLMouth', 'DLHair', 'DLBG']
            # auxiliary nets for loss calculation
            if self.opt.chamfer_loss:
                self.auxiliary_model_names += ['DT1', 'DT2']
                self.auxiliary_model_names += ['Line1', 'Line2']
            if self.opt.continuity_loss:
                self.auxiliary_model_names += ['Regressor']
        else:  # during test time, only load Gs
            self.model_names = ['G']
            if self.opt.test_continuity_loss:
                self.auxiliary_model_names += ['Regressor']
        if self.opt.use_local:
            self.model_names += ['GLEyel', 'GLEyer', 'GLNose', 'GLMouth', 'GLHair', 'GLBG', 'GCombine']
            self.auxiliary_model_names += ['CLm', 'CLh']
            # auxiliary nets for local output refinement
            if self.opt.nose_ae:
                self.auxiliary_model_names += ['AE']
            if self.opt.others_ae:
                self.auxiliary_model_names += ['AEel', 'AEer', 'AEmowhite', 'AEmoblack']
        print('model_names', self.model_names)
        print('auxiliary_model_names', self.auxiliary_model_names)
        # load/define networks
        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                      opt.nnG)
        print('netG', opt.netG)

        if self.isTrain:
            use_sigmoid = opt.no_lsgan
            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                          self.gpu_ids)
            print('netD', opt.netD, opt.n_layers_D)
            if self.opt.discriminator_local:
                self.netDLEyel = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                                   opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                                   self.gpu_ids)
                self.netDLEyer = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                                   opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                                   self.gpu_ids)
                self.netDLNose = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                                   opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                                   self.gpu_ids)
                self.netDLMouth = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                                    opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                                    self.gpu_ids)
                self.netDLHair = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                                   opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                                   self.gpu_ids)
                self.netDLBG = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                                 opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain,
                                                 self.gpu_ids)

        if self.opt.use_local:
            netlocal1 = 'partunet' if self.opt.use_resnet == 0 else 'resnet_nblocks'
            netlocal2 = 'partunet2' if self.opt.use_resnet == 0 else 'resnet_6blocks'
            netlocal2_style = 'partunet2style' if self.opt.use_resnet == 0 else 'resnet_style2_6blocks'
            self.netGLEyel = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
                                               not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
            self.netGLEyer = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
                                               not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
            self.netGLNose = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
                                               not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
            self.netGLMouth = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
                                                not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
            self.netGLHair = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal2_style, opt.norm,
                                               not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=4,
                                               extra_channel=3)
            self.netGLBG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal2, opt.norm,
                                             not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=4)
            # by default combiner_type is combiner, which uses resnet
            print('combiner_type', self.opt.combiner_type)
            self.netGCombine = networks.define_G(2 * opt.output_nc, opt.output_nc, opt.ngf, self.opt.combiner_type,
                                                 opt.norm,
                                                 not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, 2)
            # auxiliary classifiers for mouth and hair
            ratio = self.opt.fineSize / 256
            self.MOUTH_H = int(self.opt.MOUTH_H * ratio)
            self.MOUTH_W = int(self.opt.MOUTH_W * ratio)
            self.netCLm = networks.define_G(opt.input_nc, 2, opt.ngf, 'classifier', opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                            nnG=3, ae_h=self.MOUTH_H, ae_w=self.MOUTH_W)
            self.netCLh = networks.define_G(opt.input_nc, 3, opt.ngf, 'classifier', opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                            nnG=opt.nnG_hairc, ae_h=opt.fineSize, ae_w=opt.fineSize)

        if self.isTrain:
            self.fake_AB_pool = ImagePool(opt.pool_size)
            # define loss functions
            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan).to(self.device)
            self.criterionL1 = torch.nn.L1Loss()

            # initialize optimizers
            self.optimizers = []
            if not self.opt.use_local:
                print('G_params 1 components')
                self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))
            else:
                G_params = list(self.netG.parameters()) + list(self.netGLEyel.parameters()) + list(
                    self.netGLEyer.parameters()) + list(self.netGLNose.parameters()) + list(
                    self.netGLMouth.parameters()) + list(self.netGCombine.parameters()) + list(
                    self.netGLHair.parameters()) + list(self.netGLBG.parameters())
                print('G_params 8 components')
                self.optimizer_G = torch.optim.Adam(G_params,
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))

            if not self.opt.discriminator_local:
                print('D_params 1 components')
                self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))
            else:  # self.opt.discriminator_local == True
                D_params = list(self.netD.parameters()) + list(self.netDLEyel.parameters()) + list(
                    self.netDLEyer.parameters()) + list(self.netDLNose.parameters()) + list(
                    self.netDLMouth.parameters()) + list(self.netDLHair.parameters()) + list(self.netDLBG.parameters())
                print('D_params 7 components')
                self.optimizer_D = torch.optim.Adam(D_params,
                                                    lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

        # ==================================auxiliary nets (loaded, parameters fixed)=============================
        if self.opt.use_local and self.opt.nose_ae:
            ratio = self.opt.fineSize / 256
            NOSE_H = self.opt.NOSE_H * ratio
            NOSE_W = self.opt.NOSE_W * ratio
            self.netAE = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
                                           not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                           latent_dim=self.opt.ae_latentno, ae_h=NOSE_H, ae_w=NOSE_W)
            self.set_requires_grad(self.netAE, False)
        if self.opt.use_local and self.opt.others_ae:
            ratio = self.opt.fineSize / 256
            EYE_H = self.opt.EYE_H * ratio
            EYE_W = self.opt.EYE_W * ratio
            MOUTH_H = self.opt.MOUTH_H * ratio
            MOUTH_W = self.opt.MOUTH_W * ratio
            self.netAEel = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
                                             not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                             latent_dim=self.opt.ae_latenteye, ae_h=EYE_H, ae_w=EYE_W)
            self.netAEer = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
                                             not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                             latent_dim=self.opt.ae_latenteye, ae_h=EYE_H, ae_w=EYE_W)
            self.netAEmowhite = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
                                                  not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                                  latent_dim=self.opt.ae_latentmo, ae_h=MOUTH_H, ae_w=MOUTH_W)
            self.netAEmoblack = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
                                                  not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                                  latent_dim=self.opt.ae_latentmo, ae_h=MOUTH_H, ae_w=MOUTH_W)
            self.set_requires_grad(self.netAEel, False)
            self.set_requires_grad(self.netAEer, False)
            self.set_requires_grad(self.netAEmowhite, False)
            self.set_requires_grad(self.netAEmoblack, False)

        if self.isTrain and self.opt.continuity_loss:
            self.nc = 1
            self.netRegressor = networks.define_G(self.nc, 1, opt.ngf, 'regressor', opt.norm,
                                                  not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p,
                                                  nnG=opt.regarch)
            self.set_requires_grad(self.netRegressor, False)

        if self.isTrain and self.opt.chamfer_loss:
            self.nc = 1
            self.netDT1 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_dt, opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
            self.netDT2 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_dt, opt.norm,
                                            not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
            self.set_requires_grad(self.netDT1, False)
            self.set_requires_grad(self.netDT2, False)
            self.netLine1 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_line, opt.norm,
                                              not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
            self.netLine2 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_line, opt.norm,
                                              not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
            self.set_requires_grad(self.netLine1, False)
            self.set_requires_grad(self.netLine2, False)

        # ==================================for test (nets loaded, parameters fixed)=============================
        if not self.isTrain and self.opt.test_continuity_loss:
            self.nc = 1
            self.netRegressor = networks.define_G(self.nc, 1, opt.ngf, 'regressor', opt.norm,
                                                  not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
                                                  nnG=opt.regarch)
            self.set_requires_grad(self.netRegressor, False)

    def set_input(self, input):
        AtoB = self.opt.which_direction == 'AtoB'
        self.real_A = input['A' if AtoB else 'B'].to(self.device)
        self.real_B = input['B' if AtoB else 'A'].to(self.device)
        self.image_paths = input['A_paths' if AtoB else 'B_paths']
        self.batch_size = len(self.image_paths)
        if self.opt.use_local:
            self.real_A_eyel = input['eyel_A'].to(self.device)
            self.real_A_eyer = input['eyer_A'].to(self.device)
            self.real_A_nose = input['nose_A'].to(self.device)
            self.real_A_mouth = input['mouth_A'].to(self.device)
            self.real_B_eyel = input['eyel_B'].to(self.device)
            self.real_B_eyer = input['eyer_B'].to(self.device)
            self.real_B_nose = input['nose_B'].to(self.device)
            self.real_B_mouth = input['mouth_B'].to(self.device)
            if self.opt.region_enm in [0, 1]:
                self.center = input['center']
            if self.opt.soft_border:
                self.softel = input['soft_eyel_mask'].to(self.device)
                self.softer = input['soft_eyer_mask'].to(self.device)
                self.softno = input['soft_nose_mask'].to(self.device)
                self.softmo = input['soft_mouth_mask'].to(self.device)
            if self.opt.compactmask:
                self.cmask = input['cmask'].to(self.device)
                self.cmask1 = self.cmask * 2 - 1  # [0,1]->[-1,1]
                self.cmaskel = input['cmaskel'].to(self.device)
                self.cmask1el = self.cmaskel * 2 - 1
                self.cmasker = input['cmasker'].to(self.device)
                self.cmask1er = self.cmasker * 2 - 1
                self.cmaskmo = input['cmaskmo'].to(self.device)
                self.cmask1mo = self.cmaskmo * 2 - 1
            self.real_A_hair = input['hair_A'].to(self.device)
            self.real_B_hair = input['hair_B'].to(self.device)
            self.mask = input['mask'].to(self.device)  # mask for non-eyes,nose,mouth
            self.mask2 = input['mask2'].to(self.device)  # mask for non-bg
            self.real_A_bg = input['bg_A'].to(self.device)
            self.real_B_bg = input['bg_B'].to(self.device)
        if (self.isTrain and self.opt.chamfer_loss):
            self.dt1gt = input['dt1gt'].to(self.device)
            self.dt2gt = input['dt2gt'].to(self.device)
        if self.isTrain and self.opt.emphasis_conti_face:
            self.face_mask = input['face_mask'].cuda(self.gpu_ids_p[0])

    def getonehot(self, outputs, classes):
        [maxv, index] = torch.max(outputs, 1)
        y = torch.unsqueeze(index, 1)
        onehot = torch.FloatTensor(self.batch_size, classes).to(self.device)
        onehot.zero_()
        onehot.scatter_(1, y, 1)
        return onehot

    def forward(self):
        if not self.opt.use_local:
            self.fake_B = self.netG(self.real_A)
        else:
            self.fake_B0 = self.netG(self.real_A)
            # EYES, MOUTH
            outputs1 = self.netCLm(self.real_A_mouth)
            onehot1 = self.getonehot(outputs1, 2)

            if not self.opt.others_ae:
                fake_B_eyel = self.netGLEyel(self.real_A_eyel)
                fake_B_eyer = self.netGLEyer(self.real_A_eyer)
                fake_B_mouth = self.netGLMouth(self.real_A_mouth)
            else:  # use AE that only constains compact region, need cmask!
                self.fake_B_eyel1 = self.netGLEyel(self.real_A_eyel)
                self.fake_B_eyer1 = self.netGLEyer(self.real_A_eyer)
                self.fake_B_mouth1 = self.netGLMouth(self.real_A_mouth)
                self.fake_B_eyel2, _ = self.netAEel(self.fake_B_eyel1)
                self.fake_B_eyer2, _ = self.netAEer(self.fake_B_eyer1)
                # USE 2 AEs
                self.fake_B_mouth2 = torch.FloatTensor(self.batch_size, self.opt.output_nc, self.MOUTH_H,
                                                       self.MOUTH_W).to(self.device)
                for i in range(self.batch_size):
                    if onehot1[i][0] == 1:
                        self.fake_B_mouth2[i], _ = self.netAEmowhite(self.fake_B_mouth1[i].unsqueeze(0))
                        # print('AEmowhite')
                    elif onehot1[i][1] == 1:
                        self.fake_B_mouth2[i], _ = self.netAEmoblack(self.fake_B_mouth1[i].unsqueeze(0))
                        # print('AEmoblack')
                fake_B_eyel = self.add_with_mask(self.fake_B_eyel2, self.fake_B_eyel1, self.cmaskel)
                fake_B_eyer = self.add_with_mask(self.fake_B_eyer2, self.fake_B_eyer1, self.cmasker)
                fake_B_mouth = self.add_with_mask(self.fake_B_mouth2, self.fake_B_mouth1, self.cmaskmo)
            # NOSE
            if not self.opt.nose_ae:
                fake_B_nose = self.netGLNose(self.real_A_nose)
            else:  # use AE that only constains compact region, need cmask!
                self.fake_B_nose1 = self.netGLNose(self.real_A_nose)
                self.fake_B_nose2, _ = self.netAE(self.fake_B_nose1)
                fake_B_nose = self.add_with_mask(self.fake_B_nose2, self.fake_B_nose1, self.cmask)

            # for visuals and later local loss
            if self.opt.region_enm in [0, 1]:
                self.fake_B_nose = fake_B_nose
                self.fake_B_eyel = fake_B_eyel
                self.fake_B_eyer = fake_B_eyer
                self.fake_B_mouth = fake_B_mouth
                # for soft border of 4 rectangle facial feature
                if self.opt.region_enm == 0 and self.opt.soft_border:
                    self.fake_B_nose = self.masked(fake_B_nose, self.softno)
                    self.fake_B_eyel = self.masked(fake_B_eyel, self.softel)
                    self.fake_B_eyer = self.masked(fake_B_eyer, self.softer)
                    self.fake_B_mouth = self.masked(fake_B_mouth, self.softmo)
            elif self.opt.region_enm in [2]:  # need to multiply cmask
                self.fake_B_nose = self.masked(fake_B_nose, self.cmask)
                self.fake_B_eyel = self.masked(fake_B_eyel, self.cmaskel)
                self.fake_B_eyer = self.masked(fake_B_eyer, self.cmasker)
                self.fake_B_mouth = self.masked(fake_B_mouth, self.cmaskmo)

            # HAIR, BG AND PARTCOMBINE
            outputs2 = self.netCLh(self.real_A_hair)
            onehot2 = self.getonehot(outputs2, 3)

            if not self.isTrain:
                opt = self.opt
                if opt.imagefolder == 'images':
                    file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch),
                                             'styleonehot.txt')
                else:
                    file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch),
                                             opt.imagefolder, 'styleonehot.txt')
                message = '%s [%d %d] [%d %d %d]' % (self.image_paths[0], onehot1[0][0], onehot1[0][1],
                                                     onehot2[0][0], onehot2[0][1], onehot2[0][2])
                with open(file_name, 'a+') as s_file:
                    s_file.write(message)
                    s_file.write('\n')

            fake_B_hair = self.netGLHair(self.real_A_hair, onehot2)
            fake_B_bg = self.netGLBG(self.real_A_bg)
            self.fake_B_hair = self.masked(fake_B_hair, self.mask * self.mask2)
            self.fake_B_bg = self.masked(fake_B_bg, self.inverse_mask(self.mask2))
            if not self.opt.compactmask:
                self.fake_B1 = self.partCombiner2_bg(fake_B_eyel, fake_B_eyer, fake_B_nose, fake_B_mouth, fake_B_hair,
                                                     fake_B_bg, self.mask * self.mask2, self.inverse_mask(self.mask2),
                                                     self.opt.comb_op)
            else:
                self.fake_B1 = self.partCombiner2_bg(fake_B_eyel, fake_B_eyer, fake_B_nose, fake_B_mouth, fake_B_hair,
                                                     fake_B_bg, self.mask * self.mask2, self.inverse_mask(self.mask2),
                                                     self.opt.comb_op, self.opt.region_enm, self.cmaskel, self.cmasker,
                                                     self.cmask, self.cmaskmo)

            self.fake_B = self.netGCombine(torch.cat([self.fake_B0, self.fake_B1], 1))

            # for AE visuals
            if self.opt.region_enm in [0, 1]:
                if self.opt.nose_ae:
                    self.fake_B_nose_v = padpart(self.fake_B_nose, 'nose', self.center, self.opt, self.device)
                    self.fake_B_nose_v1 = padpart(self.fake_B_nose1, 'nose', self.center, self.opt, self.device)
                    self.fake_B_nose_v2 = padpart(self.fake_B_nose2, 'nose', self.center, self.opt, self.device)
                    self.cmask1no = padpart(self.cmask1, 'nose', self.center, self.opt, self.device)
                if self.opt.others_ae:
                    self.fake_B_eyel_v = padpart(self.fake_B_eyel, 'eyel', self.center, self.opt, self.device)
                    self.fake_B_eyel_v1 = padpart(self.fake_B_eyel1, 'eyel', self.center, self.opt, self.device)
                    self.fake_B_eyel_v2 = padpart(self.fake_B_eyel2, 'eyel', self.center, self.opt, self.device)
                    self.cmask1el = padpart(self.cmask1el, 'eyel', self.center, self.opt, self.device)
                    self.fake_B_eyer_v = padpart(self.fake_B_eyer, 'eyer', self.center, self.opt, self.device)
                    self.fake_B_eyer_v1 = padpart(self.fake_B_eyer1, 'eyer', self.center, self.opt, self.device)
                    self.fake_B_eyer_v2 = padpart(self.fake_B_eyer2, 'eyer', self.center, self.opt, self.device)
                    self.cmask1er = padpart(self.cmask1er, 'eyer', self.center, self.opt, self.device)
                    self.fake_B_mouth_v = padpart(self.fake_B_mouth, 'mouth', self.center, self.opt, self.device)
                    self.fake_B_mouth_v1 = padpart(self.fake_B_mouth1, 'mouth', self.center, self.opt, self.device)
                    self.fake_B_mouth_v2 = padpart(self.fake_B_mouth2, 'mouth', self.center, self.opt, self.device)
                    self.cmask1mo = padpart(self.cmask1mo, 'mouth', self.center, self.opt, self.device)

            if not self.isTrain and self.opt.test_continuity_loss:
                self.ContinuityForTest(real=1)

    def backward_D(self):
        # Fake
        # stop backprop to the generator by detaching fake_B
        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
        # print('fake_AB', fake_AB.shape) # (1,4,512,512)
        pred_fake = self.netD(fake_AB.detach())  # by detach, not affect G's gradient
        self.loss_D_fake = self.criterionGAN(pred_fake, False)
        if self.opt.discriminator_local:
            fake_AB_parts = self.getLocalParts(fake_AB)
            local_names = ['DLEyel', 'DLEyer', 'DLNose', 'DLMouth', 'DLHair', 'DLBG']
            self.loss_D_fake_local = 0
            for i in range(len(fake_AB_parts)):
                net = getattr(self, 'net' + local_names[i])
                pred_fake_tmp = net(fake_AB_parts[i].detach())
                addw = self.getaddw(local_names[i])
                self.loss_D_fake_local = self.loss_D_fake_local + self.criterionGAN(pred_fake_tmp, False) * addw
            self.loss_D_fake = self.loss_D_fake + self.loss_D_fake_local

        # Real
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        pred_real = self.netD(real_AB)
        self.loss_D_real = self.criterionGAN(pred_real, True)
        if self.opt.discriminator_local:
            real_AB_parts = self.getLocalParts(real_AB)
            local_names = ['DLEyel', 'DLEyer', 'DLNose', 'DLMouth', 'DLHair', 'DLBG']
            self.loss_D_real_local = 0
            for i in range(len(real_AB_parts)):
                net = getattr(self, 'net' + local_names[i])
                pred_real_tmp = net(real_AB_parts[i])
                addw = self.getaddw(local_names[i])
                self.loss_D_real_local = self.loss_D_real_local + self.criterionGAN(pred_real_tmp, True) * addw
            self.loss_D_real = self.loss_D_real + self.loss_D_real_local

        # Combined loss
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5

        self.loss_D.backward()

    def backward_G(self):
        # First, G(A) should fake the discriminator
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB)  # (1,4,512,512)->(1,1,30,30)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
        if self.opt.discriminator_local:
            fake_AB_parts = self.getLocalParts(fake_AB)
            local_names = ['DLEyel', 'DLEyer', 'DLNose', 'DLMouth', 'DLHair', 'DLBG']
            self.loss_G_GAN_local = 0  # G_GAN_local is then added into G_GAN
            for i in range(len(fake_AB_parts)):
                net = getattr(self, 'net' + local_names[i])
                pred_fake_tmp = net(fake_AB_parts[i])
                addw = self.getaddw(local_names[i])
                self.loss_G_GAN_local = self.loss_G_GAN_local + self.criterionGAN(pred_fake_tmp, True) * addw
            if self.opt.gan_loss_strategy == 1:
                self.loss_G_GAN = (self.loss_G_GAN + self.loss_G_GAN_local) / (len(fake_AB_parts) + 1)
            elif self.opt.gan_loss_strategy == 2:
                self.loss_G_GAN_local = self.loss_G_GAN_local * 0.25
                self.loss_G_GAN = self.loss_G_GAN + self.loss_G_GAN_local

        # Second, G(A) = B
        if not self.opt.no_l1_loss:
            self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1

        if self.opt.use_local and not self.opt.no_G_local_loss:
            local_names = ['eyel', 'eyer', 'nose', 'mouth']
            self.loss_G_local = 0
            for i in range(len(local_names)):
                fakeblocal = getattr(self, 'fake_B_' + local_names[i])
                realblocal = getattr(self, 'real_B_' + local_names[i])
                addw = self.getaddw(local_names[i])
                self.loss_G_local = self.loss_G_local + self.criterionL1(fakeblocal,
                                                                         realblocal) * self.opt.lambda_local * addw
            self.loss_G_hair_local = self.criterionL1(self.fake_B_hair,
                                                      self.real_B_hair) * self.opt.lambda_local * self.opt.addw_hair
            self.loss_G_bg_local = self.criterionL1(self.fake_B_bg,
                                                    self.real_B_bg) * self.opt.lambda_local * self.opt.addw_bg

        # Third, chamfer matching (assume chamfer_2way and chamfer_only_line is true)
        if self.opt.chamfer_loss:
            if self.fake_B.shape[1] == 3:
                tmp = self.fake_B[:, 0, ...] * 0.299 + self.fake_B[:, 1, ...] * 0.587 + self.fake_B[:, 2, ...] * 0.114
                fake_B_gray = tmp.unsqueeze(1)
            else:
                fake_B_gray = self.fake_B
            if self.real_B.shape[1] == 3:
                tmp = self.real_B[:, 0, ...] * 0.299 + self.real_B[:, 1, ...] * 0.587 + self.real_B[:, 2, ...] * 0.114
                real_B_gray = tmp.unsqueeze(1)
            else:
                real_B_gray = self.real_B

            gpu_p = self.opt.gpu_ids_p[0]
            gpu = self.opt.gpu_ids[0]
            if gpu_p != gpu:
                fake_B_gray = fake_B_gray.cuda(gpu_p)
                real_B_gray = real_B_gray.cuda(gpu_p)

            # d_CM(a_i,G(p_i))
            self.dt1 = self.netDT1(fake_B_gray)
            self.dt2 = self.netDT2(fake_B_gray)
            dt1 = self.dt1 / 2.0 + 0.5  # [-1,1]->[0,1]
            dt2 = self.dt2 / 2.0 + 0.5
            if self.opt.dt_nonlinear != '':
                dt_xmax = torch.Tensor([self.opt.dt_xmax]).cuda(gpu_p)
                dt1 = nonlinearDt(dt1, self.opt.dt_nonlinear, dt_xmax)
                dt2 = nonlinearDt(dt2, self.opt.dt_nonlinear, dt_xmax)
                # print('dt1dt2',torch.min(dt1).item(),torch.max(dt1).item(),torch.min(dt2).item(),torch.max(dt2).item())

            bs = real_B_gray.shape[0]
            real_B_gray_line1 = self.netLine1(real_B_gray)
            real_B_gray_line2 = self.netLine2(real_B_gray)
            self.loss_G_chamfer = (dt1[(real_B_gray < 0) & (real_B_gray_line1 < 0)].sum() + dt2[
                (real_B_gray >= 0) & (real_B_gray_line2 >= 0)].sum()) / bs * self.opt.lambda_chamfer
            if gpu_p != gpu:
                self.loss_G_chamfer = self.loss_G_chamfer.cuda(gpu)

                # d_CM(G(p_i),a_i)
            if gpu_p != gpu:
                dt1gt = self.dt1gt.cuda(gpu_p)
                dt2gt = self.dt2gt.cuda(gpu_p)
            else:
                dt1gt = self.dt1gt
                dt2gt = self.dt2gt
            if self.opt.dt_nonlinear != '':
                dt1gt = nonlinearDt(dt1gt, self.opt.dt_nonlinear, dt_xmax)
                dt2gt = nonlinearDt(dt2gt, self.opt.dt_nonlinear, dt_xmax)
                # print('dt1gtdt2gt',torch.min(dt1gt).item(),torch.max(dt1gt).item(),torch.min(dt2gt).item(),torch.max(dt2gt).item())
            self.dt1gt = (self.dt1gt - 0.5) * 2
            self.dt2gt = (self.dt2gt - 0.5) * 2

            fake_B_gray_line1 = self.netLine1(fake_B_gray)
            fake_B_gray_line2 = self.netLine2(fake_B_gray)
            self.loss_G_chamfer2 = (dt1gt[(fake_B_gray < 0) & (fake_B_gray_line1 < 0)].sum() + dt2gt[
                (fake_B_gray >= 0) & (fake_B_gray_line2 >= 0)].sum()) / bs * self.opt.lambda_chamfer2
            if gpu_p != gpu:
                self.loss_G_chamfer2 = self.loss_G_chamfer2.cuda(gpu)

        # Fourth, line continuity loss, constrained on synthesized drawing
        if self.opt.continuity_loss:
            # Patch-based
            self.get_patches()
            self.outputs = self.netRegressor(self.fake_B_patches)
            if not self.opt.emphasis_conti_face:
                self.loss_G_continuity = (1.0 - torch.mean(self.outputs)).cuda(gpu) * self.opt.lambda_continuity
            else:
                self.loss_G_continuity = torch.mean((1.0 - self.outputs) * self.conti_weights).cuda(
                    gpu) * self.opt.lambda_continuity

        self.loss_G = self.loss_G_GAN
        if 'G_L1' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_L1
        if 'G_local' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_local
        if 'G_hair_local' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_hair_local
        if 'G_bg_local' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_bg_local
        if 'G_chamfer' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_chamfer
        if 'G_chamfer2' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_chamfer2
        if 'G_continuity' in self.loss_names:
            self.loss_G = self.loss_G + self.loss_G_continuity

        self.loss_G.backward()

    def optimize_parameters(self):
        self.forward()
        # update D
        self.set_requires_grad(self.netD, True)

        if self.opt.discriminator_local:
            self.set_requires_grad(self.netDLEyel, True)
            self.set_requires_grad(self.netDLEyer, True)
            self.set_requires_grad(self.netDLNose, True)
            self.set_requires_grad(self.netDLMouth, True)
            self.set_requires_grad(self.netDLHair, True)
            self.set_requires_grad(self.netDLBG, True)
        self.optimizer_D.zero_grad()
        self.backward_D()
        self.optimizer_D.step()

        # update G
        self.set_requires_grad(self.netD, False)
        if self.opt.discriminator_local:
            self.set_requires_grad(self.netDLEyel, False)
            self.set_requires_grad(self.netDLEyer, False)
            self.set_requires_grad(self.netDLNose, False)
            self.set_requires_grad(self.netDLMouth, False)
            self.set_requires_grad(self.netDLHair, False)
            self.set_requires_grad(self.netDLBG, False)
        self.optimizer_G.zero_grad()
        self.backward_G()
        self.optimizer_G.step()

    def get_patches(self):
        gpu_p = self.opt.gpu_ids_p[0]
        gpu = self.opt.gpu_ids[0]
        if gpu_p != gpu:
            self.fake_B = self.fake_B.cuda(gpu_p)
        # [1,1,512,512]->[bs,1,11,11]
        patches = []
        if self.isTrain and self.opt.emphasis_conti_face:
            weights = []
            W2 = int(W / 2)
        t = np.random.randint(res, size=2)
        for i in range(aa):
            for j in range(aa):
                p = self.fake_B[:, :, t[0] + i * W:t[0] + (i + 1) * W, t[1] + j * W:t[1] + (j + 1) * W]
                whitenum = torch.sum(p >= 0.0)
                # if whitenum < 5 or whitenum > W*W-5:
                if whitenum < 1 or whitenum > W * W - 1:
                    continue
                patches.append(p)
                if self.isTrain and self.opt.emphasis_conti_face:
                    weights.append(self.face_mask[:, :, t[0] + i * W + W2, t[1] + j * W + W2])
        self.fake_B_patches = torch.cat(patches, dim=0)
        if self.isTrain and self.opt.emphasis_conti_face:
            self.conti_weights = torch.cat(weights, dim=0) + 1  # 0->1,1->2

    def get_patches_real(self):
        # [1,1,512,512]->[bs,1,11,11]
        patches = []
        t = np.random.randint(res, size=2)
        for i in range(aa):
            for j in range(aa):
                p = self.real_B[:, :, t[0] + i * W:t[0] + (i + 1) * W, t[1] + j * W:t[1] + (j + 1) * W]
                whitenum = torch.sum(p >= 0.0)
                # if whitenum < 5 or whitenum > W*W-5:
                if whitenum < 1 or whitenum > W * W - 1:
                    continue
                patches.append(p)
        self.real_B_patches = torch.cat(patches, dim=0)