Update app.py

app.py

@@ -1,47 +1,93 @@
 import streamlit as st
 import numpy as np
-import matplotlib.pyplot as plt
-from model import load_model, generate_trajectory
-from environment import RobotEnv
-
-# Set Streamlit title
-st.title("6-DOF Robot Motion Planning (Diffusion Model)")
-
-# Load model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-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)
+import pybullet as p
+import pybullet_data
+import time
+import tensorflow as tf
+
+# Streamlit app
+st.title("6-DOF Robotic Arm Path Finder with Pretrained Model")
+
+# User inputs
+start_position = st.text_input("Enter start position (x, y, z):", "0, 0, 0")
+target_position = st.text_input("Enter target position (x, y, z):", "1, 1, 1")
+obstacles = st.text_area("Enter obstacle positions (x, y, z, radius):", "0.5, 0.5, 0.5, 0.2")
+
+# Parse inputs
+start = np.array([float(x) for x in start_position.split(",")])
+target = np.array([float(x) for x in target_position.split(",")])
+obstacle_list = [tuple(map(float, line.split(","))) for line in obstacles.split("\n")]
+
+# Load pretrained model (replace with your model)
+def load_pretrained_model():
+    # Example: Load a simple neural network
+    model = tf.keras.Sequential([
+        tf.keras.layers.Dense(64, activation='relu', input_shape=(6,)),  # Input: start + target
+        tf.keras.layers.Dense(64, activation='relu'),
+        tf.keras.layers.Dense(3)  # Output: next waypoint
+    ])
+    # Load pretrained weights (dummy for demonstration)
+    model.load_weights("pretrained_model_weights.h5")
+    return model
+
+# Predict path using the pretrained model
+def predict_path(model, start, target, obstacles):
+    # Prepare input (start + target + obstacle info)
+    input_data = np.concatenate([start, target, np.array(obstacles).flatten()])
+    input_data = np.expand_dims(input_data, axis=0)
+
+    # Predict waypoints
+    waypoints = model.predict(input_data)
+    return waypoints
+
+# PyBullet simulation
+def visualize_path(path):
+    # Initialize PyBullet
+    p.connect(p.GUI)
+    p.setAdditionalSearchPath(pybullet_data.getDataPath())
+    p.setGravity(0, 0, -10)
+
+    # Load plane and robot arm
+    plane_id = p.loadURDF("plane.urdf")
+    robot_id = p.loadURDF("kuka_iiwa/model.urdf", [0, 0, 0])
+
+    # Add obstacles
+    for obstacle in obstacle_list:
+        obstacle_pos = obstacle[:3]
+        obstacle_radius = obstacle[3]
+        obstacle_shape = p.createCollisionShape(p.GEOM_SPHERE, radius=obstacle_radius)
+        obstacle_id = p.createMultiBody(baseMass=0, baseCollisionShapeIndex=obstacle_shape, basePosition=obstacle_pos)
+
+    # Visualize path
+    for i in range(len(path) - 1):
+        p.addUserDebugLine(path[i], path[i + 1], lineColorRGB=[1, 0, 0], lineWidth=2)
+        p.addUserDebugText(f"{path[i]}", path[i] + [0, 0, 0.2], textColorRGB=[1, 0, 0])
+
+    # Move the robot arm along the path
+    for point in path:
+        # Set target position (simplified for demonstration)
+        p.resetBasePositionAndOrientation(robot_id, point, [0, 0, 0, 1])
+        time.sleep(1)
+
+    # Keep the simulation running
+    st.write("Simulation running... Close the PyBullet window to continue.")
+    while True:
+        p.stepSimulation()
+        time.sleep(1 / 240)
+
+# Run the app
+if st.button("Find Path and Visualize"):
+    # Load pretrained model
+    model = load_pretrained_model()
+
+    # Predict path
+    path = predict_path(model, start, target, obstacle_list)
+
+    if path is not None:
+        st.success("Path found!")
+        st.write("Path:", path)
+
+        # Visualize the path in PyBullet
+        visualize_path(path)
+    else:
+        st.error("No path found.")