| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Sequence |
| 5 | from typing import Any |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | from PIL import Image |
| 10 | from transformers import ImageProcessingMixin, ProcessorMixin |
| 11 | |
| 12 | IMAGENET_MEAN = np.array([0.484375, 0.455078125, 0.40625], dtype=np.float32) |
| 13 | IMAGENET_STD = np.array([0.228515625, 0.2236328125, 0.224609375], dtype=np.float32) |
| 14 | SIGLIP_MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32) |
| 15 | SIGLIP_STD = np.array([0.5, 0.5, 0.5], dtype=np.float32) |
| 16 | |
| 17 | |
| 18 | def to_rgb_image(image: Any) -> Image.Image: |
| 19 | if isinstance(image, Image.Image): |
| 20 | return image.convert("RGB") |
| 21 | |
| 22 | if isinstance(image, torch.Tensor): |
| 23 | image = image.detach().cpu().numpy() |
| 24 | if not isinstance(image, np.ndarray): |
| 25 | raise TypeError( |
| 26 | "OpenVLA image input must be a PIL image, numpy array, or torch tensor; " |
| 27 | f"got {type(image)}" |
| 28 | ) |
| 29 | |
| 30 | if image.ndim != 3: |
| 31 | raise ValueError( |
| 32 | f"OpenVLA image input must have 3 dimensions, got shape {image.shape}" |
| 33 | ) |
| 34 | |
| 35 | if image.shape[0] in (1, 3): |
| 36 | image = np.moveaxis(image, 0, -1) |
| 37 | |
| 38 | if image.shape[-1] == 1: |
| 39 | image = np.repeat(image, 3, axis=-1) |
| 40 | elif image.shape[-1] != 3: |
| 41 | raise ValueError( |
| 42 | f"OpenVLA image input must have 1 or 3 channels, got shape {image.shape}" |
| 43 | ) |
| 44 | |
| 45 | if image.dtype != np.uint8: |
| 46 | image = image.astype(np.float32) |
| 47 | if image.max(initial=0.0) <= 1.0: |
| 48 | image = image * 255.0 |
| 49 | image = np.clip(image, 0, 255).astype(np.uint8) |
| 50 | |
| 51 | return Image.fromarray(image).convert("RGB") |
| 52 | |
| 53 | |
| 54 | def preprocess_openvla_image(image: Any, image_size: int) -> torch.Tensor: |
| 55 | rgb_image = to_rgb_image(image) |
| 56 | rgb_image = rgb_image.resize( |
| 57 | (image_size, image_size), |
| 58 | Image.Resampling.BICUBIC, |
| 59 | ) |
| 60 | |
| 61 | raw = np.asarray(rgb_image, dtype=np.float32) / 255.0 |
| 62 | dinov2_pixels = ((raw - IMAGENET_MEAN) / IMAGENET_STD).transpose(2, 0, 1) |
| 63 | siglip_pixels = ((raw - SIGLIP_MEAN) / SIGLIP_STD).transpose(2, 0, 1) |
| 64 | pixel_values = np.concatenate([dinov2_pixels, siglip_pixels], axis=0) |
| 65 | return torch.from_numpy(pixel_values) |
| 66 | |
| 67 | |
| 68 | class OpenVLAImageProcessor(ImageProcessingMixin): |
| 69 | def __init__(self, *, image_size: int) -> None: |
| 70 | self.image_size = image_size |
| 71 | |
| 72 | def __call__( |
| 73 | self, |
| 74 | images: Any | None = None, |
| 75 | **kwargs: object, |
| 76 | ) -> dict[str, object]: |
| 77 | if images is None: |
| 78 | return {} |
| 79 | if not isinstance(images, Sequence) or isinstance(images, (str, bytes)): |
| 80 | images = [images] |
| 81 | if len(images) == 0: |
| 82 | return {} |
| 83 | |
| 84 | pixel_values = torch.stack( |
| 85 | [ |
| 86 | preprocess_openvla_image(image, image_size=self.image_size) |
| 87 | for image in images |
| 88 | ], |
| 89 | dim=0, |
| 90 | ) |
| 91 | return {"pixel_values": pixel_values} |
| 92 | |
| 93 | |
| 94 | class OpenVLAProcessor(ProcessorMixin): |
| 95 | def __init__( |
| 96 | self, |
| 97 | *, |
| 98 | image_processor: OpenVLAImageProcessor, |
| 99 | tokenizer: Any, |
| 100 | ) -> None: |
| 101 | self.image_processor = image_processor |
| 102 | self.tokenizer = tokenizer |