SyntheticImageGenerator always encodes images to base64

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

@@ -14,30 +14,21 @@ class ImageFormat(Enum):
     PNG = auto()
 
 
-class Base64Encoder:
-    def __init__(self, image_format: ImageFormat = ImageFormat.PNG):
-        self.image_format = image_format
-
-    def __call__(self, image):
-        buffered = BytesIO()
-        image.save(buffered, format=self.image_format.name)
-        data = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        return f"data:image/{self.image_format.name.lower()};base64,{data}"
-
-
 class SyntheticImageGenerator:
     def __init__(
         self,
         image_width_mean: int,
         image_height_mean: int,
         image_width_stddev: int,
         image_height_stddev: int,
+        image_format: ImageFormat = ImageFormat.PNG,
         rng: Optional[np.random.Generator] = None,
     ):
         self._image_width_mean = image_width_mean
         self._image_height_mean = image_height_mean
         self._image_width_stddev = image_width_stddev
         self._image_height_stddev = image_height_stddev
+        self.image_format = image_format
         self.rng = rng or np.random.default_rng()
 
     def __iter__(self):
@@ -50,7 +41,7 @@ def _sample_random_positive_integer(self, mean: int, stddev: int) -> int:
                 break
         return n
 
-    def random_resize(self, image):
+    def _random_resize(self, image):
         width = self._sample_random_positive_integer(
             self._image_width_mean, self._image_width_stddev
         )
@@ -59,10 +50,17 @@ def random_resize(self, image):
         )
         return image.resize((width, height))
 
-    def get_next_image(self):
+    def _get_next_image(self):
         return Image.new("RGB", (100, 100), color="white")
 
-    def __next__(self):
-        image = self.get_next_image()
-        image = self.random_resize(image)
-        return image
+    def _encode(self, image):
+        buffered = BytesIO()
+        image.save(buffered, format=self.image_format.name)
+        data = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        return f"data:image/{self.image_format.name.lower()};base64,{data}"
+
+    def __next__(self) -> str:
+        image = self._get_next_image()
+        image = self._random_resize(image)
+        base64_string = self._encode(image)
+        return base64_string

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

@@ -7,13 +7,18 @@
 import pytest
 from genai_perf.exceptions import GenAIPerfException
 from genai_perf.llm_inputs.synthetic_image_generator import (
-    Base64Encoder,
     ImageFormat,
     SyntheticImageGenerator,
 )
 from PIL import Image
 
 
+def decode_image(base64_string):
+    _, data = base64_string.split(",")
+    decoded_data = base64.b64decode(data)
+    return Image.open(BytesIO(decoded_data))
+
+
 @pytest.mark.parametrize(
     "expected_image_size",
     [
@@ -30,9 +35,9 @@ def test_different_image_size(expected_image_size):
         image_height_stddev=0,
     )
 
-    image = next(sut)
+    base64_string = next(sut)
+    image = decode_image(base64_string)
 
-    assert isinstance(image, Image.Image), "generator produces unexpected type of data"
     assert image.size == expected_image_size, "image not resized to the target size"
 
 
@@ -63,10 +68,11 @@ def test_generator_deterministic():
 
 @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 = Base64Encoder(image_format=image_format)
+    IMAGE_SIZE = 100, 100
+    STDDEV = 100, 100
+    sut = SyntheticImageGenerator(*IMAGE_SIZE, *STDDEV, image_format=image_format)
 
-    base64String = sut(image)
+    base64String = next(sut)
 
     base64prefix = f"data:image/{image_format.name.lower()};base64,"
     assert base64String.startswith(base64prefix), "unexpected prefix"