@@ -65,8 +63,6 @@ def process(
# parsed_content_list = str(parsed_content_list)
return image, str(parsed_content_list)
-
-
with gr.Blocks() as demo:
gr.Markdown(MARKDOWN)
with gr.Row():
diff --git a/imgs/gradioicon.png b/imgs/gradioicon.png
new file mode 100644
index 0000000..d6b62c8
Binary files /dev/null and b/imgs/gradioicon.png differ
diff --git a/imgs/header_bar.png b/imgs/header_bar.png
new file mode 100644
index 0000000..148fbee
Binary files /dev/null and b/imgs/header_bar.png differ
diff --git a/imgs/header_bar_thin.png b/imgs/header_bar_thin.png
new file mode 100644
index 0000000..e4be8ba
Binary files /dev/null and b/imgs/header_bar_thin.png differ
diff --git a/imgs/omniboxicon.png b/imgs/omniboxicon.png
new file mode 100644
index 0000000..de05f63
Binary files /dev/null and b/imgs/omniboxicon.png differ
diff --git a/imgs/omniparsericon.png b/imgs/omniparsericon.png
new file mode 100644
index 0000000..877427a
Binary files /dev/null and b/imgs/omniparsericon.png differ
diff --git a/imgs/som_overlaid_omni.png b/imgs/som_overlaid_omni.png
new file mode 100644
index 0000000..49271bc
Binary files /dev/null and b/imgs/som_overlaid_omni.png differ
diff --git a/omniparser.py b/omniparser.py
deleted file mode 100644
index 634ae9f..0000000
--- a/omniparser.py
+++ /dev/null
@@ -1,60 +0,0 @@
-from utils import get_som_labeled_img, check_ocr_box, get_yolo_model
-import torch
-from ultralytics import YOLO
-from PIL import Image
-from typing import Dict, Tuple, List
-import io
-import base64
-
-
-config = {
- 'som_model_path': 'finetuned_icon_detect.pt',
- 'device': 'cpu',
- 'caption_model_path': 'Salesforce/blip2-opt-2.7b',
- 'draw_bbox_config': {
- 'text_scale': 0.8,
- 'text_thickness': 2,
- 'text_padding': 3,
- 'thickness': 3,
- },
- 'BOX_TRESHOLD': 0.05
-}
-
-
-class Omniparser(object):
- def __init__(self, config: Dict):
- self.config = config
-
- self.som_model = get_yolo_model(model_path=config['som_model_path'])
- # self.caption_model_processor = get_caption_model_processor(config['caption_model_path'], device=cofig['device'])
- # self.caption_model_processor['model'].to(torch.float32)
-
- def parse(self, image_path: str):
- print('Parsing image:', image_path)
- ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.9})
- text, ocr_bbox = ocr_bbox_rslt
-
- draw_bbox_config = self.config['draw_bbox_config']
- BOX_TRESHOLD = self.config['BOX_TRESHOLD']
- dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_path, self.som_model, BOX_TRESHOLD = BOX_TRESHOLD, output_coord_in_ratio=False, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=None, ocr_text=text,use_local_semantics=False)
-
- image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
- # formating output
- return_list = [{'from': 'omniparser', 'shape': {'x':coord[0], 'y':coord[1], 'width':coord[2], 'height':coord[3]},
- 'text': parsed_content_list[i].split(': ')[1], 'type':'text'} for i, (k, coord) in enumerate(label_coordinates.items()) if i < len(parsed_content_list)]
- return_list.extend(
- [{'from': 'omniparser', 'shape': {'x':coord[0], 'y':coord[1], 'width':coord[2], 'height':coord[3]},
- 'text': 'None', 'type':'icon'} for i, (k, coord) in enumerate(label_coordinates.items()) if i >= len(parsed_content_list)]
- )
-
- return [image, return_list]
-
-parser = Omniparser(config)
-image_path = 'examples/pc_1.png'
-
-# time the parser
-import time
-s = time.time()
-image, parsed_content_list = parser.parse(image_path)
-device = config['device']
-print(f'Time taken for Omniparser on {device}:', time.time() - s)
diff --git a/omnitool/gradio/.gitignore b/omnitool/gradio/.gitignore
new file mode 100644
index 0000000..c036379
--- /dev/null
+++ b/omnitool/gradio/.gitignore
@@ -0,0 +1 @@
+tmp/
\ No newline at end of file
diff --git a/demo/__init__.py b/omnitool/gradio/__init__.py
similarity index 100%
rename from demo/__init__.py
rename to omnitool/gradio/__init__.py
diff --git a/demo/gui_agent/anthropic_agent.py b/omnitool/gradio/agent/anthropic_agent.py
similarity index 73%
rename from demo/gui_agent/anthropic_agent.py
rename to omnitool/gradio/agent/anthropic_agent.py
index 09d857e..b1c744e 100644
--- a/demo/gui_agent/anthropic_agent.py
+++ b/omnitool/gradio/agent/anthropic_agent.py
@@ -1,207 +1,162 @@
-"""
-Agentic sampling loop that calls the Anthropic API and local implenmentation of anthropic-defined computer use tools.
-"""
-import asyncio
-import platform
-from collections.abc import Callable
-from datetime import datetime
-from enum import StrEnum
-from typing import Any, cast
-
-from anthropic import Anthropic, AnthropicBedrock, AnthropicVertex, APIResponse
-from anthropic.types import (
- ToolResultBlockParam,
-)
-from anthropic.types.beta import (
- BetaContentBlock,
- BetaContentBlockParam,
- BetaImageBlockParam,
- BetaMessage,
- BetaMessageParam,
- BetaTextBlockParam,
- BetaToolResultBlockParam,
-)
-from anthropic.types import TextBlock
-from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock
-
-from tools import BashTool, ComputerTool, EditTool, ToolCollection, ToolResult
-
-from PIL import Image
-from io import BytesIO
-import gradio as gr
-from typing import Dict
-
-
-BETA_FLAG = "computer-use-2024-10-22"
-
-
-class APIProvider(StrEnum):
- ANTHROPIC = "anthropic"
- BEDROCK = "bedrock"
- VERTEX = "vertex"
-
-
-PROVIDER_TO_DEFAULT_MODEL_NAME: dict[APIProvider, str] = {
- APIProvider.ANTHROPIC: "claude-3-5-sonnet-20241022",
- APIProvider.BEDROCK: "anthropic.claude-3-5-sonnet-20241022-v2:0",
- APIProvider.VERTEX: "claude-3-5-sonnet-v2@20241022",
-}
-
-
-# Check OS
-
-SYSTEM_PROMPT = f"""
-* You are utilizing a Windows system with internet access.
-* The current date is {datetime.today().strftime('%A, %B %d, %Y')}.
-
-"""
-
-
-class AnthropicActor:
- def __init__(
- self,
- model: str,
- provider: APIProvider,
- system_prompt_suffix: str,
- api_key: str,
- api_response_callback: Callable[[APIResponse[BetaMessage]], None],
- max_tokens: int = 4096,
- only_n_most_recent_images: int | None = None,
- selected_screen: int = 0,
- print_usage: bool = True,
- ):
- self.model = model
- self.provider = provider
- self.system_prompt_suffix = system_prompt_suffix
- self.api_key = api_key
- self.api_response_callback = api_response_callback
- self.max_tokens = max_tokens
- self.only_n_most_recent_images = only_n_most_recent_images
- self.selected_screen = selected_screen
-
- self.tool_collection = ToolCollection(
- ComputerTool(selected_screen=selected_screen),
- BashTool(),
- EditTool(),
- )
-
- self.system = (
- f"{SYSTEM_PROMPT}{' ' + system_prompt_suffix if system_prompt_suffix else ''}"
- )
-
- self.total_token_usage = 0
- self.total_cost = 0
- self.print_usage = print_usage
-
- # Instantiate the appropriate API client based on the provider
- if provider == APIProvider.ANTHROPIC:
- self.client = Anthropic(api_key=api_key)
- elif provider == APIProvider.VERTEX:
- self.client = AnthropicVertex()
- elif provider == APIProvider.BEDROCK:
- self.client = AnthropicBedrock()
-
- def __call__(
- self,
- *,
- messages: list[BetaMessageParam]
- ):
- """
- Generate a response given history messages.
- """
- if self.only_n_most_recent_images:
- _maybe_filter_to_n_most_recent_images(messages, self.only_n_most_recent_images)
-
- # Call the API synchronously
- raw_response = self.client.beta.messages.with_raw_response.create(
- max_tokens=self.max_tokens,
- messages=messages,
- model=self.model,
- system=self.system,
- tools=self.tool_collection.to_params(),
- betas=["computer-use-2024-10-22"],
- )
-
- self.api_response_callback(cast(APIResponse[BetaMessage], raw_response))
-
- response = raw_response.parse()
- print(f"AnthropicActor response: {response}")
-
- self.total_token_usage += response.usage.input_tokens + response.usage.output_tokens
- self.total_cost += (response.usage.input_tokens * 3 / 1000000 + response.usage.output_tokens * 15 / 1000000)
-
- if self.print_usage:
- print(f"Claude total token usage so far: {self.total_token_usage}, total cost so far: $USD{self.total_cost}")
-
- return response
-
-
-def _maybe_filter_to_n_most_recent_images(
- messages: list[BetaMessageParam],
- images_to_keep: int,
- min_removal_threshold: int = 10,
-):
- """
- With the assumption that images are screenshots that are of diminishing value as
- the conversation progresses, remove all but the final `images_to_keep` tool_result
- images in place, with a chunk of min_removal_threshold to reduce the amount we
- break the implicit prompt cache.
- """
- if images_to_keep is None:
- return messages
-
- tool_result_blocks = cast(
- list[ToolResultBlockParam],
- [
- item
- for message in messages
- for item in (
- message["content"] if isinstance(message["content"], list) else []
- )
- if isinstance(item, dict) and item.get("type") == "tool_result"
- ],
- )
-
- total_images = sum(
- 1
- for tool_result in tool_result_blocks
- for content in tool_result.get("content", [])
- if isinstance(content, dict) and content.get("type") == "image"
- )
-
- images_to_remove = total_images - images_to_keep
- # for better cache behavior, we want to remove in chunks
- images_to_remove -= images_to_remove % min_removal_threshold
-
- for tool_result in tool_result_blocks:
- if isinstance(tool_result.get("content"), list):
- new_content = []
- for content in tool_result.get("content", []):
- if isinstance(content, dict) and content.get("type") == "image":
- if images_to_remove > 0:
- images_to_remove -= 1
- continue
- new_content.append(content)
- tool_result["content"] = new_content
-
-
-
-if __name__ == "__main__":
- pass
- # client = Anthropic(api_key="")
- # response = client.beta.messages.with_raw_response.create(
- # max_tokens=4096,
- # model="claude-3-5-sonnet-20241022",
- # system=SYSTEM_PROMPT,
- # # tools=ToolCollection(
- # # ComputerTool(selected_screen=0),
- # # BashTool(),
- # # EditTool(),
- # # ).to_params(),
- # betas=["computer-use-2024-10-22"],
- # messages=[
- # {"role": "user", "content": "click on (199, 199)."}
- # ],
- # )
-
- # print(f"AnthropicActor response: {response.parse().usage.input_tokens+response.parse().usage.output_tokens}")
\ No newline at end of file
+"""
+Agentic sampling loop that calls the Anthropic API and local implenmentation of anthropic-defined computer use tools.
+"""
+import asyncio
+import platform
+from collections.abc import Callable
+from datetime import datetime
+from enum import StrEnum
+from typing import Any, cast
+
+from anthropic import Anthropic, AnthropicBedrock, AnthropicVertex, APIResponse
+from anthropic.types import (
+ ToolResultBlockParam,
+)
+from anthropic.types.beta import (
+ BetaContentBlock,
+ BetaContentBlockParam,
+ BetaImageBlockParam,
+ BetaMessage,
+ BetaMessageParam,
+ BetaTextBlockParam,
+ BetaToolResultBlockParam,
+)
+from anthropic.types import TextBlock
+from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock
+
+from tools import ComputerTool, ToolCollection, ToolResult
+
+from PIL import Image
+from io import BytesIO
+import gradio as gr
+from typing import Dict
+
+BETA_FLAG = "computer-use-2024-10-22"
+
+class APIProvider(StrEnum):
+ ANTHROPIC = "anthropic"
+ BEDROCK = "bedrock"
+ VERTEX = "vertex"
+
+SYSTEM_PROMPT = f"""
+* You are utilizing a Windows system with internet access.
+* The current date is {datetime.today().strftime('%A, %B %d, %Y')}.
+
+"""
+
+class AnthropicActor:
+ def __init__(
+ self,
+ model: str,
+ provider: APIProvider,
+ api_key: str,
+ api_response_callback: Callable[[APIResponse[BetaMessage]], None],
+ max_tokens: int = 4096,
+ only_n_most_recent_images: int | None = None,
+ print_usage: bool = True,
+ ):
+ self.model = model
+ self.provider = provider
+ self.api_key = api_key
+ self.api_response_callback = api_response_callback
+ self.max_tokens = max_tokens
+ self.only_n_most_recent_images = only_n_most_recent_images
+
+ self.tool_collection = ToolCollection(ComputerTool())
+
+ self.system = SYSTEM_PROMPT
+
+ self.total_token_usage = 0
+ self.total_cost = 0
+ self.print_usage = print_usage
+
+ # Instantiate the appropriate API client based on the provider
+ if provider == APIProvider.ANTHROPIC:
+ self.client = Anthropic(api_key=api_key)
+ elif provider == APIProvider.VERTEX:
+ self.client = AnthropicVertex()
+ elif provider == APIProvider.BEDROCK:
+ self.client = AnthropicBedrock()
+
+ def __call__(
+ self,
+ *,
+ messages: list[BetaMessageParam]
+ ):
+ """
+ Generate a response given history messages.
+ """
+ if self.only_n_most_recent_images:
+ _maybe_filter_to_n_most_recent_images(messages, self.only_n_most_recent_images)
+
+ # Call the API synchronously
+ raw_response = self.client.beta.messages.with_raw_response.create(
+ max_tokens=self.max_tokens,
+ messages=messages,
+ model=self.model,
+ system=self.system,
+ tools=self.tool_collection.to_params(),
+ betas=["computer-use-2024-10-22"],
+ )
+
+ self.api_response_callback(cast(APIResponse[BetaMessage], raw_response))
+
+ response = raw_response.parse()
+ print(f"AnthropicActor response: {response}")
+
+ self.total_token_usage += response.usage.input_tokens + response.usage.output_tokens
+ self.total_cost += (response.usage.input_tokens * 3 / 1000000 + response.usage.output_tokens * 15 / 1000000)
+
+ if self.print_usage:
+ print(f"Claude total token usage so far: {self.total_token_usage}, total cost so far: $USD{self.total_cost}")
+
+ return response
+
+
+def _maybe_filter_to_n_most_recent_images(
+ messages: list[BetaMessageParam],
+ images_to_keep: int,
+ min_removal_threshold: int = 10,
+):
+ """
+ With the assumption that images are screenshots that are of diminishing value as
+ the conversation progresses, remove all but the final `images_to_keep` tool_result
+ images in place, with a chunk of min_removal_threshold to reduce the amount we
+ break the implicit prompt cache.
+ """
+ if images_to_keep is None:
+ return messages
+
+ tool_result_blocks = cast(
+ list[ToolResultBlockParam],
+ [
+ item
+ for message in messages
+ for item in (
+ message["content"] if isinstance(message["content"], list) else []
+ )
+ if isinstance(item, dict) and item.get("type") == "tool_result"
+ ],
+ )
+
+ total_images = sum(
+ 1
+ for tool_result in tool_result_blocks
+ for content in tool_result.get("content", [])
+ if isinstance(content, dict) and content.get("type") == "image"
+ )
+
+ images_to_remove = total_images - images_to_keep
+ # for better cache behavior, we want to remove in chunks
+ images_to_remove -= images_to_remove % min_removal_threshold
+
+ for tool_result in tool_result_blocks:
+ if isinstance(tool_result.get("content"), list):
+ new_content = []
+ for content in tool_result.get("content", []):
+ if isinstance(content, dict) and content.get("type") == "image":
+ if images_to_remove > 0:
+ images_to_remove -= 1
+ continue
+ new_content.append(content)
+ tool_result["content"] = new_content
\ No newline at end of file
diff --git a/omnitool/gradio/agent/llm_utils/groqclient.py b/omnitool/gradio/agent/llm_utils/groqclient.py
new file mode 100644
index 0000000..c31a502
--- /dev/null
+++ b/omnitool/gradio/agent/llm_utils/groqclient.py
@@ -0,0 +1,59 @@
+from groq import Groq
+import os
+from .utils import is_image_path
+
+def run_groq_interleaved(messages: list, system: str, model_name: str, api_key: str, max_tokens=256, temperature=0.6):
+ """
+ Run a chat completion through Groq's API, ignoring any images in the messages.
+ """
+ api_key = api_key or os.environ.get("GROQ_API_KEY")
+ if not api_key:
+ raise ValueError("GROQ_API_KEY is not set")
+
+ client = Groq(api_key=api_key)
+ # avoid using system messages for R1
+ final_messages = [{"role": "user", "content": system}]
+
+ if isinstance(messages, list):
+ for item in messages:
+ if isinstance(item, dict):
+ # For dict items, concatenate all text content, ignoring images
+ text_contents = []
+ for cnt in item["content"]:
+ if isinstance(cnt, str):
+ if not is_image_path(cnt): # Skip image paths
+ text_contents.append(cnt)
+ else:
+ text_contents.append(str(cnt))
+
+ if text_contents: # Only add if there's text content
+ message = {"role": "user", "content": " ".join(text_contents)}
+ final_messages.append(message)
+ else: # str
+ message = {"role": "user", "content": item}
+ final_messages.append(message)
+
+ elif isinstance(messages, str):
+ final_messages.append({"role": "user", "content": messages})
+
+ try:
+ completion = client.chat.completions.create(
+ model="deepseek-r1-distill-llama-70b",
+ messages=final_messages,
+ temperature=0.6,
+ max_completion_tokens=max_tokens,
+ top_p=0.95,
+ stream=False,
+ reasoning_format="raw"
+ )
+
+ response = completion.choices[0].message.content
+ final_answer = response.split('\n')[-1] if '' in response else response
+ final_answer = final_answer.replace("
", "")
+ token_usage = completion.usage.total_tokens
+
+ return final_answer, token_usage
+ except Exception as e:
+ print(f"Error in interleaved Groq: {e}")
+
+ return str(e), 0
\ No newline at end of file
diff --git a/omnitool/gradio/agent/llm_utils/oaiclient.py b/omnitool/gradio/agent/llm_utils/oaiclient.py
new file mode 100644
index 0000000..ad42110
--- /dev/null
+++ b/omnitool/gradio/agent/llm_utils/oaiclient.py
@@ -0,0 +1,62 @@
+import os
+import logging
+import base64
+import requests
+from .utils import is_image_path, encode_image
+
+def run_oai_interleaved(messages: list, system: str, model_name: str, api_key: str, max_tokens=256, temperature=0, provider_base_url: str = "https://api.openai.com/v1"):
+ headers = {"Content-Type": "application/json",
+ "Authorization": f"Bearer {api_key}"}
+ final_messages = [{"role": "system", "content": system}]
+
+ if type(messages) == list:
+ for item in messages:
+ contents = []
+ if isinstance(item, dict):
+ for cnt in item["content"]:
+ if isinstance(cnt, str):
+ if is_image_path(cnt) and 'o3-mini' not in model_name:
+ # 03 mini does not support images
+ base64_image = encode_image(cnt)
+ content = {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+ else:
+ content = {"type": "text", "text": cnt}
+ else:
+ # in this case it is a text block from anthropic
+ content = {"type": "text", "text": str(cnt)}
+
+ contents.append(content)
+
+ message = {"role": 'user', "content": contents}
+ else: # str
+ contents.append({"type": "text", "text": item})
+ message = {"role": "user", "content": contents}
+
+ final_messages.append(message)
+
+
+ elif isinstance(messages, str):
+ final_messages = [{"role": "user", "content": messages}]
+
+ payload = {
+ "model": model_name,
+ "messages": final_messages,
+ }
+ if 'o1' in model_name or 'o3-mini' in model_name:
+ payload['reasoning_effort'] = 'low'
+ payload['max_completion_tokens'] = max_tokens
+ else:
+ payload['max_tokens'] = max_tokens
+
+ response = requests.post(
+ f"{provider_base_url}/chat/completions", headers=headers, json=payload
+ )
+
+
+ try:
+ text = response.json()['choices'][0]['message']['content']
+ token_usage = int(response.json()['usage']['total_tokens'])
+ return text, token_usage
+ except Exception as e:
+ print(f"Error in interleaved openAI: {e}. This may due to your invalid API key. Please check the response: {response.json()} ")
+ return response.json()
\ No newline at end of file
diff --git a/omnitool/gradio/agent/llm_utils/omniparserclient.py b/omnitool/gradio/agent/llm_utils/omniparserclient.py
new file mode 100644
index 0000000..e90ddef
--- /dev/null
+++ b/omnitool/gradio/agent/llm_utils/omniparserclient.py
@@ -0,0 +1,44 @@
+import requests
+import base64
+from pathlib import Path
+from tools.screen_capture import get_screenshot
+from agent.llm_utils.utils import encode_image
+
+OUTPUT_DIR = "./tmp/outputs"
+
+class OmniParserClient:
+ def __init__(self,
+ url: str) -> None:
+ self.url = url
+
+ def __call__(self,):
+ screenshot, screenshot_path = get_screenshot()
+ screenshot_path = str(screenshot_path)
+ image_base64 = encode_image(screenshot_path)
+ response = requests.post(self.url, json={"base64_image": image_base64})
+ response_json = response.json()
+ print('omniparser latency:', response_json['latency'])
+
+ som_image_data = base64.b64decode(response_json['som_image_base64'])
+ screenshot_path_uuid = Path(screenshot_path).stem.replace("screenshot_", "")
+ som_screenshot_path = f"{OUTPUT_DIR}/screenshot_som_{screenshot_path_uuid}.png"
+ with open(som_screenshot_path, "wb") as f:
+ f.write(som_image_data)
+
+ response_json['width'] = screenshot.size[0]
+ response_json['height'] = screenshot.size[1]
+ response_json['original_screenshot_base64'] = image_base64
+ response_json['screenshot_uuid'] = screenshot_path_uuid
+ response_json = self.reformat_messages(response_json)
+ return response_json
+
+ def reformat_messages(self, response_json: dict):
+ screen_info = ""
+ for idx, element in enumerate(response_json["parsed_content_list"]):
+ element['idx'] = idx
+ if element['type'] == 'text':
+ screen_info += f'ID: {idx}, Text: {element["content"]}\n'
+ elif element['type'] == 'icon':
+ screen_info += f'ID: {idx}, Icon: {element["content"]}\n'
+ response_json['screen_info'] = screen_info
+ return response_json
\ No newline at end of file
diff --git a/omnitool/gradio/agent/llm_utils/utils.py b/omnitool/gradio/agent/llm_utils/utils.py
new file mode 100644
index 0000000..12ab36e
--- /dev/null
+++ b/omnitool/gradio/agent/llm_utils/utils.py
@@ -0,0 +1,13 @@
+import base64
+
+def is_image_path(text):
+ image_extensions = (".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".tif")
+ if text.endswith(image_extensions):
+ return True
+ else:
+ return False
+
+def encode_image(image_path):
+ """Encode image file to base64."""
+ with open(image_path, "rb") as image_file:
+ return base64.b64encode(image_file.read()).decode("utf-8")
\ No newline at end of file
diff --git a/omnitool/gradio/agent/vlm_agent.py b/omnitool/gradio/agent/vlm_agent.py
new file mode 100644
index 0000000..7599c39
--- /dev/null
+++ b/omnitool/gradio/agent/vlm_agent.py
@@ -0,0 +1,338 @@
+import json
+from collections.abc import Callable
+from typing import cast, Callable
+import uuid
+from PIL import Image, ImageDraw
+import base64
+from io import BytesIO
+
+from anthropic import APIResponse
+from anthropic.types import ToolResultBlockParam
+from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock, BetaMessageParam, BetaUsage
+
+from agent.llm_utils.oaiclient import run_oai_interleaved
+from agent.llm_utils.groqclient import run_groq_interleaved
+from agent.llm_utils.utils import is_image_path
+import time
+import re
+
+OUTPUT_DIR = "./tmp/outputs"
+
+def extract_data(input_string, data_type):
+ # Regular expression to extract content starting from '```python' until the end if there are no closing backticks
+ pattern = f"```{data_type}" + r"(.*?)(```|$)"
+ # Extract content
+ # re.DOTALL allows '.' to match newlines as well
+ matches = re.findall(pattern, input_string, re.DOTALL)
+ # Return the first match if exists, trimming whitespace and ignoring potential closing backticks
+ return matches[0][0].strip() if matches else input_string
+
+class VLMAgent:
+ def __init__(
+ self,
+ model: str,
+ provider: str,
+ api_key: str,
+ output_callback: Callable,
+ api_response_callback: Callable,
+ max_tokens: int = 4096,
+ only_n_most_recent_images: int | None = None,
+ print_usage: bool = True,
+ ):
+ if model == "omniparser + gpt-4o":
+ self.model = "gpt-4o-2024-11-20"
+ elif model == "omniparser + R1":
+ self.model = "deepseek-r1-distill-llama-70b"
+ elif model == "omniparser + qwen2.5vl":
+ self.model = "qwen2.5-vl-72b-instruct"
+ elif model == "omniparser + o1":
+ self.model = "o1"
+ elif model == "omniparser + o3-mini":
+ self.model = "o3-mini"
+ else:
+ raise ValueError(f"Model {model} not supported")
+
+
+ self.provider = provider
+ self.api_key = api_key
+ self.api_response_callback = api_response_callback
+ self.max_tokens = max_tokens
+ self.only_n_most_recent_images = only_n_most_recent_images
+ self.output_callback = output_callback
+
+ self.print_usage = print_usage
+ self.total_token_usage = 0
+ self.total_cost = 0
+ self.step_count = 0
+
+ self.system = ''
+
+ def __call__(self, messages: list, parsed_screen: list[str, list, dict]):
+ self.step_count += 1
+ image_base64 = parsed_screen['original_screenshot_base64']
+ latency_omniparser = parsed_screen['latency']
+ self.output_callback(f'-- Step {self.step_count}: --', sender="bot")
+ screen_info = str(parsed_screen['screen_info'])
+ screenshot_uuid = parsed_screen['screenshot_uuid']
+ screen_width, screen_height = parsed_screen['width'], parsed_screen['height']
+
+ boxids_and_labels = parsed_screen["screen_info"]
+ system = self._get_system_prompt(boxids_and_labels)
+
+ # drop looping actions msg, byte image etc
+ planner_messages = messages
+ _remove_som_images(planner_messages)
+ _maybe_filter_to_n_most_recent_images(planner_messages, self.only_n_most_recent_images)
+
+ if isinstance(planner_messages[-1], dict):
+ if not isinstance(planner_messages[-1]["content"], list):
+ planner_messages[-1]["content"] = [planner_messages[-1]["content"]]
+ planner_messages[-1]["content"].append(f"{OUTPUT_DIR}/screenshot_{screenshot_uuid}.png")
+ planner_messages[-1]["content"].append(f"{OUTPUT_DIR}/screenshot_som_{screenshot_uuid}.png")
+
+ start = time.time()
+ if "gpt" in self.model or "o1" in self.model or "o3-mini" in self.model:
+ vlm_response, token_usage = run_oai_interleaved(
+ messages=planner_messages,
+ system=system,
+ model_name=self.model,
+ api_key=self.api_key,
+ max_tokens=self.max_tokens,
+ provider_base_url="https://api.openai.com/v1",
+ temperature=0,
+ )
+ print(f"oai token usage: {token_usage}")
+ self.total_token_usage += token_usage
+ if 'gpt' in self.model:
+ self.total_cost += (token_usage * 2.5 / 1000000) # https://openai.com/api/pricing/
+ elif 'o1' in self.model:
+ self.total_cost += (token_usage * 15 / 1000000) # https://openai.com/api/pricing/
+ elif 'o3-mini' in self.model:
+ self.total_cost += (token_usage * 1.1 / 1000000) # https://openai.com/api/pricing/
+ elif "r1" in self.model:
+ vlm_response, token_usage = run_groq_interleaved(
+ messages=planner_messages,
+ system=system,
+ model_name=self.model,
+ api_key=self.api_key,
+ max_tokens=self.max_tokens,
+ )
+ print(f"groq token usage: {token_usage}")
+ self.total_token_usage += token_usage
+ self.total_cost += (token_usage * 0.99 / 1000000)
+ elif "qwen" in self.model:
+ vlm_response, token_usage = run_oai_interleaved(
+ messages=planner_messages,
+ system=system,
+ model_name=self.model,
+ api_key=self.api_key,
+ max_tokens=min(2048, self.max_tokens),
+ provider_base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
+ temperature=0,
+ )
+ print(f"qwen token usage: {token_usage}")
+ self.total_token_usage += token_usage
+ self.total_cost += (token_usage * 2.2 / 1000000) # https://help.aliyun.com/zh/model-studio/getting-started/models?spm=a2c4g.11186623.0.0.74b04823CGnPv7#fe96cfb1a422a
+ else:
+ raise ValueError(f"Model {self.model} not supported")
+ latency_vlm = time.time() - start
+ self.output_callback(f"LLM: {latency_vlm:.2f}s, OmniParser: {latency_omniparser:.2f}s", sender="bot")
+
+ print(f"{vlm_response}")
+
+ if self.print_usage:
+ print(f"Total token so far: {self.total_token_usage}. Total cost so far: $USD{self.total_cost:.5f}")
+
+ vlm_response_json = extract_data(vlm_response, "json")
+ vlm_response_json = json.loads(vlm_response_json)
+
+ img_to_show_base64 = parsed_screen["som_image_base64"]
+ if "Box ID" in vlm_response_json:
+ bbox = parsed_screen["parsed_content_list"][int(vlm_response_json["Box ID"])]["bbox"]
+ vlm_response_json["box_centroid_coordinate"] = [int((bbox[0] + bbox[2]) / 2 * screen_width), int((bbox[1] + bbox[3]) / 2 * screen_height)]
+ img_to_show_data = base64.b64decode(img_to_show_base64)
+ img_to_show = Image.open(BytesIO(img_to_show_data))
+
+ draw = ImageDraw.Draw(img_to_show)
+ x, y = vlm_response_json["box_centroid_coordinate"]
+ radius = 10
+ draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill='red')
+ draw.ellipse((x - radius*3, y - radius*3, x + radius*3, y + radius*3), fill=None, outline='red', width=2)
+
+ buffered = BytesIO()
+ img_to_show.save(buffered, format="PNG")
+ img_to_show_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+ self.output_callback(f'
', sender="bot")
+ self.output_callback(
+ f'
'
+ f' Parsed Screen elemetns by OmniParser
'
+ f' {screen_info}'
+ f' ',
+ sender="bot"
+ )
+ vlm_plan_str = ""
+ for key, value in vlm_response_json.items():
+ if key == "Reasoning":
+ vlm_plan_str += f'{value}'
+ else:
+ vlm_plan_str += f'\n{key}: {value}'
+
+ # construct the response so that anthropicExcutor can execute the tool
+ response_content = [BetaTextBlock(text=vlm_plan_str, type='text')]
+ if 'box_centroid_coordinate' in vlm_response_json:
+ move_cursor_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}',
+ input={'action': 'mouse_move', 'coordinate': vlm_response_json["box_centroid_coordinate"]},
+ name='computer', type='tool_use')
+ response_content.append(move_cursor_block)
+
+ if vlm_response_json["Next Action"] == "None":
+ print("Task paused/completed.")
+ elif vlm_response_json["Next Action"] == "type":
+ sim_content_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}',
+ input={'action': vlm_response_json["Next Action"], 'text': vlm_response_json["value"]},
+ name='computer', type='tool_use')
+ response_content.append(sim_content_block)
+ else:
+ sim_content_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}',
+ input={'action': vlm_response_json["Next Action"]},
+ name='computer', type='tool_use')
+ response_content.append(sim_content_block)
+ response_message = BetaMessage(id=f'toolu_{uuid.uuid4()}', content=response_content, model='', role='assistant', type='message', stop_reason='tool_use', usage=BetaUsage(input_tokens=0, output_tokens=0))
+ return response_message, vlm_response_json
+
+ def _api_response_callback(self, response: APIResponse):
+ self.api_response_callback(response)
+
+ def _get_system_prompt(self, screen_info: str = ""):
+ main_section = f"""
+You are using a Windows device.
+You are able to use a mouse and keyboard to interact with the computer based on the given task and screenshot.
+You can only interact with the desktop GUI (no terminal or application menu access).
+
+You may be given some history plan and actions, this is the response from the previous loop.
+You should carefully consider your plan base on the task, screenshot, and history actions.
+
+Here is the list of all detected bounding boxes by IDs on the screen and their description:{screen_info}
+
+Your available "Next Action" only include:
+- type: types a string of text.
+- left_click: move mouse to box id and left clicks.
+- right_click: move mouse to box id and right clicks.
+- double_click: move mouse to box id and double clicks.
+- hover: move mouse to box id.
+- scroll_up: scrolls the screen up.
+- scroll_down: scrolls the screen down.
+- wait: waits for 1 second for the device to load or respond.
+
+Based on the visual information from the screenshot image and the detected bounding boxes, please determine the next action, the Box ID you should operate on (if action is not 'type', 'hover', 'scroll_up', 'scroll_down', 'wait'), and the value (if the action is 'type') in order to complete the task.
+
+Output format:
+```json
+{{
+ "Reasoning": str, # describe what is in the current screen, taking into account the history, then describe your step-by-step thoughts on how to achieve the task, choose one action from available actions at a time.
+ "Next Action": "action_type, action description" | "None" # one action at a time, describe it in short and precisely.
+ "Box ID": n,
+ "value": "xxx" # only provide value field if the action is type, else don't include value key
+}}
+```
+
+One Example:
+```json
+{{
+ "Reasoning": "The current screen shows google result of amazon, in previous action I have searched amazon on google. Then I need to click on the first search results to go to amazon.com.",
+ "Next Action": "left_click",
+ "Box ID": m
+}}
+```
+
+Another Example:
+```json
+{{
+ "Reasoning": "The current screen shows the front page of amazon. There is no previous action. Therefore I need to type "Apple watch" in the search bar.",
+ "Next Action": "type",
+ "Box ID": n,
+ "value": "Apple watch"
+}}
+```
+
+IMPORTANT NOTES:
+1. You should only give a single action at a time.
+
+"""
+ thinking_model = "r1" in self.model
+ if not thinking_model:
+ main_section += """
+2. You should give an analysis to the current screen, and reflect on what has been done by looking at the history, then describe your step-by-step thoughts on how to achieve the task.
+
+"""
+ else:
+ main_section += """
+2. In
XML tags give an analysis to the current screen, and reflect on what has been done by looking at the history, then describe your step-by-step thoughts on how to achieve the task. In )})
',
- sender="bot")
- yield None
-
- output_callback(f"{colorful_text_vlm} sending action to {colorful_text_showui}:\n{next_action}", sender="bot")
-
- actor_response = actor(messages=next_action)
- yield actor_response
-
- for message, tool_result_content in executor(actor_response, messages):
- time.sleep(1)
- yield message
-
- # since showui executor has no feedback for now, we use "actor_response" to represent its response
- # update messages for the next loop
- messages.append({"role": "user",
- "content": ["History plan:\n" + str(json.loads(vlm_response)) +
- "\nHistory actions:\n" + str(actor_response["content"])]})
- print(f"End of loop. Messages: {str(messages)[:100000]}. Total cost: $USD{planner.total_cost:.5f}")
\ No newline at end of file
diff --git a/demo/omniparser_agent/vlm_agent.py b/demo/omniparser_agent/vlm_agent.py
deleted file mode 100644
index 53ba46a..0000000
--- a/demo/omniparser_agent/vlm_agent.py
+++ /dev/null
@@ -1,417 +0,0 @@
-import json
-import asyncio
-import platform
-from collections.abc import Callable
-from datetime import datetime
-from enum import StrEnum
-from typing import Any, cast, Dict, Callable
-import uuid
-import requests
-from PIL import Image, ImageDraw
-import base64
-from io import BytesIO
-
-from anthropic import Anthropic, AnthropicBedrock, AnthropicVertex, APIResponse
-from anthropic.types import TextBlock, ToolResultBlockParam
-from anthropic.types.beta import BetaMessage, BetaTextBlock, BetaToolUseBlock, BetaMessageParam, BetaUsage
-
-from tools.screen_capture import get_screenshot
-from gui_agent.llm_utils.oai import run_oai_interleaved, encode_image
-from gui_agent.llm_utils.qwen import run_qwen
-from gui_agent.llm_utils.llm_utils import extract_data
-from tools.colorful_text import colorful_text_showui, colorful_text_vlm
-
-
-SYSTEM_PROMPT = f"""
',
- sender="bot")
- self.output_callback(f'Set of Marks Screenshot for {colorful_text_vlm}:\n
', sender="bot")
-
-
- if isinstance(planner_messages[-1], dict):
- if not isinstance(planner_messages[-1]["content"], list):
- planner_messages[-1]["content"] = [planner_messages[-1]["content"]]
- planner_messages[-1]["content"].append(screenshot_path)
- planner_messages[-1]["content"].append(som_screenshot_path)
-
- print(f"Sending messages to VLMPlanner : {planner_messages}")
-
- if "gpt" in self.model:
- vlm_response, token_usage = run_oai_interleaved(
- messages=planner_messages,
- system=system,
- llm=self.model,
- api_key=self.api_key,
- max_tokens=self.max_tokens,
- temperature=0,
- )
- print(f"oai token usage: {token_usage}")
- self.total_token_usage += token_usage
- self.total_cost += (token_usage * 0.15 / 1000000) # https://openai.com/api/pricing/
-
- elif "qwen" in self.model:
- vlm_response, token_usage = run_qwen(
- messages=planner_messages,
- system=system,
- llm=self.model,
- api_key=self.api_key,
- max_tokens=self.max_tokens,
- temperature=0,
- )
- print(f"qwen token usage: {token_usage}")
- self.total_token_usage += token_usage
- self.total_cost += (token_usage * 0.02 / 7.25 / 1000) # 1USD=7.25CNY, https://help.aliyun.com/zh/dashscope/developer-reference/tongyi-qianwen-vl-plus-api
- elif "phi" in self.model:
- pass # TODO
- else:
- raise ValueError(f"Model {self.model} not supported")
-
- print(f"VLMPlanner response: {vlm_response}")
-
- if self.print_usage:
- print(f"VLMPlanner total token usage so far: {self.total_token_usage}. Total cost so far: $USD{self.total_cost:.5f}")
-
- vlm_response_json = extract_data(vlm_response, "json")
- vlm_response_json = json.loads(vlm_response_json)
-
- # map "box_id" to "idx" in parsed_screen, and output the xy coordinate of bbox
- # TODO add try except for the case when "box_id" is not in the response
- # if 'Box ID' in vlm_response_json:
- try:
- bbox = parsed_screen["parsed_content_list"][int(vlm_response_json["Box ID"])]["bbox"]
- vlm_response_json["coordinate"] = [int((bbox[0] + bbox[2]) / 2 * screen_width), int((bbox[1] + bbox[3]) / 2 * screen_height)]
- # draw a circle on the screenshot image to indicate the action
- self.draw_action(vlm_response_json, image_base64)
- except:
- print("No Box ID in the response.")
-
-
- # vlm_plan_str = '\n'.join([f'{key}: {value}' for key, value in json.loads(response).items()])
- vlm_plan_str = ""
- for key, value in vlm_response_json.items():
- if key == "Reasoning":
- vlm_plan_str += f'{value}'
- else:
- vlm_plan_str += f'\n{key}: {value}'
-
- # self.output_callback(f"{colorful_text_vlm}:\n{vlm_plan_str}", sender="bot")
-
- # construct the response so that anthropicExcutor can execute the tool
- analysis = BetaTextBlock(text=vlm_plan_str, type='text')
- if 'coordinate' in vlm_response_json:
- move_cursor_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}',
- input={'action': 'mouse_move', 'coordinate': vlm_response_json["coordinate"]},
- name='computer', type='tool_use')
- response_content = [analysis, move_cursor_block]
- else:
- response_content = [analysis]
- if vlm_response_json["Next Action"] == "type":
- click_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}', input={'action': 'left_click'}, name='computer', type='tool_use')
- sim_content_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}',
- input={'action': vlm_response_json["Next Action"], 'text': vlm_response_json["value"]},
- name='computer', type='tool_use')
- response_content.extend([click_block, sim_content_block])
- elif vlm_response_json["Next Action"] == "None":
- print("Task paused/completed.")
- else:
- sim_content_block = BetaToolUseBlock(id=f'toolu_{uuid.uuid4()}',
- input={'action': vlm_response_json["Next Action"]},
- name='computer', type='tool_use')
- response_content.append(sim_content_block)
-
- response = BetaMessage(id=f'toolu_{uuid.uuid4()}', content=response_content, model='', role='assistant', type='message', stop_reason='tool_use', usage=BetaUsage(input_tokens=0, output_tokens=0))
- return response, vlm_response_json
-
-
- def _api_response_callback(self, response: APIResponse):
- self.api_response_callback(response)
-
-
- def reformat_messages(self, messages: list):
- pass
-
- def _get_system_prompt(self, screen_info: str = ""):
- datetime_str = datetime.now().strftime("%A, %B %d, %Y")
- os_name = platform.system()
- return f"""
-You are using an {os_name} device.
-You are able to use a mouse and keyboard to interact with the computer based on the given task and screenshot.
-You can only interact with the desktop GUI (no terminal or application menu access).
-
-You may be given some history plan and actions, this is the response from the previous loop.
-You should carefully consider your plan base on the task, screenshot, and history actions.
-
-Here is the list of all detected bounding boxes by IDs on the screen and their description:{screen_info}
-
-Your available "Next Action" only include:
-- type: type a string of text.
-- left_click: Describe the ui element to be clicked.
-- enter: Press an enter key.
-- escape: Press an ESCAPE key.
-- hover: Describe the ui element to be hovered.
-- scroll: Scroll the screen, you must specify up or down.
-- press: Describe the ui element to be pressed.
-
-Based on the visual information from the screenshot image and the detected bounding boxes, please determine the next action, the Box ID you should operate on, and the value (if the action is 'type') in order to complete the task.
-
-Output format:
-```json
-{{
- "Reasoning": str, # describe what is in the current screen, taking into account the history, then describe your step-by-step thoughts on how to achieve the task, choose one action from available actions at a time.
- "Next Action": "action_type, action description" | "None" # one action at a time, describe it in short and precisely.
- 'Box ID': n,
- 'value': "xxx" # if the action is type, you should provide the text to type.
-}}
-```
-
-One Example:
-```json
-{{
- "Reasoning": "The current screen shows google result of amazon, in previous action I have searched amazon on google. Then I need to click on the first search results to go to amazon.com.",
- "Next Action": "left_click",
- 'Box ID': m,
-}}
-```
-
-Another Example:
-```json
-{{
- "Reasoning": "The current screen shows the front page of amazon. There is no previous action. Therefore I need to type "Apple watch" in the search bar.",
- "Next Action": "type",
- 'Box ID': n,
- 'value': "Apple watch"
-}}
-```
-
-IMPORTANT NOTES:
-1. You should only give a single action at a time.
-2. You should give an analysis to the current screen, and reflect on what has been done by looking at the history, then describe your step-by-step thoughts on how to achieve the task.
-3. Attach the next action prediction in the "Next Action".
-4. You should not include other actions, such as keyboard shortcuts.
-5. When the task is completed, you should say "Next Action": "None" in the json field.
-"""
- def draw_action(self, vlm_response_json, image_base64):
- # draw a circle using the coordinate in parsed_screen['som_image_base64']
- image_data = base64.b64decode(image_base64)
- image = Image.open(BytesIO(image_data))
-
- draw = ImageDraw.Draw(image)
- x, y = vlm_response_json["coordinate"]
- radius = 10
- draw.ellipse((x - radius, y - radius, x + radius, y + radius), outline='red', width=3)
- buffered = BytesIO()
- image.save('demo.png')
- image.save(buffered, format="PNG")
- image_with_circle_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
- self.output_callback(f'Action performed on the Screenshot (red circle), for {colorful_text_vlm}:\n
', sender="bot")
-
-
-def _keep_latest_images(messages):
- for i in range(len(messages)-1):
- if isinstance(messages[i]["content"], list):
- for cnt in messages[i]["content"]:
- if isinstance(cnt, str):
- if cnt.endswith((".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".tif")):
- messages[i]["content"].remove(cnt)
- return messages
-
-
-def _maybe_filter_to_n_most_recent_images(
- messages: list[BetaMessageParam],
- images_to_keep: int,
- min_removal_threshold: int = 10,
-):
- """
- With the assumption that images are screenshots that are of diminishing value as
- the conversation progresses, remove all but the final `images_to_keep` tool_result
- images in place, with a chunk of min_removal_threshold to reduce the amount we
- break the implicit prompt cache.
- """
- if images_to_keep is None:
- return messages
-
- tool_result_blocks = cast(
- list[ToolResultBlockParam],
- [
- item
- for message in messages
- for item in (
- message["content"] if isinstance(message["content"], list) else []
- )
- if isinstance(item, dict) and item.get("type") == "tool_result"
- ],
- )
-
- total_images = sum(
- 1
- for tool_result in tool_result_blocks
- for content in tool_result.get("content", [])
- if isinstance(content, dict) and content.get("type") == "image"
- )
-
- images_to_remove = total_images - images_to_keep
- # for better cache behavior, we want to remove in chunks
- images_to_remove -= images_to_remove % min_removal_threshold
-
- for tool_result in tool_result_blocks:
- if isinstance(tool_result.get("content"), list):
- new_content = []
- for content in tool_result.get("content", []):
- if isinstance(content, dict) and content.get("type") == "image":
- if images_to_remove > 0:
- images_to_remove -= 1
- continue
- new_content.append(content)
- tool_result["content"] = new_content
-
-
-def _message_filter_callback(messages):
- filtered_list = []
- try:
- for msg in messages:
- if msg.get('role') in ['user']:
- if not isinstance(msg["content"], list):
- msg["content"] = [msg["content"]]
- if isinstance(msg["content"][0], TextBlock):
- filtered_list.append(str(msg["content"][0].text)) # User message
- elif isinstance(msg["content"][0], str):
- filtered_list.append(msg["content"][0]) # User message
- else:
- print("[_message_filter_callback]: drop message", msg)
- continue
-
- # elif msg.get('role') in ['assistant']:
- # if isinstance(msg["content"][0], TextBlock):
- # msg["content"][0] = str(msg["content"][0].text)
- # elif isinstance(msg["content"][0], BetaTextBlock):
- # msg["content"][0] = str(msg["content"][0].text)
- # elif isinstance(msg["content"][0], BetaToolUseBlock):
- # msg["content"][0] = str(msg['content'][0].input)
- # elif isinstance(msg["content"][0], Dict) and msg["content"][0]["content"][-1]["type"] == "image":
- # msg["content"][0] = f'
'
- # else:
- # print("[_message_filter_callback]: drop message", msg)
- # continue
- # filtered_list.append(msg["content"][0]) # User message
-
- else:
- print("[_message_filter_callback]: drop message", msg)
- continue
-
- except Exception as e:
- print("[_message_filter_callback]: error", e)
-
- return filtered_list
\ No newline at end of file
diff --git a/demo/remote_request.py b/demo/remote_request.py
deleted file mode 100644
index d0bc1c3..0000000
--- a/demo/remote_request.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# uvicorn remote_request:app --host 0.0.0.0 --port 8000 --reload
-
-import sys
-import os
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from utils import get_som_labeled_img, check_ocr_box, get_caption_model_processor, get_yolo_model
-import torch
-from PIL import Image
-from typing import Dict, Tuple, List
-import base64
-
-
-config = {
- 'som_model_path': '../weights/icon_detect_v1_5/model_v1_5.pt',
- 'device': 'cpu',
- 'caption_model_name': 'florence2',
- 'caption_model_path': '../weights/icon_caption_florence',
- 'BOX_TRESHOLD': 0.05
-}
-
-
-class Omniparser(object):
- def __init__(self, config: Dict):
- self.config = config
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
- self.som_model = get_yolo_model(model_path=config['som_model_path'])
- self.caption_model_processor = get_caption_model_processor(model_name=config['caption_model_name'], model_name_or_path=config['caption_model_path'], device=device)
- print('Omniparser initialized!!!')
-
- def parse(self, image_base64: str):
- image_path = '../imgs/demo_image.jpg'
- with open(image_path, "wb") as fh:
- fh.write(base64.b64decode(image_base64))
- print('Parsing image:', image_path)
-
- image = Image.open(image_path)
- print('image size:', image.size)
-
- box_overlay_ratio = max(image.size) / 3200
- draw_bbox_config = {
- 'text_scale': 0.8 * box_overlay_ratio,
- 'text_thickness': max(int(2 * box_overlay_ratio), 1),
- 'text_padding': max(int(3 * box_overlay_ratio), 1),
- 'thickness': max(int(3 * box_overlay_ratio), 1),
- }
- BOX_TRESHOLD = config['BOX_TRESHOLD']
-
- ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.8}, use_paddleocr=True)
- text, ocr_bbox = ocr_bbox_rslt
- dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_path, self.som_model, BOX_TRESHOLD = BOX_TRESHOLD, output_coord_in_ratio=True, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=self.caption_model_processor, ocr_text=text,use_local_semantics=True, iou_threshold=0.7, scale_img=False, batch_size=128)
- with open('../imgs/demo_image_som.jpg', "wb") as fh:
- fh.write(base64.b64decode(dino_labled_img))
-
- return dino_labled_img, parsed_content_list
-
-
-from fastapi import FastAPI
-from pydantic import BaseModel
-
-app = FastAPI()
-
-class Item(BaseModel):
- base64_image: str
- prompt: str
-
-Omniparser = Omniparser(config)
-
-@app.post("/send_text/")
-async def send_text(item: Item):
- print('start parsing...')
- import time
- start = time.time()
- dino_labled_img, parsed_content_list = Omniparser.parse(item.base64_image)
- latency = time.time() - start
- print('time:', latency)
- return {"som_image_base64": dino_labled_img, "parsed_content_list": parsed_content_list, 'latency': latency}
\ No newline at end of file
diff --git a/demo/tmp/outputs/screenshot.png b/demo/tmp/outputs/screenshot.png
deleted file mode 100644
index eca7672..0000000
Binary files a/demo/tmp/outputs/screenshot.png and /dev/null differ
diff --git a/demo/tools/bash.py b/demo/tools/bash.py
deleted file mode 100644
index 56eb912..0000000
--- a/demo/tools/bash.py
+++ /dev/null
@@ -1,136 +0,0 @@
-import asyncio
-import os
-from typing import ClassVar, Literal
-
-from anthropic.types.beta import BetaToolBash20241022Param
-
-from .base import BaseAnthropicTool, CLIResult, ToolError, ToolResult
-
-
-class _BashSession:
- """A session of a bash shell."""
-
- _started: bool
- _process: asyncio.subprocess.Process
-
- command: str = "/bin/bash"
- _output_delay: float = 0.2 # seconds
- _timeout: float = 120.0 # seconds
- _sentinel: str = "<