File size: 4,724 Bytes
59524ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import torch
from PIL import Image
from conversation import conv_templates
from builder import load_pretrained_model
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}") |