55import matplotlib .animation as animation
66import time
77from collections import deque
8+ from ...state .agent import AgentEnum
9+ import numpy as np
810
911class MPLVisualization (Component ):
1012 """Runs a matplotlib visualization at 10Hz.
@@ -21,8 +23,12 @@ def __init__(self, rate : float = 10.0, save_as : str = None):
2123 self .axs = None
2224 self .tstart = 0
2325 self .plot_t_range = 10
24- self .plot_values = {}
25- self .plot_events = {}
26+
27+ # Separate vehicle and pedestrian tracking
28+ self .vehicle_plot_values = {}
29+ self .pedestrian_plot_values = {}
30+ self .vehicle_plot_events = {}
31+ self .pedestrian_plot_events = {}
2632
2733 def rate (self ) -> float :
2834 return self ._rate
@@ -41,7 +47,7 @@ def initialize(self):
4147 self .writer .setup (plt .gcf (), self .save_as , dpi = 100 )
4248 plt .ion ()
4349 # to run GUI event loop
44- self .fig ,self .axs = plt .subplots (1 ,2 ,figsize = (12 ,6 ))
50+ self .fig ,self .axs = plt .subplots (1 ,3 ,figsize = (18 ,6 ))
4551 self .fig .canvas .mpl_connect ('close_event' , self .on_close )
4652 plt .show (block = False )
4753 self .tstart = time .time ()
@@ -53,49 +59,76 @@ def on_close(self,event):
5359 def debug (self , source , item , value ):
5460 t = time .time () - self .tstart
5561 item = source + '.' + item
56- if item not in self .plot_values :
57- self .plot_values [item ] = deque ()
58- plot = self .plot_values [item ]
59- self .plot_values [item ].append ((t ,value ))
62+ # Determine which plot dict to use based on source
63+ if source .startswith ('ped_' ):
64+ target_dict = self .pedestrian_plot_values
65+ else :
66+ target_dict = self .vehicle_plot_values
67+
68+ if item not in target_dict :
69+ target_dict [item ] = deque ()
70+ plot = target_dict [item ]
71+ plot .append ((t ,value ))
6072 while t - plot [0 ][0 ] > self .plot_t_range :
6173 plot .popleft ()
6274
6375 def debug_event (self , source , event ):
6476 t = time .time () - self .tstart
6577 event = source + '.' + event
66- if event not in self .plot_events :
67- self .plot_events [event ] = deque ()
68- plot = self .plot_events [event ]
78+ target_dict = self .pedestrian_plot_events if source .startswith ('ped_' ) else self .vehicle_plot_events
79+
80+ if event not in target_dict :
81+ target_dict [event ] = deque ()
82+ plot = target_dict [event ]
6983 plot .append (t )
7084 while t - plot [0 ] > self .plot_t_range :
7185 plot .popleft ()
7286
7387 def update (self , state ):
7488 if not plt .fignum_exists (self .fig .number ):
75- #plot closed
7689 return
7790 self .num_updates += 1
78- self .debug ("vehicle" ,"velocity" ,state .vehicle .v )
79- self .debug ("vehicle" ,"front wheel angle" ,state .vehicle .front_wheel_angle )
91+
92+ # Vehicle metrics
93+ self .debug ("vehicle" , "velocity" , state .vehicle .v )
94+ self .debug ("vehicle" , "front_wheel_angle" , state .vehicle .front_wheel_angle )
95+
96+ # Pedestrian metrics
97+ for agent_id , agent in state .agents .items ():
98+ if agent .type == AgentEnum .PEDESTRIAN :
99+ # self.debug(f"ped_{agent_id}", "x", agent.pose.x)
100+ # self.debug(f"ped_{agent_id}", "y", agent.pose.y)
101+ self .debug (f"ped_{ agent_id } " , "velocity" , np .linalg .norm (agent .velocity )) # Magnitude of resultant velocity
102+ self .debug (f"ped_{ agent_id } " , "yaw_rate" , agent .yaw_rate )
103+
80104 time_str = time .strftime ("%Y-%m-%d %H:%M:%S" , time .localtime (state .t ))
81105 #frame=ObjectFrameEnum.CURRENT
82106 #state = state.to_frame(frame)
83107 xrange = [state .vehicle .pose .x - 10 , state .vehicle .pose .x + 10 ]
84108 yrange = [state .vehicle .pose .y - 10 , state .vehicle .pose .y + 10 ]
85109 #plot main visualization
86110 mpl_visualization .plot (state ,title = "Scene %d at %s" % (self .num_updates ,time_str ),xrange = xrange ,yrange = yrange ,show = False ,ax = self .axs [0 ])
87- #plot figure on axs[1]
111+
112+ # Vehicle plot (axs[1])
88113 self .axs [1 ].clear ()
89- for k ,v in self .plot_values .items ():
114+ for k ,v in self .vehicle_plot_values .items ():
90115 t = [x [0 ] for x in v ]
91116 y = [x [1 ] for x in v ]
92117 self .axs [1 ].plot (t ,y ,label = k )
93- for i ,(k ,v ) in enumerate (self .plot_events .items ()):
94- for t in v :
95- self .axs [1 ].axvline (x = t ,linestyle = '--' ,color = 'C' + str (i ),label = k )
118+ self .axs [1 ].set_title ('Vehicle Metrics' )
96119 self .axs [1 ].set_xlabel ('Time (s)' )
97120 self .axs [1 ].legend ()
98121
122+ # Pedestrian plot (axs[2])
123+ self .axs [2 ].clear ()
124+ for k ,v in self .pedestrian_plot_values .items ():
125+ t = [x [0 ] for x in v ]
126+ y = [x [1 ] for x in v ]
127+ self .axs [2 ].plot (t ,y ,label = k )
128+ self .axs [2 ].set_title ('Pedestrian Metrics' )
129+ self .axs [2 ].set_xlabel ('Time (s)' )
130+ self .axs [2 ].legend ()
131+
99132 self .fig .canvas .draw_idle ()
100133 self .fig .canvas .flush_events ()
101134
0 commit comments