extract_specific_png.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. import os
  2. import argparse
  3. import glob
  4. def main():
  5. parser = argparse.ArgumentParser(description="Extract specific samples indices")
  6. parser.add_argument('--folder', type=str, help='folder with all rawls files', required=True)
  7. parser.add_argument('--index', type=str, help='current rawls image index', required=True)
  8. parser.add_argument('--nsamples', type=str, help='expected nsamples for image', required=True)
  9. parser.add_argument('--output', type=str, help='folder with all png files', required=True)
  10. args = parser.parse_args()
  11. p_folder = args.folder
  12. p_output = args.output
  13. p_index = args.index
  14. p_samples = args.nsamples
  15. expected_index = str(p_index)
  16. while len(expected_index) < 6:
  17. expected_index = "0" + expected_index
  18. output_index = ""
  19. while len(output_index) < 6:
  20. output_index = "0" + output_index
  21. images_path = glob.glob(f"{p_folder}/**/**/*{expected_index}.png")
  22. for img in sorted(images_path):
  23. # replace expected Samples value
  24. img_data = img.split('-')
  25. img_data[-2] = "S" + p_samples
  26. img_data[-1] = output_index + ".png"
  27. output_path = '-'.join(img_data)
  28. output_path = output_path.replace(p_folder, p_output)
  29. output_folder, _ = os.path.split(output_path)
  30. if not os.path.exists(output_folder):
  31. os.makedirs(output_folder)
  32. if not os.path.exists(output_path):
  33. os.system(f'cp {img} {output_path}')
  34. else:
  35. print(f'{output_path} already exists')
  36. if __name__ == "__main__":
  37. main()