image.py 5.44 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
from io import BytesIO
from pathlib import Path
6

7
import pybase64
8
9
10
import torch
from PIL import Image

11
12
from vllm.logger import init_logger

13
from .base import MediaIO, MediaWithBytes
14

15
16
logger = init_logger(__file__)

17

18
19
20
def rescale_image_size(
    image: Image.Image, size_factor: float, transpose: int = -1
) -> Image.Image:
21
22
23
24
25
26
27
    """Rescale the dimensions of an image by a constant factor."""
    new_width = int(image.width * size_factor)
    new_height = int(image.height * size_factor)
    image = image.resize((new_width, new_height))
    if transpose >= 0:
        image = image.transpose(Image.Transpose(transpose))
    return image
28
29


30
def rgba_to_rgb(
31
    image: Image.Image,
32
    background_color: tuple[int, int, int] | list[int] = (255, 255, 255),
33
) -> Image.Image:
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
    """Convert an RGBA image to RGB with filled background color."""
    assert image.mode == "RGBA"
    converted = Image.new("RGB", image.size, background_color)
    converted.paste(image, mask=image.split()[3])  # 3 is the alpha channel
    return converted


def convert_image_mode(image: Image.Image, to_mode: str):
    if image.mode == to_mode:
        return image
    elif image.mode == "RGBA" and to_mode == "RGB":
        return rgba_to_rgb(image)
    else:
        return image.convert(to_mode)


50
class ImageMediaIO(MediaIO[Image.Image]):
51
    def __init__(self, image_mode: str = "RGB", **kwargs) -> None:
52
53
54
        super().__init__()

        self.image_mode = image_mode
55
56
57
58
59
60
        # `kwargs` contains custom arguments from
        # --media-io-kwargs for this modality.
        # They can be passed to the underlying
        # media loaders (e.g. custom implementations)
        # for flexible control.
        self.kwargs = kwargs
61

62
63
        # Extract RGBA background color from kwargs if provided
        # Default to white background for backward compatibility
64
        rgba_bg = kwargs.get("rgba_background_color", (255, 255, 255))
65
66
67
68
69
        # Convert list to tuple for consistency
        if isinstance(rgba_bg, list):
            rgba_bg = tuple(rgba_bg)

        # Validate rgba_background_color format
70
71
72
73
74
        if not (
            isinstance(rgba_bg, tuple)
            and len(rgba_bg) == 3
            and all(isinstance(c, int) and 0 <= c <= 255 for c in rgba_bg)
        ):
75
76
            raise ValueError(
                "rgba_background_color must be a list or tuple of 3 integers "
77
78
                "in the range [0, 255]."
            )
79
80
        self.rgba_background_color = rgba_bg

81
82
83
    def _convert_image_mode(
        self, image: Image.Image | MediaWithBytes[Image.Image]
    ) -> Image.Image:
84
        """Convert image mode with custom background color."""
85
86
        if isinstance(image, MediaWithBytes):
            image = image.media
87
88
89
90
91
92
93
        if image.mode == self.image_mode:
            return image
        elif image.mode == "RGBA" and self.image_mode == "RGB":
            return rgba_to_rgb(image, self.rgba_background_color)
        else:
            return convert_image_mode(image, self.image_mode)

94
    def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]:
95
        image = Image.open(BytesIO(data))
96
        return MediaWithBytes(self._convert_image_mode(image), data)
97

98
    def load_base64(self, media_type: str, data: str) -> MediaWithBytes[Image.Image]:
99
        return self.load_bytes(pybase64.b64decode(data, validate=True))
100

101
102
103
104
105
    def load_file(self, filepath: Path) -> MediaWithBytes[Image.Image]:
        with open(filepath, "rb") as f:
            data = f.read()
        image = Image.open(BytesIO(data))
        return MediaWithBytes(self._convert_image_mode(image), data)
106
107
108
109
110

    def encode_base64(
        self,
        media: Image.Image,
        *,
111
        image_format: str | None = None,
112
    ) -> str:
113
114
115
116
117
118
119
120
121
        if image_format is None:
            logger.warning_once(
                "The default format of `ImageMediaIO.encode_base64` will be changed "
                'from "JPEG" to "PNG" in v0.15 to avoid lossy compression. '
                "To continue using the old default, "
                'pass `format="JPEG"` explicitly to silence this warning.'
            )
            image_format = "JPEG"

122
123
124
        image = media

        with BytesIO() as buffer:
125
            image = self._convert_image_mode(image)
126
127
128
            image.save(buffer, image_format)
            data = buffer.getvalue()

129
        return pybase64.b64encode(data).decode("utf-8")
130
131
132
133
134
135
136
137


class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
    def __init__(self) -> None:
        super().__init__()

    def load_bytes(self, data: bytes) -> torch.Tensor:
        buffer = BytesIO(data)
138
139
140
141
142
        # Enable sparse tensor integrity checks to prevent out-of-bounds
        # writes from maliciously crafted tensors
        with torch.sparse.check_sparse_tensor_invariants():
            tensor = torch.load(buffer, weights_only=True)
            return tensor.to_dense()
143
144

    def load_base64(self, media_type: str, data: str) -> torch.Tensor:
145
        return self.load_bytes(pybase64.b64decode(data, validate=True))
146
147

    def load_file(self, filepath: Path) -> torch.Tensor:
148
149
150
151
152
        # Enable sparse tensor integrity checks to prevent out-of-bounds
        # writes from maliciously crafted tensors
        with torch.sparse.check_sparse_tensor_invariants():
            tensor = torch.load(filepath, weights_only=True)
            return tensor.to_dense()
153
154

    def encode_base64(self, media: torch.Tensor) -> str:
155
        return pybase64.b64encode(media.numpy()).decode("utf-8")