diff --git a/lama_cleaner/helper.py b/lama_cleaner/helper.py index f48cecf..e8a5f68 100644 --- a/lama_cleaner/helper.py +++ b/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 diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index 8fa9f82..9451349 100644 --- a/lama_cleaner/server.py +++ b/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}", diff --git a/lama_cleaner/tests/icc_profile_test.jpg b/lama_cleaner/tests/icc_profile_test.jpg new file mode 100644 index 0000000..b603ef9 Binary files /dev/null and b/lama_cleaner/tests/icc_profile_test.jpg differ diff --git a/lama_cleaner/tests/icc_profile_test.png b/lama_cleaner/tests/icc_profile_test.png new file mode 100644 index 0000000..90d18ac Binary files /dev/null and b/lama_cleaner/tests/icc_profile_test.png differ diff --git a/lama_cleaner/tests/pnginfo_test.png b/lama_cleaner/tests/png_parameter_test.png similarity index 100% rename from lama_cleaner/tests/pnginfo_test.png rename to lama_cleaner/tests/png_parameter_test.png diff --git a/lama_cleaner/tests/test_save_exif.py b/lama_cleaner/tests/test_save_exif.py index 5921256..36f48a5 100644 --- a/lama_cleaner/tests/test_save_exif.py +++ b/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)