def __getitem__(self, idx):
if isinstance(idx, slice):
start, stop, step = idx.start, idx.stop, idx.step
if start is None:
start = 0
if stop is None:
stop = len(self.poses)
if step is None:
step = 1
return [
self.__getitem__(i)
for i in range(start, stop, step)
if i < len(self.poses)
]
pos, rot, label = self.poses[idx]
new_state = AgentState()
new_state.position = pos
new_state.rotation = rot
self.sim.agents[0].set_state(new_state)
obs = self.sim.get_sensor_observations()
sample = {
out_name: obs[self.out_name_to_sensor_name[out_name]]
for out_name in self.output
}
sample["label"] = self.label_map[label]
return sample
def initialize_agent(self, agent_id, initial_state=None):
agent = self.get_agent(agent_id=agent_id)
if initial_state is None:
initial_state = AgentState()
initial_state.position = self._sim.pathfinder.get_random_navigable_point(
)
initial_state.rotation = utils.quat_from_angle_axis(
np.random.uniform(0, 2.0 * np.pi), np.array([0, 1, 0]))
agent.set_state(initial_state)
self._last_state = agent.state
return agent
def initialize_agent(self, agent_id, initial_state=None):
agent = self.get_agent(agent_id=agent_id)
if initial_state is None:
initial_state = AgentState()
if self.pathfinder.is_loaded:
initial_state.position = self.pathfinder.get_random_navigable_point()
initial_state.rotation = quat_from_angle_axis(
self.random.uniform_float(0, 2.0 * np.pi), np.array([0, 1, 0])
)
agent.set_state(initial_state, is_initial=True)
self._last_state = agent.state
return agent
def __getitem__(self, idx):
if isinstance(idx, slice):
start, stop, step = idx.start, idx.stop, idx.step
if start is None:
start = 0
if stop is None:
stop = len(self.mode_to_data[self.mode])
if step is None:
step = 1
return [
self.__getitem__(i)
for i in range(start, stop, step)
if i < len(self.mode_to_data[self.mode])
]
mymode = self.mode.lower()
if self.use_caching:
cache_entry = (idx, mymode)
if cache_entry in self.cache:
return self.cache[cache_entry]
poses = self.mode_to_data[mymode]
pos, rot, label, fp = poses[idx]
# Only switch scene if it is different from the last one accessed
if fp != self.cur_fp:
self.sim.reconfigure(self._config_sim(fp, self.img_size))
self.cur_fp = fp
new_state = AgentState()
new_state.position = pos
new_state.rotation = rot
self.sim.agents[0].set_state(new_state)
obs = self.sim.get_sensor_observations()
sample = {
out_name: obs[self.out_name_to_sensor_name[out_name]]
for out_name in self.output
}
sample["label"] = label
if self.use_caching:
self.cache.add(cache_entry, sample)
return sample