import utils
import argparse
import numpy as np
import hparams_config
import efficientdet_arch
import tensorflow.compat.v1 as tf


def test(imgs, model_name, img_size=512):
    raw_images, images, scales = [], [], []
    for f in tf.io.gfile.glob(image_path_pattern):
        image = Image.open(f)
        raw_images.append(image)
        image, scale = image_preprocess(image, image_size)
        images.append(image)
        scales.append(scale)
    with tf.Session() as sess:
        X = tf.placeholder(tf.float32, shape=(1, img_size, img_size, 3))
        class_outputs, box_outputs = efficientdet_arch.efficientdet(X, model_name=model_name)
        sess.run(tf.global_variables_initializer())
        if tf.io.gfile.isdir(model_name):
            model_name = tf.train.latest_checkpoint(model_name)

        var_dict = utils.get_ema_vars()
        tf.train.get_or_create_global_step()
        sess.run(tf.global_variables_initializer())
        saver = tf.train.Saver(var_dict, max_to_keep=1)
        saver.restore(sess, model_name)
        import time
        times = []
        for img in imgs:
            img = img[np.newaxis, ...]
            start = time.time()
            sess.run(class_outputs, feed_dict={X: img})
            spent = time.time() - start
            print(spent)
            times.append(spent)
        print('mean time of 99 times: %.4f' % np.array(times[1:]).mean())


if __name__ == '__main__':
    parser = argparse.ArgumentParser('test')
    parser.add_argument('--model', type=str, default='0')
    args = parser.parse_args()
    model = 'efficientdet-d%s' % args.model
    img_size = hparams_config.efficientdet_model_param_dict[model]['image_size']
    test(, model, img_size=img_size)