Update app.py

app.py

@@ -1,34 +1,47 @@
+import streamlit as st
+import numpy as np
+import matplotlib.pyplot as plt
+from model import load_model, generate_trajectory
+from environment import RobotEnv
 
-import torch
-from diffusion_policy.model.diffusion_policy import DiffusionPolicy
-from diffusion_policy.env.robot_env import RobotEnv  # Custom environment for robot motion
+# Set Streamlit title
+st.title("6-DOF Robot Motion Planning (Diffusion Model)")
 
-# Check if CUDA is available
+# Load model
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
-
-print("Using device:", device)
-
-# Load a pretrained diffusion model for robot motion
-model_path = "diffusion-policy/pretrained/6dof_robot"  # Example path, replace with actual model
-model = DiffusionPolicy.load_from_checkpoint(model_path, map_location=device)
-model.to(device)
-model.eval()  # Set model to evaluation mode
-
-# Define the robot environment
-env = RobotEnv(num_dof=6)  # 6-DOF robot environment
-start_state = env.get_random_start_state()  # Get a random starting position
-goal_state = env.get_random_goal_state()  # Get a random goal position
-
-# Generate trajectory using diffusion model
-num_steps = 50
-trajectory = torch.zeros((num_steps, 6), device=device, dtype=dtype)  # Store trajectory
-
-with torch.no_grad():  # No gradients needed for inference
-    state = start_state.clone().to(device)
-    for t in range(num_steps):
-        action = model.predict(state)  # Predict the next action using the diffusion model
-        state = env.step(state, action)  # Move the robot in simulation
-        trajectory[t] = state
-
-print("Generated 6-DOF trajectory:\n", trajectory.cpu().numpy())
+model = load_model(device)
+
+# Create robot environment
+env = RobotEnv()
+
+# User input for start and goal state
+st.sidebar.header("Set Start & Goal Positions")
+start_state = env.get_random_start_state()
+goal_state = env.get_random_goal_state()
+
+start_state = st.sidebar.text_input("Start State (comma-separated)", ",".join(map(str, start_state)))
+goal_state = st.sidebar.text_input("Goal State (comma-separated)", ",".join(map(str, goal_state)))
+
+# Convert input to NumPy array
+try:
+    start_state = np.array([float(x) for x in start_state.split(",")])
+    goal_state = np.array([float(x) for x in goal_state.split(",")])
+except ValueError:
+    st.sidebar.error("Enter valid numbers separated by commas.")
+    start_state = env.get_random_start_state()
+    goal_state = env.get_random_goal_state()
+
+# Generate trajectory button
+if st.button("Generate Trajectory"):
+    trajectory = generate_trajectory(model, start_state, goal_state)
+
+    # Plot the trajectory
+    fig, ax = plt.subplots()
+    for i in range(6):  # 6-DOF joints
+        ax.plot(trajectory[:, i], label=f"Joint {i+1}")
+
+    ax.set_title("Generated 6-DOF Robot Trajectory")
+    ax.set_xlabel("Time Steps")
+    ax.set_ylabel("Joint Angles")
+    ax.legend()
+    st.pyplot(fig)