test_save_exif.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. import io
  2. import tempfile
  3. from pathlib import Path
  4. from typing import List
  5. from PIL import Image
  6. from sorawm.iopaint.helper import load_img, pil_to_bytes
  7. current_dir = Path(__file__).parent.absolute().resolve()
  8. def print_exif(exif):
  9. for k, v in exif.items():
  10. print(f"{k}: {v}")
  11. def extra_info(img_p: Path):
  12. ext = img_p.suffix.strip(".")
  13. img_bytes = img_p.read_bytes()
  14. np_img, _, infos = load_img(img_bytes, False, True)
  15. res_pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, infos=infos)
  16. res_img = Image.open(io.BytesIO(res_pil_bytes))
  17. return infos, res_img.info, res_pil_bytes
  18. def assert_keys(keys: List[str], infos, res_infos):
  19. for k in keys:
  20. assert k in infos
  21. assert k in res_infos
  22. assert infos[k] == res_infos[k]
  23. def run_test(file_path, keys):
  24. infos, res_infos, res_pil_bytes = extra_info(file_path)
  25. assert_keys(keys, infos, res_infos)
  26. with tempfile.NamedTemporaryFile("wb", suffix=file_path.suffix) as temp_file:
  27. temp_file.write(res_pil_bytes)
  28. temp_file.flush()
  29. infos, res_infos, res_pil_bytes = extra_info(Path(temp_file.name))
  30. assert_keys(keys, infos, res_infos)
  31. def test_png_icc_profile_png():
  32. run_test(current_dir / "icc_profile_test.png", ["icc_profile", "exif"])
  33. def test_png_icc_profile_jpeg():
  34. run_test(current_dir / "icc_profile_test.jpg", ["icc_profile", "exif"])
  35. def test_jpeg():
  36. jpg_img_p = current_dir / "bunny.jpeg"
  37. run_test(jpg_img_p, ["dpi", "exif"])
  38. def test_png_parameter():
  39. jpg_img_p = current_dir / "png_parameter_test.png"
  40. run_test(jpg_img_p, ["parameters"])