Update inference.py

inference.py

@@ -1,340 +1,98 @@
-import torch
+import subprocess
+import os
+import subprocess
 from PIL import Image
-from conversation import conv_templates
-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
+import re
+import json
 
-
-# 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):
+def process_inference_results(results):
     """
-    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.
+    Process the inference results by:
+    1. Adding bounding boxes on the image based on the coordinates in 'text'.
+    2. Extracting and returning the text prompt.
+    
+    :param results: List of inference results with bounding boxes in 'text'.
+    :return: (image, text)
     """
-    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)
+    processed_images = []
+    extracted_texts = []
 
-    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
+    for result in results:
+        image_path = result['image_path']
+        img = Image.open(image_path).convert("RGB")
 
-    if is_pad:
-        scale_w = target_width / original_width
-        scale_h = target_height / original_height
+        # this no more than extracts bounding box coordinates from the 'text'
+        bbox_str = re.search(r'\[\[([0-9,\s]+)\]\]', result['text'])
+        if bbox_str:
+            bbox = [int(coord) for coord in bbox_str.group(1).split(',')]
+            x1, y1, x2, y2 = bbox
 
-        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)
+            # Draw the bounding box on the image (optional if needed later)
+            # draw = ImageDraw.Draw(img)
+            # draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
 
-        # Resize the image
-        resized_image = image.resize((new_width, new_height))
+        extracted_texts.append(result['text'])
 
-        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))
+        processed_images.append(img)
 
-    return new_image
+    return processed_images[0], extracted_texts[0]
     
-def process_anyres_image(image, processor, grid_pinpoints, image_process_func: Optional[Callable] = None):
+def inference_and_run(image_path, prompt, conv_mode="ferret_gemma_instruct", model_path="jadechoghari/Ferret-UI-Gemma2b", box=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.
+    Run the inference and capture the errors for debugging.
     """
-    if type(grid_pinpoints) is list:
-        possible_resolutions = grid_pinpoints
-    else:
-        possible_resolutions = ast.literal_eval(grid_pinpoints)
+    data_input = [{
+        "id": 0,
+        "image": os.path.basename(image_path),
+        "image_h": Image.open(image_path).height,
+        "image_w": Image.open(image_path).width,
+        "conversations": [{"from": "human", "value": f"<image>\n{prompt}"}]
+    }]
     
-    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):
-    """
-    Generates a region mask based on provided coordinates.
-    Handles both point and box input.
-    """
-    if mask is not None:
-        assert mask.shape[0] == raw_w and mask.shape[1] == raw_h
-    coor_mask = np.zeros((raw_w, raw_h))
-
-    # if it's a point (2 coordinates)
-    if len(coor) == 2:
-        span = 5  # Define the span for the point
-        x_min = max(0, coor[0] - span)
-        x_max = min(raw_w, coor[0] + span + 1)
-        y_min = max(0, coor[1] - span)
-        y_max = min(raw_h, coor[1] + span + 1)
-        coor_mask[int(x_min):int(x_max), int(y_min):int(y_max)] = 1
-        assert (coor_mask == 1).any(), f"coor: {coor}, raw_w: {raw_w}, raw_h: {raw_h}"
-
-    # if it's a box (4 coordinates)
-    elif len(coor) == 4:
-        coor_mask[coor[0]:coor[2]+1, coor[1]:coor[3]+1] = 1
-        if mask is not None:
-            coor_mask = coor_mask * mask
-
-    # convert to torch tensor and ensure it contains non-zero values
-    coor_mask = torch.from_numpy(coor_mask)
-    assert len(coor_mask.nonzero()) != 0, "Generated mask is empty :("
-
-
-    return coor_mask
-
-
-def infer_single_prompt(image_path, prompt, model_path, region=None, model_name="ferret_gemma", conv_mode="ferret_gemma_instruct", add_region_feature=False):
-    img = Image.open(image_path).convert('RGB')
-
-    # this loads the model, image processor and tokenizer
-    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
-    # define the image size required by clip
-    image_size = {"height": 336, "width": 336}
-
-    if "<image>" in prompt:
-            prompt = prompt.split('\n')[1]
-
-    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()
-
-    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()
+    if box:
+        data_input[0]["box_x1y1x2y2"] = [[box]]
 
+    with open("eval.json", "w") as json_file:
+        json.dump(data_input, json_file)
     
-
-    # 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
-        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"[{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()
-
+    print("eval.json file created successfully.")
     
-        
-    # generate model output
-    with torch.inference_mode():
-        # Use region_masks in model's forward call
-        model.orig_forward = model.forward
-        model.forward = partial(
-            model.orig_forward,
-            region_masks=region_masks
-        )
-        # explcit add of attention mask
-        output_ids = model.generate(
-            input_ids,
-            images=images,
-            max_new_tokens=1024,
-            num_beams=1,
-            region_masks=region_masks,  # pass the region mask to the model
-            image_sizes=[img.size]
-        )
-        model.forward = model.orig_forward
-
-    # we decode the output
-    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
-    return output_text.strip()
+    cmd = [
+        "python", "-m", "model_UI", 
+        "--model_path", model_path,
+        "--data_path", "eval.json", 
+        "--image_path", ".", 
+        "--answers_file", "eval_output.jsonl", 
+        "--num_beam", "1", 
+        "--max_new_tokens", "1024",
+        "--conv_mode", conv_mode
+    ]
+
+    if box:
+        cmd.extend(["--region_format", "box", "--add_region_feature"])
+
+        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+        print(f"Subprocess output:\n{result.stdout}")
+        print(f"Subprocess error (if any):\n{result.stderr}")
+        print(f"Inference completed. Output written to eval_output.jsonl")
+
+        output_folder = 'eval_output.jsonl'
+        if os.path.exists(output_folder):
+            json_files = [f for f in os.listdir(output_folder) if f.endswith(".jsonl")]
+            if json_files:
+                output_file_path = os.path.join(output_folder, json_files[0])
+                with open(output_file_path, "r") as output_file:
+                    results = [json.loads(line) for line in output_file]
+                
+                return process_inference_results(results)
+            else:
+                print("No output JSONL files found.")
+                return None, None
+        else:
+            print("Output folder not found.")
+            return None, None
 
-# 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.
-    """
-    if region is not None:
-        add_region_feature=True
-    if task in box_in_tasks and region is None:
-        raise ValueError(f"Task {task} requires a bounding box region.")
-    
-    if task in box_in_tasks:
-        print(f"Processing {task} with bounding box region.")
-        return infer_single_prompt(image_path, prompt, model_path, region, add_region_feature=add_region_feature)
-    
-    elif task in box_out_tasks:
-        print(f"Processing {task} without bounding box region.")
-        return infer_single_prompt(image_path, prompt, model_path)
-    
-    elif task in no_box_tasks:
-        print(f"Processing {task} without image or bounding box.")
-        return infer_single_prompt(image_path, prompt, model_path)
-    
-    else:
-        raise ValueError(f"Unknown task type: {task}")
+    except subprocess.CalledProcessError as e:
+        print(f"Error occurred during inference:\n{e}")
+        print(f"Subprocess output:\n{e.output}")
+        return None, None