diff --git a/pyrit/converter/base_image_to_image_converter.py b/pyrit/converter/base_image_to_image_converter.py index 94ad9b891b..1a2b8d4ae1 100644 --- a/pyrit/converter/base_image_to_image_converter.py +++ b/pyrit/converter/base_image_to_image_converter.py @@ -105,6 +105,21 @@ def _prepare_image_for_jpeg(self, image: Image.Image) -> Image.Image: return background return image.convert("RGB") + @staticmethod + def _convert_to_rgb_or_rgba(image: Image.Image) -> Image.Image: + """ + Convert an image to RGB, or RGBA when it has transparency, unless it already is one of them. + + Args: + image (PIL.Image.Image): The image to convert. + + Returns: + PIL.Image.Image: The image in RGB or RGBA mode. + """ + if image.mode in ("RGB", "RGBA"): + return image + return image.convert("RGBA" if image.has_transparency_data else "RGB") + def _transform_image(self, image: Image.Image, original_format: str) -> tuple[BytesIO, str]: """ Resolve format, prepare the image, apply the transform, and save to a buffer. diff --git a/pyrit/converter/image_color_saturation_converter.py b/pyrit/converter/image_color_saturation_converter.py index 239edcc7bc..b54cf78a13 100644 --- a/pyrit/converter/image_color_saturation_converter.py +++ b/pyrit/converter/image_color_saturation_converter.py @@ -77,4 +77,7 @@ def _apply_transform(self, image: Image.Image) -> Image.Image: Returns: PIL.Image.Image: The adjusted image. """ + # ImageEnhance.Color rejects palette, bilevel and 16-bit images. + if image.mode not in ("L", "LA"): + image = self._convert_to_rgb_or_rgba(image) return ImageEnhance.Color(image).enhance(self._level) diff --git a/pyrit/converter/image_rotation_converter.py b/pyrit/converter/image_rotation_converter.py index 8c7efc0514..d855a3745c 100644 --- a/pyrit/converter/image_rotation_converter.py +++ b/pyrit/converter/image_rotation_converter.py @@ -86,4 +86,6 @@ def _apply_transform(self, image: Image.Image) -> Image.Image: Returns: PIL.Image.Image: The rotated image. """ + # The fill color is an RGB tuple, which grayscale and other single-band modes reject. + image = self._convert_to_rgb_or_rgba(image) return image.rotate(self._angle, expand=True, fillcolor=self._fill_color) diff --git a/tests/unit/converter/test_image_color_saturation_converter.py b/tests/unit/converter/test_image_color_saturation_converter.py index 5c9876fdc4..96247c2e2b 100644 --- a/tests/unit/converter/test_image_color_saturation_converter.py +++ b/tests/unit/converter/test_image_color_saturation_converter.py @@ -220,3 +220,47 @@ def test_invalid_level_non_finite(level: float) -> None: """Non-finite saturation levels should fail during converter construction.""" with pytest.raises(ValueError, match="Level must be finite"): ImageColorSaturationConverter(level=level) + + +@pytest.mark.parametrize("mode", ["P", "1", "I;16", "L", "LA"]) +def test_image_color_saturation_converter_handles_non_rgb_png_modes(mode: str) -> None: + converter = ImageColorSaturationConverter(level=0.5) + + adjusted_io, output_format = converter._transform_image(Image.new(mode, (40, 30)), "PNG") + + assert output_format == "PNG" + assert Image.open(adjusted_io).size == (40, 30) + + +@pytest.mark.parametrize( + "mode, transparency, expected_alpha", + [ + ("P", 0, [0, 255, 255]), + ("P", bytes([0, 128, 255]), [0, 128, 255]), + ("LA", None, [0, 128, 255]), + ("RGBA", None, [0, 128, 255]), + ], +) +@pytest.mark.parametrize("level", [0.0, 0.5, 2.0]) +def test_image_color_saturation_converter_preserves_png_alpha( + *, mode: str, transparency: int | bytes | None, expected_alpha: list[int], level: float +) -> None: + image = Image.new(mode, (3, 1)) + if mode == "P": + image.putpalette([200, 40, 20, 40, 200, 20, 40, 20, 200]) + image.putdata([0, 1, 2]) + image.info["transparency"] = transparency + else: + image.putdata([(60, alpha) if mode == "LA" else (200, 40, 20, alpha) for alpha in expected_alpha]) + image_bytes = BytesIO() + image.save(image_bytes, format="PNG") + converter = ImageColorSaturationConverter(level=level) + + with Image.open(image_bytes) as source: + adjusted_io, output_format = converter._transform_image(source, "PNG") + + assert output_format == "PNG" + with Image.open(adjusted_io) as adjusted: + assert adjusted.size == (3, 1) + assert adjusted.mode in ("LA", "RGBA") + assert adjusted.getchannel("A").tobytes() == bytes(expected_alpha) diff --git a/tests/unit/converter/test_image_rotation_converter.py b/tests/unit/converter/test_image_rotation_converter.py index 705e921f2f..0d8664e060 100644 --- a/tests/unit/converter/test_image_rotation_converter.py +++ b/tests/unit/converter/test_image_rotation_converter.py @@ -270,3 +270,48 @@ def test_invalid_angle_non_finite(angle: float) -> None: """Non-finite rotation angles should fail during converter construction.""" with pytest.raises(ValueError, match="Angle must be finite"): ImageRotationConverter(angle=angle) + + +@pytest.mark.parametrize("mode", ["L", "LA", "1", "I;16", "P"]) +def test_image_rotation_converter_handles_non_rgb_png_modes(mode: str) -> None: + converter = ImageRotationConverter(angle=45.0, fill_color=(255, 0, 0)) + + rotated_io, output_format = converter._transform_image(Image.new(mode, (40, 30)), "PNG") + rotated = Image.open(rotated_io) + + assert output_format == "PNG" + assert rotated.mode in ("RGB", "RGBA") + assert rotated.getpixel((0, 0))[:3] == (255, 0, 0) + + +@pytest.mark.parametrize( + "mode, transparency, expected_alpha", + [ + ("P", 0, [0, 255, 255]), + ("P", bytes([0, 128, 255]), [0, 128, 255]), + ("LA", None, [0, 128, 255]), + ("RGBA", None, [0, 128, 255]), + ], +) +def test_image_rotation_converter_preserves_png_alpha( + *, mode: str, transparency: int | bytes | None, expected_alpha: list[int] +) -> None: + image = Image.new(mode, (3, 1)) + if mode == "P": + image.putpalette([200, 40, 20, 40, 200, 20, 40, 20, 200]) + image.putdata([0, 1, 2]) + image.info["transparency"] = transparency + else: + image.putdata([(60, alpha) if mode == "LA" else (200, 40, 20, alpha) for alpha in expected_alpha]) + image_bytes = BytesIO() + image.save(image_bytes, format="PNG") + converter = ImageRotationConverter(angle=90.0) + + with Image.open(image_bytes) as source: + rotated_io, output_format = converter._transform_image(source, "PNG") + + assert output_format == "PNG" + with Image.open(rotated_io) as rotated: + assert rotated.size == (1, 3) + assert rotated.mode == "RGBA" + assert rotated.getchannel("A").tobytes() == bytes(expected_alpha[::-1])