File size: 3,286 Bytes
5502d37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | 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"]
|