| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import struct |
| 4 | from io import BytesIO |
| 5 | from pathlib import Path |
| 6 | |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | from PIL import Image, ImageChops |
| 10 | |
| 11 | from vllm.multimodal.image import ( |
| 12 | _has_transparency, |
| 13 | convert_image_mode, |
| 14 | normalize_image, |
| 15 | ) |
| 16 | |
| 17 | pytestmark = pytest.mark.cpu_test |
| 18 | |
| 19 | ASSETS_DIR = Path(__file__).parent / "assets" |
| 20 | assert ASSETS_DIR.exists() |
| 21 | |
| 22 | |
| 23 | def test_rgb_to_rgb(): |
| 24 | # Start with an RGB image. |
| 25 | original_image = Image.open(ASSETS_DIR / "image1.png").convert("RGB") |
| 26 | converted_image = convert_image_mode(original_image, "RGB") |
| 27 | |
| 28 | # RGB to RGB should be a no-op. |
| 29 | diff = ImageChops.difference(original_image, converted_image) |
| 30 | assert diff.getbbox() is None |
| 31 | |
| 32 | |
| 33 | def test_rgba_to_rgb(): |
| 34 | original_image = Image.open(ASSETS_DIR / "rgba.png") |
| 35 | original_image_numpy = np.array(original_image) |
| 36 | |
| 37 | converted_image = convert_image_mode(original_image, "RGB") |
| 38 | converted_image_numpy = np.array(converted_image) |
| 39 | |
| 40 | for i in range(original_image_numpy.shape[0]): |
| 41 | for j in range(original_image_numpy.shape[1]): |
| 42 | # Verify that all transparent pixels are converted to white. |
| 43 | if original_image_numpy[i][j][3] == 0: |
| 44 | assert converted_image_numpy[i][j][0] == 255 |
| 45 | assert converted_image_numpy[i][j][1] == 255 |
| 46 | assert converted_image_numpy[i][j][2] == 255 |
| 47 | |
| 48 | |
| 49 | def test_palette_with_trns_to_rgb(): |
| 50 | """P-mode PNG with tRNS: transparent index should become white.""" |
| 51 | # Built synthetically because the existing assets are RGBA PNGs; |
| 52 | # this vulnerability only affects P/L/RGB images carrying a tRNS |
| 53 | # chunk, which uses a different transparency mechanism. |
| 54 | img = Image.new("P", (4, 4)) |
| 55 | palette = [0] * 768 |
| 56 | palette[0:3] = [255, 0, 0] |
| 57 | palette[3:6] = [0, 0, 255] |
| 58 | img.putpalette(palette) |
| 59 | img.putpixel((0, 0), 0) |
| 60 | img.putpixel((1, 0), 1) |
| 61 | img.info["transparency"] = 1 |
| 62 | |
| 63 | assert _has_transparency(img) |
| 64 | converted = convert_image_mode(img, "RGB") |
| 65 | assert converted.mode == "RGB" |
| 66 | r, g, b = converted.getpixel((1, 0)) |
| 67 | assert (r, g, b) == (255, 255, 255) |
| 68 | r, g, b = converted.getpixel((0, 0)) |
| 69 | assert (r, g, b) == (255, 0, 0) |
| 70 | |
| 71 | |
| 72 | def test_l_mode_no_trns_to_rgb(): |
| 73 | """L-mode without transparency should convert directly.""" |
| 74 | img = Image.new("L", (4, 4), 128) |
| 75 | assert not _has_transparency(img) |
| 76 | converted = convert_image_mode(img, "RGB") |
| 77 | assert converted.mode == "RGB" |
| 78 | assert converted.getpixel((0, 0)) == (128, 128, 128) |
| 79 | |
| 80 | |
| 81 | def test_exif_transpose_normalizes_orientation(): |
| 82 | """Image with EXIF orientation 3 (180-degree rotation) should be |
| 83 | normalized so pixel data matches visual display.""" |
| 84 | # Built synthetically because the existing assets are PNGs which |
| 85 | # don't carry EXIF orientation metadata; we need a JPEG with an |
| 86 | # injected EXIF orientation=3 tag. |
| 87 | img = Image.new("RGB", (2, 1)) |
| 88 | img.putpixel((0, 0), (255, 0, 0)) |
| 89 | img.putpixel((1, 0), (0, 0, 255)) |
| 90 | |
| 91 | buf = BytesIO() |
| 92 | img.save(buf, format="JPEG") |
| 93 | jpeg_bytes = buf.getvalue() |
| 94 | |
| 95 | exif_orientation_3 = ( |
| 96 | b"\xff\xe1" |
| 97 | + struct.pack(">H", 26) |
| 98 | + b"Exif\x00\x00" |
| 99 | + b"MM" |
| 100 | + b"\x00\x2a" |
| 101 | + b"\x00\x00\x00\x08" |
| 102 | + b"\x00\x01" |
| 103 | + b"\x01\x12" |
| 104 | + b"\x00\x03" |
| 105 | + b"\x00\x00\x00\x01" |
| 106 | + b"\x00\x03" |
| 107 | + b"\x00\x00" |
| 108 | + b"\x00\x00\x00\x00" |
| 109 | ) |
| 110 | |
| 111 | soi = jpeg_bytes[:2] |
| 112 | rest = jpeg_bytes[2:] |
| 113 | patched = soi + exif_orientation_3 + rest |
| 114 | |
| 115 | rotated = Image.open(BytesIO(patched)) |
| 116 | normalized = normalize_image(rotated) |
| 117 | assert normalized.size == (2, 1) |
| 118 | |
| 119 | |
| 120 | def test_normalize_image_no_exif(): |
| 121 | """Images without EXIF should pass through unchanged.""" |
| 122 | img = Image.new("RGB", (4, 4), (100, 100, 100)) |
| 123 | result = normalize_image(img) |
| 124 | assert result.size == img.size |
| 125 | assert result.mode == img.mode |