generate_meta_info.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. import argparse
  2. import glob
  3. import os
  4. def main(args):
  5. txt_file = open(args.meta_info, 'w')
  6. for folder, root in zip(args.input, args.root):
  7. img_paths = sorted(glob.glob(os.path.join(folder, '*')))
  8. for img_path in img_paths:
  9. img_name = os.path.relpath(img_path, root)
  10. print(img_name)
  11. txt_file.write(f'{img_name}\n')
  12. if __name__ == '__main__':
  13. """Generate meta info (txt file) for only Ground-Truth images.
  14. It can also generate meta info from several folders into one txt file.
  15. """
  16. parser = argparse.ArgumentParser()
  17. parser.add_argument(
  18. '--input',
  19. nargs='+',
  20. default=['datasets/DF2K/DF2K_HR', 'datasets/DF2K/DF2K_multiscale'],
  21. help='Input folder, can be a list')
  22. parser.add_argument(
  23. '--root',
  24. nargs='+',
  25. default=['datasets/DF2K', 'datasets/DF2K'],
  26. help='Folder root, should have the length as input folders')
  27. parser.add_argument(
  28. '--meta_info',
  29. type=str,
  30. default='datasets/DF2K/meta_info/meta_info_DF2Kmultiscale.txt',
  31. help='txt path for meta info')
  32. args = parser.parse_args()
  33. assert len(args.input) == len(args.root), ('Input folder and folder root should have the same length, but got '
  34. f'{len(args.input)} and {len(args.root)}.')
  35. main(args)