import os
import sys
import numpy as np
from tqdm import tqdm
from preprocess import preprocess


def main(data_dir, npy_dir):
    img_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if '.jpg' in f]
    img_files.sort()

    if not os.path.isdir(npy_dir):
        os.mkdir(npy_dir)

    print("INFO: Convert jpg to npy file.")
    with open(os.path.join(npy_dir, 'quant.txt'), 'w') as f:
        # 遍历需要处理的图片文件路径
        for index, img_file in tqdm(enumerate(img_files[:50])):
            # 调用预处理，得到预处理后的 numpy 数组
            img = preprocess(img_file)
            # 文件名可以随便取，但是需要保持唯一
            npy_name = f"{index}.npy"
            np.save(os.path.join(npy_dir, npy_name), img)
            # 文件名写入 quant.txt 文件中，后续量化校准会根据此文件读取使用
            f.write(npy_name + '\n')


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])
