wehub-resource-sync
126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import struct
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image, ImageChops
|
|
|
|
from vllm.multimodal.image import (
|
|
_has_transparency,
|
|
convert_image_mode,
|
|
normalize_image,
|
|
)
|
|
|
|
pytestmark = pytest.mark.cpu_test
|
|
|
|
ASSETS_DIR = Path(__file__).parent / "assets"
|
|
assert ASSETS_DIR.exists()
|
|
|
|
|
|
def test_rgb_to_rgb():
|
|
# Start with an RGB image.
|
|
original_image = Image.open(ASSETS_DIR / "image1.png").convert("RGB")
|
|
converted_image = convert_image_mode(original_image, "RGB")
|
|
|
|
# RGB to RGB should be a no-op.
|
|
diff = ImageChops.difference(original_image, converted_image)
|
|
assert diff.getbbox() is None
|
|
|
|
|
|
def test_rgba_to_rgb():
|
|
original_image = Image.open(ASSETS_DIR / "rgba.png")
|
|
original_image_numpy = np.array(original_image)
|
|
|
|
converted_image = convert_image_mode(original_image, "RGB")
|
|
converted_image_numpy = np.array(converted_image)
|
|
|
|
for i in range(original_image_numpy.shape[0]):
|
|
for j in range(original_image_numpy.shape[1]):
|
|
# Verify that all transparent pixels are converted to white.
|
|
if original_image_numpy[i][j][3] == 0:
|
|
assert converted_image_numpy[i][j][0] == 255
|
|
assert converted_image_numpy[i][j][1] == 255
|
|
assert converted_image_numpy[i][j][2] == 255
|
|
|
|
|
|
def test_palette_with_trns_to_rgb():
|
|
"""P-mode PNG with tRNS: transparent index should become white."""
|
|
# Built synthetically because the existing assets are RGBA PNGs;
|
|
# this vulnerability only affects P/L/RGB images carrying a tRNS
|
|
# chunk, which uses a different transparency mechanism.
|
|
img = Image.new("P", (4, 4))
|
|
palette = [0] * 768
|
|
palette[0:3] = [255, 0, 0]
|
|
palette[3:6] = [0, 0, 255]
|
|
img.putpalette(palette)
|
|
img.putpixel((0, 0), 0)
|
|
img.putpixel((1, 0), 1)
|
|
img.info["transparency"] = 1
|
|
|
|
assert _has_transparency(img)
|
|
converted = convert_image_mode(img, "RGB")
|
|
assert converted.mode == "RGB"
|
|
r, g, b = converted.getpixel((1, 0))
|
|
assert (r, g, b) == (255, 255, 255)
|
|
r, g, b = converted.getpixel((0, 0))
|
|
assert (r, g, b) == (255, 0, 0)
|
|
|
|
|
|
def test_l_mode_no_trns_to_rgb():
|
|
"""L-mode without transparency should convert directly."""
|
|
img = Image.new("L", (4, 4), 128)
|
|
assert not _has_transparency(img)
|
|
converted = convert_image_mode(img, "RGB")
|
|
assert converted.mode == "RGB"
|
|
assert converted.getpixel((0, 0)) == (128, 128, 128)
|
|
|
|
|
|
def test_exif_transpose_normalizes_orientation():
|
|
"""Image with EXIF orientation 3 (180-degree rotation) should be
|
|
normalized so pixel data matches visual display."""
|
|
# Built synthetically because the existing assets are PNGs which
|
|
# don't carry EXIF orientation metadata; we need a JPEG with an
|
|
# injected EXIF orientation=3 tag.
|
|
img = Image.new("RGB", (2, 1))
|
|
img.putpixel((0, 0), (255, 0, 0))
|
|
img.putpixel((1, 0), (0, 0, 255))
|
|
|
|
buf = BytesIO()
|
|
img.save(buf, format="JPEG")
|
|
jpeg_bytes = buf.getvalue()
|
|
|
|
exif_orientation_3 = (
|
|
b"\xff\xe1"
|
|
+ struct.pack(">H", 26)
|
|
+ b"Exif\x00\x00"
|
|
+ b"MM"
|
|
+ b"\x00\x2a"
|
|
+ b"\x00\x00\x00\x08"
|
|
+ b"\x00\x01"
|
|
+ b"\x01\x12"
|
|
+ b"\x00\x03"
|
|
+ b"\x00\x00\x00\x01"
|
|
+ b"\x00\x03"
|
|
+ b"\x00\x00"
|
|
+ b"\x00\x00\x00\x00"
|
|
)
|
|
|
|
soi = jpeg_bytes[:2]
|
|
rest = jpeg_bytes[2:]
|
|
patched = soi + exif_orientation_3 + rest
|
|
|
|
rotated = Image.open(BytesIO(patched))
|
|
normalized = normalize_image(rotated)
|
|
assert normalized.size == (2, 1)
|
|
|
|
|
|
def test_normalize_image_no_exif():
|
|
"""Images without EXIF should pass through unchanged."""
|
|
img = Image.new("RGB", (4, 4), (100, 100, 100))
|
|
result = normalize_image(img)
|
|
assert result.size == img.size
|
|
assert result.mode == img.mode
|