def main(_):
if not os.path.exists(FLAGS.checkpoint_dir):
os.makedirs(FLAGS.checkpoint_dir)
if not os.path.exists(FLAGS.results_dir):
os.makedirs(FLAGS.results_dir)
# gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
run_config = tf.ConfigProto()
run_config.gpu_options.allow_growth = True
with tf.Session(config = run_config) as sess:
dcgan = DCGAN(
sess = sess,
crop = FLAGS.crop,
checkpoint_dir = FLAGS.checkpoint_dir,
results_dir = FLAGS.results_dir)
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")
dcgan.anomaly_detection(FLAGS)
