From 9587d27805d87c54174cf2e02e18750ca5f861c1 Mon Sep 17 00:00:00 2001
From: Adam <adam@ineedroot.com>
Date: Wed, 26 Mar 2025 20:48:27 +0000
Subject: [PATCH] Incorperated rsi visualiser code.

---
 src/RSIPI/kuka_visualizer.py | 93 ++++++++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)
 create mode 100644 src/RSIPI/kuka_visualizer.py

diff --git a/src/RSIPI/kuka_visualizer.py b/src/RSIPI/kuka_visualizer.py
new file mode 100644
index 0000000..fa2f1ea
--- /dev/null
+++ b/src/RSIPI/kuka_visualizer.py
@@ -0,0 +1,93 @@
+import pandas as pd
+import matplotlib.pyplot as plt
+import argparse
+import os
+
+
+class KukaRSIVisualizer:
+    def __init__(self, csv_file):
+        self.csv_file = csv_file
+        if not os.path.exists(csv_file):
+            raise FileNotFoundError(f"CSV file {csv_file} not found.")
+        self.df = pd.read_csv(csv_file)
+
+    def plot_trajectory(self, save_path=None):
+        """Plots the robot's trajectory in 3D space."""
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection='3d')
+
+        ax.plot(self.df["RIst.X"], self.df["RIst.Y"], self.df["RIst.Z"], label="Actual Trajectory", linestyle='-')
+
+        if "RSol.X" in self.df.columns:
+            ax.plot(self.df["RSol.X"], self.df["RSol.Y"], self.df["RSol.Z"], label="Planned Trajectory", linestyle='--')
+
+        ax.set_xlabel("X Position")
+        ax.set_ylabel("Y Position")
+        ax.set_zlabel("Z Position")
+        ax.set_title("Robot Trajectory")
+        ax.legend()
+
+        if save_path:
+            plt.savefig(save_path)
+        plt.show()
+
+    def plot_joint_positions(self, save_path=None):
+        """Plots joint positions over time."""
+        plt.figure()
+        time_series = range(len(self.df))
+
+        for col in ["AIPos.A1", "AIPos.A2", "AIPos.A3", "AIPos.A4", "AIPos.A5", "AIPos.A6"]:
+            if col in self.df.columns:
+                plt.plot(time_series, self.df[col], label=col)
+
+        plt.xlabel("Time Steps")
+        plt.ylabel("Joint Position (Degrees)")
+        plt.title("Joint Positions Over Time")
+        plt.legend()
+
+        if save_path:
+            plt.savefig(save_path)
+        plt.show()
+
+    def plot_force_trends(self, save_path=None):
+        """Plots force trends if force data is logged."""
+        force_columns = ["PosCorr.X", "PosCorr.Y", "PosCorr.Z"]
+        plt.figure()
+        time_series = range(len(self.df))
+
+        for col in force_columns:
+            if col in self.df.columns:
+                plt.plot(time_series, self.df[col], label=col)
+
+        plt.xlabel("Time Steps")
+        plt.ylabel("Force Correction (N)")
+        plt.title("Force Trends Over Time")
+        plt.legend()
+
+        if save_path:
+            plt.savefig(save_path)
+        plt.show()
+
+    def export_graphs(self, export_dir="exports"):
+        """Exports all graphs to PNG/PDF."""
+        os.makedirs(export_dir, exist_ok=True)
+        self.plot_trajectory(save_path=os.path.join(export_dir, "trajectory.png"))
+        self.plot_joint_positions(save_path=os.path.join(export_dir, "joint_positions.png"))
+        self.plot_force_trends(save_path=os.path.join(export_dir, "force_trends.png"))
+        print(f"Graphs exported to {export_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Visualize RSI data logs.")
+    parser.add_argument("csv_file", type=str, help="Path to the RSI CSV log file.")
+    parser.add_argument("--export", action="store_true", help="Export graphs as PNG/PDF.")
+
+    args = parser.parse_args()
+    visualizer = KukaRSIVisualizer(args.csv_file)
+
+    visualizer.plot_trajectory()
+    visualizer.plot_joint_positions()
+    visualizer.plot_force_trends()
+
+    if args.export:
+        visualizer.export_graphs()