IOPaint/lama_cleaner/tests/test_save_exif.py

53 lines
1.4 KiB
Python
Raw Normal View History

2023-02-06 15:00:47 +01:00
import io
2023-02-07 14:00:19 +01:00
from pathlib import Path
2023-12-29 02:55:47 +01:00
from typing import List
2023-02-06 15:00:47 +01:00
from PIL import Image
2023-05-07 10:58:55 +02:00
from lama_cleaner.helper import pil_to_bytes, load_img
2023-02-06 15:00:47 +01:00
2023-02-07 14:00:19 +01:00
current_dir = Path(__file__).parent.absolute().resolve()
2023-02-06 15:00:47 +01:00
def print_exif(exif):
for k, v in exif.items():
print(f"{k}: {v}")
2023-12-29 02:55:47 +01:00
def extra_info(img_p: Path):
2023-05-07 10:58:55 +02:00
ext = img_p.suffix.strip(".")
img_bytes = img_p.read_bytes()
2023-12-29 02:55:47 +01:00
np_img, _, infos = load_img(img_bytes, False, True)
pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, infos=infos)
res_img = Image.open(io.BytesIO(pil_bytes))
return infos, res_img.info
2023-02-06 15:00:47 +01:00
2023-12-29 02:55:47 +01:00
def assert_keys(keys: List[str], infos, res_infos):
for k in keys:
assert k in infos
assert k in res_infos
assert infos[k] == res_infos[k]
2023-02-06 15:00:47 +01:00
2023-12-29 02:55:47 +01:00
def test_png_icc_profile_png():
infos, res_infos = extra_info(current_dir / "icc_profile_test.png")
assert_keys(["icc_profile", "exif"], infos, res_infos)
2023-02-06 15:00:47 +01:00
2023-12-29 02:55:47 +01:00
def test_png_icc_profile_jpeg():
infos, res_infos = extra_info(current_dir / "icc_profile_test.jpg")
assert_keys(["icc_profile", "exif"], infos, res_infos)
2023-02-06 15:00:47 +01:00
2023-05-07 10:58:55 +02:00
def test_jpeg():
jpg_img_p = current_dir / "bunny.jpeg"
2023-12-29 02:55:47 +01:00
infos, res_infos = extra_info(jpg_img_p)
assert_keys(["dpi", "exif"], infos, res_infos)
def test_png_parameter():
jpg_img_p = current_dir / "png_parameter_test.png"
infos, res_infos = extra_info(jpg_img_p)
assert_keys(["parameters"], infos, res_infos)