Update app.py

app.py

@@ -1,93 +1,67 @@
-import streamlit as st
 import numpy as np
-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.")
+import matplotlib.pyplot as plt
+import gradio as gr
+from scipy.spatial.distance import euclidean
+from scipy.optimize import minimize
+
+# Define the robot's arm lengths (6-DOF)
+L = np.array([1, 1, 1, 1, 1, 1])  # Lengths of each arm segment
+
+# Forward kinematics: Compute the end-effector position given joint angles
+def forward_kinematics(theta):
+    theta = np.array(theta)
+    x = L[0] * np.cos(theta[0]) + L[1] * np.cos(theta[0] + theta[1]) + L[2] * np.cos(theta[0] + theta[1] + theta[2]) + \
+         L[3] * np.cos(theta[0] + theta[1] + theta[2] + theta[3]) + L[4] * np.cos(theta[0] + theta[1] + theta[2] + theta[3] + theta[4]) + \
+         L[5] * np.cos(theta[0] + theta[1] + theta[2] + theta[3] + theta[4] + theta[5])
+    y = L[0] * np.sin(theta[0]) + L[1] * np.sin(theta[0] + theta[1]) + L[2] * np.sin(theta[0] + theta[1] + theta[2]) + \
+         L[3] * np.sin(theta[0] + theta[1] + theta[2] + theta[3]) + L[4] * np.sin(theta[0] + theta[1] + theta[2] + theta[3] + theta[4]) + \
+         L[5] * np.sin(theta[0] + theta[1] + theta[2] + theta[3] + theta[4] + theta[5])
+    return np.array([x, y])
+
+# Objective function: Minimize the distance between the end-effector and the target
+def objective_function(theta, target):
+    end_effector = forward_kinematics(theta)
+    return euclidean(end_effector, target)
+
+# Find the shortest path (joint angles) to reach the target
+def find_shortest_path(initial_angles, target):
+    result = minimize(objective_function, initial_angles, args=(target,), method='BFGS')
+    return result.x
+
+# Gradio interface
+def robot_pathfinder(target_x, target_y):
+    target = np.array([target_x, target_y])
+    initial_angles = np.array([0, 0, 0, 0, 0, 0])  # Initial joint angles
+    optimal_angles = find_shortest_path(initial_angles, target)
+    
+    # Plot the robot arm
+    fig, ax = plt.subplots()
+    ax.set_xlim(-5, 5)
+    ax.set_ylim(-5, 5)
+    ax.set_aspect('equal')
+    
+    # Draw the robot arm
+    x = [0]
+    y = [0]
+    for i in range(len(L)):
+        x.append(x[-1] + L[i] * np.cos(np.sum(optimal_angles[:i+1])))
+        y.append(y[-1] + L[i] * np.sin(np.sum(optimal_angles[:i+1])))
+    ax.plot(x, y, 'o-', lw=2, markersize=10)
+    
+    # Mark the target
+    ax.plot(target[0], target[1], 'rx', markersize=10)
+    
+    plt.title("6-DOF Robot Arm Pathfinding")
+    return fig
+
+# Gradio app
+iface = gr.Interface(
+    fn=robot_pathfinder,
+    inputs=["number", "number"],
+    outputs="plot",
+    live=True,
+    title="6-DOF Robot Shortest Path Finder",
+    description="Enter target coordinates (x, y) to find the shortest path for a 6-DOF robot arm."
+)
+
+iface.launch()