StackBots/utils/plot_servo.py

62 lines
1.7 KiB
Python

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import argparse
import sys
def plot_servo_data(csv_file):
# Read CSV file
try:
df = pd.read_csv(csv_file)
except FileNotFoundError:
print(f"Error: File '{csv_file}' not found")
sys.exit(1)
except Exception as e:
print(f"Error reading CSV file: {e}")
sys.exit(1)
# Preprocess data: subtract 500 and multiply by 0.5
for servo in ['Servo 1', 'Servo 2', 'Servo 3']:
df[servo] = (df[servo] - 500) * 0.25
# Set plot style
sns.set_style("whitegrid")
plt.figure(figsize=(12, 6))
# Plot servo angles
sns.lineplot(data=df, x='Step', y='Servo 1', label='Servo 1')
sns.lineplot(data=df, x='Step', y='Servo 2', label='Servo 2')
sns.lineplot(data=df, x='Step', y='Servo 3', label='Servo 3')
# Set title and labels
plt.title('Servo Angles vs Time (Normalized)')
plt.xlabel('Time (seconds)')
plt.ylabel('Angle (degrees)')
# Add legend
plt.legend(title='Servo ID')
# Optimize display
plt.tight_layout()
# Generate output filename based on input
output_file = csv_file.rsplit('.', 1)[0] + '_normalized_plot.png'
# Save plot
plt.savefig(output_file, dpi=300, bbox_inches='tight')
print(f"Plot saved as: {output_file}")
def main():
# Set up argument parser
parser = argparse.ArgumentParser(description='Plot servo angles from CSV file')
parser.add_argument('csv_file', help='Path to the CSV file containing servo data')
# Parse arguments
args = parser.parse_args()
# Generate plot
plot_servo_data(args.csv_file)
if __name__ == "__main__":
main()