Skip to content

Commit

Permalink
Fixed agent properties addition in scene plotter.
Browse files Browse the repository at this point in the history
  • Loading branch information
KieranRatcliffeInvertedAI committed Aug 1, 2024
1 parent 3bb9f77 commit 9ec2ea5
Showing 1 changed file with 19 additions and 9 deletions.
28 changes: 19 additions & 9 deletions invertedai/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -865,7 +865,7 @@ def reset_recording(self):
self.traffic_lights_history = []
self.agent_properties = None
self.conditional_agents = []
self.agent_attributes = None
self.agent_properties = None
self.agent_face_colors = None
self.agent_edge_colors = None

Expand Down Expand Up @@ -893,7 +893,8 @@ def record_step(
def plot_scene(
self,
agent_states: List[AgentState],
agent_attributes: List[AgentAttributes],
agent_attributes: Optional[List[AgentAttributes]] = None,
agent_properties: Optional[List[AgentProperties]] = None,
traffic_light_states: Optional[Dict[int, TrafficLightState]] = None,
conditional_agents: Optional[List[int]] = None,
ax: Optional[Axes] = None,
Expand All @@ -913,6 +914,9 @@ def plot_scene(
agent_attributes:
Static attributes of the agent, which don’t change over the course of a simulation. We assume every agent is a rectangle obeying a kinematic
bicycle model.
agent_properties:
Static attributes of the agent (with the AgentProperties data type), which don’t change over the course of a simulation. We assume every
agent is a rectangle obeying a kinematic bicycle model.
traffic_light_states:
Optional parameter containing the state of the traffic lights to be visualized in the image. This parameter should only be used if the
corresponding map contains traffic light static actors.
Expand All @@ -934,16 +938,22 @@ def plot_scene(
index ID. A value of None in this list will use the default color. This value gets overwritten by the conditional agent color.
"""
assert (agent_attributes is not None) ^ (agent_properties is not None), \
"Either agent_attributes or agent_properties is populated. Populating both or neither field is invalid."

if agent_attributes is not None:
agent_properties = [convert_attributes_to_properties(attr) for attr in agent_attributes]

self.initialize_recording(
agent_states,
agent_attributes,
agent_states=agent_states,
agent_properties=agent_properties,
traffic_light_states=traffic_light_states,
conditional_agents=conditional_agents
)

self._validate_agent_style_data(
agent_face_colors,
agent_edge_colors
agent_face_colors=agent_face_colors,
agent_edge_colors=agent_edge_colors
)

self._plot_frame(
Expand Down Expand Up @@ -1033,12 +1043,12 @@ def _plot_frame(self, idx, ax=None, numbers=None, direction_vec=False,
self._update_frame_to(idx)

def _validate_agent_style_data(self,agent_face_colors,agent_edge_colors):
if self.agent_attributes is not None:
if self.agent_properties is not None:
if agent_face_colors is not None:
if len(agent_face_colors) != len(self.agent_attributes):
if len(agent_face_colors) != len(self.agent_properties):
raise Exception("Number of agent face colors does not match number of agents.")
if agent_edge_colors is not None:
if len(agent_edge_colors) != len(self.agent_attributes):
if len(agent_edge_colors) != len(self.agent_properties):
raise Exception("Number of agent edge colors does not match number of agents.")

self.agent_face_colors = agent_face_colors
Expand Down

0 comments on commit 9ec2ea5

Please sign in to comment.