Update inference.py

inference.py

@@ -1,19 +1,200 @@
 import torch
 from PIL import Image
 from conversation import conv_templates
-from builder import load_pretrained_model  # Assuming this is your custom model loader
+from builder import load_pretrained_model
 from functools import partial
+from typing import Optional, Callable
+import ast
+import math
 import numpy as np
 DEFAULT_REGION_FEA_TOKEN = "<region_fea>"
 DEFAULT_IMAGE_TOKEN = "<image>"
 DEFAULT_IM_START_TOKEN = "<im_start>"
 DEFAULT_IM_END_TOKEN = "<im_end>"
+VOCAB_IMAGE_W = 1000  # 224
+VOCAB_IMAGE_H = 1000  # 224
+IMAGE_TOKEN_INDEX = -200
+
 
 # define the task categories
 box_in_tasks = ['widgetcaptions', 'taperception', 'ocr', 'icon_recognition', 'widget_classification', 'example_0']
 box_out_tasks = ['widget_listing', 'find_text', 'find_icons', 'find_widget', 'conversation_interaction']
 no_box_tasks = ['screen2words', 'detailed_description', 'conversation_perception', 'gpt4']
 
+def get_bbox_coor(box, ratio_w, ratio_h):
+    return box[0] * ratio_w, box[1] * ratio_h, box[2] * ratio_w, box[3] * ratio_h
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    if '<image>' in prompt:
+        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+        input_ids = []
+        for i, chunk in enumerate(prompt_chunks):
+            input_ids.extend(chunk)
+            if i < len(prompt_chunks) - 1:
+                input_ids.append(image_token_index)
+    else:
+        input_ids = tokenizer(prompt).input_ids
+    # if return_tensors == 'pt':
+    #     import torch
+    #     input_ids = torch.tensor(input_ids).unsqueeze(0)
+    
+    return input_ids
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif 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 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 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 resize_and_pad_image(image, target_resolution, is_pad=False):
+    """
+    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
+
+    if is_pad:
+        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:
+        new_image = image.resize((target_width, target_height))
+
+    return new_image
+    
+def process_anyres_image(image, processor, grid_pinpoints, image_process_func: Optional[Callable] = None):
+    """
+    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:
+        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)
+
+    # FIXME: not sure if do_pad or undo_pad may affect the referring side 
+    image_padded = resize_and_pad_image(image, best_resolution, is_pad=False)
+
+    patches = divide_to_patches(image_padded, processor.crop_size['height'])
+
+    if image_process_func:
+        resized_image_h, resized_image_w = image_process_func.keywords['size']
+        image_original_resize = image.resize((resized_image_w, resized_image_h))
+        image_patches = [image_original_resize] + patches
+        image_patches = [image_process_func(image_patch)['pixel_values'][0]
+                        for image_patch in image_patches]
+    else:
+        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, return_tensors='pt')['pixel_values'][0]
+                        for image_patch in image_patches]
+
+    return torch.stack(image_patches, dim=0)
+
+
+def process_images(images, image_processor, model_cfg, image_process_func: Optional[Callable] = None):
+    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, return_tensors='pt')['pixel_values'][0]
+            new_images.append(image)
+    elif image_aspect_ratio == "anyres":
+        # image_processor(images, return_tensors='pt', do_resize=True, do_center_crop=False, size=[image_h, image_w])['pixel_values']
+        for image in images:
+            image = process_anyres_image(image, image_processor, model_cfg.image_grid_pinpoints, image_process_func=image_process_func)
+            new_images.append(image)
+    else:
+        return image_processor(images, return_tensors='pt')['pixel_values']
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
 # function to generate the mask
 def generate_mask_for_feature(coor, raw_w, raw_h, mask=None):
     """
@@ -56,42 +237,58 @@ def infer_single_prompt(image_path, prompt, model_path, region=None, model_name=
     # define the image size required by clip
     image_size = {"height": 336, "width": 336}
 
-    # process the image
-    image_tensor = image_processor.preprocess(
-        img,
-        return_tensors='pt',
-        do_resize=True,
-        do_center_crop=False,
-        size=(image_size['height'], image_size['width'])
-    )['pixel_values'][0].unsqueeze(0)
+    if "<image>" in prompt:
+            prompt = prompt.split('\n')[1]
 
-    image_tensor = image_tensor.half().cuda()
+    if model.config.mm_use_im_start_end:
+        prompt = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + prompt
+    else:
+        prompt = DEFAULT_IMAGE_TOKEN + '\n' + prompt
 
     # generate the prompt per template requirement
     conv = conv_templates[conv_mode].copy()
     conv.append_message(conv.roles[0], prompt)
     conv.append_message(conv.roles[1], None)
     prompt_input = conv.get_prompt()
-    
-    # add the special tokens
-    prompt_input = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + prompt_input
+
+    input_ids = tokenizer(prompt_input, return_tensors='pt')['input_ids'].cuda()
+
+    # raw_w, raw_h = img.size # check if shouldnt be width and height
+    raw_w = image_size["width"]
+    raw_h = image_size["height"]
+    if model.config.image_aspect_ratio == "square_nocrop":
+            image_tensor = image_processor.preprocess(img, return_tensors='pt', do_resize=True, 
+                                                  do_center_crop=False, size=[raw_h, raw_w])['pixel_values'][0]
+    elif model.config.image_aspect_ratio == "anyres":
+        image_process_func = partial(image_processor.preprocess, return_tensors='pt', do_resize=True, do_center_crop=False, size=[raw_h, raw_h])
+        image_tensor = process_images([img], image_processor, model.config, image_process_func=image_process_func)[0]
+    else:
+        image_tensor = process_images([img], image_processor, model.config)[0]
+
+    images = image_tensor.unsqueeze(0).to(torch.float16).cuda()
 
     
 
     # region mask logic (if region is provided)
     region_masks = None
     if add_region_feature and region is not None:
+        # box_in is true
         raw_w, raw_h = img.size
-        region_masks = generate_mask_for_feature(region, raw_w, raw_h).unsqueeze(0).cuda().half()
+        ratio_w = VOCAB_IMAGE_W * 1.0 / raw_w
+        ratio_h = VOCAB_IMAGE_H * 1.0 / raw_h
+        # preprocess the region
+        box_x1, box_y1, box_x2, box_y2 = region
+        box_x1_textvocab, box_y1_textvocab, box_x2_textvocab, box_y2_textvocab = get_bbox_coor(box=region, ratio_h=ratio_h, ratio_w=ratio_w)
+        region_coordinate_raw = [box_x1, box_y1, box_x2, box_y2]
+
+        region_masks = generate_mask_for_feature(region_coordinate_raw, raw_w, raw_h).unsqueeze(0).cuda().half()
         region_masks = [[region_mask_i.cuda().half() for region_mask_i in region_masks]]
-        prompt_input = prompt_input.replace("<bbox_location0>", f"[{region[0]}, {region[1]}, {region[2]}, {region[3]}] {DEFAULT_REGION_FEA_TOKEN}")
+        prompt_input = prompt_input.replace("<bbox_location0>", f"[{box_x1_textvocab}, {box_y1_textvocab}, {box_x2_textvocab}, {box_y2_textvocab}] {DEFAULT_REGION_FEA_TOKEN}")
         
     # tokenize prompt
     # input_ids = tokenizer(prompt_input, return_tensors='pt')['input_ids'].cuda()
 
-    inputs = tokenizer(prompt_input, return_tensors='pt', padding=True)
-    input_ids = inputs['input_ids'].cuda()
-    attention_mask = inputs['attention_mask'].cuda() 
+    
         
     # generate model output
     with torch.inference_mode():
@@ -104,8 +301,7 @@ def infer_single_prompt(image_path, prompt, model_path, region=None, model_name=
         # explcit add of attention mask
         output_ids = model.generate(
             input_ids,
-            images=image_tensor,
-            attention_mask=attention_mask, 
+            images=images,
             max_new_tokens=1024,
             num_beams=1,
             region_masks=region_masks,  # pass the region mask to the model
@@ -119,7 +315,6 @@ def infer_single_prompt(image_path, prompt, model_path, region=None, model_name=
 
 # We also define a task-specific inference function
 def infer_ui_task(image_path, prompt, model_path, task, region=None, add_region_feature=False):
-    # region = torch.tensor(region).cuda()
     """
     Handles task types: box_in_tasks, box_out_tasks, no_box_tasks.
     """
@@ -141,4 +336,4 @@ def infer_ui_task(image_path, prompt, model_path, task, region=None, add_region_
         return infer_single_prompt(image_path, prompt, model_path)
     
     else:
-        raise ValueError(f"Unknown task type: {task}")
+        raise ValueError(f"Unknown task type: {task}")