check_indices.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. # main import
  2. import os
  3. import argparse
  4. import shutil
  5. # parameters
  6. scene_image_quality_separator = '_'
  7. scene_image_extension = '.png'
  8. def get_scene_image_quality(img_path):
  9. # if path getting last element (image name) and extract quality
  10. img_postfix = img_path.split('/')[-1].split(scene_image_quality_separator)[-1]
  11. img_quality = img_postfix.replace(scene_image_extension, '')
  12. return int(img_quality)
  13. def rename_folder_images(folder, output, expected):
  14. folders = os.listdir(folder)
  15. for folder_name in folders:
  16. folder_path = os.path.join(folder, folder_name)
  17. output_folder_path = os.path.join(output, folder_name)
  18. images = sorted(os.listdir(folder_path))
  19. last_index = get_scene_image_quality(images[-1])
  20. if last_index != expected:
  21. print('Update images indices for %s' % folder_path)
  22. if not os.path.exists(output_folder_path):
  23. os.makedirs(output_folder_path)
  24. for img in images:
  25. img_path = os.path.join(folder_path, img)
  26. current_quality = get_scene_image_quality(img_path)
  27. img_prefix_split = img_path.split('/')[-1].split(scene_image_quality_separator)
  28. del img_prefix_split[-1]
  29. img_prefix = "_".join(img_prefix_split)
  30. index_str = str(current_quality * 20)
  31. while len(index_str) < 5:
  32. index_str = "0" + index_str
  33. img_output_name = img_prefix + '_' + index_str + '.png'
  34. img_output_path = os.path.join(output_folder_path, img_output_name)
  35. shutil.copy2(img_path, img_output_path)
  36. else:
  37. print('Max expected found for', folder_path, '(no need to update)')
  38. def main():
  39. parser = argparse.ArgumentParser(description="rename image with correct indices")
  40. parser.add_argument('--folder', type=str, help="folder with HD images", required=True)
  41. parser.add_argument('--output', type=str, help="output folder", required=True)
  42. parser.add_argument('--expected', type=int, help="max expected index", required=True)
  43. args = parser.parse_args()
  44. p_folder = args.folder
  45. p_output = args.output
  46. p_expected = args.expected
  47. rename_folder_images(p_folder, p_output, p_expected)
  48. if __name__ == "__main__":
  49. main()