def build_model(sess, embedding_dim, batch_size):
model = CondGAN(
lr_imsize=cfg.TEST.LR_IMSIZE,
hr_lr_ratio=int(cfg.TEST.HR_IMSIZE / cfg.TEST.LR_IMSIZE))
embeddings = tf.placeholder(
tf.float32, [batch_size, embedding_dim],
name='conditional_embeddings')
with pt.defaults_scope(phase=pt.Phase.test):
with tf.variable_scope("g_net"):
c = sample_encoded_context(embeddings, model)
z = tf.random_normal([batch_size, cfg.Z_DIM])
fake_images = model.get_generator(tf.concat(1, [c, z]))
with tf.variable_scope("hr_g_net"):
hr_c = sample_encoded_context(embeddings, model)
hr_fake_images = model.hr_get_generator(fake_images, hr_c)
ckt_path = cfg.TEST.PRETRAINED_MODEL
if ckt_path.find('.ckpt') != -1:
print("Reading model parameters from %s" % ckt_path)
saver = tf.train.Saver(tf.all_variables())
saver.restore(sess, ckt_path)
else:
print("Input a valid model path.")
return embeddings, fake_images, hr_fake_images
def build_model(sess, embedding_dim, batch_size):
model = CondGAN(
lr_imsize=cfg.TEST.LR_IMSIZE,
hr_lr_ratio=int(cfg.TEST.HR_IMSIZE/cfg.TEST.LR_IMSIZE))
embeddings = tf.placeholder(
tf.float32, [batch_size, embedding_dim],
name='conditional_embeddings')
with pt.defaults_scope(phase=pt.Phase.test):
with tf.variable_scope("g_net"):
c = sample_encoded_context(embeddings, model)
z = tf.random_normal([batch_size, cfg.Z_DIM])
fake_images = model.get_generator(tf.concat(1, [c, z]))
with tf.variable_scope("hr_g_net"):
hr_c = sample_encoded_context(embeddings, model)
hr_fake_images = model.hr_get_generator(fake_images, hr_c)
ckt_path = cfg.TEST.PRETRAINED_MODEL
if ckt_path.find('.ckpt') != -1:
print("Reading model parameters from %s" % ckt_path)
saver = tf.train.Saver(tf.all_variables())
saver.restore(sess, ckt_path)
else:
print("Input a valid model path.")
return embeddings, fake_images, hr_fake_images
datadir = 'Data/%s' % cfg.DATASET_NAME
dataset = TextDataset(datadir, cfg.EMBEDDING_TYPE, 4)
filename_test = '%s/test' % (datadir)
dataset.test = dataset.get_data(filename_test)
if cfg.TRAIN.FLAG:
filename_train = '%s/train' % (datadir)
dataset.train = dataset.get_data(filename_train)
ckt_logs_dir = "ckt_logs/%s/%s_%s" % \
(cfg.DATASET_NAME, cfg.CONFIG_NAME, timestamp)
mkdir_p(ckt_logs_dir)
else:
s_tmp = cfg.TRAIN.PRETRAINED_MODEL
ckt_logs_dir = s_tmp[:s_tmp.find('.ckpt')]
model = CondGAN(lr_imsize=int(dataset.image_shape[0] /
dataset.hr_lr_ratio),
hr_lr_ratio=dataset.hr_lr_ratio)
algo = CondGANTrainer(model=model,
dataset=dataset,
ckt_logs_dir=ckt_logs_dir)
if cfg.TRAIN.FLAG:
algo.train()
else:
''' For every input text embedding/sentence in the
training and test datasets, generate cfg.TRAIN.NUM_COPY
images with randomness from noise z and conditioning augmentation.'''
algo.evaluate()
datadir = '%s/%s' % (args.data_path, cfg.DATASET_NAME)
dataset = datastore.create(datadir, cfg)
print('Using dataset:')
print(dataset)
dataset.test = dataset.get_data(os.path.join(datadir, 'test'))
if cfg.TRAIN.FLAG:
dataset.train = dataset.get_data(os.path.join(datadir, 'train'))
ckt_logs_dir = "ckt_logs/%s/%s_%s" % \
(cfg.DATASET_NAME, cfg.CONFIG_NAME, timestamp)
mkdir_p(ckt_logs_dir)
else:
s_tmp = cfg.TRAIN.PRETRAINED_MODEL
ckt_logs_dir = s_tmp[:s_tmp.find('.ckpt')]
model = CondGAN(dataset.lr_imsize, dataset.hr_lr_ratio)
algo = CondGANTrainer(model=model,
dataset=dataset,
ckt_logs_dir=ckt_logs_dir)
if cfg.TRAIN.FLAG:
algo.train()
else:
''' For every input text embedding/sentence in the
training and test datasets, generate cfg.TRAIN.NUM_COPY
images with randomness from noise z and conditioning augmentation.'''
algo.evaluate()
