initial setup

app.py

@@ -1,7 +1,143 @@
 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
+from io import BytesIO
+import re
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+models = {
+    "OS-Copilot/OS-Atlas-Base-7B": Qwen2VLForConditionalGeneration.from_pretrained("OS-Copilot/OS-Atlas-Base-7B", torch_dtype="auto", device_map="auto"),
+}
+
+processors = {
+    "OS-Copilot/OS-Atlas-Base-7B": AutoProcessor.from_pretrained("OS-Copilot/OS-Atlas-Base-7B")
+}
+
+
+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
+
+
+@spaces.GPU
+def run_example(image, text_input, system_prompt, model_id="OS-Copilot/OS-Atlas-Base-7B"):
+    model = models[model_id].eval()
+    processor = processors[model_id]
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": f"data:image;base64,{image_to_base64(image)}"},
+                {"type": "text", "text": text_input},
+            ],
+        }
+    ]
+
+    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)
+
+    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 output_text, boxes, draw_bounding_boxes(image, scaled_boxes)
+
+css = """
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
+default_system_prompt = "You are a helpfull assistant to detect objects in images. When asked to detect elements based on a description you return bounding boxes for all elements in the form of [xmin, ymin, xmax, ymax] whith the values beeing scaled to 1000 by 1000 pixels. When there are more than one result, answer with a list of bounding boxes in the form of [[xmin, ymin, xmax, ymax], [xmin, ymin, xmax, ymax], ...]."
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(
+    """
+    # OS-Atlas Demo
+    """)
+    with gr.Tab(label="OS-Atlas Input"):
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Image(label="Input Image", type="pil")
+                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="OS-Copilot/OS-Atlas-Base-7B")
+                system_prompt = gr.Textbox(label="System Prompt", value=default_system_prompt)
+                text_input = gr.Textbox(label="User Prompt")
+                submit_btn = gr.Button(value="Submit")
+            with gr.Column():
+                model_output_text = gr.Textbox(label="Model Output Text")
+                parsed_boxes = gr.Textbox(label="Parsed Boxes")
+                annotated_image = gr.Image(label="Annotated Image")
+
+        gr.Examples(
+            examples=[
+                ["assets/image1.jpg", "detect goats", default_system_prompt],
+                ["assets/image2.jpg", "detect blue button", default_system_prompt],
+                ["assets/image3.jpg", "detect person on bike", default_system_prompt],
+            ],
+            inputs=[input_img, text_input, system_prompt],
+            outputs=[model_output_text, parsed_boxes, annotated_image],
+            fn=run_example,
+            cache_examples=True,
+            label="Try examples"
+        )
+
+        submit_btn.click(run_example, [input_img, text_input, system_prompt, model_selector], [model_output_text, parsed_boxes, annotated_image])
+
+demo.launch(debug=True)

requirements.txt

@@ -0,0 +1,8 @@
+numpy==1.24.4
+Pillow==10.3.0
+Requests==2.31.0
+torch
+torchvision
+transformers
+accelerate==0.30.0
+qwen-vl-utils