junyangwang0410/Mobile-Agent
9
1import io2import os3import shutil4import base645import gradio as gr6from PIL import Image, ImageDraw7 8from MobileAgent.text_localization import ocr9from MobileAgent.icon_localization import det10from MobileAgent.local_server import mobile_agent_infer11 12from modelscope import snapshot_download13from modelscope.pipelines import pipeline14from modelscope.utils.constant import Tasks15 16 17chatbot_css = """18<style>19.chat-container {20 display: flex;21 flex-direction: column;22 overflow-y: auto;23 max-height: 630px;24 margin: 10px;25}26.user-message, .bot-message {27 margin: 5px;28 padding: 10px;29 border-radius: 10px;30}31.user-message {32 text-align: right;33 background-color: #7B68EE;34 color: white;35 align-self: flex-end;36}37.bot-message {38 text-align: left;39 background-color: #ADD8E6;40 color: black;41 align-self: flex-start;42}43.user-image {44 text-align: right;45 align-self: flex-end;46 max-width: 150px;47 max-height: 300px;48}49.bot-image {50 text-align: left;51 align-self: flex-start;52 max-width: 200px;53 max-height: 400px;54}55</style>56"""57 58 59temp_file = "temp"60screenshot = "screenshot"61cache = "cache"62if not os.path.exists(temp_file):63 os.mkdir(temp_file)64if not os.path.exists(screenshot):65 os.mkdir(screenshot)66if not os.path.exists(cache):67 os.mkdir(cache)68 69 70groundingdino_dir = snapshot_download('AI-ModelScope/GroundingDINO', revision='v1.0.0')71groundingdino_model = pipeline('grounding-dino-task', model=groundingdino_dir)72ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_resnet18_ocr-detection-line-level_damo')73ocr_recognition = pipeline(Tasks.ocr_recognition, model='damo/cv_convnextTiny_ocr-recognition-document_damo')74 75 76def encode_image(image_path):77 with open(image_path, "rb") as image_file:78 return base64.b64encode(image_file.read()).decode('utf-8')79 80 81def get_all_files_in_folder(folder_path):82 file_list = []83 for file_name in os.listdir(folder_path):84 file_list.append(file_name)85 return file_list86 87 88def crop(image, box, i):89 image = Image.open(image)90 x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])91 if x1 >= x2-10 or y1 >= y2-10:92 return93 cropped_image = image.crop((x1, y1, x2, y2))94 cropped_image.save(f"./temp/{i}.png", format="PNG")95 96 97def merge_text_blocks(text_list, coordinates_list):98 merged_text_blocks = []99 merged_coordinates = []100 101 sorted_indices = sorted(range(len(coordinates_list)), key=lambda k: (coordinates_list[k][1], coordinates_list[k][0]))102 sorted_text_list = [text_list[i] for i in sorted_indices]103 sorted_coordinates_list = [coordinates_list[i] for i in sorted_indices]104 105 num_blocks = len(sorted_text_list)106 merge = [False] * num_blocks107 108 for i in range(num_blocks):109 if merge[i]:110 continue111 112 anchor = i113 114 group_text = [sorted_text_list[anchor]]115 group_coordinates = [sorted_coordinates_list[anchor]]116 117 for j in range(i+1, num_blocks):118 if merge[j]:119 continue120 121 if abs(sorted_coordinates_list[anchor][0] - sorted_coordinates_list[j][0]) < 10 and \122 sorted_coordinates_list[j][1] - sorted_coordinates_list[anchor][3] >= -10 and sorted_coordinates_list[j][1] - sorted_coordinates_list[anchor][3] < 30 and \123 abs(sorted_coordinates_list[anchor][3] - sorted_coordinates_list[anchor][1] - (sorted_coordinates_list[j][3] - sorted_coordinates_list[j][1])) < 10:124 group_text.append(sorted_text_list[j])125 group_coordinates.append(sorted_coordinates_list[j])126 merge[anchor] = True127 anchor = j128 merge[anchor] = True129 130 merged_text = "\n".join(group_text)131 min_x1 = min(group_coordinates, key=lambda x: x[0])[0]132 min_y1 = min(group_coordinates, key=lambda x: x[1])[1]133 max_x2 = max(group_coordinates, key=lambda x: x[2])[2]134 max_y2 = max(group_coordinates, key=lambda x: x[3])[3]135 136 merged_text_blocks.append(merged_text)137 merged_coordinates.append([min_x1, min_y1, max_x2, max_y2])138 139 return merged_text_blocks, merged_coordinates140 141 142def get_perception_infos(screenshot_file):143 width, height = Image.open(screenshot_file).size144 145 text, coordinates = ocr(screenshot_file, ocr_detection, ocr_recognition)146 text, coordinates = merge_text_blocks(text, coordinates)147 148 perception_infos = []149 for i in range(len(coordinates)):150 perception_info = {"text": "text: " + text[i], "coordinates": coordinates[i]}151 perception_infos.append(perception_info)152 153 coordinates = det(screenshot_file, "icon", groundingdino_model)154 155 for i in range(len(coordinates)):156 perception_info = {"text": "icon", "coordinates": coordinates[i]}157 perception_infos.append(perception_info)158 159 image_box = []160 image_id = []161 for i in range(len(perception_infos)):162 if perception_infos[i]['text'] == 'icon':163 image_box.append(perception_infos[i]['coordinates'])164 image_id.append(i)165 166 for i in range(len(image_box)):167 crop(screenshot_file, image_box[i], image_id[i])168 169 images = get_all_files_in_folder(temp_file)170 if len(images) > 0:171 images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))172 image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]173 icon_map = {}174 prompt = 'This image is an icon from a phone screen. Please briefly describe the shape and color of this icon in one sentence.'175 176 string_image = []177 for i in range(len(images)):178 image_path = os.path.join(temp_file, images[i])179 string_image.append({"image_name": images[i], "image_file": encode_image(image_path)})180 query_data = {"task": "caption", "images": string_image, "query": prompt}181 response_query = mobile_agent_infer(query_data)182 icon_map = response_query["icon_map"]183 184 for i, j in zip(image_id, range(1, len(image_id)+1)):185 if icon_map.get(str(j)):186 perception_infos[i]['text'] = "icon: " + icon_map[str(j)]187 188 for i in range(len(perception_infos)):189 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)]190 191 return perception_infos, width, height192 193 194def image_to_base64(image):195 buffered = io.BytesIO()196 image.save(buffered, format="PNG")197 img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")198 img_html = f'<img src="data:image/png;base64,{img_str}" />'199 return img_html200 201 202def chatbot(image, instruction, add_info, history, chat_log):203 if history == {}:204 thought_history = []205 summary_history = []206 action_history = []207 summary = ""208 action = ""209 completed_requirements = ""210 memory = ""211 insight = ""212 error_flag = False213 user_msg = "<div class='user-message'>{}</div>".format(instruction)214 else:215 thought_history = history["thought_history"]216 summary_history = history["summary_history"]217 action_history = history["action_history"]218 summary = history["summary"]219 action = history["action"]220 completed_requirements = history["completed_requirements"]221 memory = history["memory"][0]222 insight = history["insight"]223 error_flag = history["error_flag"]224 user_msg = "<div class='user-message'>{}</div>".format("I have uploaded the screenshot. Please continue operating.")225 226 images = get_all_files_in_folder(cache)227 if len(images) > 0 and len(images) <= 100:228 images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))229 image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]230 cur_image_id = image_id[-1] + 1231 elif len(images) > 100:232 images = sorted(images, key=lambda x: int(x.split('/')[-1].split('.')[0]))233 image_id = [int(image.split('/')[-1].split('.')[0]) for image in images]234 cur_image_id = image_id[-1] + 1235 os.remove(os.path.join(cache, str(image_id[0])+".png"))236 else:237 cur_image_id = 1238 239 image.save(os.path.join(cache, str(cur_image_id) + ".png"), format="PNG")240 screenshot_file = os.path.join(cache, str(cur_image_id) + ".png")241 perception_infos, width, height = get_perception_infos(screenshot_file)242 shutil.rmtree(temp_file)243 os.mkdir(temp_file)244 245 local_screenshot_file = encode_image(screenshot_file)246 query_data = {247 "task": "decision",248 "screenshot_file": local_screenshot_file,249 "instruction": instruction,250 "perception_infos": perception_infos,251 "width": width,252 "height": height,253 "summary_history": summary_history,254 "action_history": action_history,255 "summary": summary,256 "action": action,257 "add_info": add_info,258 "error_flag": error_flag,259 "completed_requirements": completed_requirements,260 "memory": memory,261 "memory_switch": True,262 "insight": insight263 }264 265 response_query = mobile_agent_infer(query_data)266 output_action = response_query["decision"]267 output_memory = response_query["memory"]268 if output_action == "No token":269 bot_response = ["<div class='bot-message'>{}</div>".format("Sorry, the resources can be exhausted today.")]270 chat_html = "<div class='chat-container'>{}</div>".format("".join(bot_response))271 return chatbot_css + chat_html, history, chat_log272 273 thought = output_action.split("### Thought ###")[-1].split("### Action ###")[0].replace("\n", " ").replace(":", "").replace(" ", " ").strip()274 summary = output_action.split("### Operation ###")[-1].replace("\n", " ").replace(" ", " ").strip()275 action = output_action.split("### Action ###")[-1].split("### Operation ###")[0].replace("\n", " ").replace(" ", " ").strip()276 277 output_memory = output_memory.split("### Important content ###")[-1].split("\n\n")[0].strip() + "\n"278 if "None" not in output_memory and output_memory not in memory:279 memory += output_memory280 281 if "Open app" in action:282 bot_response = "Please click the red circle and upload the current screenshot again."283 app_name = action.split("(")[-1].split(")")[0]284 text, coordinate = ocr(screenshot_file, ocr_detection, ocr_recognition)285 for ti in range(len(text)):286 if app_name == text[ti]:287 name_coordinate = [int((coordinate[ti][0] + coordinate[ti][2])/2), int((coordinate[ti][1] + coordinate[ti][3])/2)]288 x, y = name_coordinate[0], name_coordinate[1]289 radius = 75290 draw = ImageDraw.Draw(image)291 draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline='red', width=10)292 break293 294 elif "Tap" in action:295 bot_response = "Please click the red circle and upload the current screenshot again."296 coordinate = action.split("(")[-1].split(")")[0].split(", ")297 x, y = int(coordinate[0]), int(coordinate[1])298 radius = 75299 draw = ImageDraw.Draw(image)300 draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline='red', width=10)301 302 elif "Swipe" in action:303 bot_response = "Please slide from red circle to blue circle and upload the current screenshot again."304 coordinate1 = action.split("Swipe (")[-1].split("), (")[0].split(", ")305 coordinate2 = action.split("), (")[-1].split(")")[0].split(", ")306 x1, y1 = int(coordinate1[0]), int(coordinate1[1])307 x2, y2 = int(coordinate2[0]), int(coordinate2[1])308 radius = 75309 draw = ImageDraw.Draw(image)310 draw.ellipse([x1 - radius, y1 - radius, x1 + radius, y1 + radius], outline='red', width=10)311 draw.ellipse([x2 - radius, y2 - radius, x2 + radius, y2 + radius], outline='blue', width=10)312 313 elif "Type" in action:314 if "(text)" not in action:315 text = action.split("(")[-1].split(")")[0]316 else:317 text = action.split(" \"")[-1].split("\"")[0]318 bot_response = f"Please type the \"{text}\" and upload the current screenshot again."319 320 elif "Back" in action:321 bot_response = f"Please back to previous page and upload the current screenshot again."322 323 elif "Home" in action:324 bot_response = f"Please back to home page and upload the current screenshot again."325 326 elif "Stop" in action:327 bot_response = f"Task completed."328 329 bot_text1 = "<div class='bot-message'>{}</div>".format("### Decision ###")330 bot_thought = "<div class='bot-message'>{}</div>".format("Thought: " + thought)331 bot_action = "<div class='bot-message'>{}</div>".format("Action: " + action)332 bot_operation = "<div class='bot-message'>{}</div>".format("Operation: " + summary)333 bot_text2 = "<div class='bot-message'>{}</div>".format("### Memory ###")334 bot_memory = "<div class='bot-message'>{}</div>".format(output_memory)335 bot_response = "<div class='bot-message'>{}</div>".format(bot_response)336 if image is not None:337 bot_img_html = image_to_base64(image)338 bot_response = "<div class='bot-image'>{}</div>".format(bot_img_html) + bot_response339 340 chat_log.append(user_msg)341 342 thought_history.append(thought)343 summary_history.append(summary)344 action_history.append(action)345 346 history["thought_history"] = thought_history347 history["summary_history"] = summary_history348 history["action_history"] = action_history349 history["summary"] = summary350 history["action"] = action351 history["memory"] = memory,352 history["memory_switch"] = True,353 history["insight"] = insight354 history["error_flag"] = error_flag355 356 query_data = {357 "task": "planning",358 "instruction": instruction,359 "thought_history": thought_history,360 "summary_history": summary_history,361 "action_history": action_history,362 "completed_requirements": "",363 "add_info": add_info364 }365 366 response_query = mobile_agent_infer(query_data)367 output_planning = response_query["planning"]368 if output_planning == "No token":369 bot_response = ["<div class='bot-message'>{}</div>".format("Sorry, the resources can be exhausted today.")]370 chat_html = "<div class='chat-container'>{}</div>".format("".join(bot_response))371 return chatbot_css + chat_html, history, chat_log372 373 output_planning = output_planning.split("### Completed contents ###")[-1].replace("\n", " ").strip()374 history["completed_requirements"] = output_planning375 376 bot_text3 = "<div class='bot-message'>{}</div>".format("### Planning ###")377 output_planning = "<div class='bot-message'>{}</div>".format(output_planning)378 379 chat_log.append(bot_text3)380 chat_log.append(output_planning)381 chat_log.append(bot_text1)382 chat_log.append(bot_thought)383 chat_log.append(bot_action)384 chat_log.append(bot_operation)385 chat_log.append(bot_text2)386 chat_log.append(bot_memory)387 chat_log.append(bot_response)388 389 chat_html = "<div class='chat-container'>{}</div>".format("".join(chat_log))390 391 return chatbot_css + chat_html, history, chat_log392 393 394def lock_input(instruction):395 return gr.update(value=instruction, interactive=False), gr.update(value=None)396 397 398def reset_demo():399 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>", {}, []400 401 402tos_markdown = ("""<div style="display:flex; gap: 0.25rem;" align="center">403 <a href='https://github.com/X-PLUG/MobileAgent'><img src='https://img.shields.io/badge/Github-Code-blue'></a>404 <a href="https://arxiv.org/abs/2406.01014"><img src="https://img.shields.io/badge/Arxiv-2406.01014-red"></a>405 <a href='https://github.com/X-PLUG/MobileAgent/stargazers'><img src='https://img.shields.io/github/stars/X-PLUG/MobileAgent.svg?style=social'></a>406</div>407If you like our project, please give us a star ✨ on Github for latest update.408 409**Terms of use**4101. Input your instruction in \"Instruction\", for example \"Turn on the dark mode\".4112. You can input helpful operation knowledge in \"Knowledge\".4123. 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.4134. The 5 cases in \"Examples\" are a complete flow. Click and submit from top to bottom to experience.4145. Due to limited resources, each operation may take a long time, please be patient and wait.415 416**使用说明**4171. 在“Instruction”中输入你的指令，例如“打开深色模式”。4182. 你可以在“Knowledge”中输入帮助性的操作知识。4193. 点击“Submit”来获得操作。你需要根据输出来操作手机，并且上传操作后的截图。4204. “Example”中的5个例子是一个任务。从上到下点击它们并且点击“Submit”来体验。4215. 由于资源有限，每次操作的时间会比较长，请耐心等待。""")422 423title_markdowm = ("""# Mobile-Agent-v2: Mobile Device Operation Assistant with Effective Navigation via Multi-Agent Collaboration""")424 425instruction_input = gr.Textbox(label="Instruction", placeholder="Input your instruction")426knowledge_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\"")427with gr.Blocks() as demo:428 history_state = gr.State(value={})429 history_output = gr.State(value=[])430 with gr.Row():431 gr.Markdown(title_markdowm)432 with gr.Row():433 with gr.Column(scale=5):434 gr.Markdown(tos_markdown)435 with gr.Row():436 image_input = gr.Image(label="Screenshot", type="pil", height=550, width=230)437 gr.Examples(examples=[438 ["./example/example_1.jpg", "Turn on the dark mode"],439 ["./example/example_2.jpg", "Turn on the dark mode"],440 ["./example/example_3.jpg", "Turn on the dark mode"],441 ["./example/example_4.jpg", "Turn on the dark mode"],442 ["./example/example_5.jpg", "Turn on the dark mode"],443 ], inputs=[image_input, instruction_input, knowledge_input])444 445 with gr.Column(scale=6):446 instruction_input.render()447 knowledge_input.render()448 with gr.Row():449 start_button = gr.Button("Submit")450 clear_button = gr.Button("Clear")451 output_component = gr.HTML(label="Chat history", value="<div class='chat-container'></div>")452 453 start_button.click(454 fn=lambda image, instruction, add_info, history, output: chatbot(image, instruction, add_info, history, output),455 inputs=[image_input, instruction_input, knowledge_input, history_state, history_output],456 outputs=[output_component, history_state, history_output]457 )458 459 clear_button.click(460 fn=reset_demo,461 inputs=[],462 outputs=[instruction_input, knowledge_input, output_component, history_state, history_output]463 )464 465demo.queue().launch(share=True)