Upload any_res.py

src/data/any_res.py

@@ -0,0 +1,257 @@
+import base64
+import torch
+import math
+import ast
+from PIL import Image
+from io import BytesIO
+
+
+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 select_best_resolution_v2(original_size, possible_resolutions):
+    """
+    Selects the best resolution from a list of possible resolutions based on the original size and aspect ratio.
+
+    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
+    original_aspect_ratio = original_height / original_width
+    original_area = original_width * original_height
+    best_fit = None
+    min_aspect_ratio_diff = float('inf')
+    min_area_ratio = float('inf')
+
+    for width, height in possible_resolutions:
+        aspect_ratio = height / width
+        area = width * height
+        aspect_ratio_diff = max(aspect_ratio, original_aspect_ratio) / min(aspect_ratio, original_aspect_ratio)
+        area_ratio = max(area, original_area) / min(area, original_area)
+
+        if aspect_ratio_diff < min_aspect_ratio_diff or (aspect_ratio_diff == min_aspect_ratio_diff and area_ratio < min_area_ratio):
+            min_aspect_ratio_diff = aspect_ratio_diff
+            min_area_ratio = area_ratio
+            best_fit = (width, height)
+
+    return best_fit
+
+
+def resize_and_pad_image(image, target_resolution, keep_ratio=False):
+    """
+    Resize and pad an image to a target resolution
+
+    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
+
+    if keep_ratio:
+        # maintaining aspect ratio
+        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))
+    else:
+        # not maintaining aspect ratio
+        new_image = image.resize((target_width, target_height))
+        
+    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)
+    width1, height1 = select_best_resolution(image_size, possible_resolutions)
+    width2, height2 = select_best_resolution_v2(image_size, possible_resolutions)
+    if width1*height1 > width2*height2:
+        width, height = width2, height2
+    else:
+        width, height = width1, height1
+    return width // patch_size, height // patch_size
+
+
+def process_anyres_image(image, image_transform, grid_pinpoints, base_image_size):
+    """
+    Process an image with variable resolutions.
+
+    Args:
+        image (PIL.Image.Image): The input image to be processed.
+        image_transform: The image processor object.
+        grid_pinpoints (str): A string representation of a list of possible resolutions.
+
+    Returns:
+        torch.Tensor: A tensor 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)
+    width1, height1 = select_best_resolution(image.size, possible_resolutions)
+    width2, height2 = select_best_resolution_v2(image.size, possible_resolutions)
+    if width1*height1 > width2*height2:
+        width, height = width2, height2
+    else:
+        width, height = width1, height1
+    best_resolution = [width, height]
+
+    image_padded = resize_and_pad_image(image, best_resolution)
+
+    patches = divide_to_patches(image_padded, base_image_size)
+
+    image_original_resize = image.resize((base_image_size, base_image_size))
+
+    image_patches =  patches + [image_original_resize] # add the original image as the last patch
+    image_patches = [image_transform(image_patch)
+                     for image_patch in image_patches]
+    
+    patch_grid = (best_resolution[0]//base_image_size, best_resolution[1]//base_image_size)
+    x_index = (torch.arange(patch_grid[0]).repeat(patch_grid[1], 1)  + 0.5)/patch_grid[0]
+    y_index = (torch.arange(patch_grid[1]).unsqueeze(1).repeat(1, patch_grid[0]) + 0.5)/patch_grid[1]
+    patch_pos = torch.stack([x_index, y_index], dim=-1).flatten(0, 1) # h*w, 2
+
+    origin_pos = torch.tensor([[0.5, 0.5]])
+    patch_pos  = torch.cat([patch_pos, origin_pos], dim=0) # h*w+1, 2
+
+    return torch.stack(image_patches, dim=0), patch_pos
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def anyres_data_collate(batch, tokenizer, dataset_name=None):
+    results = {}
+    keys = batch[0].keys()
+
+    for key in keys:
+        cur = [batch[i][key] for i in range(len(batch)) if batch[i][key] is not None]
+        if len(cur) == 0:
+            results[key] = None
+        elif isinstance(cur[0], torch.Tensor):
+            if key in ['embeds_gen_mask', 'embeds_cmp_mask', 'images', 'images_patch_length', 'patch_position', 'image_size']:
+                results[key] = torch.cat(cur, dim=0)
+            else:
+                if key in ['input_ids']:
+                    results[key] = torch.nn.utils.rnn.pad_sequence(cur, batch_first=True, padding_value=tokenizer.pad_token_id)
+                elif key in ['attention_mask']:
+                    results[key] = torch.nn.utils.rnn.pad_sequence(cur, batch_first=True, padding_value=0)
+                elif key in ['labels']:
+                    results[key] = torch.nn.utils.rnn.pad_sequence(cur, batch_first=True, padding_value=-100)
+                elif key in ['ids_gen_mask', 'ids_cmp_mask']:
+                    results[key] = torch.nn.utils.rnn.pad_sequence(cur, batch_first=True, padding_value=False)
+
+                else:
+                    results[key] = torch.stack(cur, dim=0)
+        else:
+            results[key] = cur
+
+    results['dataset_name'] = dataset_name
+
+    return results
+
+
+def anyres_data_collate_old(batch, dataset_name=None):
+    results = {}
+    keys = batch[0].keys()
+
+    for key in keys:
+        cur = [batch[i][key] for i in range(len(batch)) if batch[i][key] is not None]
+        if len(cur) == 0:
+            results[key] = None
+        elif isinstance(cur[0], torch.Tensor):
+            if key in ['embeds_gen_mask', 'embeds_cmp_mask', 'images', 'images_patch_length', 'patch_position', 'image_size']:
+                results[key] = torch.cat(cur, dim=0)
+            else:
+                results[key] = torch.stack(cur, dim=0)
+        else:
+            results[key] = cur
+
+    results['dataset_name'] = dataset_name
+
+    return results