File size: 5,682 Bytes
59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac d347b1f 59524ac 6e3e9da 59524ac d347b1f 59524ac d347b1f |
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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
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
DEFAULT_REGION_FEA_TOKEN = "<region_fea>"
DEFAULT_IMAGE_TOKEN = "<image>"
DEFAULT_IM_START_TOKEN = "<im_start>"
DEFAULT_IM_END_TOKEN = "<im_end>"
# 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_llama", conv_mode="ferret_llama_3", 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}
# 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()
# add the special tokens
prompt_input = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + prompt_input
# region mask logic (if region is provided)
region_masks = None
if add_region_feature and 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_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}")
# 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():
# 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=image_tensor,
attention_mask=attention_mask,
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, 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}")
|