Create app.py

app.py

@@ -0,0 +1,286 @@
+'''
+author:        caishaofei <[email protected]>
+date:          2024-09-20 20:10:44
+Copyright © Team CraftJarvis All rights reserved
+'''
+import re
+import os
+import cv2
+import time
+from pathlib import Path
+import argparse
+import requests
+import gradio as gr
+import torch
+import numpy as np
+from io import BytesIO
+from PIL import Image, ImageDraw
+from rocket.arm.sessions import Session, Pointer
+
+COLORS = [
+    (255, 0, 0), (0, 255, 0), (0, 0, 255), 
+    (255, 255, 0), (255, 0, 255), (0, 255, 255),
+    (255, 255, 255), (0, 0, 0), (128, 128, 128),
+    (128, 0, 0), (128, 128, 0), (0, 128, 0),
+    (128, 0, 128), (0, 128, 128), (0, 0, 128),
+]
+
+SEGMENT_MAPPING = {
+    "Hunt": 0, "Use": 3, "Mine": 2, "Interact": 3, "Craft": 4, "Switch": 5, "Approach": 6
+}
+
+NOOP_ACTION = {
+    "back": 0,
+    "drop": 0,
+    "forward": 0,
+    "hotbar.1": 0,
+    "hotbar.2": 0,
+    "hotbar.3": 0,
+    "hotbar.4": 0,
+    "hotbar.5": 0,
+    "hotbar.6": 0,
+    "hotbar.7": 0,
+    "hotbar.8": 0,
+    "hotbar.9": 0,
+    "inventory": 0,
+    "jump": 0,
+    "left": 0,
+    "right": 0,
+    "sneak": 0,
+    "sprint": 0,
+    "camera": np.array([0, 0]),
+    "attack": 0,
+    "use": 0,
+}
+
+def reset_fn(env_name, session):
+    image = session.reset(env_name)
+    return image, session
+
+def step_fn(act_key, session):
+    action = NOOP_ACTION.copy()
+    if act_key != "null":
+        action[act_key] = 1
+    image = session.step(action)
+    return image, session
+
+def loop_step_fn(steps, session):
+    for i in range(steps):
+        image = session.step()
+        status = f"Running Agent `Rocket` steps: {i+1}/{steps}. "
+        yield image, session.num_steps, status, session
+
+def clear_memory_fn(session):
+    image = session.current_image
+    session.clear_agent_memory()
+    return image, "0", session
+
+def get_points_with_draw(image, label, session, evt: gr.SelectData):
+    points = session.points
+    point_label = session.points_label
+    x, y = evt.index[0], evt.index[1]
+    point_radius, point_color = 5, (0, 255, 0) if label == 'Add Points' else (255, 0, 0)
+    points.append([x, y])
+    point_label.append(1 if label == 'Add Points' else 0)
+    cv2.circle(image, (x, y), point_radius, point_color, -1)
+    return image, session
+
+def clear_points_fn(session):
+    session.clear_points()
+    return session.current_image, session
+
+def segment_fn(session):
+    if len(session.points) == 0:
+        return session.current_image, session
+    session.segment()
+    image = session.apply_mask()
+    return image, session
+
+def clear_segment_fn(session):
+    session.clear_obj_mask()
+    session.tracking_flag = False
+    return session.current_image, False, session
+
+def set_tracking_mode(tracking_flag, session):
+    session.tracking_flag = tracking_flag
+    return session
+
+def set_segment_type(segment_type, session):
+    session.segment_type = segment_type
+    return session
+
+def play_fn(session):
+    image = session.step()
+    return image, session
+
+memory_length = gr.Textbox(value="0", interactive=False, show_label=False)
+
+def make_video_fn(session, make_video, save_video, progress=gr.Progress()):
+    images = session.image_history
+    if len(images) == 0:
+        return session, make_video, save_video
+    filepath = "rocket.mp4"
+    h, w = images[0].shape[:2]
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video = cv2.VideoWriter(filepath, fourcc, 20.0, (w, h))
+    for image in progress.tqdm(images):
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        video.write(image)
+    video.release()
+    session.image_history = []
+    return session, gr.Button("Make Video", visible=False), gr.DownloadButton("Download!", value=filepath, visible=True)
+
+def save_video_fn(session, make_video, save_video):
+    return session, gr.Button("Make Video", visible=True), gr.DownloadButton("Download!", visible=False)
+
+def choose_sam_fn(sam_choice, session):
+    session.sam_choice = sam_choice
+    session.load_sam()
+    return session
+
+def molmo_fn(molmo_text, molmo_session, rocket_session, display_image):
+    image = rocket_session.current_image.copy()
+    points = molmo_session.gen_point(image=image, prompt=molmo_text)
+    molmo_result = molmo_session.molmo_result
+    for x, y in points:
+        x, y = int(x), int(y)
+        point_radius, point_color = 5, (0, 255, 0) 
+        rocket_session.points.append([x, y])
+        rocket_session.points_label.append(1)
+        cv2.circle(display_image, (x, y), point_radius, point_color, -1)
+    return molmo_result, display_image
+
+def extract_points(data):
+    # 匹配 x 和 y 坐标的值，支持 <points> 和 <point> 标签
+    pattern = r'x\d?="([-+]?\d*\.\d+|\d+)" y\d?="([-+]?\d*\.\d+|\d+)"'
+    points = re.findall(pattern, data)
+    # 将提取到的坐标转换为浮点数
+    points = [(float(x)/100*640, float(y)/100*360) for x, y in points]
+    return points
+
+def draw_gradio_components(args):
+
+    with gr.Blocks() as demo:
+        
+        gr.Markdown(
+            """
+            # Welcome to Explore ROCKET-1 in Minecraft!!
+            ## Please follow next steps to interact with the agent:
+            1. Reset the environment by selecting an environment name.
+            2. Select a SAM2 checkpoint to load.
+            3. Use your mouse to add or remove points on the image. 
+            4. Select the segment type you want to perform.
+            5. Enable `tracking` mode if you want to track objects while stepping actions. 
+            6. Click `New Segment` to segment the image based on the points you added. 
+            7. Call the agent by clicking `Call Rocket` to run the agent for a certain number of steps. 
+            ## Hints:
+            1. You can use the `Make Video` button to generate a video of the agent's actions.
+            2. You can use the `Clear Memory` button to clear the ROCKET-1's memory. 
+            3. You can use the `Clear Segment` button to clear SAM's memory. 
+            4. You can use the `Manually Step` button to manually step the agent. 
+            """
+        )
+        
+        rocket_session = gr.State(Session(
+            sam_path=args.sam_path,
+        ))
+        molmo_session = gr.State(Pointer(
+            model_id="molmo-72b-0924", 
+            model_url="http://172.17.30.127:8000/v1", 
+        ))
+        with gr.Row():
+            
+            with gr.Column(scale=2):
+                # start_image = Image.open("start.png").resize((640, 360))
+                start_image = np.zeros((360, 640, 3), dtype=np.uint8)
+                
+                with gr.Group():
+                    display_image = gr.Image(
+                        value=np.array(start_image), 
+                        interactive=False, 
+                        show_label=False, 
+                        label="Real-time Environment Observation", 
+                        streaming=True
+                    )
+                    display_status = gr.Textbox("Status Bar", interactive=False, show_label=False)
+            
+            with gr.Column(scale=1):
+                
+                sam_choice = gr.Radio(
+                    choices=["large", "base", "small", "tiny"],
+                    value="base",
+                    label="Select SAM2 checkpoint",
+                )
+                sam_choice.select(fn=choose_sam_fn, inputs=[sam_choice, rocket_session], outputs=[rocket_session], show_progress=False)
+                
+                with gr.Group():
+                    add_or_remove = gr.Radio(
+                        choices=["Add Points", "Remove Areas"], 
+                        value="Add Points", 
+                        label="Use you mouse to add or remove points",
+                    )
+                    clear_points_btn = gr.Button("Clear Points")
+                    clear_points_btn.click(clear_points_fn, inputs=[rocket_session], outputs=[display_image, rocket_session], show_progress=True)
+                
+                with gr.Group():
+                    segment_type = gr.Radio(
+                        choices=["Approach", "Interact", "Hunt", "Mine", "Craft", "Switch"],
+                        value="Approach", 
+                        label="What do you want with this segment?",
+                    )
+                    track_flag = gr.Checkbox(True, label="Enable tracking objects while steping actions")
+                    track_flag.select(fn=set_tracking_mode, inputs=[track_flag, rocket_session], outputs=[rocket_session], show_progress=False)
+                    with gr.Group(), gr.Row():
+                        new_segment_btn = gr.Button("New Segment")
+                        clear_segment_btn = gr.Button("Clear Segment")
+                        new_segment_btn.click(segment_fn, inputs=[rocket_session], outputs=[display_image, rocket_session], show_progress=True)
+                        clear_segment_btn.click(clear_segment_fn, inputs=[rocket_session], outputs=[display_image, track_flag, rocket_session], show_progress=True)
+
+            display_image.select(get_points_with_draw, inputs=[display_image, add_or_remove, rocket_session], outputs=[display_image, rocket_session])
+            segment_type.select(set_segment_type, inputs=[segment_type, rocket_session], outputs=[rocket_session], show_progress=False)
+
+        with gr.Row():
+            with gr.Group():
+                env_list = [f"rocket/{x.stem}" for x in Path("../env_configs/rocket").glob("*.yaml") if 'base' not in x.name != 'base']
+                env_name = gr.Dropdown(env_list, multiselect=False, min_width=200, show_label=False, label="Env Name")
+                reset_btn = gr.Button("Reset Environment")
+                reset_btn.click(fn=reset_fn, inputs=[env_name, rocket_session], outputs=[display_image, rocket_session], show_progress=True)
+
+            with gr.Group():
+                action_list = [x for x in NOOP_ACTION.keys()]
+                act_key = gr.Dropdown(action_list, multiselect=False, min_width=200, show_label=False, label="Action")
+                step_btn = gr.Button("Manually Step")
+                step_btn.click(fn=step_fn, inputs=[act_key, rocket_session], outputs=[display_image, rocket_session], show_progress=False)
+
+            with gr.Group():
+                steps = gr.Slider(1, 600, 30, 1, label="Steps", show_label=False)
+                play_btn = gr.Button("Call Rocket")
+                play_btn.click(fn=loop_step_fn, inputs=[steps, rocket_session], outputs=[display_image, memory_length, display_status, rocket_session], show_progress=False)
+
+            with gr.Group():
+                memory_length.render()
+                clear_states_btn = gr.Button("Clear Memory")
+                clear_states_btn.click(fn=clear_memory_fn, inputs=rocket_session, outputs=[display_image, memory_length, rocket_session], show_progress=False)
+            
+            make_video_btn = gr.Button("Make Video")
+            save_video_btn = gr.DownloadButton("Download!!", visible=False)
+            make_video_btn.click(make_video_fn, inputs=[rocket_session, make_video_btn, save_video_btn], outputs=[rocket_session, make_video_btn, save_video_btn], show_progress=False)
+            save_video_btn.click(save_video_fn, inputs=[rocket_session, make_video_btn, save_video_btn], outputs=[rocket_session, make_video_btn, save_video_btn], show_progress=False)
+        with gr.Row():
+            with gr.Group():
+                molmo_text = gr.Textbox("pinpoint the", label="Molmo Text", show_label=True, min_width=200)
+                molmo_btn = gr.Button("Generate")
+                output_text = gr.Textbox("", label="Molmo Output", show_label=False, min_width=200)
+                molmo_btn.click(molmo_fn, inputs=[molmo_text, molmo_session, rocket_session, display_image],outputs=[output_text, display_image],show_progress=False)
+
+        demo.queue()
+        demo.launch(share=False,server_port=args.port)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--sam-path", type=str, default="/app/ROCKET-1/rocket/realtime_sam/checkpoints")
+    parser.add_argument("--molmo-id", type=str, default="molmo-72b-0924")
+    parser.add_argument("--molmo-url", type=str, default="http://127.0.0.1:8000/v1")
+    args = parser.parse_args()
+    draw_gradio_components(args)