File size: 4,629 Bytes
39e9e02 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import re
import math
import torch
import string
import numpy as np
import pandas as pd
from PIL import Image
import torch.distributed as dist
import torchvision.transforms as T
from transformers import AutoModel, AutoTokenizer
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
def round_by_factor(number: int, factor: int) -> int:
"""Returns the closest integer to 'number' that is divisible by 'factor'."""
return round(number / factor) * factor
def ceil_by_factor(number: int, factor: int) -> int:
"""Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
return math.ceil(number / factor) * factor
def floor_by_factor(number: int, factor: int) -> int:
"""Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
return math.floor(number / factor) * factor
# copy from https://github.com/QwenLM/Qwen2.5-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py#L60
def smart_resize(
height: int, width: int, factor: int = 32, min_pixels: int = 65536, max_pixels: int = 4194304
) -> tuple[int, int]:
"""
Rescales the image so that the following conditions are met:
1. Both dimensions (height and width) are divisible by 'factor'.
2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
3. The aspect ratio of the image is maintained as closely as possible.
"""
if max(height, width) / min(height, width) > 200:
raise ValueError(
f"absolute aspect ratio must be smaller than {200}, got {max(height, width) / min(height, width)}"
)
h_bar = max(factor, round_by_factor(height, factor))
w_bar = max(factor, round_by_factor(width, factor))
if h_bar * w_bar > max_pixels:
beta = math.sqrt((height * width) / max_pixels)
h_bar = max(factor, floor_by_factor(height / beta, factor))
w_bar = max(factor, floor_by_factor(width / beta, factor))
elif h_bar * w_bar < min_pixels:
beta = math.sqrt(min_pixels / (height * width))
h_bar = ceil_by_factor(height * beta, factor)
w_bar = ceil_by_factor(width * beta, factor)
return h_bar, w_bar
def dynamic_preprocess_native_resolution(
image, size_factor=32, min_pixels=65536, max_pixels=4194304, **kwargs
):
width, height = image.size
resized_height, resized_width = smart_resize(
height,
width,
factor=size_factor,
min_pixels=min_pixels,
max_pixels=max_pixels,
)
image = image.resize((resized_width, resized_height))
return image
def preprocess_pixel_values(pixel_values, patch_size=16):
c, h, w = pixel_values.shape
grid_h = h // patch_size
grid_w = w // patch_size
flatten_pixel_values = (
pixel_values.view(c, grid_h, patch_size, grid_w, patch_size)
.permute(1, 3, 0, 2, 4) # [grid_h, grid_w, c, patch_size, patch_size]
.reshape(grid_h * grid_w, c * patch_size ** 2)
)
grid_hw = torch.tensor([[grid_h, grid_w]]).to(device=pixel_values.device)
return flatten_pixel_values, grid_hw
def load_image_native(
image, patch_size=16, downsample_ratio=0.5, min_pixels=65536, max_pixels=4194304, upscale=False
):
"""
Load and preprocess an image file, converting it to RGB mode,
resizing, normalizing, and optionally adding a thumbnail version.
"""
if not isinstance(image, Image.Image):
image = Image.open(image)
if image.mode == "RGBA":
bg_color = get_contrasting_background(image)
if bg_color:
background = Image.new("RGB", image.size, bg_color)
background.paste(image, mask=image.split()[3])
image = background.convert("RGB")
else:
image = image.convert("RGB")
else:
image = image.convert("RGB")
if upscale:
image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
transform = T.Compose(
[
T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
T.ToTensor(),
T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
]
)
new_image = dynamic_preprocess_native_resolution(
image, size_factor=int(patch_size // downsample_ratio), min_pixels=min_pixels, max_pixels=max_pixels
)
pixel_values, grid_hw = preprocess_pixel_values(transform(new_image).to(torch.float32), patch_size=patch_size)
# print(f'Transfer image_size from ({image.height, image.width}) to ({new_image.height, new_image.width})')
return pixel_values, grid_hw
|