def test2(status='X2Y'):
    start = time.time()
    config = tf.ConfigProto(allow_soft_placement=True)
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        content_img = cv2.resize(getImg(FLAGS.input),
                                 (FLAGS.img_size, FLAGS.img_size))
        content = np.expand_dims(encode(content_img), 0)
        network = CycleGAN(FLAGS.batch_size,
                           FLAGS.ngf,
                           FLAGS.img_size,
                           is_training=True,
                           use_E=FLAGS.USE_E,
                           Norm=FLAGS.Norm)
        network.test(content.shape)
        sess.run(tf.global_variables_initializer())
        var_list = [
            var for var in tf.global_variables() if 'generator_net' in var.name
            or 'discriminator_net' in var.name or 'edge_net' in var.name
        ]
        saver = tf.train.Saver(var_list=var_list)
        saver.restore(sess, FLAGS.checkpoint)
        print("restored all")
        s = time.time()
        if status == 'X2Y':
            out = sess.run(network.Ygenerated,
                           feed_dict={network.testA: content})
        else:
            out = sess.run(network.Xgenerated,
                           feed_dict={network.testB: content})
        # result = np.clip(out[0], 0, 255).astype(np.uint8)
        print("Transform in {} s".format((time.time() - s)))
        saveImg(out[0], FLAGS.output)
        print("Finished all process in {} s".format(time.time() - start))
def generateBatch(files, batch_shape):
    batch = np.zeros(batch_shape, dtype=np.float32)
    while True:
        try:
            choosed = random.sample(files, batch_shape[0])
            for i, s in enumerate(choosed):
                batch[i] = augmentor(resizeTo(getImg(s), 800, 1800))
                batch[i] = encode(batch[i])
            yield batch
        except:
            continue
def generateBatch(files, batch_shape):
    batch = np.zeros(batch_shape, dtype=np.float32)
    while True:
        try:
            choosed = random.sample(files, batch_shape[0])
            for i, s in enumerate(choosed):
                batch[i] = imgRandomCrop(s, 256, 256, FLAGS.img_size)
                batch[i] = encode(batch[i])
            yield batch
        except:
            continue
def test1(dir, status='X2Y'):
    capture = cv2.VideoCapture(dir)
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, FLAGS.img_size)
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, FLAGS.img_size)
    config = tf.ConfigProto(allow_soft_placement=True)
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        _, content_img = capture.read()
        content = np.expand_dims(encode(content_img), 0)
        network = CycleGAN(FLAGS.batch_size,
                           FLAGS.ngf,
                           FLAGS.img_size,
                           is_training=False,
                           Norm=FLAGS.Norm)
        network.test(content.shape)
        sess.run(tf.global_variables_initializer())
        var_list = [
            var for var in tf.global_variables() if 'generator_net' in var.name
            or 'discriminator_net' in var.name or 'edge_net' in var.name
        ]
        saver = tf.train.Saver(var_list=var_list)
        saver.restore(sess, FLAGS.checkpoint)
        print("restored all")
        while True:
            if status == 'X2Y':
                out = sess.run(network.Ygenerated,
                               feed_dict={network.testA: content})
            else:
                out = sess.run(network.Xgenerated,
                               feed_dict={network.testB: content})
            # result = np.clip(out, 0, 255).astype(np.uint8)
            cv2.imshow("test", result)
            c = cv2.waitKey(1)
            if c == 27:
                break
            _, content_img = capture.read()
            content = np.expand_dims(encode(content_img), 0)
        capture.release()
        cv2.destroyAllWindows()