From d6ba5b5f367672edd3d406140e201c47fcd0ee80 Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Wed, 16 Oct 2024 07:54:16 +1100 Subject: [PATCH 1/9] created image_overlay_converter.py --- .../image_overlay_converter.py | 99 +++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 pyrit/prompt_converter/image_overlay_converter.py diff --git a/pyrit/prompt_converter/image_overlay_converter.py b/pyrit/prompt_converter/image_overlay_converter.py new file mode 100644 index 000000000..c8c22edef --- /dev/null +++ b/pyrit/prompt_converter/image_overlay_converter.py @@ -0,0 +1,99 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import logging +import base64 +from typing import Optional + +from PIL import Image +from io import BytesIO + +from pyrit.models import data_serializer_factory +from pyrit.models import PromptDataType +from pyrit.prompt_converter import PromptConverter, ConverterResult +from pyrit.memory import MemoryInterface, DuckDBMemory + +logger = logging.getLogger(__name__) + + +class ImageOverlayConverter(PromptConverter): + """ + A converter that takes in a base image, and a secondary image to embed within the main image. + + Args: + base_image_path (str): File path of the base image + secondary_image_path (str): File path of the second image to embed with the base image + x_pos (int, optional): X coordinate to place second image on the base image (0 is left most). Defaults to 0. + y_pos (int, optional): Y coordinate to place second image on the base image (0 is upper most). Defaults to 0. + memory: (memory, optional): Memory to store the chat messages. DuckDBMemory will be used by default. + """ + + def __init__( + self, + base_image_path: str, + secondary_image_path: str, + x_pos: Optional[int] = 0, + y_pos: Optional[int] = 0, + memory: Optional[MemoryInterface] = None, + ): + if not base_image_path: + raise ValueError("Please provide valid base image path") + if not secondary_image_path: + raise ValueError("Please provide valid overlay image path") + + self._base_image_path = base_image_path + self._secondary_image_path = secondary_image_path + self._x_pos = x_pos + self._y_pos = y_pos + self._memory = memory or DuckDBMemory() + + def _add_overlay_image(self) -> Image.Image: + """ + Embed the second image onto the base image + + Returns: + Image.Image: The combined image with overlay. + """ + # Open the images + base_image = Image.open(self._base_image_path) + second_image = Image.open(self._secondary_image_path) + + # Make a copy of base image, in case the user want to keep their input images unchanged + copied_base_image = base_image.copy() + + # Paste the second image onto the base image + copied_base_image.paste(second_image, (self._x_pos, self._y_pos), second_image) + + return copied_base_image + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image") -> ConverterResult: + """ + Converter the base image to embed the second image onto it. + + Args: + prompt (str): Not used in this converter but may be required by the interface + input_type (PromptDataType): type of data, should be image + Returns: + ConverterResult: converted image with file path + """ + if not self.input_supported(input_type): + raise ValueError("Input type not supported") + + img_serializer = data_serializer_factory(value=self._base_image_path, data_type="image_path", memory=self._memory) + + # Add overlay to the image + updated_img = self._add_overlay_image() + + # Encode the image using base64 and return ConverterResult + # I took the following code from add_image_text_converter.py @author rdheekonda + image_bytes = BytesIO() + mime_type = img_serializer.get_mime_type(self._base_image_path) + image_type = mime_type.split("/")[-1] + updated_img.save(image_bytes, format=image_type) + image_str = base64.b64encode(image_bytes.getvalue()) + # Save image as generated UUID filename + await img_serializer.save_b64_image(data=image_str) + return ConverterResult(output_text=str(img_serializer.value), output_type="image_path") + + def input_supported(self, input_type: PromptDataType) -> bool: + return input_type == "image" From 947ef15c0eb2de2185913bcdeb4ec21668853bce Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Wed, 16 Oct 2024 07:55:03 +1100 Subject: [PATCH 2/9] update __init__.py to add the new converter --- pyrit/prompt_converter/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyrit/prompt_converter/__init__.py b/pyrit/prompt_converter/__init__.py index a232e5dd1..d04b8681a 100644 --- a/pyrit/prompt_converter/__init__.py +++ b/pyrit/prompt_converter/__init__.py @@ -43,6 +43,7 @@ from pyrit.prompt_converter.unicode_sub_converter import UnicodeSubstitutionConverter from pyrit.prompt_converter.url_converter import UrlConverter from pyrit.prompt_converter.variation_converter import VariationConverter +from pyrit.prompt_converter.image_overlay_converter import ImageOverlayConverter __all__ = [ @@ -86,4 +87,5 @@ "UnicodeSubstitutionConverter", "UrlConverter", "VariationConverter", + "ImageOverlayConverter", ] From 7a7bfad288171b3292311ec47c41ee1403242879 Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Wed, 16 Oct 2024 08:28:47 +1100 Subject: [PATCH 3/9] minor change on orders --- pyrit/prompt_converter/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyrit/prompt_converter/__init__.py b/pyrit/prompt_converter/__init__.py index d04b8681a..81eea213b 100644 --- a/pyrit/prompt_converter/__init__.py +++ b/pyrit/prompt_converter/__init__.py @@ -25,6 +25,7 @@ FuzzerSimilarConverter, ) from pyrit.prompt_converter.human_in_the_loop_converter import HumanInTheLoopConverter +from pyrit.prompt_converter.image_overlay_converter import ImageOverlayConverter from pyrit.prompt_converter.leetspeak_converter import LeetspeakConverter from pyrit.prompt_converter.morse_converter import MorseConverter from pyrit.prompt_converter.noise_converter import NoiseConverter @@ -43,7 +44,7 @@ from pyrit.prompt_converter.unicode_sub_converter import UnicodeSubstitutionConverter from pyrit.prompt_converter.url_converter import UrlConverter from pyrit.prompt_converter.variation_converter import VariationConverter -from pyrit.prompt_converter.image_overlay_converter import ImageOverlayConverter + __all__ = [ @@ -67,6 +68,7 @@ "FuzzerShortenConverter", "FuzzerSimilarConverter", "HumanInTheLoopConverter", + "ImageOverlayConverter", "LeetspeakConverter", "LLMGenericTextConverter", "MorseConverter", @@ -87,5 +89,4 @@ "UnicodeSubstitutionConverter", "UrlConverter", "VariationConverter", - "ImageOverlayConverter", ] From d2107474e92c24fcabaa688d13a8b1830927b53d Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Fri, 18 Oct 2024 14:14:25 +1100 Subject: [PATCH 4/9] update on image_overlay_converter.py, to make the whole class more consistent with other converters --- .../image_overlay_converter.py | 30 ++++++++----------- 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/pyrit/prompt_converter/image_overlay_converter.py b/pyrit/prompt_converter/image_overlay_converter.py index c8c22edef..1dc8affc4 100644 --- a/pyrit/prompt_converter/image_overlay_converter.py +++ b/pyrit/prompt_converter/image_overlay_converter.py @@ -22,7 +22,6 @@ class ImageOverlayConverter(PromptConverter): Args: base_image_path (str): File path of the base image - secondary_image_path (str): File path of the second image to embed with the base image x_pos (int, optional): X coordinate to place second image on the base image (0 is left most). Defaults to 0. y_pos (int, optional): Y coordinate to place second image on the base image (0 is upper most). Defaults to 0. memory: (memory, optional): Memory to store the chat messages. DuckDBMemory will be used by default. @@ -31,47 +30,42 @@ class ImageOverlayConverter(PromptConverter): def __init__( self, base_image_path: str, - secondary_image_path: str, x_pos: Optional[int] = 0, y_pos: Optional[int] = 0, memory: Optional[MemoryInterface] = None, ): - if not base_image_path: - raise ValueError("Please provide valid base image path") - if not secondary_image_path: - raise ValueError("Please provide valid overlay image path") self._base_image_path = base_image_path - self._secondary_image_path = secondary_image_path self._x_pos = x_pos self._y_pos = y_pos self._memory = memory or DuckDBMemory() - def _add_overlay_image(self) -> Image.Image: + def _add_overlay_image(self, overlay_image: Image.Image) -> Image.Image: """ Embed the second image onto the base image Returns: Image.Image: The combined image with overlay. """ + if not overlay_image: + raise ValueError("Please provide a valid image") # Open the images base_image = Image.open(self._base_image_path) - second_image = Image.open(self._secondary_image_path) # Make a copy of base image, in case the user want to keep their input images unchanged copied_base_image = base_image.copy() # Paste the second image onto the base image - copied_base_image.paste(second_image, (self._x_pos, self._y_pos), second_image) + copied_base_image.paste(overlay_image, (self._x_pos, self._y_pos), overlay_image) return copied_base_image - async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image") -> ConverterResult: + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: """ Converter the base image to embed the second image onto it. Args: - prompt (str): Not used in this converter but may be required by the interface + prompt (str): The filename of the base image input_type (PromptDataType): type of data, should be image Returns: ConverterResult: converted image with file path @@ -79,15 +73,17 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "imag if not self.input_supported(input_type): raise ValueError("Input type not supported") - img_serializer = data_serializer_factory(value=self._base_image_path, data_type="image_path", memory=self._memory) + img_serializer = data_serializer_factory(value=prompt, data_type="image_path", memory=self._memory) + second_img_bytes = await img_serializer.read_data() + second_img = Image.open(BytesIO(second_img_bytes)) - # Add overlay to the image - updated_img = self._add_overlay_image() + # Add overlay to the base image + updated_img = self._add_overlay_image(second_img) # Encode the image using base64 and return ConverterResult # I took the following code from add_image_text_converter.py @author rdheekonda image_bytes = BytesIO() - mime_type = img_serializer.get_mime_type(self._base_image_path) + mime_type = img_serializer.get_mime_type(prompt) image_type = mime_type.split("/")[-1] updated_img.save(image_bytes, format=image_type) image_str = base64.b64encode(image_bytes.getvalue()) @@ -96,4 +92,4 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "imag return ConverterResult(output_text=str(img_serializer.value), output_type="image_path") def input_supported(self, input_type: PromptDataType) -> bool: - return input_type == "image" + return input_type == "image_path" From 09efd5a303787641fd62583aec50691f265cee86 Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Sun, 20 Oct 2024 14:45:22 +1100 Subject: [PATCH 5/9] add some test cases of image_overlay_converter.py to git --- .../converter/test_image_overlay_converter.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 tests/converter/test_image_overlay_converter.py diff --git a/tests/converter/test_image_overlay_converter.py b/tests/converter/test_image_overlay_converter.py new file mode 100644 index 000000000..c3aa94732 --- /dev/null +++ b/tests/converter/test_image_overlay_converter.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +import os + +import pytest +from PIL import Image + +from pyrit.prompt_converter import ImageOverlayConverter + +from io import BytesIO + + +@pytest.fixture +def base_image_path(): + img = Image.new("RGB", (100, 100), color=(255, 255, 255)) + img.save("base_test.png") + return "base_test.png" + + +@pytest.fixture +def overlay_image_byte(): + img = Image.new("RGBA", (20, 20), color=(125, 125, 125, 125)) + img_bytes = BytesIO() + img.save(img_bytes, format="PNG") + img_bytes = img_bytes.getvalue() + return img_bytes + + +def test_image_overlay_converter_initialization(base_image_path): + converter = ImageOverlayConverter( + base_image_path=base_image_path, x_pos=10, y_pos=10, memory=None + ) + assert converter._base_image_path == "base_test.png" + assert converter._x_pos == 10 + assert converter._y_pos == 10 + os.remove("base_test.png") + + +def test_image_overlay_converter_invalid_image(): + with pytest.raises(ValueError): + ImageOverlayConverter(base_image_path="") + + +def test_image_overlay_converter_add_overlay_image(base_image_path, overlay_image_byte): + converter = ImageOverlayConverter(base_image_path=base_image_path) + base_image = Image.open(base_image_path) + overlay_image = Image.open(overlay_image_byte) + pixels_before = list(base_image.getdata()) + + # Adding overlay image + updated_image = converter._add_overlay_image(overlay_image) + pixels_after = list(updated_image.getdata()) + + assert updated_image is not None + # Check for pixels changes + assert pixels_before != pixels_after + os.remove("base_test.png") + + +def test_image_overlay_converter_input_supported(base_image_path): + converter = ImageOverlayConverter(base_image_path=base_image_path) + assert converter.input_supported("image_path") is True + assert converter.input_supported("text") is False + os.remove("base_test.png") \ No newline at end of file From ec3b110bd59c067ef61f9538d99c34857920df12 Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Sun, 20 Oct 2024 14:45:47 +1100 Subject: [PATCH 6/9] update the image_overlay_converter.py --- .../image_overlay_converter.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/pyrit/prompt_converter/image_overlay_converter.py b/pyrit/prompt_converter/image_overlay_converter.py index 1dc8affc4..32e24d72e 100644 --- a/pyrit/prompt_converter/image_overlay_converter.py +++ b/pyrit/prompt_converter/image_overlay_converter.py @@ -28,12 +28,14 @@ class ImageOverlayConverter(PromptConverter): """ def __init__( - self, - base_image_path: str, - x_pos: Optional[int] = 0, - y_pos: Optional[int] = 0, - memory: Optional[MemoryInterface] = None, + self, + base_image_path: str, + x_pos: Optional[int] = 0, + y_pos: Optional[int] = 0, + memory: Optional[MemoryInterface] = None, ): + if not base_image_path: + raise ValueError("Please provide valid image path") self._base_image_path = base_image_path self._x_pos = x_pos @@ -44,6 +46,9 @@ def _add_overlay_image(self, overlay_image: Image.Image) -> Image.Image: """ Embed the second image onto the base image + Args: + overlay_image(Image.Image): The second image to lay on the base one. + Returns: Image.Image: The combined image with overlay. """ @@ -65,8 +70,9 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "imag Converter the base image to embed the second image onto it. Args: - prompt (str): The filename of the base image - input_type (PromptDataType): type of data, should be image + prompt (str): The filename of the second image + input_type (PromptDataType): type of data, should be image_path + Returns: ConverterResult: converted image with file path """ From 3af29ea5f622fc6197c48627569e9a3b32aab94f Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Sun, 27 Oct 2024 13:31:39 +1100 Subject: [PATCH 7/9] update --- pyrit/prompt_converter/image_overlay_converter.py | 13 +++++-------- tests/converter/test_image_overlay_converter.py | 7 ------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/pyrit/prompt_converter/image_overlay_converter.py b/pyrit/prompt_converter/image_overlay_converter.py index 32e24d72e..d14c4abb2 100644 --- a/pyrit/prompt_converter/image_overlay_converter.py +++ b/pyrit/prompt_converter/image_overlay_converter.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import logging import base64 from typing import Optional @@ -13,8 +12,6 @@ from pyrit.prompt_converter import PromptConverter, ConverterResult from pyrit.memory import MemoryInterface, DuckDBMemory -logger = logging.getLogger(__name__) - class ImageOverlayConverter(PromptConverter): """ @@ -28,11 +25,11 @@ class ImageOverlayConverter(PromptConverter): """ def __init__( - self, - base_image_path: str, - x_pos: Optional[int] = 0, - y_pos: Optional[int] = 0, - memory: Optional[MemoryInterface] = None, + self, + base_image_path: str, + x_pos: Optional[int] = 0, + y_pos: Optional[int] = 0, + memory: Optional[MemoryInterface] = None, ): if not base_image_path: raise ValueError("Please provide valid image path") diff --git a/tests/converter/test_image_overlay_converter.py b/tests/converter/test_image_overlay_converter.py index c3aa94732..a28818e3c 100644 --- a/tests/converter/test_image_overlay_converter.py +++ b/tests/converter/test_image_overlay_converter.py @@ -54,11 +54,4 @@ def test_image_overlay_converter_add_overlay_image(base_image_path, overlay_imag assert updated_image is not None # Check for pixels changes assert pixels_before != pixels_after - os.remove("base_test.png") - - -def test_image_overlay_converter_input_supported(base_image_path): - converter = ImageOverlayConverter(base_image_path=base_image_path) - assert converter.input_supported("image_path") is True - assert converter.input_supported("text") is False os.remove("base_test.png") \ No newline at end of file From e09e20ad9e8ceb55bb91c375aef0b74051cad6a3 Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Mon, 11 Nov 2024 18:28:19 +1100 Subject: [PATCH 8/9] update on image_overlay_converter.py to fix the issue from PR comments --- .../image_overlay_converter.py | 61 ++++++++++--------- 1 file changed, 33 insertions(+), 28 deletions(-) diff --git a/pyrit/prompt_converter/image_overlay_converter.py b/pyrit/prompt_converter/image_overlay_converter.py index d14c4abb2..40f81cbb6 100644 --- a/pyrit/prompt_converter/image_overlay_converter.py +++ b/pyrit/prompt_converter/image_overlay_converter.py @@ -6,6 +6,7 @@ from PIL import Image from io import BytesIO +from pathlib import Path from pyrit.models import data_serializer_factory from pyrit.models import PromptDataType @@ -34,33 +35,37 @@ def __init__( if not base_image_path: raise ValueError("Please provide valid image path") + file_from_path = Path(base_image_path) + if not file_from_path.is_file(): + raise ValueError("File does not exist") + + if x_pos < 0 or y_pos < 0: + raise ValueError("Position is out of boundary ") + self._base_image_path = base_image_path self._x_pos = x_pos self._y_pos = y_pos self._memory = memory or DuckDBMemory() - def _add_overlay_image(self, overlay_image: Image.Image) -> Image.Image: + def _add_overlay_image(self, overlay_image_path: str) -> Image.Image: """ Embed the second image onto the base image Args: - overlay_image(Image.Image): The second image to lay on the base one. + overlay_image_path(str): The Path of second image Returns: Image.Image: The combined image with overlay. """ - if not overlay_image: - raise ValueError("Please provide a valid image") + if not overlay_image_path: + raise ValueError("Please provide a valid image path") # Open the images - base_image = Image.open(self._base_image_path) - - # Make a copy of base image, in case the user want to keep their input images unchanged - copied_base_image = base_image.copy() - - # Paste the second image onto the base image - copied_base_image.paste(overlay_image, (self._x_pos, self._y_pos), overlay_image) - - return copied_base_image + with Image.open(self._base_image_path) as base_image, Image.open(overlay_image_path) as overlay_image: + # Paste the second image onto the base image + # And make a copy of the result, so it is accessible after "with" close + base_image.paste(overlay_image, (self._x_pos, self._y_pos), overlay_image) + result_image = base_image.copy() + return result_image async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: """ @@ -76,23 +81,23 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "imag if not self.input_supported(input_type): raise ValueError("Input type not supported") - img_serializer = data_serializer_factory(value=prompt, data_type="image_path", memory=self._memory) - second_img_bytes = await img_serializer.read_data() - second_img = Image.open(BytesIO(second_img_bytes)) + second_img_from_path = Path(prompt) + if not second_img_from_path.is_file(): + raise ValueError("Overlay Image File does not exist") # Add overlay to the base image - updated_img = self._add_overlay_image(second_img) - - # Encode the image using base64 and return ConverterResult - # I took the following code from add_image_text_converter.py @author rdheekonda - image_bytes = BytesIO() - mime_type = img_serializer.get_mime_type(prompt) - image_type = mime_type.split("/")[-1] - updated_img.save(image_bytes, format=image_type) - image_str = base64.b64encode(image_bytes.getvalue()) - # Save image as generated UUID filename - await img_serializer.save_b64_image(data=image_str) - return ConverterResult(output_text=str(img_serializer.value), output_type="image_path") + updated_img = self._add_overlay_image(prompt) + + # Create a new data serializer to save the images + updated_image_serializer = data_serializer_factory( + data_type="image_path", + extension="png" # you can change the file type as you want + ) + + # Save the result to the path generated by the serializer + updated_img.save(updated_image_serializer.get_data_filename()) + + return ConverterResult(output_text=str(updated_image_serializer.value), output_type="image_path") def input_supported(self, input_type: PromptDataType) -> bool: return input_type == "image_path" From 615e7fbfbf554767270ebc5dc6edf10255021981 Mon Sep 17 00:00:00 2001 From: Hongbo Wang Date: Mon, 11 Nov 2024 18:53:07 +1100 Subject: [PATCH 9/9] update on test_image_overlay_converter.py to fix the issue from PR comments --- .../converter/test_image_overlay_converter.py | 83 +++++++++++++------ 1 file changed, 56 insertions(+), 27 deletions(-) diff --git a/tests/converter/test_image_overlay_converter.py b/tests/converter/test_image_overlay_converter.py index a28818e3c..058b4c54b 100644 --- a/tests/converter/test_image_overlay_converter.py +++ b/tests/converter/test_image_overlay_converter.py @@ -1,6 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. import os +import tempfile +from pathlib import Path +from unittest.mock import MagicMock import pytest from PIL import Image @@ -12,28 +15,44 @@ @pytest.fixture def base_image_path(): - img = Image.new("RGB", (100, 100), color=(255, 255, 255)) - img.save("base_test.png") - return "base_test.png" + # Create a temporary file with a unique name + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: + # Generate a simple image and save it to the temporary file path + img = Image.new("RGB", (100, 100), color=(255, 255, 255)) + img.save(tmp.name) + temp_path = tmp.name # Store the temporary file path + + yield temp_path # Provide the path to the test + + # Cleanup after the test + os.remove(temp_path) @pytest.fixture -def overlay_image_byte(): - img = Image.new("RGBA", (20, 20), color=(125, 125, 125, 125)) - img_bytes = BytesIO() - img.save(img_bytes, format="PNG") - img_bytes = img_bytes.getvalue() - return img_bytes +def overlay_image_path(): + # Create a temporary file with a unique name for overlay image + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: + # Generate a simple image and save it to the temporary file path + img = Image.new("RGB", (10, 10), color=(100, 100, 100)) + img.save(tmp.name) + temp_path = tmp.name # Store the temporary file path + + yield temp_path # Provide the path to the test + + # Cleanup after the test + os.remove(temp_path) def test_image_overlay_converter_initialization(base_image_path): + # use MagicMock for memory + memory_mock = MagicMock() + converter = ImageOverlayConverter( - base_image_path=base_image_path, x_pos=10, y_pos=10, memory=None + base_image_path=base_image_path, x_pos=10, y_pos=15, memory=memory_mock ) - assert converter._base_image_path == "base_test.png" - assert converter._x_pos == 10 - assert converter._y_pos == 10 - os.remove("base_test.png") + assert converter._base_image_path == base_image_path, " Base image path should be initialized" + assert converter._x_pos == 10, "X position should be 10" + assert converter._y_pos == 15, "Y position should be 15" def test_image_overlay_converter_invalid_image(): @@ -41,17 +60,27 @@ def test_image_overlay_converter_invalid_image(): ImageOverlayConverter(base_image_path="") -def test_image_overlay_converter_add_overlay_image(base_image_path, overlay_image_byte): +@pytest.mark.asyncio +async def test_image_overlay_converter_convert_async(base_image_path, overlay_image_path): + # Initialize the converter with the base image path converter = ImageOverlayConverter(base_image_path=base_image_path) - base_image = Image.open(base_image_path) - overlay_image = Image.open(overlay_image_byte) - pixels_before = list(base_image.getdata()) - - # Adding overlay image - updated_image = converter._add_overlay_image(overlay_image) - pixels_after = list(updated_image.getdata()) - - assert updated_image is not None - # Check for pixels changes - assert pixels_before != pixels_after - os.remove("base_test.png") \ No newline at end of file + + # Call the async `convert_async` method with the overlay image path as the prompt + result = await converter.convert_async(prompt=overlay_image_path, input_type="image_path") + + # Verify that the result contains a valid file path in `output_text` + assert isinstance(result.output_text, str), "The result should contain a file path as output_text." + output_path = Path(result.output_text) + + # Check that the output path exists and is a file + assert output_path.is_file(), "The output image file should exist." + + # Open the result image and verify its properties + with Image.open(output_path) as img: + # Ensure the output image dimensions match the base image + assert img.size == Image.open(base_image_path).size, "The output image size should match the base image size." + # Check if the image mode is appropriate + assert img.mode in ["RGB", "RGBA"], "The output image mode should be RGB or RGBA." + + # Cleanup after test + output_path.unlink()