| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from pathlib import Path |
| 4 | |
| 5 | import numpy as np |
| 6 | import pytest |
| 7 | from PIL import Image |
| 8 | |
| 9 | from vllm.multimodal.media import ImageMediaIO |
| 10 | |
| 11 | pytestmark = pytest.mark.cpu_test |
| 12 | |
| 13 | ASSETS_DIR = Path(__file__).parent.parent / "assets" |
| 14 | assert ASSETS_DIR.exists() |
| 15 | |
| 16 | |
| 17 | def test_image_media_io_rgba_custom_background(tmp_path): |
| 18 | """Test RGBA to RGB conversion with custom background colors.""" |
| 19 | # Create a simple RGBA image with transparent and opaque pixels |
| 20 | rgba_image = Image.new("RGBA", (10, 10), (255, 0, 0, 255)) # Red with full opacity |
| 21 | |
| 22 | # Make top-left quadrant transparent |
| 23 | for i in range(5): |
| 24 | for j in range(5): |
| 25 | rgba_image.putpixel((i, j), (0, 0, 0, 0)) # Fully transparent |
| 26 | |
| 27 | # Save the test image to tmp_path |
| 28 | test_image_path = tmp_path / "test_rgba.png" |
| 29 | rgba_image.save(test_image_path) |
| 30 | |
| 31 | # Test 1: Default white background (backward compatibility) |
| 32 | image_io_default = ImageMediaIO() |
| 33 | converted_default = image_io_default.load_file(test_image_path) |
| 34 | default_numpy = np.array(converted_default) |
| 35 | |
| 36 | # Check transparent pixels are white |
| 37 | assert default_numpy[0][0][0] == 255 # R |
| 38 | assert default_numpy[0][0][1] == 255 # G |
| 39 | assert default_numpy[0][0][2] == 255 # B |
| 40 | # Check opaque pixels remain red |
| 41 | assert default_numpy[5][5][0] == 255 # R |
| 42 | assert default_numpy[5][5][1] == 0 # G |
| 43 | assert default_numpy[5][5][2] == 0 # B |
| 44 | |
| 45 | # Test 2: Custom black background via kwargs |
| 46 | image_io_black = ImageMediaIO(rgba_background_color=(0, 0, 0)) |
| 47 | converted_black = image_io_black.load_file(test_image_path) |
| 48 | black_numpy = np.array(converted_black) |
| 49 | |
| 50 | # Check transparent pixels are black |
| 51 | assert black_numpy[0][0][0] == 0 # R |
| 52 | assert black_numpy[0][0][1] == 0 # G |
| 53 | assert black_numpy[0][0][2] == 0 # B |
| 54 | # Check opaque pixels remain red |
| 55 | assert black_numpy[5][5][0] == 255 # R |
| 56 | assert black_numpy[5][5][1] == 0 # G |
| 57 | assert black_numpy[5][5][2] == 0 # B |
| 58 | |
| 59 | # Test 3: Custom blue background via kwargs (as list) |
| 60 | image_io_blue = ImageMediaIO(rgba_background_color=[0, 0, 255]) |
| 61 | converted_blue = image_io_blue.load_file(test_image_path) |
| 62 | blue_numpy = np.array(converted_blue) |
| 63 | |
| 64 | # Check transparent pixels are blue |
| 65 | assert blue_numpy[0][0][0] == 0 # R |
| 66 | assert blue_numpy[0][0][1] == 0 # G |
| 67 | assert blue_numpy[0][0][2] == 255 # B |
| 68 | |
| 69 | # Test 4: Test with load_bytes method |
| 70 | with open(test_image_path, "rb") as f: |
| 71 | image_data = f.read() |
| 72 | |
| 73 | image_io_green = ImageMediaIO(rgba_background_color=(0, 255, 0)) |
| 74 | converted_green = image_io_green.load_bytes(image_data) |
| 75 | green_numpy = np.array(converted_green) |
| 76 | |
| 77 | # Check transparent pixels are green |
| 78 | assert green_numpy[0][0][0] == 0 # R |
| 79 | assert green_numpy[0][0][1] == 255 # G |
| 80 | assert green_numpy[0][0][2] == 0 # B |
| 81 | |
| 82 | |
| 83 | def test_image_media_io_no_mode_conversion(tmp_path): |
| 84 | """image_mode=None skips conversion and preserves the original mode.""" |
| 85 | # RGBA image: opaque black pixel on a fully transparent background |
| 86 | rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0)) |
| 87 | rgba_image.putpixel((5, 5), (0, 0, 0, 255)) |
| 88 | test_image_path = tmp_path / "test_rgba.png" |
| 89 | rgba_image.save(test_image_path) |
| 90 | |
| 91 | # Default behavior: RGBA is composited onto a white background |
| 92 | image_io_default = ImageMediaIO() |
| 93 | converted_default = image_io_default.load_file(test_image_path) |
| 94 | assert converted_default.media.mode == "RGB" |
| 95 | assert converted_default.media.getpixel((0, 0)) == (255, 255, 255) |
| 96 | assert converted_default.media.getpixel((5, 5)) == (0, 0, 0) |
| 97 | |
| 98 | # image_mode=None: original mode and alpha channel are preserved |
| 99 | image_io_keep = ImageMediaIO(image_mode=None) |
| 100 | converted_keep = image_io_keep.load_file(test_image_path) |
| 101 | assert converted_keep.media.mode == "RGBA" |
| 102 | assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0) |
| 103 | assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255) |
| 104 | |
| 105 | |
| 106 | def test_image_media_io_rgba_background_color_validation(): |
| 107 | """Test that invalid rgba_background_color values are properly rejected.""" |
| 108 | # Test invalid types |
| 109 | with pytest.raises( |
| 110 | ValueError, match="rgba_background_color must be a list or tuple" |
| 111 | ): |
| 112 | ImageMediaIO(rgba_background_color="255,255,255") |
| 113 | |
| 114 | with pytest.raises( |
| 115 | ValueError, match="rgba_background_color must be a list or tuple" |
| 116 | ): |
| 117 | ImageMediaIO(rgba_background_color=255) |
| 118 | |
| 119 | # Test wrong number of elements |
| 120 | with pytest.raises( |
| 121 | ValueError, match="rgba_background_color must be a list or tuple" |
| 122 | ): |
| 123 | ImageMediaIO(rgba_background_color=(255, 255)) |
| 124 | |
| 125 | with pytest.raises( |
| 126 | ValueError, match="rgba_background_color must be a list or tuple" |
| 127 | ): |
| 128 | ImageMediaIO(rgba_background_color=(255, 255, 255, 255)) |
| 129 | |
| 130 | # Test non-integer values |
| 131 | with pytest.raises( |
| 132 | ValueError, match="rgba_background_color must be a list or tuple" |
| 133 | ): |
| 134 | ImageMediaIO(rgba_background_color=(255.0, 255.0, 255.0)) |
| 135 | |
| 136 | with pytest.raises( |
| 137 | ValueError, match="rgba_background_color must be a list or tuple" |
| 138 | ): |
| 139 | ImageMediaIO(rgba_background_color=(255, "255", 255)) |
| 140 | |
| 141 | # Test out of range values |
| 142 | with pytest.raises( |
| 143 | ValueError, match="rgba_background_color must be a list or tuple" |
| 144 | ): |
| 145 | ImageMediaIO(rgba_background_color=(256, 255, 255)) |
| 146 | |
| 147 | with pytest.raises( |
| 148 | ValueError, match="rgba_background_color must be a list or tuple" |
| 149 | ): |
| 150 | ImageMediaIO(rgba_background_color=(255, -1, 255)) |
| 151 | |
| 152 | # Test that valid values work |
| 153 | ImageMediaIO(rgba_background_color=(0, 0, 0)) # Should not raise |
| 154 | ImageMediaIO(rgba_background_color=[255, 255, 255]) # Should not raise |
| 155 | ImageMediaIO(rgba_background_color=(128, 128, 128)) # Should not raise |
| 156 | |
| 157 | |
| 158 | def test_image_media_io_load_bytes(tmp_path): |
| 159 | """Test load_bytes with valid and invalid image data.""" |
| 160 | # Save a valid RGB image to use as source bytes |
| 161 | valid_image = Image.new("RGB", (8, 8), (100, 150, 200)) |
| 162 | valid_path = tmp_path / "valid.png" |
| 163 | valid_image.save(valid_path) |
| 164 | |
| 165 | valid_data = valid_path.read_bytes() |
| 166 | |
| 167 | # Test 1: Valid image bytes load successfully and are fully decoded |
| 168 | image_io = ImageMediaIO() |
| 169 | result = image_io.load_bytes(valid_data) |
| 170 | |
| 171 | # Check the returned media is a properly loaded image |
| 172 | assert isinstance(result.media, Image.Image) |
| 173 | assert result.media.size == (8, 8) |
| 174 | assert result.media.getpixel((0, 0)) == (100, 150, 200) |
| 175 | |
| 176 | # Test 2: Garbage bytes raise ValueError |
| 177 | with pytest.raises(ValueError, match="Failed to load image"): |
| 178 | image_io.load_bytes(b"not an image") |
| 179 | |
| 180 | # Test 3: Truncated PNG header raises ValueError |
| 181 | with pytest.raises(ValueError, match="Failed to load image"): |
| 182 | image_io.load_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 10) |
| 183 | |
| 184 | # Test 4: Real PNG truncated mid-stream raises ValueError |
| 185 | with pytest.raises(ValueError, match="Failed to load image"): |
| 186 | image_io.load_bytes(valid_data[: len(valid_data) // 2]) |
| 187 | |
| 188 | # Test 5: Empty bytes raise ValueError |
| 189 | with pytest.raises(ValueError, match="Failed to load image"): |
| 190 | image_io.load_bytes(b"") |
| 191 | |
| 192 | |
| 193 | def test_image_media_io_load_file(tmp_path): |
| 194 | """Test load_file with valid and invalid image files.""" |
| 195 | # Save a valid RGB image to disk |
| 196 | valid_image = Image.new("RGB", (4, 4), (10, 20, 30)) |
| 197 | valid_path = tmp_path / "valid.png" |
| 198 | valid_image.save(valid_path) |
| 199 | |
| 200 | # Test 1: Valid image file loads successfully and is fully decoded |
| 201 | image_io = ImageMediaIO() |
| 202 | result = image_io.load_file(valid_path) |
| 203 | |
| 204 | # Check the returned media is a properly loaded image |
| 205 | assert isinstance(result.media, Image.Image) |
| 206 | assert result.media.size == (4, 4) |
| 207 | assert result.media.getpixel((0, 0)) == (10, 20, 30) |
| 208 | |
| 209 | # Test 2: File with garbage content raises ValueError |
| 210 | bad_file = tmp_path / "bad.png" |
| 211 | bad_file.write_bytes(b"this is not an image") |
| 212 | |
| 213 | with pytest.raises(ValueError, match="Failed to load image"): |
| 214 | image_io.load_file(bad_file) |
| 215 | |
| 216 | # Test 3: File with truncated PNG header raises ValueError |
| 217 | truncated_file = tmp_path / "truncated.png" |
| 218 | truncated_file.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 10) |
| 219 | |
| 220 | with pytest.raises(ValueError, match="Failed to load image"): |
| 221 | image_io.load_file(truncated_file) |
| 222 | |
| 223 | # Test 4: Real PNG file truncated mid-stream raises ValueError |
| 224 | valid_data = valid_path.read_bytes() |
| 225 | truncated_real_file = tmp_path / "truncated_real.png" |
| 226 | truncated_real_file.write_bytes(valid_data[: len(valid_data) // 2]) |
| 227 | |
| 228 | with pytest.raises(ValueError, match="Failed to load image"): |
| 229 | image_io.load_file(truncated_real_file) |
| 230 | |
| 231 | |
| 232 | def test_image_pixel_limit_respected(): |
| 233 | """A small image within the pixel limit loads successfully.""" |
| 234 | import vllm.envs as envs |
| 235 | |
| 236 | image = Image.new("RGB", (100, 100), (255, 0, 0)) |
| 237 | from io import BytesIO |
| 238 | |
| 239 | buf = BytesIO() |
| 240 | image.save(buf, format="PNG") |
| 241 | data = buf.getvalue() |
| 242 | |
| 243 | assert envs.VLLM_MAX_IMAGE_PIXELS >= 100 * 100 |
| 244 | |
| 245 | image_io = ImageMediaIO() |
| 246 | result = image_io.load_bytes(data) |
| 247 | assert result.media.size == (100, 100) |
| 248 | |
| 249 | |
| 250 | def test_image_pixel_limit_rejected(monkeypatch): |
| 251 | """An image exceeding the pixel limit is rejected before raster decode.""" |
| 252 | import vllm.envs as envs |
| 253 | |
| 254 | monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) |
| 255 | |
| 256 | image = Image.new("RGB", (20, 20), (0, 255, 0)) |
| 257 | from io import BytesIO |
| 258 | |
| 259 | buf = BytesIO() |
| 260 | image.save(buf, format="PNG") |
| 261 | data = buf.getvalue() |
| 262 | |
| 263 | image_io = ImageMediaIO() |
| 264 | with pytest.raises(ValueError, match="exceed"): |
| 265 | image_io.load_bytes(data) |
| 266 | |
| 267 | |
| 268 | def test_image_pixel_limit_disabled(monkeypatch): |
| 269 | """Setting VLLM_MAX_IMAGE_PIXELS=0 disables the pixel limit.""" |
| 270 | import vllm.envs as envs |
| 271 | |
| 272 | monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 0) |
| 273 | |
| 274 | image = Image.new("RGB", (1000, 1000), (0, 0, 255)) |
| 275 | from io import BytesIO |
| 276 | |
| 277 | buf = BytesIO() |
| 278 | image.save(buf, format="PNG") |
| 279 | data = buf.getvalue() |
| 280 | |
| 281 | image_io = ImageMediaIO() |
| 282 | result = image_io.load_bytes(data) |
| 283 | assert result.media.size == (1000, 1000) |