test input

app.py

@@ -1,7 +1,5 @@
 import gradio as gr
 import spaces
-#from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
-# from qwen_vl_utils import process_vision_info
 import torch
 import base64
 from PIL import Image, ImageDraw
@@ -17,14 +15,12 @@ from transformers import AutoModelForCausalLM
 
 
 models = {
-    #"OS-Copilot/OS-Atlas-Base-7B": Qwen2VLForConditionalGeneration.from_pretrained("OS-Copilot/OS-Atlas-Base-7B", torch_dtype="auto", device_map="auto"),
     "deepseek-ai/deepseek-vl2-tiny": AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-vl2-tiny", trust_remote_code=True),
     #"deepseek-ai/deepseek-vl2-small": AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-vl2-small", trust_remote_code=True),
     #"deepseek-ai/deepseek-vl2": AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-vl2", trust_remote_code=True)
 }
 
 processors = {
-    #"OS-Copilot/OS-Atlas-Base-7B": AutoProcessor.from_pretrained("OS-Copilot/OS-Atlas-Base-7B")
     "deepseek-ai/deepseek-vl2-tiny": DeepseekVLV2Processor.from_pretrained("deepseek-ai/deepseek-vl2-tiny",),
     #"deepseek-ai/deepseek-vl2-small": DeepseekVLV2Processor.from_pretrained("deepseek-ai/deepseek-vl2-small",),
     #"deepseek-ai/deepseek-vl2": DeepseekVLV2Processor.from_pretrained("deepseek-ai/deepseek-vl2",),
@@ -110,69 +106,21 @@ def deepseek(image, text_input, model_id):
     det_pattern = r"<\|det\|>\[\[(.+)]]<\|\/det\|>"
 
     det_content = re.search(det_pattern, answer).group(1)
-    bbox = [ int(v.strip()) / 999 for v in det_content.split(",")]
-    w, h = image.size
-    bbox = [[bbox[0] * w, bbox[1] * h, bbox[2] * w, bbox[3] * h]] 
+    if det_content is None:
+        return text_input, [], image
     
-    #scaled_boxes = rescale_bounding_boxes([bbox], image.width, image.height)
-    return text_input, bbox, draw_bounding_boxes(image, bbox)
+    bbox = [ int(v.strip()) for v in det_content.split(",")]
+    #w, h = image.size
+    #bbox = [[bbox[0] * w, bbox[1] * h, bbox[2] * w, bbox[3] * h]] 
+    
+    scaled_boxes = rescale_bounding_boxes([scaled_boxes], image.width, image.height)
+    return text_input, scaled_boxes, draw_bounding_boxes(image, scaled_boxes)
 
 
 @spaces.GPU
 def run_example(image, text_input, model_id="OS-Copilot/OS-Atlas-Base-7B"):
-
     return deepseek(image, text_input, model_id)
     
-
-def run_example_old(image, text_input, model_id="OS-Copilot/OS-Atlas-Base-7B"):
-    model = models[model_id].eval()
-    processor = processors[model_id]
-    prompt = f"In this UI screenshot, what is the position of the element corresponding to the command \"{text_input}\" (with bbox)?"
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": f"data:image;base64,{image_to_base64(image)}"},
-                {"type": "text", "text": prompt},
-            ],
-        }
-    ]
-
-    text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    )
-    inputs = inputs.to("cuda")
-
-    generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [
-        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=False, clean_up_tokenization_spaces=False
-    )
-    print(output_text)
-    text = output_text[0]
-
-    object_ref_pattern = r"<\|object_ref_start\|>(.*?)<\|object_ref_end\|>"
-    box_pattern = r"<\|box_start\|>(.*?)<\|box_end\|>"
-
-    object_ref = re.search(object_ref_pattern, text).group(1)
-    box_content = re.search(box_pattern, text).group(1)
-
-    boxes = [tuple(map(int, pair.strip("()").split(','))) for pair in box_content.split("),(")]
-    boxes = [[boxes[0][0], boxes[0][1], boxes[1][0], boxes[1][1]]]
-
-    scaled_boxes = rescale_bounding_boxes(boxes, image.width, image.height)
-    return object_ref, scaled_boxes, draw_bounding_boxes(image, scaled_boxes)
-
 css = """
   #output {
     height: 500px;