From 4aa0fa00eed77f87e80a0ec3e2954457bb0a519a Mon Sep 17 00:00:00 2001 From: Shuyu Chen Date: Mon, 22 Jun 2020 23:21:04 +0800 Subject: [PATCH] Add files via upload --- models/AE_Model.py | 208 +++++++++++++++++++ models/Combine_Model.py | 96 +++++++++ models/networks.py | 435 +++++++++++++++++++++++++++++++++++++++ options/AE_face.py | 42 ++++ options/parts_combine.py | 52 +++++ 5 files changed, 833 insertions(+) create mode 100644 models/AE_Model.py create mode 100644 models/Combine_Model.py create mode 100644 models/networks.py create mode 100644 options/AE_face.py create mode 100644 options/parts_combine.py diff --git a/models/AE_Model.py b/models/AE_Model.py new file mode 100644 index 0000000..bfa0648 --- /dev/null +++ b/models/AE_Model.py @@ -0,0 +1,208 @@ +# this is model for fix netP and mask only left eye, right eye, nose and skin +# original loss + +import random +import numpy as np +import os +from . import networks +# from scipy.ndimage import median_filter +import jittor as jt +from jittor import init +from jittor import nn +import jittor.transform as transform +# from pdb import set_trace as st +# import heapq +from numpy.linalg import solve +import time + + +class AE_Model(nn.Module): + def name(self): + return self.name + + def init_loss_filter(self): + flags = (True,True,True) + def loss_filter(kl_loss,l2_image,l2_mask): + return [l for (l,f) in zip((kl_loss,l2_image,l2_mask),flags) if f] + + return loss_filter + + def initialize(self, opt): + # assert opt.vae_encoder == True + self.opt = opt + self.save_dir = os.path.join(opt.param, opt.name) + + self.name = 'AE_Model' + # BaseModel.initialize(self, opt) + + input_nc = opt.input_nc + + self.output_nc = opt.output_nc + self.input_nc = input_nc + + self.model_partial_name = opt.partial + ##### define networks + # Generator network + netG_input_nc = input_nc + + self.net_encoder = networks.define_part_encoder(model=self.model_partial_name, + input_nc = opt.input_nc, + norm=opt.norm, + latent_dim = opt.latant_dim) + + self.net_decoder = networks.define_part_decoder(model=self.model_partial_name, + output_nc = opt.input_nc, + norm=opt.norm, + latent_dim = opt.latant_dim) + + self.load_network(self.net_encoder, 'encoder_'+self.model_partial_name, 'latest', '') + self.load_network(self.net_decoder, 'decoder_'+self.model_partial_name+'_image', 'latest', '') + + # use for test + self.feature_list_male = np.fromfile(os.path.join(opt.param, opt.name) + '/man_' + opt.partial + '_feature.bin', dtype=np.float32) + self.feature_list_male.shape = 6247, 512 + + + #girl + self.feature_list_female = np.fromfile(os.path.join(opt.param, opt.name) + '/female_' + opt.partial + '_feature.bin', dtype=np.float32) + self.feature_list_female.shape = 11456, 512 + self.feature_list = [self.feature_list_male,self.feature_list_female] + + def get_latent(self, input_image): + + input_image = (input_image-127.5)/127.5 + input_image = np.expand_dims(input_image, axis=2) + input_image = input_image.transpose(2,0,1) + input_image = np.expand_dims(input_image, axis=0) + input_image = input_image.astype('float32') + input_image = transform.to_tensor(jt.array(input_image)) + # print(input_image.shape) + mus_mouth = self.net_encoder(input_image) + + return mus_mouth + + def get_image(self, latent_vec): + + # return self.net_decoder(latent_vec) + + fakes = self.net_decoder(latent_vec) + fakes = (fakes[0,:,:,:].numpy()+1)/2 + + fakes = np.transpose(fakes, (1, 2, 0)) * 255.0 + fakes = np.clip(fakes, 0, 255) + + return fakes.astype(np.uint8) + + def get_inter(self, input_image, nearnN=3, sex=1,w_c=1,random_=-1): + generated_f = self.get_latent(input_image) + generated_f = generated_f.numpy() + + feature_list = self.feature_list[sex] + list_len = jt.array([feature_list.shape[0]]) + # a = jt.random((n,3)) + b = jt.code([1, nearnN], + "int32", [jt.array(feature_list),jt.array(generated_f), list_len], + cpu_header="#include ", + cpu_src=""" + using namespace std; + auto n=out_shape0, k=out_shape1; + int N=@in2(0); + + // 使用openmp实现自动并行化 + // 存储k近邻的距离和下标 + vector> id(N); + #pragma omp parallel for + for (int j=0; j0.5: + w_i = 0.5 + + # print(i) + mus_vec = jt.unsqueeze(mus_mouth[[i],:],1) + + fake_image = self.net_decoder(jt.array(mus_vec)) + # fake_image = fake_image[[0],:,:,:] + if i==0: + fakes = (1-fake_image)/2* w_i + else: + fakes = fakes + (1-fake_image)/2 * w_i + + fakes = 1-fakes + + fakes = fakes[0,:,:,:].detach().numpy() + + fakes = np.transpose(fakes, (1, 2, 0)) * 255.0 + fakes = np.clip(fakes, 0, 255) + + return fakes.astype(np.uint8) + + # helper loading function that can be used by subclasses + def load_network(self, network, network_label, epoch_label, save_dir='', save_path=''): + save_filename = '%s_net_%s.pkl' % (epoch_label, network_label) + if not save_dir: + save_dir = self.save_dir + save_path = os.path.join(save_dir, save_filename) + print("load_path",save_path) + if not os.path.isfile(save_path): + print('%s not exists yet!' % save_path) + else: + network.load(save_path) \ No newline at end of file diff --git a/models/Combine_Model.py b/models/Combine_Model.py new file mode 100644 index 0000000..439f10d --- /dev/null +++ b/models/Combine_Model.py @@ -0,0 +1,96 @@ +import numpy as np +import os +from . import networks + +import jittor as jt +from jittor import init +from jittor import nn +import jittor.transform as transform + +class Combine_Model(nn.Module): + def name(self): + return 'Combine_Model' + + def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss): + flags = (True, use_gan_feat_loss, use_vgg_loss, True, True) + def loss_filter(g_gan, g_gan_feat, g_vgg, d_real, d_fake): + return [l for (l,f) in zip((g_gan,g_gan_feat,g_vgg,d_real,d_fake),flags) if f] + return loss_filter + + def initialize(self, opt): + self.opt = opt + self.save_dir = os.path.join(opt.param, opt.name) + # BaseModel.initialize(self, opt) + input_nc = opt.input_nc + + ##### define networks + # Generator network + self.part = {'': (0, 0, 512), + 'eye1': (108, 156, 128), + 'eye2': (255, 156, 128), + 'nose': (182, 232, 160), + 'mouth': (169, 301, 192)} + + + self.Decoder_Part = {} + + for key in self.part.keys(): + self.Decoder_Part[key] = networks.define_feature_decoder(model=key, + output_nc = 32, norm=opt.norm, + latent_dim = opt.latant_dim) + + self.netG = networks.define_G(opt.num_inter_channels, opt.output_nc, opt.ngf, + opt.n_downsample_global, opt.n_blocks_global, opt.norm) + + self.load_network(self.netG, 'G', opt.which_epoch, '') + + for key in self.part.keys(): + self.load_network(self.Decoder_Part[key], 'DE_'+key, opt.which_epoch, '') + + def inference(self, part_v, image=None): + + eye1_code = part_v['eye1'] + eye2_code = part_v['eye2'] + nose_code = part_v['nose'] + mouth_code = part_v['mouth'] + bg_code = part_v[''] + + eye1_r_feature = self.Decoder_Part['eye1'](eye1_code) + eye2_r_feature = self.Decoder_Part['eye2'](eye2_code) + nose_r_feature = self.Decoder_Part['nose'](nose_code) + mouth_r_feature = self.Decoder_Part['mouth'](mouth_code) + bg_r_feature = self.Decoder_Part[''](bg_code) + + bg_r_feature[:, :, 301:301 + 192, 169:169 + 192] = mouth_r_feature + bg_r_feature[:, :, 232:232 + 160 - 36, 182:182 + 160] = nose_r_feature[:, :, :-36, :] + bg_r_feature[:, :, 156:156 + 128, 108:108 + 128] = eye1_r_feature + bg_r_feature[:, :, 156:156 + 128, 255:255 + 128] = eye2_r_feature + + input_concat = bg_r_feature + + fake_image = self.netG(input_concat) + + # fakes = fake_image.detach().numpy() + fakes = fake_image[0, :, :, :].detach().numpy() + + fakes = (np.transpose(fakes, (1, 2, 0)) + 1) / 2.0 * 255.0 + fakes = np.clip(fakes, 0, 255) + + return fakes.astype(np.uint8) + + # helper loading function that can be used by subclasses + def load_network(self, network, network_label, epoch_label, save_dir='', save_path=''): + save_filename = '%s_net_%s.pkl' % (epoch_label, network_label) + if not save_dir: + save_dir = self.save_dir + save_path = os.path.join(save_dir, save_filename) + print("load_path",save_path) + if not os.path.isfile(save_path): + print('%s not exists yet!' % save_path) + else: + network.load(save_path) + +class InferenceModel(Combine_Model): + def forward(self, inp): + label, image = inp + return self.inference(label, image) diff --git a/models/networks.py b/models/networks.py new file mode 100644 index 0000000..7059411 --- /dev/null +++ b/models/networks.py @@ -0,0 +1,435 @@ +import jittor as jt +from jittor import init +from jittor import nn +from jittor import models +import functools +import numpy as np + + +############################################################################### +# Functions +############################################################################### + +def weights_init_normal(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + jt.init.gauss_(m.weight, 0.0, 0.02) + elif classname.find("BatchNorm") != -1: + jt.init.gauss_(m.weight, 1.0, 0.02) + jt.init.constant_(m.bias, 0.0) + +def get_norm_layer(norm_type='instance'): + if (norm_type == 'batch'): + norm_layer = nn.BatchNorm + elif (norm_type == 'instance'): + norm_layer = nn.InstanceNorm2d + else: + raise NotImplementedError(('normalization layer [%s] is not found' % norm_type)) + return norm_layer + +class MSELoss: + def __init__(self): + pass + + def __call__(self, output, target): + from jittor.nn import mse_loss + return mse_loss(output, target) + +class BCELoss: + def __init__(self): + pass + + def __call__(self, output, target): + from jittor.nn import bce_loss + return bce_loss(output, target) + +############################ +# Model function +############################ +def define_part_encoder(model='mouth', norm='instance', input_nc=1, latent_dim=512): + norm_layer = get_norm_layer(norm_type=norm) + image_size = 512 + if 'eye' in model: + image_size = 128 + elif 'mouth' in model: + image_size = 192 + elif 'nose' in model: + image_size = 160 + elif 'face' in model: + image_size = 512 + else: + print("Whole Image !!") + + net_encoder = EncoderGenerator_Res(norm_layer,image_size,input_nc, latent_dim) # input longsize 256 to 512*4*4 + print("net_encoder of part "+model+" is:",image_size) + + return net_encoder + +def define_part_decoder(model='mouth', norm='instance', output_nc=1, latent_dim=512): + norm_layer = get_norm_layer(norm_type=norm) + + image_size = 512 + if 'eye' in model: + image_size = 128 + elif 'mouth' in model: + image_size = 192 + elif 'nose' in model: + image_size = 160 + else: + print("Whole Image !!") + + net_decoder = DecoderGenerator_image_Res(norm_layer,image_size,output_nc, latent_dim) # input longsize 256 to 512*4*4 + + print("net_decoder to image of part "+model+" is:",image_size) + + return net_decoder + + +def define_feature_decoder(model='mouth', norm='instance', output_nc=1, latent_dim=512): + norm_layer = get_norm_layer(norm_type=norm) + + image_size = 512 + if 'eye' in model: + image_size = 128 + elif 'mouth' in model: + image_size = 192 + elif 'nose' in model: + image_size = 160 + else: + print("Whole Image !!") + + net_decoder = DecoderGenerator_feature_Res(norm_layer,image_size,output_nc, latent_dim) # input longsize 256 to 512*4*4 + + print("net_decoder to image of part "+model+" is:",image_size) + # print(net_decoder) + + return net_decoder +def define_G(input_nc, output_nc, ngf, n_downsample_global=3, n_blocks_global=9, norm='instance'): + norm_layer = get_norm_layer(norm_type=norm) + netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer) + return netG + + +############################################################################## +# Losses +############################################################################## + + + +# Define a resnet block +class ResnetBlock(nn.Module): + def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False): + super(ResnetBlock, self).__init__() + self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout) + + def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout): + conv_block = [] + p = 0 + if (padding_type == 'reflect'): + conv_block += [nn.ReflectionPad2d(1)] + elif (padding_type == 'replicate'): + conv_block += [nn.ReplicationPad2d(1)] + elif (padding_type == 'zero'): + p = 1 + else: + raise NotImplementedError(('padding [%s] is not implemented' % padding_type)) + conv_block += [nn.Conv(dim, dim, 3, padding=p), norm_layer(dim), activation] + if use_dropout: + conv_block += [nn.Dropout(0.5)] + + p = 0 + if (padding_type == 'reflect'): + conv_block += [nn.ReflectionPad2d(1)] + elif (padding_type == 'replicate'): + conv_block += [nn.ReplicationPad2d(1)] + elif (padding_type == 'zero'): + p = 1 + else: + raise NotImplementedError(('padding [%s] is not implemented' % padding_type)) + conv_block += [nn.Conv(dim, dim, 3, padding=p), norm_layer(dim)] + return nn.Sequential(*conv_block) + + def execute(self, x): + # print(x.shape) + out = (x + self.conv_block(x)) + return out + +class EncoderGenerator_Res(nn.Module): + """docstring for EncoderGenerator""" + def __init__(self, norm_layer, image_size, input_nc, latent_dim=512): + super( EncoderGenerator_Res, self).__init__() + layers_list = [] + + latent_size = int(image_size/32) + longsize = 512*latent_size*latent_size + self.longsize = longsize + # print(image_size,latent_size, longsize) + + activation = nn.ReLU() + padding_type='reflect' + norm_layer=nn.BatchNorm + + # encode + layers_list.append(EncoderBlock(channel_in=input_nc, channel_out=32, kernel_size=4, padding=1, stride=2)) # 176 176 + + dim_size = 32 + for i in range(4): + layers_list.append(ResnetBlock(dim_size, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) + layers_list.append(EncoderBlock(channel_in=dim_size, channel_out=dim_size*2, kernel_size=4, padding=1, stride=2)) + dim_size *= 2 + + layers_list.append(ResnetBlock(512, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) + + # final shape Bx256*7*6 + self.conv = nn.Sequential(*layers_list) + self.fc_mu = nn.Sequential(nn.Linear(in_features=longsize, out_features=latent_dim))#, + + # self.fc_var = nn.Sequential(nn.Linear(in_features=longsize, out_features=latent_dim))#, + + for m in self.modules(): + weights_init_normal(m) + + def execute(self, ten): + # ten = ten[:,:,:] + # ten2 = jt.reshape(ten,[ten.size()[0],-1]) + # print(ten.shape, ten2.shape) + ten = self.conv(ten) + ten = jt.reshape(ten,[ten.size()[0],-1]) + # print(ten.shape,self.longsize) + mu = self.fc_mu(ten) + # logvar = self.fc_var(ten) + return mu#,logvar + +class DecoderGenerator_image_Res(nn.Module): + def __init__(self, norm_layer, image_size, output_nc, latent_dim=512): + super(DecoderGenerator_image_Res, self).__init__() + # start from B*1024 + latent_size = int(image_size/32) + self.latent_size = latent_size + longsize = 512*latent_size*latent_size + + activation = nn.ReLU() + padding_type='reflect' + norm_layer=nn.BatchNorm + + self.fc = nn.Sequential(nn.Linear(in_features=latent_dim, out_features=longsize)) + layers_list = [] + + layers_list.append(ResnetBlock(512, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + dim_size = 256 + for i in range(4): + layers_list.append(DecoderBlock(channel_in=dim_size*2, channel_out=dim_size, kernel_size=4, padding=1, stride=2, output_padding=0)) #latent*2 + layers_list.append(ResnetBlock(dim_size, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) + dim_size = int(dim_size/2) + + layers_list.append(DecoderBlock(channel_in=32, channel_out=32, kernel_size=4, padding=1, stride=2, output_padding=0)) #352 352 + layers_list.append(ResnetBlock(32, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + # layers_list.append(DecoderBlock(channel_in=64, channel_out=64, kernel_size=4, padding=1, stride=2, output_padding=0)) #96*160 + layers_list.append(nn.ReflectionPad2d(2)) + layers_list.append(nn.Conv(32,output_nc,kernel_size=5,padding=0)) + + self.conv = nn.Sequential(*layers_list) + + for m in self.modules(): + weights_init_normal(m) + + def execute(self, ten): + # print("in DecoderGenerator, print some shape ") + # print(ten.size()) + ten = self.fc(ten) + # print(ten.size()) + ten = jt.reshape(ten,(ten.size()[0],512, self.latent_size, self.latent_size)) + # print(ten.size()) + ten = self.conv(ten) + + return ten + + # def __call__(self, *args, **kwargs): + # return super(DecoderGenerator_image_Res, self).__call__(*args, **kwargs) + + +class DecoderGenerator_feature_Res(nn.Module): + def __init__(self, norm_layer, image_size, output_nc, latent_dim=512): + super(DecoderGenerator_feature_Res, self).__init__() + # start from B*1024 + latent_size = int(image_size/32) + self.latent_size = latent_size + longsize = 512*latent_size*latent_size + + activation = nn.ReLU() + padding_type='reflect' + norm_layer=nn.BatchNorm + + self.fc = nn.Sequential(nn.Linear(in_features=latent_dim, out_features=longsize)) + layers_list = [] + + layers_list.append(ResnetBlock(512, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + layers_list.append(DecoderBlock(channel_in=512, channel_out=256, kernel_size=4, padding=1, stride=2, output_padding=0)) #22 22 + layers_list.append(ResnetBlock(256, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + layers_list.append(DecoderBlock(channel_in=256, channel_out=256, kernel_size=4, padding=1, stride=2, output_padding=0)) #44 44 + layers_list.append(ResnetBlock(256, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + layers_list.append(DecoderBlock(channel_in=256, channel_out=128, kernel_size=4, padding=1, stride=2, output_padding=0)) #88 88 + layers_list.append(ResnetBlock(128, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + layers_list.append(DecoderBlock(channel_in=128, channel_out=64, kernel_size=4, padding=1, stride=2, output_padding=0)) #176 176 + layers_list.append(ResnetBlock(64, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + + layers_list.append(DecoderBlock(channel_in=64, channel_out=64, kernel_size=4, padding=1, stride=2, output_padding=0)) #352 352 + layers_list.append(ResnetBlock(64, padding_type=padding_type, activation=activation, norm_layer=norm_layer)) # 176 176 + + # layers_list.append(DecoderBlock(channel_in=64, channel_out=64, kernel_size=4, padding=1, stride=2, output_padding=0)) #96*160 + layers_list.append(nn.ReflectionPad2d(2)) + layers_list.append(nn.Conv(64,output_nc,kernel_size=5,padding=0)) + + self.conv = nn.Sequential(*layers_list) + + for m in self.modules(): + weights_init_normal(m) + + def execute(self, ten): + # print("in DecoderGenerator, print some shape ") + # print(ten.size()) + ten = self.fc(ten) + # print(ten.size()) + ten = jt.reshape(ten,(ten.size()[0],512, self.latent_size, self.latent_size)) + # print(ten.size()) + ten = self.conv(ten) + + return ten + + +# decoder block (used in the decoder) +class DecoderBlock(nn.Module): + + def __init__(self, channel_in, channel_out, kernel_size=4, padding=1, stride=2, output_padding=0, norelu=False): + super(DecoderBlock, self).__init__() + layers_list = [] + layers_list.append(nn.ConvTranspose(channel_in, channel_out, kernel_size, padding=padding, stride=stride, output_padding=output_padding)) + layers_list.append(nn.BatchNorm(channel_out, momentum=0.9)) + if (norelu == False): + layers_list.append(nn.LeakyReLU(1)) + self.conv = nn.Sequential(*layers_list) + + def execute(self, ten): + ten = self.conv(ten) + return ten + +# encoder block (used in encoder and discriminator) +class EncoderBlock(nn.Module): + def __init__(self, channel_in, channel_out, kernel_size=7, padding=3, stride=4): + super(EncoderBlock, self).__init__() + # convolution to halve the dimensions + self.conv = nn.Conv(channel_in, channel_out, kernel_size, padding=padding, stride=stride) + self.bn = nn.BatchNorm(channel_out, momentum=0.9) + self.relu = nn.LeakyReLU(1) + + def execute(self, ten, out=False, t=False): + # print('ten',ten.shape) + # here we want to be able to take an intermediate output for reconstruction error + if out: + ten = self.conv(ten) + ten_out = ten + ten = self.bn(ten) + ten = self.relu(ten) + return (ten, ten_out) + else: + ten = self.conv(ten) + ten = self.bn(ten) + # print(ten.shape) + ten = self.relu(ten) + return ten + +class GlobalGenerator(nn.Module): + + def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm, padding_type='reflect'): + assert (n_blocks >= 0) + super(GlobalGenerator, self).__init__() + activation = nn.ReLU() + + model = [nn.ReflectionPad2d(3), nn.Conv(input_nc, ngf, 7, padding=0), norm_layer(ngf), activation] + ### downsample + for i in range(n_downsampling): + mult = (2 ** i) + model += [nn.Conv((ngf * mult), ((ngf * mult) * 2), 3, stride=2, padding=1), norm_layer(((ngf * mult) * 2)), activation] + + ### resnet blocks + mult = (2 ** n_downsampling) + for i in range(n_blocks): + model += [ResnetBlock((ngf * mult), padding_type=padding_type, activation=activation, norm_layer=norm_layer)] + + ### upsample + for i in range(n_downsampling): + mult = (2 ** (n_downsampling - i)) + model += [nn.ConvTranspose((ngf * mult), int(((ngf * mult) / 2)), 3, stride=2, padding=1, output_padding=1), norm_layer(int(((ngf * mult) / 2))), activation] + model += [nn.ReflectionPad2d(3), nn.Conv(ngf, output_nc, 7, padding=0), nn.Tanh()] + self.model = nn.Sequential(*model) + + for m in self.modules(): + weights_init_normal(m) + + def execute(self, input): + return self.model(input) + +############################################################################## +# Losses +############################################################################## + +class ToTensor: + def __init__(self): + pass + + def __call__(self, img): + from jittor.transform import to_tensor + return to_tensor(img) + +class GANLoss(nn.Module): + + def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0, tensor=ToTensor): + super(GANLoss, self).__init__() + self.real_label = target_real_label + self.fake_label = target_fake_label + self.real_label_var = None + self.fake_label_var = None + self.Tensor = tensor + # if use_lsgan: + self.loss = MSELoss() + # else: + # self.loss = BCELoss() + + def get_target_tensor(self, input, target_is_real): + target_tensor = None + # print("Input",input.shape) + if target_is_real: + create_label = ((self.real_label_var is None) or (self.real_label_var.numel() != input.numel())) + if create_label: + real_tensor = jt.transform.to_tensor(jt.ones(input.shape)) + # real_tensor = self.Tensor(input.shape).fill_(self.real_label) + # self.real_label_var = Variable(real_tensor, requires_grad=False) + self.real_label_var = real_tensor.stop_grad() + target_tensor = self.real_label_var + else: + create_label = ((self.fake_label_var is None) or (self.fake_label_var.numel() != input.numel())) + if create_label: + + fake_tensor = jt.transform.to_tensor(jt.zeros(input.shape)) + self.fake_label_var = fake_tensor.stop_grad() + target_tensor = self.fake_label_var + return target_tensor + + def __call__(self, input, target_is_real): + if isinstance(input[0], list): + loss = 0 + for input_i in input: + pred = input_i[(- 1)] + target_tensor = self.get_target_tensor(pred, target_is_real) + loss += self.loss(pred, target_tensor) + return loss + else: + target_tensor = self.get_target_tensor(input[(- 1)], target_is_real) + return self.loss(input[(- 1)], target_tensor) diff --git a/options/AE_face.py b/options/AE_face.py new file mode 100644 index 0000000..08cdead --- /dev/null +++ b/options/AE_face.py @@ -0,0 +1,42 @@ +import argparse +import os + +class wholeOptions(): + def __init__(self): + self.parser = argparse.ArgumentParser() + self.initialized = False + + def initialize(self): + # experiment specifics + self.parser.add_argument('--name', type=str, default='AE_whole', + help='name of the experiment. It decides where to store samples and models') + self.parser.add_argument('--partial', type=str, default='', help='facial part mouth/eye1/eye2/nose/') + self.parser.add_argument('--latant_dim', type=int, default=512, help='vae latent dim') + + + self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU') + self.parser.add_argument('--model', type=str, default='ae', help='which model to use') + + self.parser.add_argument('--input_nc', type=int, default=1, help='# of input image channels') + self.parser.add_argument('--output_nc', type=int, default=1, help='# of output image channels') + self.parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization') + + self.parser.add_argument('--param', type=str, default='./Params', help='models are saved here') + + self.initialized = True + + def parse(self, save=True): + if not self.initialized: + self.initialize() + self.opt = self.parser.parse_args() + self.opt.isTrain = False + + str_ids = self.opt.gpu_ids.split(',') + self.opt.gpu_ids = [] + for str_id in str_ids: + id = int(str_id) + if id >= 0: + self.opt.gpu_ids.append(id) + args = vars(self.opt) + + return self.opt diff --git a/options/parts_combine.py b/options/parts_combine.py new file mode 100644 index 0000000..77ebb94 --- /dev/null +++ b/options/parts_combine.py @@ -0,0 +1,52 @@ +import argparse +import os + +class CombineOptions(): + def __init__(self): + self.parser = argparse.ArgumentParser() + self.initialized = False + def initialize(self): + self.parser.add_argument('--name', type=str, default='Combine', + help='name of folder') + self.parser.add_argument('--param', type=str, default='./Params', help='models are saved here') + self.parser.add_argument('--model', type=str, default='pix2pixHD', help='which model to use') + self.parser.add_argument('--norm', type=str, default='instance', + help='instance normalization or batch normalization') + self.parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator') + self.parser.add_argument('--local_rank', type=int, default=0, help='local rank for distributed training') + + # input/output sizes + self.parser.add_argument('--loadSize', type=int, default=512, help='scale images to this size') + self.parser.add_argument('--fineSize', type=int, default=512, help='then crop to this size') + self.parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels') + self.parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels') + + # for generator + self.parser.add_argument('--ngf', type=int, default=56, help='# of gen filters in first conv layer') + self.parser.add_argument('--n_downsample_global', type=int, default=4, + help='number of downsampling layers in netG') + self.parser.add_argument('--n_blocks_global', type=int, default=9, + help='number of residual blocks in the global generator network') + self.parser.add_argument('--n_blocks_local', type=int, default=3, + help='number of residual blocks in the local enhancer network') + self.parser.add_argument('--n_local_enhancers', type=int, default=1, help='number of local enhancers to use') + self.parser.add_argument('--niter_fix_global', type=int, default=0, + help='number of epochs that we only train the outmost local enhancer') + + self.parser.add_argument('--n_downsample_E', type=int, default=4, help='# of downsampling layers in encoder') + self.parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer') + + self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model') + self.parser.add_argument('--latant_dim', type=int, default=512, help='vae latent dim') + self.parser.add_argument('--num_inter_channels', default=32, help='# of channels of the intermediate feature maps') + + self.isTrain = False + + def parse(self, save=True): + if not self.initialized: + self.initialize() + self.opt = self.parser.parse_args() + + args = vars(self.opt) + + return self.opt \ No newline at end of file