76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
|
import base64
|
||
|
|
import mimetypes
|
||
|
|
from tempfile import NamedTemporaryFile
|
||
|
|
from typing import Dict, Tuple
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
import pytest_asyncio
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from vllm.multimodal.utils import ImageFetchAiohttp
|
||
|
|
|
||
|
|
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||
|
|
TEST_IMAGE_URLS = [
|
||
|
|
"https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",
|
||
|
|
"https://upload.wikimedia.org/wikipedia/commons/f/fa/Grayscale_8bits_palette_sample_image.png",
|
||
|
|
"https://upload.wikimedia.org/wikipedia/commons/thumb/9/91/Venn_diagram_rgb.svg/1280px-Venn_diagram_rgb.svg.png",
|
||
|
|
"https://upload.wikimedia.org/wikipedia/commons/0/0b/RGBA_comp.png",
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest_asyncio.fixture(scope="session")
|
||
|
|
async def url_images() -> Dict[str, Image.Image]:
|
||
|
|
return {
|
||
|
|
image_url: await ImageFetchAiohttp.fetch_image(image_url)
|
||
|
|
for image_url in TEST_IMAGE_URLS
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def get_supported_suffixes() -> Tuple[str, ...]:
|
||
|
|
# We should at least test the file types mentioned in GPT-4 with Vision
|
||
|
|
OPENAI_SUPPORTED_SUFFIXES = ('.png', '.jpeg', '.jpg', '.webp', '.gif')
|
||
|
|
|
||
|
|
# Additional file types that are supported by us
|
||
|
|
EXTRA_SUPPORTED_SUFFIXES = ('.bmp', '.tiff')
|
||
|
|
|
||
|
|
return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES
|
||
|
|
|
||
|
|
|
||
|
|
def _image_equals(a: Image.Image, b: Image.Image) -> bool:
|
||
|
|
return (np.asarray(a) == np.asarray(b.convert(a.mode))).all()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
@pytest.mark.parametrize("image_url", TEST_IMAGE_URLS)
|
||
|
|
@pytest.mark.parametrize("suffix", get_supported_suffixes())
|
||
|
|
async def test_fetch_image_base64(url_images: Dict[str, Image.Image],
|
||
|
|
image_url: str, suffix: str):
|
||
|
|
url_image = url_images[image_url]
|
||
|
|
|
||
|
|
try:
|
||
|
|
mime_type = Image.MIME[Image.registered_extensions()[suffix]]
|
||
|
|
except KeyError:
|
||
|
|
try:
|
||
|
|
mime_type = mimetypes.types_map[suffix]
|
||
|
|
except KeyError:
|
||
|
|
pytest.skip('No MIME type')
|
||
|
|
|
||
|
|
with NamedTemporaryFile(suffix=suffix) as f:
|
||
|
|
try:
|
||
|
|
url_image.save(f.name)
|
||
|
|
except Exception as e:
|
||
|
|
if e.args[0] == 'cannot write mode RGBA as JPEG':
|
||
|
|
pytest.skip('Conversion not supported')
|
||
|
|
|
||
|
|
raise
|
||
|
|
|
||
|
|
base64_image = base64.b64encode(f.read()).decode("utf-8")
|
||
|
|
data_url = f"data:{mime_type};base64,{base64_image}"
|
||
|
|
|
||
|
|
data_image = await ImageFetchAiohttp.fetch_image(data_url)
|
||
|
|
if _image_equals(url_image, Image.open(f)):
|
||
|
|
assert _image_equals(url_image, data_image)
|
||
|
|
else:
|
||
|
|
pass # Lossy format; only check that image can be opened
|