def main(_):
pp = pprint.PrettyPrinter()
pp.pprint(flags.FLAGS.__flags)
if FLAGS.input_width is None:
FLAGS.input_width = FLAGS.input_height
if FLAGS.output_width is None:
FLAGS.output_width = FLAGS.output_height
if not os.path.exists(FLAGS.checkpoint_dir):
os.makedirs(FLAGS.checkpoint_dir)
if not os.path.exists(FLAGS.sample_dir):
os.makedirs(FLAGS.sample_dir)
#gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
#run_config = tf.compat.v1.ConfigProto
#run_config.gpu_options.allow_growth = True
#run_config.gpu_options.visible_device_list = ''
with tf.compat.v1.Session() as sess:
dcgan = DCGAN(
sess,
input_width=FLAGS.input_width,
input_height=FLAGS.input_height,
output_width=FLAGS.output_width,
output_height=FLAGS.output_height,
batch_size=FLAGS.batch_size,
sample_num=FLAGS.sample_num,
dataset_name=FLAGS.dataset,
input_fname_pattern=FLAGS.input_fname_pattern,
crop=FLAGS.crop,
checkpoint_dir=FLAGS.checkpoint_dir,
data_dir=FLAGS.data_dir)
model_vars = tf.trainable_variables()
slim.model_analyzer.analyze_vars(model_vars, print_info=True)
if FLAGS.train:
dcgan.train(FLAGS)
else:
if not dcgan.load(FLAGS.checkpoint_dir)[0]:
raise Exception("[!] Train a model first, then run test mode")
if FLAGS.predict:
dcgan.predict(FLAGS.predict_dataset)
else:
dcgan.test()
def predict(latent_dim, height, width, channels):
random_latent_vectors = np.random.normal(size=(100, latent_dim))
dcgan = DCGAN(latent_dim, height, width, channels)
dcgan.load_weights('gan_epoch20.h5')
# dcgan.load_weights('gan_epoch1.h5')
generated_images = dcgan.predict(random_latent_vectors)
for i, generated_image in enumerate(generated_images):
img = image.array_to_img(denormalize(generated_image), scale=False)
img.save(os.path.join('generated', str(i) + '.png'))
