nvidia/Eagle2-2B
343.3k
1# --------------------------------------------------------2# NVIDIA3# Copyright (c) 2025 NVIDIA4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6 7# copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/image_processing_llava_onevision_fast.py8from typing import List, Optional, Union9 10from transformers.image_processing_utils import BatchFeature, get_patch_output_size, select_best_resolution11from transformers.image_processing_utils_fast import (12 BASE_IMAGE_PROCESSOR_FAST_DOCSTRING,13 BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS,14 BaseImageProcessorFast,15 DefaultFastImageProcessorKwargs,16 divide_to_patches,17 group_images_by_shape,18 reorder_images,19)20from transformers.image_utils import (21 OPENAI_CLIP_MEAN,22 OPENAI_CLIP_STD,23 IMAGENET_STANDARD_MEAN, # 0.5, 0.5, 0.524 IMAGENET_STANDARD_STD, # 0.5, 0.5, 0.525 ChannelDimension,26 ImageInput,27 VideoInput,28 PILImageResampling,29 SizeDict,30 get_image_size,31 make_flat_list_of_images,32 make_batched_videos,33 validate_kwargs34)35from transformers.processing_utils import Unpack36from transformers.utils import TensorType, add_start_docstrings, is_torch_available, is_torchvision_v2_available37 38 39if is_torch_available():40 import torch41if is_torchvision_v2_available():42 from transformers.image_utils import pil_torch_interpolation_mapping43 44 from torchvision.transforms.v2 import functional as F45else:46 from torchvision.transforms import functional as F47 48def crop(img: torch.Tensor, left: int, top: int, right: int, bottom: int) -> torch.Tensor:49 """Crop the given numpy array.50 51 Args:52 img (torch.Tensor): Image to be cropped. Format should be (C, H, W).53 left (int): The left coordinate of the crop box.54 top (int): The top coordinate of the crop box.55 right (int): The right coordinate of the crop box.56 bottom (int): The bottom coordinate of the crop box.57 58 Returns:59 torch.Tensor: Cropped image.60 """61 if not isinstance(img, torch.Tensor):62 raise TypeError('img should be torch.Tensor. Got {}'.format(type(img)))63 64 if img.ndim not in [2, 3]:65 raise ValueError('Image should have 2 or 3 dimensions. Got {}'.format(img.ndim))66 67 img_height = img.shape[1]68 img_width = img.shape[2]69 if top < 0 or left < 0 or bottom > img_height or right > img_width:70 raise ValueError('Crop coordinates out of bounds')71 72 if top >= bottom or left >= right:73 raise ValueError('Invalid crop coordinates')74 75 return img[:, top:bottom, left:right]76 77 78class Eagle2_5_VLFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):79 max_dynamic_tiles: Optional[int]80 min_dynamic_tiles: Optional[int]81 use_thumbnail: Optional[bool]82 pad_during_tiling: Optional[bool]83 do_pad: Optional[bool]84 85 86@add_start_docstrings(87 "Constructs a fast ConvNeXT image processor. Based on [`SiglipImageProcessor`] with incorporation of processing each video frame.",88 BASE_IMAGE_PROCESSOR_FAST_DOCSTRING,89 """90 image_grid_pinpoints (`List[List[int]]`, *optional*):91 A list of possible resolutions to use for processing high resolution images. The best resolution is selected92 based on the original size of the image. Can be overridden by `image_grid_pinpoints` in the `preprocess`93 method. Not used for processing videos.94 do_pad (`bool`, *optional*):95 Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest96 number of patches in the batch. Padding will be applied to the bottom and right with zeros.97 """,98)99class Eagle2_5_VLImageProcessorFast(BaseImageProcessorFast):100 resample = PILImageResampling.BICUBIC101 image_mean = IMAGENET_STANDARD_MEAN102 image_std = IMAGENET_STANDARD_STD103 size = {"height": 448, "width": 448}104 default_to_square = False105 crop_size = None106 do_resize = True107 do_center_crop = None108 do_rescale = True109 do_normalize = True110 do_convert_rgb = True111 do_pad = True112 max_dynamic_tiles = 12113 min_dynamic_tiles = 1114 use_thumbnail = True115 pad_during_tiling = False116 valid_kwargs = Eagle2_5_VLFastImageProcessorKwargs117 model_input_names = ["pixel_values_videos"]118 119 def __init__(self, **kwargs: Unpack[Eagle2_5_VLFastImageProcessorKwargs]):120 super().__init__(**kwargs)121 122 @add_start_docstrings(123 BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS,124 """125 max_dynamic_tiles (`int`, *optional*):126 The maximum number of dynamic tiles to use for processing high resolution images.127 min_dynamic_tiles (`int`, *optional*):128 The minimum number of dynamic tiles to use for processing high resolution images.129 use_thumbnail (`bool`, *optional*):130 Whether to use a thumbnail for processing high resolution images.131 pad_during_tiling (`bool`, *optional*):132 Whether to pad the image during tiling.133 do_pad (`bool`, *optional*):134 Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest135 number of patches in the batch. Padding will be applied to the bottom and right with zeros.136 """,137 )138 def preprocess(self, images: ImageInput, **kwargs: Unpack[Eagle2_5_VLFastImageProcessorKwargs]) -> BatchFeature:139 return super().preprocess(images, **kwargs)140 141 def _prepare_images_structure(142 self,143 images: ImageInput,144 ) -> ImageInput:145 """146 Prepare the images structure for processing.147 148 Args:149 images (`ImageInput`):150 The input images to process.151 152 Returns:153 `ImageInput`: The images with a valid nesting.154 """155 return make_flat_list_of_images(images)156 157 def _prepare_videos_structure(self, videos: VideoInput) -> VideoInput:158 return self._prepare_images_structure(videos)159 160 def _prepare_input_videos(161 self,162 videos: VideoInput,163 do_convert_rgb: Optional[bool] = None,164 input_data_format: Optional[Union[str, ChannelDimension]] = None,165 device: Optional["torch.device"] = None,166 ) -> list["torch.Tensor"]:167 """168 Prepare the input images for processing.169 """170 videos = self._prepare_videos_structure(videos)171 process_video_fn = partial(172 self._process_image,173 do_convert_rgb=do_convert_rgb,174 input_data_format=input_data_format,175 device=device,176 )177 # todo: yoni - check if we can parallelize this efficiently178 processed_videos = []179 for video in videos:180 processed_videos.append(process_video_fn(video))181 182 return processed_videos183 184 def _resize_for_patching(185 self,186 image: "torch.Tensor",187 target_resolution: tuple,188 interpolation: "F.InterpolationMode",189 input_data_format: ChannelDimension,190 ) -> "torch.Tensor":191 """192 Resizes an image to a target resolution while maintaining aspect ratio.193 194 Args:195 image ("torch.Tensor"):196 The input image.197 target_resolution (tuple):198 The target resolution (height, width) of the image.199 interpolation (`InterpolationMode`):200 Resampling filter to use if resizing the image.201 input_data_format (`ChannelDimension` or `str`):202 The channel dimension format of the input image.203 204 Returns:205 "torch.Tensor": The resized and padded image.206 """207 new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)208 209 # Resize the image210 resized_image = F.resize(image, (new_height, new_width), interpolation=interpolation)211 212 return resized_image213 214 def find_closest_aspect_ratio(self, aspect_ratio, target_ratios, width, height, image_size):215 """216 previous version mainly foucs on ratio.217 We also consider area ratio here.218 """219 best_factor = float('-inf')220 best_ratio = (1, 1)221 area = width * height222 for ratio in target_ratios:223 target_aspect_ratio = ratio[0] / ratio[1]224 ratio_diff = abs(aspect_ratio - target_aspect_ratio)225 area_ratio = (ratio[0]*ratio[1]*image_size*image_size)/ area226 """227 new area > 60% of original image area is enough.228 """229 factor_based_on_area_n_ratio = min((ratio[0]*ratio[1]*image_size*image_size)/ area, 0.6)* \230 min(target_aspect_ratio/aspect_ratio, aspect_ratio/target_aspect_ratio)231 232 if factor_based_on_area_n_ratio > best_factor:233 best_factor = factor_based_on_area_n_ratio234 best_ratio = ratio235 236 return best_ratio237 238 def _pad_for_patching(239 self, image: "torch.Tensor", target_resolution: tuple, input_data_format: ChannelDimension240 ) -> "torch.Tensor":241 """242 Pad an image to a target resolution while maintaining aspect ratio.243 """244 target_height, target_width = target_resolution245 new_height, new_width = get_patch_output_size(image, target_resolution, input_data_format)246 247 paste_x = (target_width - new_width) // 2248 paste_y = (target_height - new_height) // 2249 250 padded_image = F.pad(image, padding=[paste_x, paste_y, paste_x, paste_y])251 252 return padded_image253 254 def _get_image_patches(255 self,256 image: "torch.Tensor",257 min_num: int,258 max_num: int,259 size: tuple,260 tile_size: int,261 use_thumbnail: bool,262 interpolation: "F.InterpolationMode",263 pad_during_tiling: bool,264 ) -> List["torch.Tensor"] :265 image_size = get_image_size(image, channel_dim=ChannelDimension.FIRST)266 orig_height, orig_width = image_size267 aspect_ratio = orig_width / orig_height268 269 # calculate the existing image aspect ratio270 target_ratios = set(271 (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if272 i * j <= max_num and i * j >= min_num)273 target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])274 275 # find the closest aspect ratio to the target276 target_aspect_ratio = self.find_closest_aspect_ratio(277 aspect_ratio, target_ratios, orig_width, orig_height, tile_size)278 279 # calculate the target width and height280 target_width = tile_size * target_aspect_ratio[0]281 target_height = tile_size * target_aspect_ratio[1]282 blocks = target_aspect_ratio[0] * target_aspect_ratio[1]283 if pad_during_tiling:284 resized_image = self._resize_for_patching(285 image, (target_height, target_width), interpolation=interpolation, input_data_format=ChannelDimension.FIRST286 )287 padded_image = self._pad_for_patching(resized_image, (target_height, target_width), input_data_format=ChannelDimension.FIRST)288 image_used_to_split = padded_image289 else:290 image_used_to_split = F.resize(image, (target_height, target_width), interpolation=interpolation)291 292 processed_tiles = []293 for i in range(blocks):294 box = (295 (i % (target_width // tile_size)) * tile_size,296 (i // (target_width // tile_size)) * tile_size,297 ((i % (target_width // tile_size)) + 1) * tile_size,298 ((i // (target_width // tile_size)) + 1) * tile_size299 )300 # split the image301 split_img = crop(image_used_to_split, box[0], box[1], box[2], box[3])302 processed_tiles.append(split_img)303 assert len(processed_tiles) == blocks304 305 if use_thumbnail and len(processed_tiles) != 1:306 thumbnail_img = F.resize(image, (tile_size, tile_size), interpolation=interpolation)307 processed_tiles.append(thumbnail_img)308 309 return processed_tiles310 311 def _pad_for_batching(312 self,313 pixel_values: List["torch.Tensor"],314 ) -> List["torch.Tensor"]:315 """316 Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.317 318 Args:319 pixel_values (`List[torch.Tensor]`):320 An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`)321 322 Returns:323 List[`torch.Tensor`]: The padded images.324 """325 max_patch = max(len(x) for x in pixel_values)326 pixel_values = [327 torch.nn.functional.pad(image, pad=[0, 0, 0, 0, 0, 0, 0, max_patch - image.shape[0]])328 for image in pixel_values329 ]330 331 return pixel_values332 333 def _preprocess(334 self,335 images: List["torch.Tensor"],336 do_resize: bool,337 size: SizeDict,338 max_dynamic_tiles: int,339 min_dynamic_tiles: int,340 use_thumbnail: bool,341 pad_during_tiling: bool,342 interpolation: Optional["F.InterpolationMode"],343 do_center_crop: bool,344 crop_size: SizeDict,345 do_rescale: bool,346 rescale_factor: float,347 do_normalize: bool,348 image_mean: Optional[Union[float, List[float]]],349 image_std: Optional[Union[float, List[float]]],350 do_pad: bool,351 return_tensors: Optional[Union[str, TensorType]],352 ) -> BatchFeature:353 processed_images = []354 image_sizes = []355 # Determine the size tuple356 if size and size.height and size.width:357 size_tuple = (size.height, size.width)358 else:359 size_tuple = (size.shortest_edge, size.shortest_edge)360 361 # Determine the patch size362 if crop_size and crop_size.height:363 tile_size = crop_size.height364 elif size and size.height:365 tile_size = size.height366 else:367 tile_size = size.shortest_edge368 369 for image in images:370 image_patches = self._get_image_patches(371 image,372 min_num=min_dynamic_tiles,373 max_num=max_dynamic_tiles,374 size=size_tuple,375 tile_size=tile_size,376 use_thumbnail=use_thumbnail,377 interpolation=interpolation,378 pad_during_tiling=pad_during_tiling,379 )380 381 # Group images by size for batched processing382 processed_image_patches_grouped = {}383 grouped_image_patches, grouped_image_patches_index = group_images_by_shape(image_patches)384 385 for shape, stacked_image_patches in grouped_image_patches.items():386 if do_resize:387 stacked_image_patches = self.resize(388 image=stacked_image_patches,389 size=size,390 interpolation=interpolation,391 )392 if do_center_crop:393 stacked_image_patches = self.center_crop(stacked_image_patches, crop_size)394 # Fused rescale and normalize395 stacked_image_patches = self.rescale_and_normalize(396 stacked_image_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std397 )398 processed_image_patches_grouped[shape] = stacked_image_patches399 processed_image_patches = reorder_images(processed_image_patches_grouped, grouped_image_patches_index)400 processed_image_patches = (401 torch.stack(processed_image_patches, dim=0) if return_tensors else processed_image_patches402 )403 processed_images.append(processed_image_patches)404 image_sizes.append(get_image_size(image, ChannelDimension.FIRST))405 406 if do_pad:407 processed_images = self._pad_for_batching(processed_images)408 409 # processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images410 processed_images = torch.cat(processed_images, dim=0) if return_tensors else processed_images411 return BatchFeature(412 data={"pixel_values": processed_images, "image_sizes": image_sizes}, tensor_type=return_tensors413 )414 415 416 def preprocess(self, images: ImageInput, videos: VideoInput=None, **kwargs: Unpack[Eagle2_5_VLFastImageProcessorKwargs]) -> BatchFeature:417 validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self.valid_kwargs.__annotations__.keys())418 # Set default kwargs from self. This ensures that if a kwarg is not provided419 # by the user, it gets its default value from the instance, or is set to None.420 for kwarg_name in self.valid_kwargs.__annotations__:421 kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))422 423 # Extract parameters that are only used for preparing the input images424 do_convert_rgb = kwargs.pop("do_convert_rgb")425 input_data_format = kwargs.pop("input_data_format")426 device = kwargs.pop("device")427 # Prepare input images428 if images is not None:429 images = self._prepare_input_images(430 images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device431 )432 433 if videos is not None:434 videos = self._prepare_input_images(435 images=videos, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device436 )437 438 # Update kwargs that need further processing before being validated439 kwargs = self._further_process_kwargs(**kwargs)440 441 # Validate kwargs442 self._validate_preprocess_kwargs(**kwargs)443 444 # torch resize uses interpolation instead of resample445 resample = kwargs.pop("resample")446 kwargs["interpolation"] = (447 pil_torch_interpolation_mapping[resample] if isinstance(resample, (PILImageResampling, int)) else resample448 )449 450 # Pop kwargs that are not needed in _preprocess451 kwargs.pop("default_to_square")452 kwargs.pop("data_format")453 if images is not None:454 return self._preprocess(images, **kwargs)455 elif videos is not None:456 return self._preprocess(videos, **kwargs)457 458__all__ = ["Eagle2_5_VLImageProcessorFast"]459