fix icc_profile

lama_cleaner/helper.py

@@ -135,31 +135,27 @@ def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes:
     return image_bytes
 
 
-def pil_to_bytes(pil_img, ext: str, quality: int = 95, exif_infos={}) -> bytes:
+def pil_to_bytes(pil_img, ext: str, quality: int = 95, infos={}) -> bytes:
     with io.BytesIO() as output:
-        kwargs = {k: v for k, v in exif_infos.items() if v is not None}
-        if ext == "png" and "parameters" in kwargs:
+        kwargs = {k: v for k, v in infos.items() if v is not None}
+        if ext == 'jpg':
+            ext = 'jpeg'
+        if "png" == ext.lower() and "parameters" in kwargs:
             pnginfo_data = PngImagePlugin.PngInfo()
             pnginfo_data.add_text("parameters", kwargs["parameters"])
             kwargs["pnginfo"] = pnginfo_data
 
-        pil_img.save(
-            output,
-            format=ext,
-            quality=quality,
-            **kwargs,
-        )
+        pil_img.save(output, format=ext, quality=quality, **kwargs)
         image_bytes = output.getvalue()
     return image_bytes
 
 
-def load_img(img_bytes, gray: bool = False, return_exif: bool = False):
+def load_img(img_bytes, gray: bool = False, return_info: bool = False):
     alpha_channel = None
     image = Image.open(io.BytesIO(img_bytes))
 
-    if return_exif:
-        info = image.info or {}
-        exif_infos = {"exif": image.getexif(), "parameters": info.get("parameters")}
+    if return_info:
+        infos = image.info
 
     try:
         image = ImageOps.exif_transpose(image)
@@ -178,8 +174,8 @@ def load_img(img_bytes, gray: bool = False, return_exif: bool = False):
             image = image.convert("RGB")
             np_img = np.array(image)
 
-    if return_exif:
-        return np_img, alpha_channel, exif_infos
+    if return_info:
+        return np_img, alpha_channel, infos
     return np_img, alpha_channel
 
 

lama_cleaner/server.py

@@ -150,7 +150,7 @@ def save_image():
     origin_image_bytes = input["image"].read()  # RGB
     # ext = get_image_ext(origin_image_bytes)
     ext = "png"
-    image, alpha_channel, exif_infos = load_img(origin_image_bytes, return_exif=True)
+    image, alpha_channel, infos = load_img(origin_image_bytes, return_info=True)
     save_path = (global_config.output_dir / filename).with_suffix(f".{ext}")
 
     if alpha_channel is not None:
@@ -166,7 +166,7 @@ def save_image():
         pil_image,
         ext,
         quality=global_config.image_quality,
-        exif_infos=exif_infos,
+        infos=infos,
     )
     try:
         with open(save_path, "wb") as fw:
@@ -230,7 +230,7 @@ def process():
     input = request.files
     # RGB
     origin_image_bytes = input["image"].read()
-    image, alpha_channel, exif_infos = load_img(origin_image_bytes, return_exif=True)
+    image, alpha_channel, exif_infos = load_img(origin_image_bytes, return_info=True)
 
     mask, _ = load_img(input["mask"].read(), gray=True)
     mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
@@ -337,7 +337,7 @@ def process():
             Image.fromarray(res_np_img),
             ext,
             quality=global_config.image_quality,
-            exif_infos=exif_infos,
+            infos=exif_infos,
         )
     )
 
@@ -363,8 +363,8 @@ def run_plugin():
         return "Plugin not found", 500
 
     origin_image_bytes = files["image"].read()  # RGB
-    rgb_np_img, alpha_channel, exif_infos = load_img(
-        origin_image_bytes, return_exif=True
+    rgb_np_img, alpha_channel, infos = load_img(
+        origin_image_bytes, return_info=True
     )
 
     start = time.time()
@@ -416,7 +416,7 @@ def run_plugin():
                     Image.fromarray(rgb_res),
                     ext,
                     quality=global_config.image_quality,
-                    exif_infos=exif_infos,
+                    infos=infos,
                 )
             ),
             mimetype=f"image/{ext}",

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)