| 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 |
| |
