Spaces:
Runtime error
Runtime error
| 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) |