1234567891011121314151617181920212223242526 |
- 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}')
|