| | from typing import Dict, List, Any |
| | from peft import AutoPeftModelForCausalLM |
| | import transformers |
| | import os |
| | import tempfile |
| | from PIL import Image, ImageDraw |
| | from io import BytesIO |
| | import base64, json |
| |
|
| | COORDINATE_PROMPT = 'In this UI screenshot, what is the position of the element corresponding to the command \"{command}\" (with point)?' |
| |
|
| | PARTITION_PROMPT = 'In this UI screenshot, what is the partition of the element corresponding to the command \"{command}\" (with quadrant number)?' |
| |
|
| | class EndpointHandler(): |
| | def __init__(self, path=""): |
| | self.model = transformers.AutoModelForCausalLM.from_pretrained( |
| | path, |
| | device_map="cuda", |
| | trust_remote_code=True, |
| | fp16=True).eval() |
| | self.tokenizer = transformers.AutoTokenizer.from_pretrained( |
| | path, |
| | cache_dir=None, |
| | model_max_length=2048, |
| | padding_side="right", |
| | use_fast=False, |
| | trust_remote_code=True, |
| | ) |
| | self.tokenizer.pad_token_id = self.tokenizer.eod_id |
| | return |
| | |
| | def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
| | """ |
| | data args: |
| | image (:obj: `str`) |
| | task (:obj: `str`) |
| | k (:obj: `str`) |
| | context (:obj: 'str') |
| | kwargs |
| | Return: |
| | A :obj:`list` | `dict`: will be serialized and returned |
| | """ |
| | |
| | with tempfile.TemporaryDirectory() as temp_dir: |
| | image = os.path.join(temp_dir, "image.png") |
| | img = Image.open(BytesIO(base64.b64decode(data["inputs"]["image"]))) |
| | img.save(image) |
| | command = data["inputs"]["task"] |
| | K = int(data["inputs"]["k"]) |
| | keep_context = bool(data["inputs"]["context"]) |
| |
|
| | print(image) |
| | print(command) |
| | print(K) |
| | print(keep_context) |
| |
|
| | images = [image] |
| | partitions = [] |
| |
|
| | for k in range(K): |
| | query = self.tokenizer.from_list_format(([{ 'image': context_image } for context_image in images] if keep_context else [{'image': image}]) + |
| | [{'text': PARTITION_PROMPT.format(command=command)}]) |
| | response, _ = self.model.chat(self.tokenizer, query=query, history=None) |
| |
|
| | partition = int(response.split(" ")[-1]) |
| | partitions.append(partition) |
| |
|
| | |
| | with Image.open(image) as img: |
| | width, height = img.size |
| | if partition == 1: |
| | img = img.crop((width // 2, 0, width, height // 2)) |
| | elif partition == 2: |
| | img = img.crop((0, 0, width // 2, height // 2)) |
| | elif partition == 3: |
| | img = img.crop((0, height // 2, width // 2, height)) |
| | elif partition == 4: |
| | img = img.crop((width // 2, height // 2, width, height)) |
| | |
| | new_path = os.path.join(temp_dir, f"partition{k}.png") |
| | img.save(new_path) |
| | image = new_path |
| | images.append(image) |
| | |
| | query = self.tokenizer.from_list_format(([{ 'image': context_image } for context_image in images] if keep_context else [{'image': image}]) + |
| | [{'text': COORDINATE_PROMPT.format(command=command)}]) |
| | response, _ = self.model.chat(self.tokenizer, query=query, history=None) |
| | print("Coordinate Response:", response) |
| |
|
| | x = float(response.split(",")[0].split("(")[1]) |
| | y = float(response.split(",")[1].split(")")[0]) |
| |
|
| | for partition in partitions[::-1]: |
| | if partition == 1: |
| | x = x/2 + 0.5 |
| | y = y/2 |
| | elif partition == 2: |
| | x = x/2 |
| | y = y/2 |
| | elif partition == 3: |
| | x = x/2 |
| | y = y/2 + 0.5 |
| | elif partition == 4: |
| | x = x/2 + 0.5 |
| | y = y/2 + 0.5 |
| | print("rescaled point:", x, y) |
| |
|
| | response = {} |
| | response['x'] = x |
| | response['y'] = y |
| | return response |
| | |
