| import math |
|
|
| from transformers.image_utils import SizeDict |
| from transformers.models.idefics3.image_processing_idefics3 import ( |
| MAX_IMAGE_SIZE, |
| Idefics3ImageProcessor, |
| Idefics3ImageProcessorKwargs, |
| _resize_output_size_rescale_to_max_len, |
| _resize_output_size_scale_below_upper_bound, |
| ) |
|
|
|
|
| def get_capped_longest_edge(height: int, width: int, longest_edge: int, tile_size: int, max_image_tiles: int) -> int: |
| for edge in range(longest_edge, tile_size - 1, -tile_size): |
| resized_height, resized_width = _resize_output_size_rescale_to_max_len(height, width, max_len=edge) |
| resized_height, resized_width = _resize_output_size_scale_below_upper_bound( |
| resized_height, resized_width, max_len=MAX_IMAGE_SIZE |
| ) |
| aspect_ratio = resized_width / resized_height |
| if resized_width >= resized_height: |
| cols = math.ceil(resized_width / tile_size) |
| rows = math.ceil(int(cols * tile_size / aspect_ratio) / tile_size) |
| else: |
| rows = math.ceil(resized_height / tile_size) |
| cols = math.ceil(int(rows * tile_size * aspect_ratio) / tile_size) |
| if rows * cols <= max_image_tiles: |
| return edge |
| return tile_size |
|
|
|
|
| class UltraVBertImageProcessorKwargs(Idefics3ImageProcessorKwargs, total=False): |
| max_image_tiles: int | None |
|
|
|
|
| class UltraVBertImageProcessor(Idefics3ImageProcessor): |
| max_image_tiles = None |
| valid_kwargs = UltraVBertImageProcessorKwargs |
|
|
| def _preprocess(self, images, *, do_resize, size, resample, do_image_splitting, max_image_size, |
| max_image_tiles=None, **kwargs): |
| if do_resize and do_image_splitting and max_image_tiles is not None: |
| tile_size = max_image_size["longest_edge"] |
| images = [ |
| [ |
| self.resize( |
| image, |
| SizeDict( |
| longest_edge=get_capped_longest_edge( |
| *image.shape[-2:], size.longest_edge, tile_size, max_image_tiles |
| ) |
| ), |
| resample=resample, |
| ) |
| for image in batch_images |
| ] |
| for batch_images in images |
| ] |
| do_resize = False |
| return super()._preprocess(images, do_resize=do_resize, size=size, resample=resample, |
| do_image_splitting=do_image_splitting, max_image_size=max_image_size, **kwargs) |
|
|
| def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None): |
| images_kwargs = dict(images_kwargs or {}) |
| max_image_tiles = images_kwargs.get("max_image_tiles", self.max_image_tiles) |
| if max_image_tiles is not None: |
| size = images_kwargs.get("size", self.size) |
| tile_size = images_kwargs.get("max_image_size", self.max_image_size)["longest_edge"] |
| longest_edge = get_capped_longest_edge(height, width, size["longest_edge"], tile_size, max_image_tiles) |
| images_kwargs["size"] = {"longest_edge": longest_edge} |
| return super().get_number_of_image_patches(height, width, images_kwargs) |
|
|
|
|
| __all__ = ["UltraVBertImageProcessor"] |
|
|
