File size: 1,637 Bytes
cbbdd92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import io
import tempfile
from pathlib import Path
from typing import List
from PIL import Image
from iopaint.helper import pil_to_bytes, load_img
current_dir = Path(__file__).parent.absolute().resolve()
def print_exif(exif):
for k, v in exif.items():
print(f"{k}: {v}")
def extra_info(img_p: Path):
ext = img_p.suffix.strip(".")
img_bytes = img_p.read_bytes()
np_img, _, infos = load_img(img_bytes, False, True)
res_pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, infos=infos)
res_img = Image.open(io.BytesIO(res_pil_bytes))
return infos, res_img.info, res_pil_bytes
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]
def run_test(file_path, keys):
infos, res_infos, res_pil_bytes = extra_info(file_path)
assert_keys(keys, infos, res_infos)
with tempfile.NamedTemporaryFile("wb", suffix=file_path.suffix) as temp_file:
temp_file.write(res_pil_bytes)
temp_file.flush()
infos, res_infos, res_pil_bytes = extra_info(Path(temp_file.name))
assert_keys(keys, infos, res_infos)
def test_png_icc_profile_png():
run_test(current_dir / "icc_profile_test.png", ["icc_profile", "exif"])
def test_png_icc_profile_jpeg():
run_test(current_dir / "icc_profile_test.jpg", ["icc_profile", "exif"])
def test_jpeg():
jpg_img_p = current_dir / "bunny.jpeg"
run_test(jpg_img_p, ["dpi", "exif"])
def test_png_parameter():
jpg_img_p = current_dir / "png_parameter_test.png"
run_test(jpg_img_p, ["parameters"])
|