Upload 2 files

handler.py

@@ -0,0 +1,115 @@
+from typing import Dict, List, Any
+from peft import AutoPeftModelForCausalLM
+import transformers
+import os
+import tempfile
+from PIL import Image, ImageDraw
+
+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 = AutoPeftModelForCausalLM.from_pretrained(
+                    path,
+                    device_map="cuda",
+                    trust_remote_code=True,
+                    fp16=True).eval()
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+                    path,
+                    cache_dir=None,
+                    model_max_length=2048,
+                    padding_side="right",
+                    use_fast=False,
+                    trust_remote_code=True,
+                )
+        tokenizer.pad_token_id = tokenizer.eod_id
+        return
+ 
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            image (:obj: `PIL.Image`)
+            task (:obj: `str`)
+            k (:obj: `str`)
+            context (:obj: 'str')
+            kwargs
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # open temp directory
+        with tempfile.TemporaryDirectory() as temp_dir:
+            image = os.path.join(temp_dir, "image.jpg")
+            data["image"].save(image)
+            img = Image.open(image)
+            command = data["task"]
+            K = int(data["k"])
+            keep_context = bool(data["context"])
+
+            print(image)
+            print(command)
+            print(k)
+            print(keep_context)
+
+            images = [image]
+            partitions = []
+
+            try: 
+                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)
+
+                    # get cropped image of the 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)
+
+            except:
+                print("Invalid response")
+                print()
+
+            response = {}
+            response['x'] = x
+            response['y'] = y
+            return response
+            

requirements.txt

@@ -0,0 +1,3 @@
+peft
+transformers
+PIL