copy_img.py 1009 B

1234567891011121314151617181920212223242526
  1. import os, random, shutil, tqdm, argparse
  2. import numpy as np
  3. PROJECT_NAME = os.environ['PROJECT_NAME']
  4. def parse_opt():
  5. parser = argparse.ArgumentParser()
  6. parser.add_argument('--path', type=str, required=True, help='base_image_path')
  7. parser.add_argument('--num', type=int, required=True, help='image num')
  8. opt = parser.parse_known_args()[0]
  9. return opt
  10. if __name__ == '__main__':
  11. opt = parse_opt()
  12. pwd, base_image_path = os.getcwd(), opt.path
  13. dest_image_path = f'{pwd}/{PROJECT_NAME}/img_train'
  14. image_num = opt.num
  15. if os.path.exists(dest_image_path):
  16. shutil.rmtree(dest_image_path)
  17. os.makedirs(dest_image_path, exist_ok=True)
  18. base_image_listdir = os.listdir(base_image_path)
  19. random.shuffle(base_image_listdir)
  20. 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}'):
  21. shutil.copy(f'{base_image_path}/{path}', f'{dest_image_path}/{path}')