fix pnginfo data

lama_cleaner/tests/test_save_exif.py

@@ -3,12 +3,9 @@ from pathlib import Path
 
 from PIL import Image
 
-from lama_cleaner.helper import pil_to_bytes
-
+from lama_cleaner.helper import pil_to_bytes, load_img
 
 current_dir = Path(__file__).parent.absolute().resolve()
-png_img_p = current_dir / "image.png"
-jpg_img_p = current_dir / "bunny.jpeg"
 
 
 def print_exif(exif):
@@ -16,27 +13,31 @@ def print_exif(exif):
         print(f"{k}: {v}")
 
 
-def test_png():
-    img = Image.open(png_img_p)
-    exif = img.getexif()
-    print_exif(exif)
+def run_test(img_p: Path):
+    print(img_p)
+    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_bytes = pil_to_bytes(img, ext="png", exif=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)
     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")
 
-    assert dict(exif) == dict(res_exif)
+
+def test_png():
+    run_test(current_dir / "image.png")
+    run_test(current_dir / "pnginfo_test.png")
 
 
 def test_jpeg():
-    img = Image.open(jpg_img_p)
-    exif = img.getexif()
-    print_exif(exif)
-
-    pil_bytes = pil_to_bytes(img, ext="jpeg", exif=exif)
-
-    res_img = Image.open(io.BytesIO(pil_bytes))
-    res_exif = res_img.getexif()
-
-    assert dict(exif) == dict(res_exif)
+    jpg_img_p = current_dir / "bunny.jpeg"
+    run_test(jpg_img_p)