generate_multiscale_DF2K.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. import argparse
  2. import glob
  3. import os
  4. from PIL import Image
  5. def main(args):
  6. # For DF2K, we consider the following three scales,
  7. # and the smallest image whose shortest edge is 400
  8. scale_list = [0.75, 0.5, 1 / 3]
  9. shortest_edge = 400
  10. path_list = sorted(glob.glob(os.path.join(args.input, '*')))
  11. for path in path_list:
  12. print(path)
  13. basename = os.path.splitext(os.path.basename(path))[0]
  14. img = Image.open(path)
  15. width, height = img.size
  16. for idx, scale in enumerate(scale_list):
  17. print(f'\t{scale:.2f}')
  18. rlt = img.resize((int(width * scale), int(height * scale)), resample=Image.LANCZOS)
  19. rlt.save(os.path.join(args.output, f'{basename}T{idx}.png'))
  20. # save the smallest image which the shortest edge is 400
  21. if width < height:
  22. ratio = height / width
  23. width = shortest_edge
  24. height = int(width * ratio)
  25. else:
  26. ratio = width / height
  27. height = shortest_edge
  28. width = int(height * ratio)
  29. rlt = img.resize((int(width), int(height)), resample=Image.LANCZOS)
  30. rlt.save(os.path.join(args.output, f'{basename}T{idx+1}.png'))
  31. if __name__ == '__main__':
  32. """Generate multi-scale versions for GT images with LANCZOS resampling.
  33. It is now used for DF2K dataset (DIV2K + Flickr 2K)
  34. """
  35. parser = argparse.ArgumentParser()
  36. parser.add_argument('--input', type=str, default='datasets/DF2K/DF2K_HR', help='Input folder')
  37. parser.add_argument('--output', type=str, default='datasets/DF2K/DF2K_multiscale', help='Output folder')
  38. args = parser.parse_args()
  39. os.makedirs(args.output, exist_ok=True)
  40. main(args)