Update app.py

app.py

@@ -1,5 +1,5 @@
 import numpy as np
-import matplotlib.pyplot as plt
+import plotly.graph_objects as go
 import gradio as gr
 from scipy.spatial.distance import euclidean
 from scipy.optimize import minimize
@@ -10,13 +10,17 @@ 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)
+    # Compute x, y, z coordinates for each joint
     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])
+    z = 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, z])
 
 # Objective function: Minimize the distance between the end-effector and the target
 def objective_function(theta, target):
@@ -29,39 +33,16 @@ def find_shortest_path(initial_angles, target):
     return result.x
 
 # Gradio interface
-def robot_pathfinder(target_x, target_y):
-    target = np.array([target_x, target_y])
+def robot_pathfinder(target_x, target_y, target_z):
+    target = np.array([target_x, target_y, target_z])
     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]
+    # Compute joint positions for visualization
+    joint_positions = []
+    joint_positions.append([0, 0, 0])  # Base position
     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()
+        x = joint_positions[-1][0] + L[i] * np.cos(np.sum(optimal_angles[:i+1]))
+        y = joint_positions[-1][1] + L[i] * np.sin(np.sum(optimal_angles[:i+1]))
+        z = joint_positions[-1][2] + L[i] * np.sin(np.sum(optimal_angles[:i+1]))
+        joint_positions