FEAT Add image overlay converter

pyrit/prompt_converter/__init__.py

@@ -26,6 +26,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.malicious_question_generator_converter import MaliciousQuestionGeneratorConverter
@@ -48,6 +49,7 @@
 from pyrit.prompt_converter.variation_converter import VariationConverter
 
 
+
 __all__ = [
     "AddImageTextConverter",
     "AddTextImageConverter",
@@ -70,6 +72,7 @@
     "FuzzerShortenConverter",
     "FuzzerSimilarConverter",
     "HumanInTheLoopConverter",
+    "ImageOverlayConverter",
     "LeetspeakConverter",
     "LLMGenericTextConverter",
     "MaliciousQuestionGeneratorConverter",

pyrit/prompt_converter/image_overlay_converter.py

@@ -0,0 +1,103 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import base64
+from typing import Optional
+
+from PIL import Image
+from io import BytesIO
+from pathlib import Path
+
+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
+
+
+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
+        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,
+            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")
+
+        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_path: str) -> Image.Image:
+        """
+        Embed the second image onto the base image
+
+        Args:
+            overlay_image_path(str): The Path of second image
+
+        Returns:
+            Image.Image: The combined image with overlay.
+        """
+        if not overlay_image_path:
+            raise ValueError("Please provide a valid image path")
+        # Open the images
+        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:
+        """
+        Converter the base image to embed the second image onto it.
+
+        Args:
+            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
+        """
+        if not self.input_supported(input_type):
+            raise ValueError("Input type not supported")
+
+        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(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"

tests/converter/test_image_overlay_converter.py

@@ -0,0 +1,86 @@
+# 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
+
+from pyrit.prompt_converter import ImageOverlayConverter
+
+from io import BytesIO
+
+
+@pytest.fixture
+def base_image_path():
+    # 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_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=15, memory=memory_mock
+    )
+    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():
+    with pytest.raises(ValueError):
+        ImageOverlayConverter(base_image_path="")
+
+
+@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)
+
+    # 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()