init

app.py

@@ -1,7 +1,177 @@
 import gradio as gr
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+from transformers.image_utils import load_image
+import torch
+import spaces
+import re
+import base64
+from PIL import Image, ImageDraw
+from io import BytesIO
+from img_utils import smart_resize
 
-def greet(name):
-    return "Hello " + name + "!!"
+SYS_PROMPT = """You are a helpful assistant.
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{{"type": "function", "function": {{"name": "computer_use", "description": "Use a mouse and keyboard to interact with a computer, and take screenshots.\n* This is an interface to a desktop GUI. You do not have access to a terminal or applications menu. You must click on desktop icons to start applications.\n* Some applications may take time to start or process actions, so you may need to wait and take successive screenshots to see the results of your actions. E.g. if you click on Firefox and a window doesn't open, try wait and taking another screenshot.\n* The screen's resolution is {width}x{height}.\n* Whenever you intend to move the cursor to click on an element like an icon, you should consult a screenshot to determine the coordinates of the element before moving the cursor.\n* If you tried clicking on a program or link but it failed to load, even after waiting, try adjusting your cursor position so that the tip of the cursor visually falls on the element that you want to click.\n* Make sure to click any buttons, links, icons, etc with the cursor tip in the center of the element. Don't click boxes on their edges unless asked.", "parameters": {{"properties": {{"action": {{"description": "The action to perform. The available actions are:\n* `key`: Performs key down presses on the arguments passed in order, then performs key releases in reverse order.\n* `type`: Type a string of text on the keyboard.\n* `mouse_move`: Move the cursor to a specified (x, y) pixel coordinate on the screen.\n* `left_click`: Click the left mouse button.\n* `left_click_drag`: Click and drag the cursor to a specified (x, y) pixel coordinate on the screen.\n* `right_click`: Click the right mouse button.\n* `middle_click`: Click the middle mouse button.\n* `double_click`: Double-click the left mouse button.\n* `scroll`: Performs a scroll of the mouse scroll wheel.\n* `wait`: Wait specified seconds for the change to happen.\n* `terminate`: Terminate the current task and report its completion status.", "enum": ["key", "type", "mouse_move", "left_click", "left_click_drag", "right_click", "middle_click", "double_click", "scroll", "wait", "terminate"], "type": "string"}}, "keys": {{"description": "Required only by `action=key`.", "type": "array"}}, "text": {{"description": "Required only by `action=type`.", "type": "string"}}, "coordinate": {{"description": "(x, y): The x (pixels from the left edge) and y (pixels from the top edge) coordinates to move the mouse to. Required only by `action=mouse_move`, `action=left_click_drag`, `action=left_click`, `action=right_click`, `action=double_click`.", "type": "array"}}, "pixels": {{"description": "The amount of scrolling to perform. Positive values scroll up, negative values scroll down. Required only by `action=scroll`.", "type": "number"}}, "time": {{"description": "The seconds to wait. Required only by `action=wait`.", "type": "number"}}, "status": {{"description": "The status of the task. Required only by `action=terminate`.", "type": "string", "enum": ["success", "failure"]}}}}, "required": ["action"], "type": "object"}}}}}}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": <function-name>, "arguments": <args-json-object>}}
+</tool_call>
+"""
+
+MODEL_ID = "xlangai/Jedi-7B-1080p"
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16
+).to("cuda").eval()
+
+def image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
+
+def draw_bounding_boxes(image, bounding_boxes, outline_color="red", line_width=2):
+    draw = ImageDraw.Draw(image)
+    for box in bounding_boxes:
+        xmin, ymin, xmax, ymax = box
+        draw.rectangle([xmin, ymin, xmax, ymax], outline=outline_color, width=line_width)
+    return image
+
+def rescale_bounding_boxes(bounding_boxes, original_width, original_height, scaled_width=1000, scaled_height=1000):
+    x_scale = original_width / scaled_width
+    y_scale = original_height / scaled_height
+    rescaled_boxes = []
+    for box in bounding_boxes:
+        xmin, ymin, xmax, ymax = box
+        rescaled_box = [
+            xmin * x_scale,
+            ymin * y_scale,
+            xmax * x_scale,
+            ymax * y_scale
+        ]
+        rescaled_boxes.append(rescaled_box)
+    return rescaled_boxes
+
+def parse_coordinates(response):
+    try:
+        action = json.loads(
+            response.split("<tool_call>\n")[1].split("\n</tool_call>")[0]
+        )
+        action_name = action["name"]
+        action_type = action["arguments"]["action"]
+
+        print(f"action_name: {action_name}, action_type: {action_type}")
+
+        if action_type == "wait":
+            return [-1, -1, -1, -1]
+
+        action_args = action["arguments"]["coordinate"]
+
+        # Return as [x1, y1, x2, y2] format for bounding box
+        return [action_args[0], action_args[1], action_args[0], action_args[1]]
+    except Exception as e:
+        print(f"Error parsing coordinates: {e}\nResponse: {response}")
+        return None
+
+@spaces.GPU
+def run_inference(image, text_input):
+    if image is None:
+        return "Please upload an image", "", None
+    
+    if not text_input:
+        text_input = "Describe this image in detail"
+
+    width, height = smart_resize(image.height, image.width)
+
+    messages = [
+        {
+            "role": "system",
+            "content": SYS_PROMPT.format(width=width, height=height)
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": text_input},
+            ],
+        }
+    ]
+    
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(
+        text=[text],
+        images=[image],
+        return_tensors="pt",
+        padding=True,
+    ).to("cuda")
+    
+    generated_ids = model.generate(**inputs, max_new_tokens=512)
+    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
+    )[0]
+    
+    try:
+        # Parse coordinates using the function from the reference code
+        coordinates = parse_coordinates(full_response)
+        
+        if coordinates is None:
+            return f"Failed to parse coordinates from response: {output_text}", "", image
+        
+        if coordinates == [-1, -1, -1, -1]:
+            return "Task deemed infeasible by the model (wait action returned)", "", image
+        
+        # Create a bounding box with a small area around the click point
+        click_x, click_y = coordinates[0], coordinates[1]
+        box_size = 20  # 20px box around the click point
+        box = [
+            click_x - box_size/2, 
+            click_y - box_size/2, 
+            click_x + box_size/2, 
+            click_y + box_size/2
+        ]
+        
+        # Draw the bounding box on the image
+        annotated_image = draw_bounding_boxes(image.copy(), [box])
+        
+        return f"Click coordinates: ({click_x}, {click_y})", str(coordinates), annotated_image
+
+css = """
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(
+    """
+    # Qwen2.5-VL Image Element Locator
+    Upload an image and provide a description of an element to locate.
+    """)
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(label="Input Image", type="pil")
+            text_input = gr.Textbox(label="User Prompt (e.g., 'select search textfield')")
+            submit_btn = gr.Button(value="Submit")
+        with gr.Column():
+            model_output_text = gr.Textbox(label="Model Output Text", lines=5)
+            model_output_box = gr.Textbox(label="Coordinates", lines=2)
+            annotated_image = gr.Image(label="Annotated Image")
+
+    submit_btn.click(run_inference, [input_img, text_input], [model_output_text, model_output_box, annotated_image])
+
+demo.launch(debug=True)

img_utils.py

@@ -0,0 +1,233 @@
+import math
+from typing import List, Union, Dict, Any
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """返回最接近 number 的且能被 factor 整除的整数"""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """返回大于等于 number 的且能被 factor 整除的整数"""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """返回小于等于 number 的且能被 factor 整除的整数"""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(height, width, factor=28, min_pixels=56 * 56, max_pixels=14 * 14 * 4 * 1280, max_long_side=8192):
+    """缩放后图片满足以下条件:
+    1. 长宽能被 factor 整除
+    2. pixels 总数被限制在 [min_pixels, max_pixels] 内
+    3. 最长边限制在 max_long_side 内
+    4. 保证其长宽比基本不变
+    """
+    if height < 2 or width < 2:
+        raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+    elif max(height, width) / min(height, width) > 200:
+        raise ValueError(f"absolute aspect ratio must be smaller than 100, got {height} / {width}")
+
+    if max(height, width) > max_long_side:
+        beta = max(height, width) / max_long_side
+        height, width = int(height / beta), int(width / beta)
+
+    h_bar = round_by_factor(height, factor)
+    w_bar = round_by_factor(width, factor)
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def update_image_size_(image_ele: dict, min_tokens=1, max_tokens=12800, merge_base=2, patch_size=14):
+    """根据 min_tokens, max_tokens 更新 image_ele 的尺寸信息
+
+    Args:
+        image_ele (dict):
+            - image_ele["image"]: str 图片路径
+            - image_ele["height"]: int 图片原始高度
+            - image_ele["width"]: int 图片原始宽度
+
+    Returns:
+        更新后的 image_ele, 新增如下 key-value pair
+        dict:
+            - image_ele["resized_height"]: int 输入到模型的真实高度
+            - image_ele["resized_width"]: int 输入到模型的真实宽度
+            - image_ele["seq_len"]: int 输入到模型所占的序列长度
+    """
+    height, width = image_ele["height"], image_ele["width"]
+    pixels_per_token = patch_size * patch_size * merge_base * merge_base
+    resized_height, resized_width = smart_resize(
+        height,
+        width,
+        factor=merge_base * patch_size,
+        min_pixels=pixels_per_token * min_tokens,
+        max_pixels=pixels_per_token * max_tokens,
+        max_long_side=50000,
+    )
+    image_ele.update(
+        {
+            "resized_height": resized_height,
+            "resized_width": resized_width,
+            "seq_len": resized_height * resized_width // pixels_per_token + 2,
+        }
+    )
+    return image_ele
+
+
+def _convert_bbox_format_from_abs_origin(bbox, image_ele: dict, *, tgt_format: str):
+    x1, y1, x2, y2 = bbox
+    if tgt_format == "abs_origin":
+        new_bbox = [int(x1), int(y1), int(x2), int(y2)]
+    elif tgt_format == "abs_resized":
+        new_bbox = [
+            int(x1 / image_ele["width"] * image_ele["resized_width"]),
+            int(y1 / image_ele["height"] * image_ele["resized_height"]),
+            int(x2 / image_ele["width"] * image_ele["resized_width"]),
+            int(y2 / image_ele["height"] * image_ele["resized_height"]),
+        ]
+    elif tgt_format == "qwen-vl":
+        new_bbox = [
+            int(x1 / image_ele["width"] * 999),
+            int(y1 / image_ele["height"] * 999),
+            int(x2 / image_ele["width"] * 999),
+            int(y2 / image_ele["height"] * 999),
+        ]
+    elif tgt_format == "rel":
+        new_bbox = [
+            float(x1 / image_ele["width"]),
+            float(y1 / image_ele["height"]),
+            float(x2 / image_ele["width"]),
+            float(y2 / image_ele["height"]),
+        ]
+    elif tgt_format == "molmo":
+        new_bbox = [
+            round(x1 / image_ele["width"] * 100, ndigits=1),
+            round(y1 / image_ele["height"] * 100, ndigits=1),
+            round(x2 / image_ele["width"] * 100, ndigits=1),
+            round(y2 / image_ele["height"] * 100, ndigits=1),
+        ]
+    else:
+        assert False, f"Unknown tgt_format: {tgt_format}"
+    return new_bbox
+
+
+def _convert_bbox_format_to_abs_origin(bbox, image_ele: dict, *, src_format: str):
+    x1, y1, x2, y2 = bbox
+    if src_format == "abs_origin":
+        new_bbox = [int(x1), int(y1), int(x2), int(y2)]
+    elif src_format == "abs_resized":
+        new_bbox = [
+            int(x1 / image_ele["resized_width"] * image_ele["width"]),
+            int(y1 / image_ele["resized_height"] * image_ele["height"]),
+            int(x2 / image_ele["resized_width"] * image_ele["width"]),
+            int(y2 / image_ele["resized_height"] * image_ele["height"]),
+        ]
+    elif src_format == "qwen-vl":
+        new_bbox = [
+            int(x1 / 999 * image_ele["width"]),
+            int(y1 / 999 * image_ele["height"]),
+            int(x2 / 999 * image_ele["width"]),
+            int(y2 / 999 * image_ele["height"]),
+        ]
+    elif src_format == "rel":
+        new_bbox = [
+            int(x1 * image_ele["width"]),
+            int(y1 * image_ele["height"]),
+            int(x2 * image_ele["width"]),
+            int(y2 * image_ele["height"]),
+        ]
+    elif src_format == "molmo":
+        new_bbox = [
+            int(x1 / 100 * image_ele["width"]),
+            int(y1 / 100 * image_ele["height"]),
+            int(x2 / 100 * image_ele["width"]),
+            int(y2 / 100 * image_ele["height"]),
+        ]
+    else:
+        assert False, f"Unknown src_format: {src_format}"
+    return new_bbox
+
+
+def convert_bbox_format(bbox, image_ele: dict, *, src_format: str, tgt_format: str):
+    bbox_abs_origin = _convert_bbox_format_to_abs_origin(bbox, image_ele, src_format=src_format)
+    bbox_tgt_format = _convert_bbox_format_from_abs_origin(bbox_abs_origin, image_ele, tgt_format=tgt_format)
+    return bbox_tgt_format
+
+
+def _convert_point_format_from_abs_origin(point, image_ele: dict, *, tgt_format: str):
+    x, y = point
+    if tgt_format == "abs_origin":
+        new_point = [int(x), int(y)]
+    elif tgt_format == "abs_resized":
+        new_point = [
+            int(x / image_ele["width"] * image_ele["resized_width"]),
+            int(y / image_ele["height"] * image_ele["resized_height"]),
+        ]
+    elif tgt_format == "qwen-vl":
+        new_point = [
+            int(x / image_ele["width"] * 999),
+            int(y / image_ele["height"] * 999),
+        ]
+    elif tgt_format == "rel":
+        new_point = [
+            float(x / image_ele["width"]),
+            float(y / image_ele["height"]),
+        ]
+    elif tgt_format == "molmo":
+        new_point = [
+            round(x / image_ele["width"] * 100, ndigits=1),
+            round(y / image_ele["height"] * 100, ndigits=1),
+        ]
+    else:
+        assert False, f"Unknown tgt_format: {tgt_format}"
+    return new_point
+
+
+def _convert_point_format_to_abs_origin(point, image_ele: dict, *, src_format: str):
+    x, y = point
+    if src_format == "abs_origin":
+        new_point = [int(x), int(y)]
+    elif src_format == "abs_resized":
+        new_point = [
+            int(x / image_ele["resized_width"] * image_ele["width"]),
+            int(y / image_ele["resized_height"] * image_ele["height"]),
+        ]
+    elif src_format == "qwen-vl":
+        new_point = [
+            int(x / 999 * image_ele["width"]),
+            int(y / 999 * image_ele["height"]),
+        ]
+    elif src_format == "rel":
+        new_point = [
+            int(x * image_ele["width"]),
+            int(y * image_ele["height"]),
+        ]
+    elif src_format == "molmo":
+        new_point = [
+            int(x / 100 * image_ele["width"]),
+            int(y / 100 * image_ele["height"]),
+        ]
+    else:
+        assert False, f"Unknown src_format: {src_format}"
+    return new_point
+
+
+def convert_point_format(point, image_ele: dict, *, src_format: str, tgt_format: str):
+    point_abs_origin = _convert_point_format_to_abs_origin(point, image_ele, src_format=src_format)
+    point_tgt_format = _convert_point_format_from_abs_origin(point_abs_origin, image_ele, tgt_format=tgt_format)
+    return point_tgt_format
+
+
+__all__ = [
+    "update_image_size_",
+    "convert_bbox_format",
+    "convert_point_format",
+]