import cv2
import numpy as np
from utility import *
from argument import *
from core import HuiMvOcr
from time import time


def main(args: "ArgType"):
    image_file_list = get_image_file_list(args.image_dir)
    engine = HuiMvOcr(args)

    # warm up 10 times
    if args.warmup:
        img = np.random.uniform(0, 255, [640, 640, 3]).astype(np.uint8)
        for i in range(10):
            engine.ocr_one(img)
    # single
    print("single start")
    total_time, res = 0, []
    for image_file in image_file_list:
        img, flag_gif = check_and_read(image_file)
        if not flag_gif:
            img = cv2.imread(image_file)  # noqa
        st = time()
        rec_res = engine.ocr_one(img, cls=True, use_space=False)
        elapse = time() - st
        total_time += elapse
        res.append(rec_res)
    print(f"single total time: {total_time}")
    for i in range(len(res)):
        print(f"file: {image_file_list[i]}")
        print("res:", res[i])

    # multi
    images = [cv2.imread(file) for file in image_file_list]  # noqa
    print("\n\nmulti start")
    st = time()
    res = engine.ocr_multi(images, cls=True, use_space=False)
    print(f"multi total time {time() - st}")
    for i in range(len(res)):
        print(f"file: {image_file_list[i]}")
        print("res:", res[i])


if __name__ == "__main__":
    main(Args())