def load_warpGAN(self):
network = WarpGAN()
network.load_model(self.warpgan_dir)
return network
def main(args):
# I/O
config_file = args.config_file
config = imp.load_source('config', config_file)
if args.name:
config.name = args.name
trainset = Dataset(config.train_dataset_path, prefix=config.data_prefix)
network = WarpGAN()
network.initialize(config, trainset.num_classes)
# Initalization for running
if config.save_model:
log_dir = utils.create_log_dir(config, config_file)
summary_writer = tf.summary.FileWriter(log_dir, network.graph)
if config.restore_model:
network.restore_model(config.restore_model, config.restore_scopes)
proc_func = lambda images: preprocess(images, config, True)
trainset.start_batch_queue(config.batch_size, proc_func=proc_func)
# Main Loop
print(
'\nStart Training\nname: {}\n# epochs: {}\nepoch_size: {}\nbatch_size: {}\n'
.format(config.name, config.num_epochs, config.epoch_size,
config.batch_size))
global_step = 0
start_time = time.time()
for epoch in range(config.num_epochs):
if epoch == 0: test(network, config, log_dir, global_step)
# Training
for step in range(config.epoch_size):
# Prepare input
learning_rate = utils.get_updated_learning_rate(
global_step, config)
batch = trainset.pop_batch_queue()
print("batch: ", step)
print("epoch: ", epoch)
wl, sm, global_step = network.train(batch['images'],
batch['labels'],
batch['is_photo'],
learning_rate,
config.keep_prob)
wl['lr'] = learning_rate
# Display
if step % config.summary_interval == 0:
duration = time.time() - start_time
start_time = time.time()
utils.display_info(epoch, step, duration, wl)
if config.save_model:
summary_writer.add_summary(sm, global_step=global_step)
# Testing
test(network, config, log_dir, global_step)
# Save the model
if config.save_model and epoch % 49 == 0:
network.save_model(log_dir, global_step)
def _build_pipe(ctx):
network = WarpGAN()
network.load_model(PARAMS.get('model_dir', './model'))
styles = np.random.normal(0., 1., (1, network.input_style.shape[1].value))
return network, styles
"--num_styles",
help="The number of images to generate with different styles",
type=int,
default=5)
parser.add_argument("--scale",
help="The path to the input directory",
type=float,
default=1.0)
parser.add_argument("--aligned",
help="Set true if the input face is already normalized",
action='store_true')
args = parser.parse_args()
if __name__ == '__main__':
network = WarpGAN()
network.load_model(args.model_dir)
img = misc.imread(args.input, mode='RGB')
if not args.aligned:
from align.detect_align import detect_align
img = detect_align(img)
img = (img - 127.5) / 128.0
images = np.tile(img[None], [args.num_styles, 1, 1, 1])
scales = args.scale * np.ones((args.num_styles))
styles = np.random.normal(
0., 1., (args.num_styles, network.input_style.shape[1].value))
from warpgan import WarpGAN model_dir = "./pretrained/warpgan_pretrained" network = WarpGAN() network.load_model(model_dir)
def init_and_load_model(model_dir):
network = WarpGAN()
network.load_model(model_dir)
return network