254 lines
8.7 KiB
Python
254 lines
8.7 KiB
Python
# Source: https://github.com/haotian-liu/LLaVA/blob/main/llava/mm_utils.py
|
|
import ast
|
|
import base64
|
|
import math
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
def select_best_resolution(original_size, possible_resolutions):
|
|
"""
|
|
Selects the best resolution from a list of possible resolutions based on the original size.
|
|
|
|
Args:
|
|
original_size (tuple): The original size of the image in the format (width, height).
|
|
possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
|
|
|
|
Returns:
|
|
tuple: The best fit resolution in the format (width, height).
|
|
"""
|
|
original_width, original_height = original_size
|
|
best_fit = None
|
|
max_effective_resolution = 0
|
|
min_wasted_resolution = float("inf")
|
|
|
|
for width, height in possible_resolutions:
|
|
scale = min(width / original_width, height / original_height)
|
|
downscaled_width, downscaled_height = int(original_width * scale), int(
|
|
original_height * scale
|
|
)
|
|
effective_resolution = min(
|
|
downscaled_width * downscaled_height, original_width * original_height
|
|
)
|
|
wasted_resolution = (width * height) - effective_resolution
|
|
|
|
if effective_resolution > max_effective_resolution or (
|
|
effective_resolution == max_effective_resolution
|
|
and wasted_resolution < min_wasted_resolution
|
|
):
|
|
max_effective_resolution = effective_resolution
|
|
min_wasted_resolution = wasted_resolution
|
|
best_fit = (width, height)
|
|
|
|
return best_fit
|
|
|
|
|
|
def resize_and_pad_image(image, target_resolution):
|
|
"""
|
|
Resize and pad an image to a target resolution while maintaining aspect ratio.
|
|
|
|
Args:
|
|
image (PIL.Image.Image): The input image.
|
|
target_resolution (tuple): The target resolution (width, height) of the image.
|
|
|
|
Returns:
|
|
PIL.Image.Image: The resized and padded image.
|
|
"""
|
|
original_width, original_height = image.size
|
|
target_width, target_height = target_resolution
|
|
|
|
scale_w = target_width / original_width
|
|
scale_h = target_height / original_height
|
|
|
|
if scale_w < scale_h:
|
|
new_width = target_width
|
|
new_height = min(math.ceil(original_height * scale_w), target_height)
|
|
else:
|
|
new_height = target_height
|
|
new_width = min(math.ceil(original_width * scale_h), target_width)
|
|
|
|
# Resize the image
|
|
resized_image = image.resize((new_width, new_height))
|
|
|
|
new_image = Image.new("RGB", (target_width, target_height), (0, 0, 0))
|
|
paste_x = (target_width - new_width) // 2
|
|
paste_y = (target_height - new_height) // 2
|
|
new_image.paste(resized_image, (paste_x, paste_y))
|
|
|
|
return new_image
|
|
|
|
|
|
def divide_to_patches(image, patch_size):
|
|
"""
|
|
Divides an image into patches of a specified size.
|
|
|
|
Args:
|
|
image (PIL.Image.Image): The input image.
|
|
patch_size (int): The size of each patch.
|
|
|
|
Returns:
|
|
list: A list of PIL.Image.Image objects representing the patches.
|
|
"""
|
|
patches = []
|
|
width, height = image.size
|
|
for i in range(0, height, patch_size):
|
|
for j in range(0, width, patch_size):
|
|
box = (j, i, j + patch_size, i + patch_size)
|
|
patch = image.crop(box)
|
|
patches.append(patch)
|
|
|
|
return patches
|
|
|
|
|
|
def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
|
|
"""
|
|
Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
|
|
|
|
Args:
|
|
image_size (tuple): The size of the input image in the format (width, height).
|
|
grid_pinpoints (str): A string representation of a list of possible resolutions.
|
|
patch_size (int): The size of each image patch.
|
|
|
|
Returns:
|
|
tuple: The shape of the image patch grid in the format (width, height).
|
|
"""
|
|
if type(grid_pinpoints) is list:
|
|
possible_resolutions = grid_pinpoints
|
|
else:
|
|
possible_resolutions = ast.literal_eval(grid_pinpoints)
|
|
width, height = select_best_resolution(image_size, possible_resolutions)
|
|
return width // patch_size, height // patch_size
|
|
|
|
|
|
def process_anyres_image(image, processor, grid_pinpoints):
|
|
"""
|
|
Process an image with variable resolutions.
|
|
|
|
Args:
|
|
image (PIL.Image.Image): The input image to be processed.
|
|
processor: The image processor object.
|
|
grid_pinpoints (str): A string representation of a list of possible resolutions.
|
|
|
|
Returns:
|
|
np.array: An np array containing the processed image patches.
|
|
"""
|
|
if type(grid_pinpoints) is list:
|
|
possible_resolutions = grid_pinpoints
|
|
else:
|
|
possible_resolutions = ast.literal_eval(grid_pinpoints)
|
|
best_resolution = select_best_resolution(image.size, possible_resolutions)
|
|
image_padded = resize_and_pad_image(image, best_resolution)
|
|
|
|
patches = divide_to_patches(image_padded, processor.crop_size["height"])
|
|
|
|
image_original_resize = image.resize(
|
|
(processor.size["shortest_edge"], processor.size["shortest_edge"])
|
|
)
|
|
|
|
image_patches = [image_original_resize] + patches
|
|
image_patches = [
|
|
processor.preprocess(image_patch)["pixel_values"][0]
|
|
for image_patch in image_patches
|
|
]
|
|
return np.stack(image_patches, axis=0)
|
|
|
|
|
|
def load_image_from_base64(image):
|
|
return Image.open(BytesIO(base64.b64decode(image)))
|
|
|
|
|
|
def expand2square(pil_img, background_color):
|
|
width, height = pil_img.size
|
|
if width == height:
|
|
return pil_img
|
|
if pil_img.mode == "L":
|
|
pil_img = pil_img.convert("RGB")
|
|
if width > height:
|
|
result = Image.new(pil_img.mode, (width, width), background_color)
|
|
result.paste(pil_img, (0, (width - height) // 2))
|
|
return result
|
|
else:
|
|
result = Image.new(pil_img.mode, (height, height), background_color)
|
|
result.paste(pil_img, ((height - width) // 2, 0))
|
|
return result
|
|
|
|
|
|
def unpad_image(tensor, original_size):
|
|
"""
|
|
Unpads a PyTorch tensor of a padded and resized image.
|
|
|
|
Args:
|
|
tensor (torch.Tensor): The image tensor, assumed to be in CxHxW format.
|
|
original_size (tuple): The original size of the image (height, width).
|
|
|
|
Returns:
|
|
torch.Tensor: The unpadded image tensor.
|
|
"""
|
|
original_width, original_height = original_size
|
|
current_height, current_width = tensor.shape[1:]
|
|
|
|
original_aspect_ratio = original_width / original_height
|
|
current_aspect_ratio = current_width / current_height
|
|
|
|
if original_aspect_ratio > current_aspect_ratio:
|
|
scale_factor = current_width / original_width
|
|
new_height = int(original_height * scale_factor)
|
|
padding = (current_height - new_height) // 2
|
|
unpadded_tensor = tensor[:, padding : current_height - padding, :]
|
|
else:
|
|
scale_factor = current_height / original_height
|
|
new_width = int(original_width * scale_factor)
|
|
padding = (current_width - new_width) // 2
|
|
unpadded_tensor = tensor[:, :, padding : current_width - padding]
|
|
|
|
return unpadded_tensor
|
|
|
|
|
|
def unpad_image_shape(current_height, current_width, original_size):
|
|
"""
|
|
Unpads a PyTorch tensor of a padded and resized image
|
|
and returns the new shape.
|
|
"""
|
|
original_width, original_height = original_size
|
|
|
|
original_aspect_ratio = original_width / original_height
|
|
current_aspect_ratio = current_width / current_height
|
|
|
|
if original_aspect_ratio > current_aspect_ratio:
|
|
scale_factor = current_width / original_width
|
|
new_height = int(original_height * scale_factor)
|
|
padding = (current_height - new_height) // 2
|
|
new_shape = (current_height - 2 * padding, current_width)
|
|
else:
|
|
scale_factor = current_height / original_height
|
|
new_width = int(original_width * scale_factor)
|
|
padding = (current_width - new_width) // 2
|
|
new_shape = (current_height, current_width - 2 * padding)
|
|
|
|
return new_shape
|
|
|
|
|
|
def process_images(images, image_processor, model_cfg):
|
|
image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
|
|
new_images = []
|
|
if image_aspect_ratio == "pad":
|
|
for image in images:
|
|
image = expand2square(
|
|
image, tuple(int(x * 255) for x in image_processor.image_mean)
|
|
)
|
|
image = image_processor.preprocess(image)["pixel_values"][0]
|
|
new_images.append(image)
|
|
elif image_aspect_ratio == "anyres":
|
|
for image in images:
|
|
image = process_anyres_image(
|
|
image, image_processor, model_cfg.image_grid_pinpoints
|
|
)
|
|
new_images.append(image)
|
|
else:
|
|
return image_processor(images)["pixel_values"]
|
|
if all(x.shape == new_images[0].shape for x in new_images):
|
|
new_images = np.stack(new_images, axis=0)
|
|
return new_images
|