fix icc_profile

lama_cleaner/tests/test_save_exif.py

@@ -1,5 +1,6 @@
 import io
 from pathlib import Path
+from typing import List
 
 from PIL import Image
 
@@ -13,31 +14,39 @@ def print_exif(exif):
         print(f"{k}: {v}")
 
 
-def run_test(img_p: Path):
-    print(img_p)
+def extra_info(img_p: Path):
     ext = img_p.suffix.strip(".")
     img_bytes = img_p.read_bytes()
-    np_img, _, exif_infos = load_img(img_bytes, False, True)
-    print(exif_infos)
-    print("Original exif_infos")
-    print_exif(exif_infos["exif"])
-
-    pil_to_bytes(Image.fromarray(np_img), ext=ext, exif_infos={})
-
-    pil_bytes = pil_to_bytes(Image.fromarray(np_img), ext=ext, exif_infos=exif_infos)
+    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))
-    print(f"Result img info: {res_img.info}")
-    res_exif = res_img.getexif()
-    print_exif(res_exif)
-    assert res_exif == exif_infos["exif"]
-    assert exif_infos["parameters"] == res_img.info.get("parameters")
+    return infos, res_img.info
 
 
-def test_png():
-    run_test(current_dir / "image.png")
-    run_test(current_dir / "pnginfo_test.png")
+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 test_png_icc_profile_png():
+    infos, res_infos = extra_info(current_dir / "icc_profile_test.png")
+    assert_keys(["icc_profile", "exif"], infos, res_infos)
+
+
+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)
 
 
 def test_jpeg():
     jpg_img_p = current_dir / "bunny.jpeg"
-    run_test(jpg_img_p)
+    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)