Create inference.py

inference.py

@@ -0,0 +1,125 @@
+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 functools import partial
+import numpy as np
+
+# 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']
+
+# 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"):
+    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 (e.g., 224x224 or 336x336)
+    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)
+
+    image_tensor = image_tensor.half().cuda()
+
+    # 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()
+
+    # tokenize prompt
+    input_ids = tokenizer(prompt_input, return_tensors='pt')['input_ids'].cuda()
+
+    # region mask logic (if region is provided)
+    region_masks = None
+    if region is not None:
+        raw_w, raw_h = img.size
+        region_masks = generate_mask_for_feature(region, raw_w, raw_h).unsqueeze(0).cuda().half()
+        region_masks = [[region_masks]]  # Wrap the mask in lists as expected by the model
+
+    # 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
+        )
+        output_ids = model.generate(
+            input_ids,
+            images=image_tensor,
+            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()
+
+# We also define a task-specific inference function
+def infer_ui_task(image_path, prompt, model_path, task, region=None):
+    """
+    Handles task types: box_in_tasks, box_out_tasks, no_box_tasks.
+    """
+    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)
+    
+    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}")