Synthetic image generator

src/c++/perf_analyzer/genai-perf/genai_perf/llm_inputs/synthetic_image_generator.py

@@ -0,0 +1,87 @@
+import base64
+import os
+from enum import Enum, auto
+from io import BytesIO
+from pathlib import Path
+from typing import Optional, Tuple, cast
+
+import numpy as np
+from genai_perf.exceptions import GenAIPerfException
+from PIL import Image
+
+
+class ImageFormat(Enum):
+    JPEG = auto()
+    PNG = auto()
+
+
+class RandomFormatBase64Encoder:
+    def __init__(self, image_formats: ImageFormat = ImageFormat.PNG):
+        self.image_formats = image_formats
+
+    def __call__(self, image):
+        choice = np.random.randint(len(self.image_formats))
+        image_format = self.image_formats[choice]
+        buffered = BytesIO()
+        image.save(buffered, format=image_format.name)
+        data = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        prefix = f"data:image/{image_format.name.lower()};base64"
+        return f"{prefix},{data}"
+
+
+def images_from_file_generator(image_path: Path):
+    if not image_path.exists():
+        raise GenAIPerfException(f"File not found: {image_path}")
+
+    image = Image.open(image_path)
+    while True:
+        yield image
+
+
+def white_images_generator():
+    white_image = Image.new("RGB", (100, 100), color="white")
+    while True:
+        yield white_image
+
+
+def build_synthetic_image_generator(
+    mean_size: Tuple[int, int],
+    dimensions_stddev: Tuple[int, int],
+    image_path: Optional[Path] = None,
+):
+    if image_path is None:
+        image_iterator = white_images_generator()
+    else:
+        image_path = cast(Path, image_path)
+        image_iterator = images_from_file_generator(image_path)
+
+    return SyntheticImageGenerator(
+        mean_size=mean_size,
+        dimensions_stddev=dimensions_stddev,
+        image_iterator=image_iterator,
+    )
+
+
+class SyntheticImageGenerator:
+    def __init__(
+        self,
+        mean_size,
+        dimensions_stddev,
+        image_iterator,
+    ):
+        self.image_iterator = image_iterator
+        self.mean_size = mean_size
+        self.dimensions_stddev = dimensions_stddev
+
+    def __iter__(self):
+        return self
+
+    def random_resize(self, image):
+        new_size = np.random.normal(self.mean_size, self.dimensions_stddev)
+        new_size = new_size.astype(int)
+        return image.resize(new_size)
+
+    def __next__(self):
+        image = next(self.image_iterator)
+        image = self.random_resize(image)
+        return image

src/c++/perf_analyzer/genai-perf/tests/test_synthetic_image_generator.py

@@ -0,0 +1,95 @@
+import base64
+from io import BytesIO, StringIO
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+from genai_perf.exceptions import GenAIPerfException
+from genai_perf.llm_inputs.synthetic_image_generator import (
+    ImageFormat,
+    RandomFormatBase64Encoder,
+    SyntheticImageGenerator,
+    images_from_file_generator,
+    white_images_generator,
+)
+from PIL import Image
+
+
+@pytest.mark.parametrize(
+    "image_size",
+    [
+        (100, 100),
+        (200, 200),
+    ],
+)
+def test_different_image_size(image_size):
+    sut = SyntheticImageGenerator(
+        mean_size=image_size,
+        dimensions_stddev=[0, 0],
+        image_iterator=white_images_generator(),
+    )
+
+    image = next(sut)
+
+    assert isinstance(image, Image.Image), "generator produces unexpected type of data"
+    assert image.size == image_size, "image not resized to the target size"
+
+
+@patch("pathlib.Path.exists", return_value=False)
+def test_images_from_file_raises_when_file_not_found(mock_exists):
+    DUMMY_PATH = Path("dummy-image.png")
+    sut = images_from_file_generator(DUMMY_PATH)
+
+    with pytest.raises(GenAIPerfException):
+        next(sut)
+
+
+DUMMY_IMAGE = Image.new("RGB", (100, 100), color="blue")
+
+
+@patch("pathlib.Path.exists", return_value=True)
+@patch(
+    "PIL.Image.open",
+    return_value=DUMMY_IMAGE,
+)
+def test_images_from_file_generates_multiple_times(mock_file, mock_exists):
+    DUMMY_PATH = Path("dummy-image.png")
+    sut = images_from_file_generator(DUMMY_PATH)
+
+    image = next(sut)
+    mock_exists.assert_called_once()
+    mock_file.assert_called_once_with(DUMMY_PATH)
+    assert image == DUMMY_IMAGE, "unexpected image produced"
+
+    image = next(sut)
+    assert image == DUMMY_IMAGE, "unexpected image produced"
+
+
+def test_white_images_generator():
+    sut = white_images_generator()
+
+    image = next(sut)
+    assert isinstance(image, Image.Image), "generator produces unexpected type of data"
+    white_pixel = np.array([[[255, 255, 255]]])
+    assert (np.array(image) == white_pixel).all(), "not all pixels are white"
+
+
+@pytest.mark.parametrize("image_format", [ImageFormat.PNG, ImageFormat.JPEG])
+def test_base64_encoding_with_different_formats(image_format):
+    image = Image.new("RGB", (100, 100))
+    sut = RandomFormatBase64Encoder(image_formats=[image_format])
+
+    base64String = sut(image)
+
+    base64prefix = f"data:image/{image_format.name.lower()};base64,"
+    assert base64String.startswith(base64prefix), "unexpected prefix"
+    data = base64String[len(base64prefix) :]
+
+    # test if generator encodes to base64
+    img_data = base64.b64decode(data)
+    img_bytes = BytesIO(img_data)
+    # test if an image is encoded
+    image = Image.open(img_bytes)
+
+    assert image.format == image_format.name