from sdsv_dewarp.api import AlignImage
import cv2
import glob
import os 
import tqdm
import time
import argparse



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input")
    parser.add_argument("--out")
    parser.add_argument("--device", type=str, default="cuda:1")

    args = parser.parse_args()
    model  = AlignImage(device=args.device)


    img_dir = args.input
    out_dir = args.out
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    img_paths  = glob.glob(img_dir + "/*")

    times = []
    for img_path in tqdm.tqdm(img_paths):
        t1 = time.time()
        img = cv2.imread(img_path)
        if img is None:
            print(img_path)
            continue

        aligned_img, is_blank, angle_align = model(img)

        times.append(time.time() - t1)

        if not is_blank:
            cv2.imwrite(os.path.join(out_dir, os.path.basename(img_path)), aligned_img)
        else:
            cv2.imwrite(os.path.join(out_dir, os.path.basename(img_path)), img)


    times = times[1:]
    print("Avg time: ", sum(times) / len(times))