import os, random, shutil, tqdm, argparse
import numpy as np
PROJECT_NAME = os.environ['PROJECT_NAME']

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--path', type=str, required=True, help='base_image_path')
    parser.add_argument('--num', type=int, required=True, help='image num')
    opt = parser.parse_known_args()[0]
    return opt

if __name__ == '__main__':
    opt = parse_opt()
    pwd, base_image_path = os.getcwd(), opt.path
    dest_image_path = f'{pwd}/{PROJECT_NAME}/img_train'
    image_num = opt.num

    if os.path.exists(dest_image_path):
        shutil.rmtree(dest_image_path)
    os.makedirs(dest_image_path, exist_ok=True)

    base_image_listdir = os.listdir(base_image_path)
    random.shuffle(base_image_listdir)

    for path in tqdm.tqdm(base_image_listdir[:image_num] if image_num > 0 else base_image_listdir, desc=f'from {base_image_path} copy to {dest_image_path}'):
        shutil.copy(f'{base_image_path}/{path}', f'{dest_image_path}/{path}')