Mobile-Agent / app.py
阳渠
Mobile-Agent-v2
1e96bca
import io
import os
import shutil
import base64
import gradio as gr
from PIL import Image, ImageDraw
from MobileAgent.text_localization import ocr
from MobileAgent.icon_localization import det
from MobileAgent.local_server import mobile_agent_infer
from modelscope import snapshot_download
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
chatbot_css = """
<style>
.chat-container {
display: flex;
flex-direction: column;
overflow-y: auto;
max-height: 630px;
margin: 10px;
}
.user-message, .bot-message {
margin: 5px;
padding: 10px;
border-radius: 10px;
}
.user-message {
text-align: right;
background-color: #7B68EE;
color: white;
align-self: flex-end;
}
.bot-message {
text-align: left;
background-color: #ADD8E6;
color: black;
align-self: flex-start;
}
.user-image {
text-align: right;
align-self: flex-end;
max-width: 150px;
max-height: 300px;
}
.bot-image {
text-align: left;
align-self: flex-start;
max-width: 200px;
max-height: 400px;
}
</style>
"""
temp_file = "temp"
screenshot = "screenshot"
cache = "cache"
if not os.path.exists(temp_file):
os.mkdir(temp_file)
if not os.path.exists(screenshot):
os.mkdir(screenshot)
if not os.path.exists(cache):
os.mkdir(cache)
groundingdino_dir = snapshot_download('AI-ModelScope/GroundingDINO', revision='v1.0.0')
groundingdino_model = pipeline('grounding-dino-task', model=groundingdino_dir)
ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_resnet18_ocr-detection-line-level_damo')
ocr_recognition = pipeline(Tasks.ocr_recognition, model='damo/cv_convnextTiny_ocr-recognition-document_damo')
def encode_image(image_path):
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode('utf-8')
def get_all_files_in_folder(folder_path):
file_list = []
for file_name in os.listdir(folder_path):
file_list.append(file_name)
return file_list
def crop(image, box, i):
image = Image.open(image)
x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
if x1 >= x2-10 or y1 >= y2-10:
return
cropped_image = image.crop((x1, y1, x2, y2))
cropped_image.save(f"./temp/{i}.png", format="PNG")
def merge_text_blocks(text_list, coordinates_list):
merged_text_blocks = []
merged_coordinates = []
sorted_indices = sorted(range(len(coordinates_list)), key=lambda k: (coordinates_list[k][1], coordinates_list[k][0]))
sorted_text_list = [text_list[i] for i in sorted_indices]
sorted_coordinates_list = [coordinates_list[i] for i in sorted_indices]
num_blocks = len(sorted_text_list)
merge = [False] * num_blocks
for i in range(num_blocks):
if merge[i]:
continue
anchor = i
group_text = [sorted_text_list[anchor]]
group_coordinates = [sorted_coordinates_list[anchor]]
for j in range(i+1, num_blocks):
if merge[j]:
continue
if abs(sorted_coordinates_list[anchor][0] - sorted_coordinates_list[j][0]) < 10 and \
sorted_coordinates_list[j][1] - sorted_coordinates_list[anchor][3] >= -10 and sorted_coordinates_list[j][1] - sorted_coordinates_list[anchor][3] < 30 and \
abs(sorted_coordinates_list[anchor][3] - sorted_coordinates_list[anchor][1] - (sorted_coordinates_list[j][3] - sorted_coordinates_list[j][1])) < 10:
group_text.append(sorted_text_list[j])
group_coordinates.append(sorted_coordinates_list[j])
merge[anchor] = True
anchor = j
merge[anchor] = True
merged_text = "\n".join(group_text)
min_x1 = min(group_coordinates, key=lambda x: x[0])[0]
min_y1 = min(group_coordinates, key=lambda x: x[1])[1]
max_x2 = max(group_coordinates, key=lambda x: x[2])[2]
max_y2 = max(group_coordinates, key=lambda x: x[3])[3]
merged_text_blocks.append(merged_text)
merged_coordinates.append([min_x1, min_y1, max_x2, max_y2])
return merged_text_blocks, merged_coordinates
def get_perception_infos(screenshot_file):
width, height = Image.open(screenshot_file).size
text, coordinates = ocr(screenshot_file, ocr_detection, ocr_recognition)
text, coordinates = merge_text_blocks(text, coordinates)
perception_infos = []
for i in range(len(coordinates)):
perception_info = {"text": "text: " + text[i], "coordinates": coordinates[i]}
perception_infos.append(perception_info)
coordinates = det(screenshot_file, "icon", groundingdino_model)
for i in range(len(coordinates)):
perception_info = {"text": "icon", "coordinates": coordinates[i]}
perception_infos.append(perception_info)
image_box = []
image_id = []
for i in range(len(perception_infos)):
if perception_infos[i]['text'] == 'icon':
image_box.append(perception_infos[i]['coordinates'])
image_id.append(i)
for i in range(len(image_box)):
crop(screenshot_file, image_box[i], image_id[i])
images = get_all_files_in_folder(temp_file)
if len(images) > 0:
images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))
image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]
icon_map = {}
prompt = 'This image is an icon from a phone screen. Please briefly describe the shape and color of this icon in one sentence.'
string_image = []
for i in range(len(images)):
image_path = os.path.join(temp_file, images[i])
string_image.append({"image_name": images[i], "image_file": encode_image(image_path)})
query_data = {"task": "caption", "images": string_image, "query": prompt}
response_query = mobile_agent_infer(query_data)
icon_map = response_query["icon_map"]
for i, j in zip(image_id, range(1, len(image_id)+1)):
if icon_map.get(str(j)):
perception_infos[i]['text'] = "icon: " + icon_map[str(j)]
for i in range(len(perception_infos)):
perception_infos[i]['coordinates'] = [int((perception_infos[i]['coordinates'][0]+perception_infos[i]['coordinates'][2])/2), int((perception_infos[i]['coordinates'][1]+perception_infos[i]['coordinates'][3])/2)]
return perception_infos, width, height
def image_to_base64(image):
buffered = io.BytesIO()
image.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
img_html = f'<img src="data:image/png;base64,{img_str}" />'
return img_html
def chatbot(image, instruction, add_info, history, chat_log):
if history == {}:
thought_history = []
summary_history = []
action_history = []
summary = ""
action = ""
completed_requirements = ""
memory = ""
insight = ""
error_flag = False
user_msg = "<div class='user-message'>{}</div>".format(instruction)
else:
thought_history = history["thought_history"]
summary_history = history["summary_history"]
action_history = history["action_history"]
summary = history["summary"]
action = history["action"]
completed_requirements = history["completed_requirements"]
memory = history["memory"][0]
insight = history["insight"]
error_flag = history["error_flag"]
user_msg = "<div class='user-message'>{}</div>".format("I have uploaded the screenshot. Please continue operating.")
images = get_all_files_in_folder(cache)
if len(images) > 0 and len(images) <= 100:
images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))
image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]
cur_image_id = image_id[-1] + 1
elif len(images) > 100:
images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))
image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]
cur_image_id = image_id[-1] + 1
os.remove(os.path.join(cache, str(image_id[0])+".png"))
else:
cur_image_id = 1
image.save(os.path.join(cache, str(cur_image_id) + ".png"), format="PNG")
screenshot_file = os.path.join(cache, str(cur_image_id) + ".png")
perception_infos, width, height = get_perception_infos(screenshot_file)
shutil.rmtree(temp_file)
os.mkdir(temp_file)
local_screenshot_file = encode_image(screenshot_file)
query_data = {
"task": "decision",
"screenshot_file": local_screenshot_file,
"instruction": instruction,
"perception_infos": perception_infos,
"width": width,
"height": height,
"summary_history": summary_history,
"action_history": action_history,
"summary": summary,
"action": action,
"add_info": add_info,
"error_flag": error_flag,
"completed_requirements": completed_requirements,
"memory": memory,
"memory_switch": True,
"insight": insight
}
response_query = mobile_agent_infer(query_data)
output_action = response_query["decision"]
output_memory = response_query["memory"]
if output_action == "No token":
bot_response = ["<div class='bot-message'>{}</div>".format("Sorry, the resources can be exhausted today.")]
chat_html = "<div class='chat-container'>{}</div>".format("".join(bot_response))
return chatbot_css + chat_html, history, chat_log
thought = output_action.split("### Thought ###")[-1].split("### Action ###")[0].replace("\n", " ").replace(":", "").replace(" ", " ").strip()
summary = output_action.split("### Operation ###")[-1].replace("\n", " ").replace(" ", " ").strip()
action = output_action.split("### Action ###")[-1].split("### Operation ###")[0].replace("\n", " ").replace(" ", " ").strip()
output_memory = output_memory.split("### Important content ###")[-1].split("\n\n")[0].strip() + "\n"
if "None" not in output_memory and output_memory not in memory:
memory += output_memory
if "Open app" in action:
bot_response = "Please click the red circle and upload the current screenshot again."
app_name = action.split("(")[-1].split(")")[0]
text, coordinate = ocr(screenshot_file, ocr_detection, ocr_recognition)
for ti in range(len(text)):
if app_name == text[ti]:
name_coordinate = [int((coordinate[ti][0] + coordinate[ti][2])/2), int((coordinate[ti][1] + coordinate[ti][3])/2)]
x, y = name_coordinate[0], name_coordinate[1]
radius = 75
draw = ImageDraw.Draw(image)
draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline='red', width=10)
break
elif "Tap" in action:
bot_response = "Please click the red circle and upload the current screenshot again."
coordinate = action.split("(")[-1].split(")")[0].split(", ")
x, y = int(coordinate[0]), int(coordinate[1])
radius = 75
draw = ImageDraw.Draw(image)
draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline='red', width=10)
elif "Swipe" in action:
bot_response = "Please slide from red circle to blue circle and upload the current screenshot again."
coordinate1 = action.split("Swipe (")[-1].split("), (")[0].split(", ")
coordinate2 = action.split("), (")[-1].split(")")[0].split(", ")
x1, y1 = int(coordinate1[0]), int(coordinate1[1])
x2, y2 = int(coordinate2[0]), int(coordinate2[1])
radius = 75
draw = ImageDraw.Draw(image)
draw.ellipse([x1 - radius, y1 - radius, x1 + radius, y1 + radius], outline='red', width=10)
draw.ellipse([x2 - radius, y2 - radius, x2 + radius, y2 + radius], outline='blue', width=10)
elif "Type" in action:
if "(text)" not in action:
text = action.split("(")[-1].split(")")[0]
else:
text = action.split(" \"")[-1].split("\"")[0]
bot_response = f"Please type the \"{text}\" and upload the current screenshot again."
elif "Back" in action:
bot_response = f"Please back to previous page and upload the current screenshot again."
elif "Home" in action:
bot_response = f"Please back to home page and upload the current screenshot again."
elif "Stop" in action:
bot_response = f"Task completed."
bot_text1 = "<div class='bot-message'>{}</div>".format("### Decision ###")
bot_thought = "<div class='bot-message'>{}</div>".format("Thought: " + thought)
bot_action = "<div class='bot-message'>{}</div>".format("Action: " + action)
bot_operation = "<div class='bot-message'>{}</div>".format("Operation: " + summary)
bot_text2 = "<div class='bot-message'>{}</div>".format("### Memory ###")
bot_memory = "<div class='bot-message'>{}</div>".format(output_memory)
bot_response = "<div class='bot-message'>{}</div>".format(bot_response)
if image is not None:
bot_img_html = image_to_base64(image)
bot_response = "<div class='bot-image'>{}</div>".format(bot_img_html) + bot_response
chat_log.append(user_msg)
thought_history.append(thought)
summary_history.append(summary)
action_history.append(action)
history["thought_history"] = thought_history
history["summary_history"] = summary_history
history["action_history"] = action_history
history["summary"] = summary
history["action"] = action
history["memory"] = memory,
history["memory_switch"] = True,
history["insight"] = insight
history["error_flag"] = error_flag
query_data = {
"task": "planning",
"instruction": instruction,
"thought_history": thought_history,
"summary_history": summary_history,
"action_history": action_history,
"completed_requirements": "",
"add_info": add_info
}
response_query = mobile_agent_infer(query_data)
output_planning = response_query["planning"]
if output_planning == "No token":
bot_response = ["<div class='bot-message'>{}</div>".format("Sorry, the resources can be exhausted today.")]
chat_html = "<div class='chat-container'>{}</div>".format("".join(bot_response))
return chatbot_css + chat_html, history, chat_log
output_planning = output_planning.split("### Completed contents ###")[-1].replace("\n", " ").strip()
history["completed_requirements"] = output_planning
bot_text3 = "<div class='bot-message'>{}</div>".format("### Planning ###")
output_planning = "<div class='bot-message'>{}</div>".format(output_planning)
chat_log.append(bot_text3)
chat_log.append(output_planning)
chat_log.append(bot_text1)
chat_log.append(bot_thought)
chat_log.append(bot_action)
chat_log.append(bot_operation)
chat_log.append(bot_text2)
chat_log.append(bot_memory)
chat_log.append(bot_response)
chat_html = "<div class='chat-container'>{}</div>".format("".join(chat_log))
return chatbot_css + chat_html, history, chat_log
def lock_input(instruction):
return gr.update(value=instruction, interactive=False), gr.update(value=None)
def reset_demo():
return gr.update(value="", interactive=True), gr.update(value="If you want to tap an icon of an app, use the action \"Open app\"", interactive=True), "<div class='chat-container'></div>", {}, []
tos_markdown = ("""<div style="display:flex; gap: 0.25rem;" align="center">
<a href='https://github.com/X-PLUG/MobileAgent'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
<a href="https://arxiv.org/abs/2406.01014"><img src="https://img.shields.io/badge/Arxiv-2406.01014-red"></a>
<a href='https://github.com/X-PLUG/MobileAgent/stargazers'><img src='https://img.shields.io/github/stars/X-PLUG/MobileAgent.svg?style=social'></a>
</div>
If you like our project, please give us a star ✨ on Github for latest update.
**Terms of use**
1. Input your instruction in \"Instruction\", for example \"Turn on the dark mode\".
2. You can input helpful operation knowledge in \"Knowledge\".
3. Click \"Submit\" to get the operation. You need to operate your mobile device according to the operation and then upload the screenshot after your operation.
4. The 5 cases in \"Examples\" are a complete flow. Click and submit from top to bottom to experience.
5. Due to limited resources, each operation may take a long time, please be patient and wait.
**使用说明**
1. 在“Instruction”中输入你的指令,例如“打开深色模式”。
2. 你可以在“Knowledge”中输入帮助性的操作知识。
3. 点击“Submit”来获得操作。你需要根据输出来操作手机,并且上传操作后的截图。
4. “Example”中的5个例子是一个任务。从上到下点击它们并且点击“Submit”来体验。
5. 由于资源有限,每次操作的时间会比较长,请耐心等待。""")
title_markdowm = ("""# Mobile-Agent-v2: Mobile Device Operation Assistant with Effective Navigation via Multi-Agent Collaboration""")
instruction_input = gr.Textbox(label="Instruction", placeholder="Input your instruction")
knowledge_input = gr.Textbox(label="Knowledge", placeholder="Input your knowledge", value="If you want to tap an icon of an app, use the action \"Open app\"")
with gr.Blocks() as demo:
history_state = gr.State(value={})
history_output = gr.State(value=[])
with gr.Row():
gr.Markdown(title_markdowm)
with gr.Row():
with gr.Column(scale=5):
gr.Markdown(tos_markdown)
with gr.Row():
image_input = gr.Image(label="Screenshot", type="pil", height=550, width=230)
gr.Examples(examples=[
["./example/example_1.jpg", "Turn on the dark mode"],
["./example/example_2.jpg", "Turn on the dark mode"],
["./example/example_3.jpg", "Turn on the dark mode"],
["./example/example_4.jpg", "Turn on the dark mode"],
["./example/example_5.jpg", "Turn on the dark mode"],
], inputs=[image_input, instruction_input, knowledge_input])
with gr.Column(scale=6):
instruction_input.render()
knowledge_input.render()
with gr.Row():
start_button = gr.Button("Submit")
clear_button = gr.Button("Clear")
output_component = gr.HTML(label="Chat history", value="<div class='chat-container'></div>")
start_button.click(
fn=lambda image, instruction, add_info, history, output: chatbot(image, instruction, add_info, history, output),
inputs=[image_input, instruction_input, knowledge_input, history_state, history_output],
outputs=[output_component, history_state, history_output]
)
clear_button.click(
fn=reset_demo,
inputs=[],
outputs=[instruction_input, knowledge_input, output_component, history_state, history_output]
)
demo.queue().launch(share=True)