def main():
if (len(sys.argv) == 1):
raise NameError('[ERROR] No dataset key')
elif (sys.argv[1] == 'lfw'):
FLAGS.updates_per_epoch = 380
FLAGS.log_interval = 120
FLAGS.out_dir = 'data/output/lfw/'
FLAGS.list_dir = 'data/imglist/lfw/'
FLAGS.pc_dir = 'data/pcomp/lfw/'
#add other datasets here
else:
raise NameError('[ERROR] Incorrect dataset key')
data_loader = lab_imageloader(FLAGS.in_dir, \
os.path.join(FLAGS.out_dir, 'images'), \
listdir=FLAGS.list_dir)
#Train colorfield VAE
graph_vae = tf.Graph()
with graph_vae.as_default():
model_colorfield = vae(FLAGS, nch=2)
dnn = network(model_colorfield, data_loader, 2, FLAGS)
latent_vars_colorfield, latent_vars_colorfield_musigma_test = \
dnn.train_vae(os.path.join(FLAGS.out_dir, 'models'), FLAGS.is_train)
np.save(os.path.join(FLAGS.out_dir, 'lv_color_train.mat'),
latent_vars_colorfield)
np.save(os.path.join(FLAGS.out_dir, 'lv_color_test.mat'),
latent_vars_colorfield_musigma_test)
def main():
if (len(sys.argv) == 1):
raise NameError('[ERROR] No dataset key')
elif (sys.argv[1] == 'lfw'):
FLAGS.updates_per_epoch = 380
FLAGS.log_interval = 120
FLAGS.out_dir = 'data/output/lfw/'
FLAGS.list_dir = 'data/imglist/lfw/'
FLAGS.pc_dir = 'data/pcomp/lfw/'
else:
raise NameError('[ERROR] Incorrect dataset key')
data_loader = lab_imageloader(FLAGS.in_dir, \
os.path.join(FLAGS.out_dir, 'images'), \
listdir=FLAGS.list_dir)
#Diverse Colorization
nmix = 8
num_batches = 31
lv_mdn_test = np.load(
os.path.join(FLAGS.out_dir, 'lv_color_mdn_test.mat.npy'))
graph_divcolor = tf.Graph()
with graph_divcolor.as_default():
model_colorfield = vae(FLAGS, nch=2, condinference_flag=True)
dnn = network(model_colorfield, data_loader, 2, FLAGS)
dnn.run_divcolor(os.path.join(FLAGS.out_dir, 'models') , \
latent_vars_colorfield_test, num_batches=num_batches)
if (FLAGS.is_run_cvae == True):
dnn.run_cvae(os.path.join(FLAGS.out_dir, 'models') , \
lv_mdn_test, num_batches=num_batches, num_repeat=8, num_cluster=5)