From b0ba3262cc0979dd49cbe68a0e0e17cedae78d60 Mon Sep 17 00:00:00 2001 From: jsong468 Date: Wed, 14 Aug 2024 16:40:16 -0700 Subject: [PATCH 1/3] created add_image_text_converter and unit tests --- pyrit/prompt_converter/__init__.py | 2 + .../add_image_text_converter.py | 147 ++++++++++++++++++ .../test_add_image_text_converter.py | 126 +++++++++++++++ 3 files changed, 275 insertions(+) create mode 100644 pyrit/prompt_converter/add_image_text_converter.py create mode 100644 tests/converter/test_add_image_text_converter.py diff --git a/pyrit/prompt_converter/__init__.py b/pyrit/prompt_converter/__init__.py index 62ffbcfdf..097ca2b46 100644 --- a/pyrit/prompt_converter/__init__.py +++ b/pyrit/prompt_converter/__init__.py @@ -4,6 +4,7 @@ from pyrit.prompt_converter.prompt_converter import ConverterResult, PromptConverter from pyrit.prompt_converter.llm_generic_text_converter import LLMGenericTextConverter +from pyrit.prompt_converter.add_image_text_converter import AddImageTextConverter from pyrit.prompt_converter.add_text_image_converter import AddTextImageConverter from pyrit.prompt_converter.ascii_art_converter import AsciiArtConverter from pyrit.prompt_converter.atbash_converter import AtbashConverter @@ -34,6 +35,7 @@ __all__ = [ + "AddImageTextConverter", "AddTextImageConverter", "AsciiArtConverter", "AtbashConverter", diff --git a/pyrit/prompt_converter/add_image_text_converter.py b/pyrit/prompt_converter/add_image_text_converter.py new file mode 100644 index 000000000..32a39d013 --- /dev/null +++ b/pyrit/prompt_converter/add_image_text_converter.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os.path +import logging +import base64 +import string +from typing import Optional + +from PIL import Image, ImageDraw, ImageFont +import textwrap +from io import BytesIO + +from pyrit.models.data_type_serializer import data_serializer_factory +from pyrit.models.prompt_request_piece import PromptDataType +from pyrit.prompt_converter import PromptConverter, ConverterResult + +logger = logging.getLogger(__name__) + + +class AddImageTextConverter(PromptConverter): + """ + Adds a string to an image and wraps the text into multiple lines if necessary. + This class is similar to AddImageTextConverter except + we pass in an image file path as an argument to the constructor as opposed to text. + + Args: + img_to_add (str): file path of image to add text to + font_name (str, optional): path of font to use. Must be a TrueType font (.ttf). Defaults to "arial.ttf". + color (tuple, optional): color to print text in, using RGB values. Defaults to (0, 0, 0). + font_size (optional, float): size of font to use. Defaults to 15. + x_pos (int, optional): x coordinate to place text in (0 is left most). Defaults to 10. + y_pos (int, optional): y coordinate to place text in (0 is upper most). Defaults to 10. + output_filename (optional, str): filename to store converted image. If not provided a unique UUID will be used + """ + + def __init__( + self, + img_to_add: str, + font_name: Optional[str] = "arial.ttf", + color: Optional[tuple[int, int, int]] = (0, 0, 0), + font_size: Optional[int] = 15, + x_pos: Optional[int] = 10, + y_pos: Optional[int] = 10, + output_filename: Optional[str] = None, + ): + if not img_to_add: + raise ValueError("Please provide valid image path") + if not os.path.exists(img_to_add): + raise FileNotFoundError("The specified image file does not exist") + if not font_name.endswith(".ttf"): + raise ValueError("The specified font must be a TrueType font with a .ttf extension") + self._img_to_add = img_to_add + self._font_name = font_name + self._font_size = font_size + self._font = self._load_font() + self._color = color + self._x_pos = x_pos + self._y_pos = y_pos + self._output_name = output_filename + + def _load_font(self): + """ + Load the font for a given font name and font size + + Returns: + ImageFont.FreeTypeFont or ImageFont.ImageFont: The loaded font object. If the specified font + cannot be loaded, the default font is returned. + + Raises: + OSError: If the font resource cannot be loaded, a warning is logged and the default font is used instead. + """ + # Try to load the specified font + try: + font = ImageFont.truetype(self._font_name, self._font_size) + except OSError: + logger.warning(f"Cannot open font resource: {self._font_name}. Using default font.") + font = ImageFont.load_default() + return font + + def _add_text_to_image(self, text: str) -> Image.Image: + """ + Adds wrapped text to the image at self._img_to_add. + + Args: + text (str): The text to add to the image. + + Returns: + Image.Image: The image with added text. + """ + if not text: + raise ValueError("Please provide valid text value") + # Open the image and create a drawing object + image = Image.open(self._img_to_add) + draw = ImageDraw.Draw(image) + + # Calculate the maximum width in pixels with margin into account + margin = 5 + max_width_pixels = image.size[0] - margin + + # Estimate the maximum chars that can fit on a line + alphabet_letters = string.ascii_letters # This gives 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ' + bbox = draw.textbbox((0, 0), alphabet_letters, font=self._font) + avg_char_width = (bbox[2] - bbox[0]) / len(alphabet_letters) + max_chars_per_line = int(max_width_pixels // avg_char_width) + + # Wrap the text + wrapped_text = textwrap.fill(text, width=max_chars_per_line) + + # Add wrapped text to image + y_offset = self._y_pos + for line in wrapped_text.split("\n"): + draw.text((self._x_pos, y_offset), line, font=self._font, fill=self._color) + bbox = draw.textbbox((self._x_pos, y_offset), line, font=self._font) + line_height = bbox[3] - bbox[1] + y_offset += line_height + + return image + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + """ + Converter that overlays input text on the img_to_add. + + Args: + prompt (str): The prompt to be added to the image. + input_type (PromptDataType): type of data + Returns: + ConverterResult: The filename of the converted image as a ConverterResult Object + """ + if not self.input_supported(input_type): + raise ValueError("Input type not supported") + + img_serializer = data_serializer_factory(value=self._img_to_add, data_type="image_path") + + # Add text to the image + updated_img = self._add_text_to_image(text=prompt) + + image_bytes = BytesIO() + mime_type = img_serializer.get_mime_type(self._img_to_add) + image_type = mime_type.split("/")[-1] + updated_img.save(image_bytes, format=image_type) + image_str = base64.b64encode(image_bytes.getvalue()) + img_serializer.save_b64_image(data=image_str, output_filename=self._output_name) + return ConverterResult(output_text=img_serializer.value, output_type="image_path") + + def input_supported(self, input_type: PromptDataType) -> bool: + return input_type == "text" diff --git a/tests/converter/test_add_image_text_converter.py b/tests/converter/test_add_image_text_converter.py new file mode 100644 index 000000000..daa8f1b7f --- /dev/null +++ b/tests/converter/test_add_image_text_converter.py @@ -0,0 +1,126 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + + +import pytest +import os + +from pyrit.prompt_converter import AddImageTextConverter, AddTextImageConverter + +from PIL import Image, ImageFont + + +@pytest.fixture +def image_text_converter_sample_image(): + img = Image.new("RGB", (100, 100), color=(125, 125, 125)) + img.save("test.png") + return "test.png" + + +def test_add_image_text_converter_initialization(image_text_converter_sample_image): + converter = AddImageTextConverter( + img_to_add=image_text_converter_sample_image, + font_name="arial.ttf", + color=(255, 255, 255), + font_size=20, + x_pos=10, + y_pos=10, + ) + assert converter._img_to_add == "test.png" + assert converter._font_name == "arial.ttf" + assert converter._color == (255, 255, 255) + assert converter._font_size == 20 + assert converter._x_pos == 10 + assert converter._y_pos == 10 + assert converter._font is not None + assert type(converter._font) is ImageFont.FreeTypeFont + os.remove("test.png") + + +def test_add_image_text_converter_invalid_font(image_text_converter_sample_image): + with pytest.raises(ValueError): + AddImageTextConverter( + img_to_add=image_text_converter_sample_image, font_name="arial.otf" + ) # Invalid font extension + os.remove("test.png") + + +def test_add_image_text_converter_null_img_to_add(): + with pytest.raises(ValueError): + AddImageTextConverter(img_to_add="", font_name="arial.ttf") + + +def test_add_image_text_converter_invalid_file_path(): + with pytest.raises(FileNotFoundError): + AddImageTextConverter(img_to_add="nonexistent_image.png", font_name="arial.ttf") + + +def test_add_image_text_converter_fallback_to_default_font(image_text_converter_sample_image, caplog): + AddImageTextConverter( + img_to_add=image_text_converter_sample_image, + font_name="nonexistent_font.ttf", + color=(255, 255, 255), + font_size=20, + x_pos=10, + y_pos=10, + ) + assert any( + record.levelname == "WARNING" and "Cannot open font resource" in record.message for record in caplog.records + ) + os.remove("test.png") + + +def test_image_text_converter_add_text_to_image(image_text_converter_sample_image): + converter = AddImageTextConverter( + img_to_add=image_text_converter_sample_image, font_name="arial.ttf", color=(255, 255, 255) + ) + image = Image.open("test.png") + pixels_before = list(image.getdata()) + updated_image = converter._add_text_to_image("Sample Text!") + pixels_after = list(updated_image.getdata()) + assert updated_image + # Check if at least one pixel changed, indicating that text was added + assert pixels_before != pixels_after + os.remove("test.png") + + +@pytest.mark.asyncio +async def test_add_image_text_converter_invalid_input_text(image_text_converter_sample_image) -> None: + converter = AddImageTextConverter(img_to_add=image_text_converter_sample_image) + with pytest.raises(ValueError): + assert await converter.convert_async(prompt="", input_type="text") # type: ignore + os.remove("test.png") + + +@pytest.mark.asyncio +async def test_add_image_text_converter_convert_async(image_text_converter_sample_image) -> None: + converter = AddImageTextConverter( + img_to_add=image_text_converter_sample_image, output_filename="sample_conv_image.png" + ) + converted_image = await converter.convert_async(prompt="Sample Text 2!", input_type="text") + assert converted_image + assert converted_image.output_text == "sample_conv_image.png" + assert converted_image.output_type == "image_path" + assert os.path.exists(converted_image.output_text) + os.remove(converted_image.output_text) + os.remove("test.png") + + +def test_text_image_converter_input_supported(image_text_converter_sample_image): + converter = AddImageTextConverter(img_to_add=image_text_converter_sample_image) + assert converter.input_supported("image_path") is False + assert converter.input_supported("text") is True + + +@pytest.mark.asyncio +async def test_add_image_text_converter_equal_to_add_text_image(image_text_converter_sample_image) -> None: + converter = AddImageTextConverter(img_to_add=image_text_converter_sample_image) + converted_image = await converter.convert_async(prompt="Sample Text!", input_type="text") + text_image_converter = AddTextImageConverter(text_to_add="Sample Text!") + converted_text_image = await text_image_converter.convert_async(prompt="test.png", input_type="image_path") + pixels_image_text = list(Image.open(converted_image.output_text).getdata()) + pixels_text_image = list(Image.open(converted_text_image.output_text).getdata()) + assert pixels_image_text == pixels_text_image + os.remove(converted_image.output_text) + os.remove("test.png") + os.remove(converted_text_image.output_text) From 782e61049a213c6dd13abdc964fbacd428bbd440 Mon Sep 17 00:00:00 2001 From: jsong468 Date: Thu, 15 Aug 2024 13:58:01 -0700 Subject: [PATCH 2/3] minor testing changes --- tests/converter/test_add_image_text_converter.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/converter/test_add_image_text_converter.py b/tests/converter/test_add_image_text_converter.py index daa8f1b7f..ed910f097 100644 --- a/tests/converter/test_add_image_text_converter.py +++ b/tests/converter/test_add_image_text_converter.py @@ -25,6 +25,7 @@ def test_add_image_text_converter_initialization(image_text_converter_sample_ima font_size=20, x_pos=10, y_pos=10, + output_filename="sample_conv_image.png", ) assert converter._img_to_add == "test.png" assert converter._font_name == "arial.ttf" @@ -34,6 +35,7 @@ def test_add_image_text_converter_initialization(image_text_converter_sample_ima assert converter._y_pos == 10 assert converter._font is not None assert type(converter._font) is ImageFont.FreeTypeFont + assert converter._output_name == "sample_conv_image.png" os.remove("test.png") @@ -97,7 +99,7 @@ async def test_add_image_text_converter_convert_async(image_text_converter_sampl converter = AddImageTextConverter( img_to_add=image_text_converter_sample_image, output_filename="sample_conv_image.png" ) - converted_image = await converter.convert_async(prompt="Sample Text 2!", input_type="text") + converted_image = await converter.convert_async(prompt="Sample Text!", input_type="text") assert converted_image assert converted_image.output_text == "sample_conv_image.png" assert converted_image.output_type == "image_path" From 46bd65618f79022d83a496e1dd0c0de567c4ea7a Mon Sep 17 00:00:00 2001 From: jsong468 Date: Thu, 15 Aug 2024 22:52:04 -0700 Subject: [PATCH 3/3] deleted redundant FileNotFound error raising --- pyrit/prompt_converter/add_image_text_converter.py | 3 --- tests/converter/test_add_image_text_converter.py | 12 +++++++----- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/pyrit/prompt_converter/add_image_text_converter.py b/pyrit/prompt_converter/add_image_text_converter.py index 32a39d013..3908b7255 100644 --- a/pyrit/prompt_converter/add_image_text_converter.py +++ b/pyrit/prompt_converter/add_image_text_converter.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import os.path import logging import base64 import string @@ -46,8 +45,6 @@ def __init__( ): if not img_to_add: raise ValueError("Please provide valid image path") - if not os.path.exists(img_to_add): - raise FileNotFoundError("The specified image file does not exist") if not font_name.endswith(".ttf"): raise ValueError("The specified font must be a TrueType font with a .ttf extension") self._img_to_add = img_to_add diff --git a/tests/converter/test_add_image_text_converter.py b/tests/converter/test_add_image_text_converter.py index ed910f097..bedea8cd7 100644 --- a/tests/converter/test_add_image_text_converter.py +++ b/tests/converter/test_add_image_text_converter.py @@ -52,11 +52,6 @@ def test_add_image_text_converter_null_img_to_add(): AddImageTextConverter(img_to_add="", font_name="arial.ttf") -def test_add_image_text_converter_invalid_file_path(): - with pytest.raises(FileNotFoundError): - AddImageTextConverter(img_to_add="nonexistent_image.png", font_name="arial.ttf") - - def test_add_image_text_converter_fallback_to_default_font(image_text_converter_sample_image, caplog): AddImageTextConverter( img_to_add=image_text_converter_sample_image, @@ -94,6 +89,13 @@ async def test_add_image_text_converter_invalid_input_text(image_text_converter_ os.remove("test.png") +@pytest.mark.asyncio +async def test_add_image_text_converter_invalid_file_path(): + converter = AddImageTextConverter(img_to_add="nonexistent_image.png", font_name="arial.ttf") + with pytest.raises(FileNotFoundError): + assert await converter.convert_async(prompt="Sample Text!", input_type="text") # type: ignore + + @pytest.mark.asyncio async def test_add_image_text_converter_convert_async(image_text_converter_sample_image) -> None: converter = AddImageTextConverter(