MiniMaxAI/MiniMax-VL-01
28633k
1from transformers.image_processing_utils import BaseImageProcessor, BatchFeature2from typing import Optional, Union, Tuple, Dict, List, Iterable3from transformers.image_transforms import to_channel_dimension_format, PaddingMode4from transformers.image_utils import ChannelDimension, to_numpy_array, make_list_of_images, get_image_size, infer_channel_dimension_format5from transformers.utils import TensorType6from PIL import Image7import numpy as np8try:9 from torchvision.transforms import InterpolationMode10 BICUBIC = InterpolationMode.BICUBIC11except ImportError:12 BICUBIC = Image.BICUBIC13 14import torch15from transformers.utils import (16 TensorType,17 is_torch_device,18 is_torch_dtype,19 requires_backends,20)21 22from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize23 24try:25 from torchvision.transforms import InterpolationMode26 BICUBIC = InterpolationMode.BICUBIC27except ImportError:28 BICUBIC = Image.BICUBIC29 30from PIL import Image31import torch32import numpy as np33import os34processor_for_vllm = int(os.getenv("PROCESSOR_FOR_VLLM", 0))35 36def select_best_resolution(original_size, possible_resolutions):37 """38 Selects the best resolution from a list of possible resolutions based on the original size.39 40 Args:41 original_size (tuple): The original size of the image in the format (width, height).42 possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].43 44 Returns:45 tuple: The best fit resolution in the format (width, height).46 """47 original_width, original_height = original_size48 best_fit = None49 max_effective_resolution = 050 min_wasted_resolution = float("inf")51 52 for width, height in possible_resolutions:53 # Calculate the downscaled size to keep the aspect ratio54 scale = min(width / original_width, height / original_height)55 downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)56 57 # Calculate effective and wasted resolutions58 effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)59 wasted_resolution = (width * height) - effective_resolution60 61 if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):62 max_effective_resolution = effective_resolution63 min_wasted_resolution = wasted_resolution64 best_fit = (width, height)65 66 return best_fit 67 68def divide_to_patches(image, patch_size):69 """70 Divides an image into patches of a specified size.71 72 Args:73 image (PIL.Image.Image): The input image.74 patch_size (int): The size of each patch.75 76 Returns:77 list: A list of PIL.Image.Image objects representing the patches.78 """79 patches = []80 width, height = image.size81 for i in range(0, height, patch_size):82 for j in range(0, width, patch_size):83 box = (j, i, j + patch_size, i + patch_size)84 patch = image.crop(box)85 patches.append(patch)86 87 return patches88 89def image_size_to_num_patches(image_size, grid_pinpoints, patch_size):90 if not isinstance(grid_pinpoints, list):91 raise TypeError("grid_pinpoints should be a list of tuples or lists")92 93 # ! VERY IMPORTANT if image_size is tensor, must convert to into tuple, otherwise it will cause wrong calculate94 if not isinstance(image_size, (list, tuple)):95 if not isinstance(image_size, (torch.Tensor, np.ndarray)):96 raise TypeError(f"image_size invalid type {type(image_size)} with value {image_size}")97 image_size = image_size.tolist()98 99 best_resolution = select_best_resolution(image_size, grid_pinpoints)100 width, height = best_resolution101 num_patches = 0102 # consider change to ceil(height/patch_size)*ceil(width/patch_size) + 1103 for i in range(0, height, patch_size):104 for j in range(0, width, patch_size):105 num_patches += 1106 # add the base patch107 num_patches += 1108 return num_patches109 110def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):111 """112 Calculate the shape of the image patch grid after the preprocessing for images of any resolution.113 114 Args:115 image_size (`tuple`):116 The size of the input image in the format (width, height).117 grid_pinpoints (`List`):118 A list containing possible resolutions. Each item in the list should be a tuple or list119 of the form `(height, width)`.120 patch_size (`int`):121 The size of each image patch.122 123 Returns:124 tuple: The shape of the image patch grid in the format (width, height).125 """126 if not isinstance(grid_pinpoints, list):127 raise TypeError("grid_pinpoints should be a list of tuples or lists")128 129 # ! VERY IMPORTANT if image_size is tensor, must convert to into tuple, otherwise it will cause wrong calculate130 if not isinstance(image_size, (list, tuple)):131 if not isinstance(image_size, (torch.Tensor, np.ndarray)):132 raise TypeError(133 f"image_size invalid type: {type(image_size)} not valid, should be either list, tuple, np.ndarray or tensor"134 )135 image_size = image_size.tolist()136 137 width, height = select_best_resolution(image_size, grid_pinpoints)138 return width // patch_size, height // patch_size139 140 141# custom transform142class KeeyRatioResize(object):143 def __init__(self, size):144 self.size = size145 146 def __call__(self, image):147 return keepratio_resize(image, self.size)148 149def keepratio_resize(image, size, return_scale=False):150 # Resize the image to keep the ratio151 w, h = image.size152 resized_w, resized_h = size153 if w / h > resized_w / resized_h:154 # resize and pad to the right and left155 new_h = int(resized_w*h/w)156 resized_image = image.resize((resized_w, new_h), Image.BICUBIC)157 158 image = Image.new('RGB', (resized_w, resized_h), (0, 0, 0))159 pad_h = (resized_h - new_h) // 2160 image.paste(resized_image, (0, pad_h))161 scale = resized_w / w162 #image.paste(resized_image, (0, 0))163 else:164 # resize and pad to the top and bottom165 new_w = int(resized_h*w/h)166 resized_image = image.resize((new_w, resized_h), Image.BICUBIC)167 image = Image.new('RGB', (resized_w, resized_h), (0, 0, 0))168 #image.paste(resized_image, (0, 0))169 pad_w = (resized_w - new_w) // 2170 image.paste(resized_image, (pad_w, 0))171 scale = resized_h / h172 if return_scale:173 return image, scale174 return image175 176def _convert_image_to_rgb(image):177 return image.convert("RGB")178 179def _transform(img_h, img_w, image_mean=(0.48145466, 0.4578275, 0.40821073), image_std=(0.26862954, 0.26130258, 0.27577711)):180 return Compose([181 # ToPILImage(),182 #RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),183 #Resize((img_h, img_w), interpolation=BICUBIC),184 _convert_image_to_rgb,185 ToTensor(),186 Normalize(image_mean, image_std),187 ])188 189 190def get_hw_multiple_of(image_size, multiple, max_size=None):191 w, h = image_size192 new_w = w if w % multiple == 0 else w + (multiple - w % multiple)193 new_h = h if h % multiple == 0 else h + (multiple - h % multiple)194 if max_size is not None:195 assert isinstance(max_size, (list, tuple)) and len(max_size) == 2196 max_w, max_h = max_size197 assert max_w % multiple == 0 and max_h % multiple == 0198 if new_w > max_w or new_h > max_h:199 # ratio = min(max_w / new_w, max_h / new_h)200 # new_w = int(new_w * ratio)201 # new_h = int(new_h * ratio)202 new_w = min((new_w * max_w) // new_w, (new_w * max_h) // new_h)203 new_h = min((new_h * max_w) // new_w, (new_h * max_h) // new_h)204 205 new_w = new_w if new_w % multiple == 0 else new_w + (multiple - new_w % multiple)206 new_h = new_h if new_h % multiple == 0 else new_h + (multiple - new_h % multiple)207 assert new_w % multiple == 0 and new_h % multiple == 0208 assert new_w <= max_w and new_h <= max_h209 return new_w, new_h210 211def resize_multiple_of(image, multiple, max_size=None):212 """213 Resize the image to the multiple of a number.214 215 Args:216 image (PIL.Image.Image): The input image.217 multiple (int): The number to which the image should be resized.218 219 Returns:220 PIL.Image.Image: The resized image.221 """222 width, height = image.size223 new_width, new_height = get_hw_multiple_of((width, height), multiple, max_size)224 return image.resize((new_width, new_height), Image.BICUBIC)225 226 227 228class CustomBatchFeature(BatchFeature):229 def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):230 """231 Convert the inner content to tensors.232 233 Args:234 tensor_type (`str` or [`~utils.TensorType`], *optional*):235 The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If236 `None`, no modification is done.237 """238 if tensor_type is None:239 return self240 241 is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)242 243 # Do the tensor conversion in batch244 for key, value in self.items():245 if key == "pixel_values":246 for i, image in enumerate(value):247 if not is_tensor(image):248 tensor = as_tensor(image)249 self[key][i] = tensor250 continue251 try:252 if not is_tensor(value):253 tensor = as_tensor(value)254 255 self[key] = tensor256 except: # noqa E722257 if key == "overflowing_values":258 raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")259 raise ValueError(260 "Unable to create tensor, you should probably activate padding "261 "with 'padding=True' to have batched tensors with the same length."262 )263 264 return self265 266 def to(self, *args, **kwargs) -> "BatchFeature":267 """268 Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in269 different `dtypes` and sending the `BatchFeature` to a different `device`.270 271 Args:272 args (`Tuple`):273 Will be passed to the `to(...)` function of the tensors.274 kwargs (`Dict`, *optional*):275 Will be passed to the `to(...)` function of the tensors.276 277 Returns:278 [`BatchFeature`]: The same instance after modification.279 """280 requires_backends(self, ["torch"])281 import torch # noqa282 283 new_data = {}284 device = kwargs.get("device")285 # Check if the args are a device or a dtype286 if device is None and len(args) > 0:287 # device should be always the first argument288 arg = args[0]289 if is_torch_dtype(arg):290 # The first argument is a dtype291 pass292 elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):293 device = arg294 else:295 # it's something else296 raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")297 # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`298 for k, v in self.items():299 if k == "pixel_values":300 new_data[k] = [v[i].to(*args, **kwargs) for i in range(len(v))]301 continue302 # check if v is a floating point303 if torch.is_floating_point(v):304 # cast and send to device305 new_data[k] = v.to(*args, **kwargs)306 elif device is not None:307 new_data[k] = v.to(device=device)308 else:309 new_data[k] = v310 self.data = new_data311 return self312 313 314def as_tensor(value):315 if isinstance(value, (list, tuple)) and len(value) > 0:316 if isinstance(value[0], np.ndarray):317 value = np.array(value)318 elif (319 isinstance(value[0], (list, tuple))320 and len(value[0]) > 0321 and isinstance(value[0][0], np.ndarray)322 ):323 value = np.array(value)324 if isinstance(value, np.ndarray):325 return torch.from_numpy(value)326 else:327 return torch.tensor(value)328 329class ImageProcessor(BaseImageProcessor):330 model_input_names = ["pixel_values"]331 332 def __init__(333 self,334 size: Optional[Union[int, Tuple[int, int], Dict[str, int]]] = None,335 image_mean: Optional[Union[float, List[float]]] = None,336 image_std: Optional[Union[float, List[float]]] = None,337 process_image_mode: Optional[str] = 'resize',338 patch_size: Optional[int] = 14,339 image_grid_pinpoints: List = None,340 **kwargs,341 ) -> None:342 super().__init__(**kwargs)343 self.size = size # (width, height)344 self.image_mean = image_mean345 self.image_std = image_std346 self.process_image_mode = process_image_mode347 image_grid_pinpoints = (348 image_grid_pinpoints349 if image_grid_pinpoints is not None350 else [[336, 672], [672, 336], [672, 672], [1008, 336], [336, 1008]]351 )352 self.image_grid_pinpoints = image_grid_pinpoints353 self.patch_size = patch_size354 355 def preprocess(self,356 images,357 return_tensors: Optional[Union[str, TensorType]] = None,358 data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,359 input_data_format: Optional[Union[str, ChannelDimension]] = None,360 **kwargs,361 ):362 if self.process_image_mode == 'resize':363 return self.resize_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)364 elif self.process_image_mode == 'anyres':365 if processor_for_vllm == 1:366 return self.anyres_for_vllm_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)367 return self.anyres_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)368 elif self.process_image_mode == 'keepratio_resize':369 return self.keepratio_resize_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)370 elif self.process_image_mode == 'dynamic_res':371 return self.dynamic_res_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)372 else:373 raise ValueError(f"Invalid process_image_mode: {self.process_image_mode}")374 375 def resize_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs):376 images = make_list_of_images(images)377 all_images = []378 for image in images:379 resized_image = image.resize(self.size, Image.BICUBIC)380 transform_img = _transform(self.size[1], self.size[0], self.image_mean, self.image_std)(resized_image)381 all_images.append(to_numpy_array(transform_img))382 383 images = [384 to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)385 for image in all_images386 ]387 388 data = {"pixel_values": images}389 return CustomBatchFeature(data=data, tensor_type=return_tensors)390 391 def keepratio_resize_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs):392 images = make_list_of_images(images)393 all_images = []394 for image in images:395 resized_image = keepratio_resize(image, self.size)396 transform_img = _transform(self.size[1], self.size[0], self.image_mean, self.image_std)(resized_image)397 all_images.append(to_numpy_array(transform_img))398 399 images = [400 to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)401 for image in all_images402 ]403 404 data = {"pixel_values": images}405 return CustomBatchFeature(data=data, tensor_type=return_tensors)406 407 def dynamic_res_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs):408 images = make_list_of_images(images)409 all_images = []410 image_sizes = []411 for image in images:412 ori_w, ori_h = image.size413 image_sizes.append([ori_h, ori_w])414 resized_image = resize_multiple_of(image, self.patch_size, max_size=self.size)415 resized_w, resized_h = resized_image.size416 transform_img = _transform(resized_h, resized_w, self.image_mean, self.image_std)(resized_image)417 all_images.append(to_numpy_array(transform_img))418 419 images = [420 as_tensor(to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format))421 for image in all_images422 ]423 424 # data = {"pixel_values": images, "image_sizes": as_tensor(image_sizes)}425 # return data426 data = {"pixel_values": images, "image_sizes": image_sizes}427 #return BatchFeature(data=data, data_format=data_format, tensor_type=return_tensors)428 429 return CustomBatchFeature(data=data, tensor_type=return_tensors)430 431 def get_image_patches(432 self,433 data: Image,434 image_grid_pinpoints,435 ):436 if not isinstance(image_grid_pinpoints, list):437 raise TypeError("grid_pinpoints must be a list of possible resolutions.")438 439 440 best_resolution = select_best_resolution(data.size, image_grid_pinpoints)441 442 resized_data, scale = keepratio_resize(data, best_resolution, return_scale=True)443 resized_data = divide_to_patches(resized_data, self.size[0])444 ori_data = data.resize(self.size, Image.BICUBIC)445 data = [ori_data] + resized_data446 return data447 448 def pad(449 self,450 image: np.ndarray,451 padding: Union[int, Tuple[int, int], Iterable[Tuple[int, int]]],452 mode: PaddingMode = PaddingMode.CONSTANT,453 constant_values: Union[float, Iterable[float]] = 0.0,454 data_format: Optional[Union[str, ChannelDimension]] = None,455 input_data_format: Optional[Union[str, ChannelDimension]] = None,456 ) -> np.ndarray:457 """458 Pads the `image` with the specified `padding` and `mode`. Padding can be in the (`height`, `width`)459 dimension of in the (`num_patches`) dimension. In the second case an iterable if tuples is expected460 as input.461 462 Args:463 image (`np.ndarray`):464 The image to pad.465 padding (`int` or `Tuple[int, int]` or `Iterable[Tuple[int, int]]`):466 Padding to apply to the edges of the height, width axes. Can be one of three formats:467 - `((before_height, after_height), (before_width, after_width))` unique pad widths for each axis.468 - `((before, after),)` yields same before and after pad for height and width.469 - `(pad,)` or int is a shortcut for before = after = pad width for all axes.470 mode (`PaddingMode`):471 The padding mode to use. Can be one of:472 - `"constant"`: pads with a constant value.473 - `"reflect"`: pads with the reflection of the vector mirrored on the first and last values of the474 vector along each axis.475 - `"replicate"`: pads with the replication of the last value on the edge of the array along each axis.476 - `"symmetric"`: pads with the reflection of the vector mirrored along the edge of the array.477 constant_values (`float` or `Iterable[float]`, *optional*):478 The value to use for the padding if `mode` is `"constant"`.479 data_format (`str` or `ChannelDimension`, *optional*):480 The channel dimension format for the output image. Can be one of:481 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.482 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.483 If unset, will use same as the input image.484 input_data_format (`str` or `ChannelDimension`, *optional*):485 The channel dimension format for the input image. Can be one of:486 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.487 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.488 If unset, will use the inferred format of the input image.489 490 Returns:491 `np.ndarray`: The padded image.492 493 """494 495 # call the general `pad` if padding on `height/width`, otherwise it's the `num_patched` dim496 if isinstance(padding, int) or len(padding) != 4:497 return pad(image, padding, mode, constant_values, data_format, input_data_format)498 499 if input_data_format is None:500 input_data_format = infer_channel_dimension_format(image)501 if mode == PaddingMode.CONSTANT:502 image = np.pad(image, padding, mode="constant", constant_values=constant_values)503 elif mode == PaddingMode.REFLECT:504 image = np.pad(image, padding, mode="reflect")505 elif mode == PaddingMode.REPLICATE:506 image = np.pad(image, padding, mode="edge")507 elif mode == PaddingMode.SYMMETRIC:508 image = np.pad(image, padding, mode="symmetric")509 else:510 raise ValueError(f"Invalid padding mode: {mode}")511 image = (512 to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image513 )514 return image515 516 def _pad_for_batching(517 self,518 pixel_values: List[np.ndarray],519 data_format: Optional[Union[str, ChannelDimension]] = None,520 input_data_format: Optional[Union[str, ChannelDimension]] = None,521 ):522 """523 Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.524 525 Args:526 pixel_values (`List[np.ndarray]`):527 An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`)528 data_format (`str` or `ChannelDimension`, *optional*):529 The channel dimension format for the output image. Can be one of:530 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.531 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.532 If unset, will use same as the input image.533 input_data_format (`str` or `ChannelDimension`, *optional*):534 The channel dimension format for the input image. Can be one of:535 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.536 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.537 If unset, will use the inferred format of the input image.538 539 Returns:540 List[`np.ndarray`]: The padded images.541 """542 max_patch = max(len(x) for x in pixel_values)543 pixel_values = [544 self.pad(545 image,546 padding=((0, max_patch - image.shape[0]), (0, 0), (0, 0), (0, 0)),547 data_format=data_format,548 input_data_format=input_data_format,549 )550 for image in pixel_values551 ]552 553 return pixel_values554 555 def anyres_for_vllm_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, do_pad: Optional[bool] = None, **kwargs):556 557 images = make_list_of_images(images)558 new_images = []559 image_sizes = []560 561 for image in images:562 ori_w, ori_h = image.size563 image_sizes.append([ori_h, ori_w])564 image_patches = self.get_image_patches(565 image,566 self.image_grid_pinpoints567 )568 all_images = []569 for image in image_patches:570 transform_img = _transform(self.size[0], self.size[1], self.image_mean, self.image_std)(image)571 img_array = to_numpy_array(transform_img)572 img_array = to_channel_dimension_format(img_array, data_format, input_channel_dim=input_data_format)573 all_images.append(img_array)574 #new_images.append(img_array)575 pixel_values = np.array(all_images)576 new_images.append(pixel_values)577 578 579 new_images = self._pad_for_batching(new_images)580 581 data = {"pixel_values": new_images, "image_sizes": image_sizes}582 return BatchFeature(data=data, tensor_type=return_tensors)583 584 585 def anyres_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, do_pad: Optional[bool] = None, **kwargs):586 587 images = make_list_of_images(images)588 new_images = []589 image_sizes = []590 591 for image in images:592 ori_w, ori_h = image.size593 image_sizes.append([ori_h, ori_w])594 image_patches = self.get_image_patches(595 image,596 self.image_grid_pinpoints597 )598 #all_images = []599 for image in image_patches:600 transform_img = _transform(self.size[0], self.size[1], self.image_mean, self.image_std)(image)601 img_array = to_numpy_array(transform_img)602 img_array = to_channel_dimension_format(img_array, data_format, input_channel_dim=input_data_format)603 #all_images.append(img_array)604 new_images.append(img_array)605 #pixel_values = np.array(all_images)606 #new_images.append(pixel_values)607 608 # if do_pad:609 # new_images = self._pad_for_batching(new_images)610 611 data = {"pixel_values": new_images, "image_sizes": image_sizes}612 return CustomBatchFeature(data=data, tensor_type=return_tensors)613 614 615 616