def test_change_state():
random_state = np.random.get_state()
np.random.seed(233)
scene_graph = habitat_sim.SceneGraph()
agent = habitat_sim.Agent(scene_graph.get_root_node().create_child())
for _ in range(100):
state = agent.state
state.position += np.random.uniform(-1, 1, size=3)
state.rotation *= quat_from_angle_axis(np.random.uniform(0, 2 * np.pi),
np.array([0.0, 1.0, 0.0]))
for v in state.sensor_states.values():
v.position += np.random.uniform(-1, 1, size=3)
v.rotation *= quat_from_angle_axis(np.random.uniform(0, 2 * np.pi),
np.array([1.0, 1.0, 1.0]))
agent.set_state(state, infer_sensor_states=False)
new_state = agent.state
_check_state_same(state, new_state)
for k, v in state.sensor_states.items():
assert k in new_state.sensor_states
_check_state_same(v, new_state.sensor_states[k])
np.random.set_state(random_state)
def place_agent(sim):
# place our agent in the scene
agent_state = habitat_sim.AgentState()
agent_state.position = [5.0, 0.0, 1.0]
agent_state.rotation = quat_from_angle_axis(
math.radians(70), np.array([0, 1.0, 0])) * quat_from_angle_axis(
math.radians(-20), np.array([1.0, 0, 0]))
agent = sim.initialize_agent(0, agent_state)
return agent.scene_node.transformation_matrix()
def get_observation(self, *args: Any, observations, episode,
**kwargs: Any):
if self.previous_episode != episode:
self.T_world_prev = SE3_Noise()
new_agent_state = self._sim.get_agent_state()
T_world_curr = SE3_Noise(new_agent_state.rotation,
new_agent_state.position)
self.T_world_prev = T_world_curr.inverse() * self.T_world_prev
noise_transform = SE3_Noise(
quat_from_angle_axis(np.random.uniform(-np.pi / 36, np.pi / 36),
geo.UP),
np.array([1.0, 0.0, 1.0]) * np.random.uniform(-0.02, 0.02),
)
T_world_curr_noisy = self.T_world_prev * (self.T_world_prev *
noise_transform).inverse()
self.T_curr_prev = T_world_curr_noisy
self.previous_episode = episode
agent_position = T_world_curr_noisy.trans
rotation_world_agent = T_world_curr_noisy.rot
goal_position = np.array(episode.goals[0].position, dtype=np.float32)
return self._compute_pointgoal(agent_position, rotation_world_agent,
goal_position)
def test_kinematics(sim):
cfg_settings = examples.settings.default_sim_settings.copy()
cfg_settings[
"scene"] = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
# enable the physics simulator: also clears available actions to no-op
cfg_settings["enable_physics"] = True
cfg_settings["depth_sensor"] = True
# test loading the physical scene
hab_cfg = examples.settings.make_cfg(cfg_settings)
sim.reconfigure(hab_cfg)
assert sim.get_physics_object_library_size() > 0
# test adding an object to the world
object_id = sim.add_object(0)
assert len(sim.get_existing_object_ids()) > 0
# test kinematics
I = np.identity(4)
# test get and set translation
sim.set_translation(np.array([0, 1.0, 0]), object_id)
assert np.allclose(sim.get_translation(object_id), np.array([0, 1.0, 0]))
# test get and set transform
sim.set_transformation(I, object_id)
assert np.allclose(sim.get_transformation(object_id), I)
# test get and set rotation
Q = quat_from_angle_axis(np.pi, np.array([0, 1.0, 0]))
expected = np.eye(4)
expected[0:3, 0:3] = quaternion.as_rotation_matrix(Q)
sim.set_rotation(quat_to_magnum(Q), object_id)
assert np.allclose(sim.get_transformation(object_id), expected)
assert np.allclose(quat_from_magnum(sim.get_rotation(object_id)), Q)
# test object removal
old_object_id = sim.remove_object(object_id)
assert len(sim.get_existing_object_ids()) == 0
assert old_object_id == object_id
object_id = sim.add_object(0)
prev_time = 0.0
for _ in range(2):
# do some kinematics here (todo: translating or rotating instead of absolute)
sim.set_translation(np.random.rand(3), object_id)
T = sim.get_transformation(object_id)
# test getting observation
obs = sim.step(
random.choice(list(hab_cfg.agents[0].action_space.keys())))
# check that time is increasing in the world
assert sim.get_world_time() > prev_time
prev_time = sim.get_world_time()
def _strafe_impl(scene_node: habitat_sim.SceneNode, forward_amount: float,
strafe_angle: float):
forward_ax = (
np.array(scene_node.absolute_transformation().rotation_scaling())
@ habitat_sim.geo.FRONT)
rotation = quat_from_angle_axis(np.deg2rad(strafe_angle),
habitat_sim.geo.UP)
move_ax = quat_rotate_vector(rotation, forward_ax)
scene_node.translate_local(move_ax * forward_amount)
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(
np.random.uniform(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 test_kinematics():
cfg_settings = examples.settings.default_sim_settings.copy()
cfg_settings[
"scene"] = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
# enable the physics simulator: also clears available actions to no-op
cfg_settings["enable_physics"] = True
cfg_settings["depth_sensor"] = True
# test loading the physical scene
hab_cfg = examples.settings.make_cfg(cfg_settings)
with habitat_sim.Simulator(hab_cfg) as sim:
obj_mgr = sim.get_object_template_manager()
assert obj_mgr.get_num_templates() > 0
# test adding an object to the world
# get handle for object 0, used to test
obj_handle_list = obj_mgr.get_template_handles("cheezit")
object_id = sim.add_object_by_handle(obj_handle_list[0])
assert len(sim.get_existing_object_ids()) > 0
# test setting the motion type
assert sim.set_object_motion_type(
habitat_sim.physics.MotionType.STATIC, object_id)
assert (sim.get_object_motion_type(object_id) ==
habitat_sim.physics.MotionType.STATIC)
assert sim.set_object_motion_type(
habitat_sim.physics.MotionType.KINEMATIC, object_id)
assert (sim.get_object_motion_type(object_id) ==
habitat_sim.physics.MotionType.KINEMATIC)
# test kinematics
I = np.identity(4)
# test get and set translation
sim.set_translation(np.array([0, 1.0, 0]), object_id)
assert np.allclose(sim.get_translation(object_id),
np.array([0, 1.0, 0]))
# test object SceneNode
object_node = sim.get_object_scene_node(object_id)
assert np.allclose(sim.get_translation(object_id),
object_node.translation)
# test get and set transform
sim.set_transformation(I, object_id)
assert np.allclose(sim.get_transformation(object_id), I)
# test get and set rotation
Q = quat_from_angle_axis(np.pi, np.array([0, 1.0, 0]))
expected = np.eye(4)
expected[0:3, 0:3] = quaternion.as_rotation_matrix(Q)
sim.set_rotation(quat_to_magnum(Q), object_id)
assert np.allclose(sim.get_transformation(object_id), expected)
assert np.allclose(quat_from_magnum(sim.get_rotation(object_id)), Q)
# test object removal
sim.remove_object(object_id)
assert len(sim.get_existing_object_ids()) == 0
obj_handle_list = obj_mgr.get_template_handles("cheezit")
object_id = sim.add_object_by_handle(obj_handle_list[0])
prev_time = 0.0
for _ in range(2):
# do some kinematics here (todo: translating or rotating instead of absolute)
sim.set_translation(np.random.rand(3), object_id)
T = sim.get_transformation(object_id) # noqa : F841
# test getting observation
sim.step(random.choice(list(
hab_cfg.agents[0].action_space.keys())))
# check that time is increasing in the world
assert sim.get_world_time() > prev_time
prev_time = sim.get_world_time()
sim.remove_object(object_id)
# test attaching/dettaching an Agent to/from physics simulation
agent_node = sim.agents[0].scene_node
obj_handle_list = obj_mgr.get_template_handles("cheezit")
object_id = sim.add_object_by_handle(obj_handle_list[0], agent_node)
sim.set_translation(np.random.rand(3), object_id)
assert np.allclose(agent_node.translation,
sim.get_translation(object_id))
sim.remove_object(object_id, False) # don't delete the agent's node
assert agent_node.translation
# test get/set RigidState
object_id = sim.add_object_by_handle(obj_handle_list[0])
targetRigidState = habitat_sim.bindings.RigidState(
mn.Quaternion(), np.array([1.0, 2.0, 3.0]))
sim.set_rigid_state(targetRigidState, object_id)
objectRigidState = sim.get_rigid_state(object_id)
assert np.allclose(objectRigidState.translation,
targetRigidState.translation)
assert objectRigidState.rotation == targetRigidState.rotation
def init_physics_test_scene(self, num_objects):
object_position = np.array(
[-0.569043, 2.04804, 13.6156]
) # above the castle table
# turn agent toward the object
print("turning agent toward the physics!")
agent_state = self._sim.get_agent(0).get_state()
agent_to_obj = object_position - agent_state.position
agent_local_forward = np.array([0, 0, -1.0])
flat_to_obj = np.array([agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
# move the agent closer to the objects if too far (this will be projected back to floor in set)
if flat_dist_to_obj > 3.0:
agent_state.position = object_position - flat_to_obj * 3.0
# unit y normal plane for rotation
det = (
flat_to_obj[0] * agent_local_forward[2]
- agent_local_forward[0] * flat_to_obj[2]
)
turn_angle = math.atan2(det, np.dot(agent_local_forward, flat_to_obj))
agent_state.rotation = quat_from_angle_axis(turn_angle, np.array([0, 1.0, 0]))
# need to move the sensors too
for sensor in agent_state.sensor_states:
agent_state.sensor_states[sensor].rotation = agent_state.rotation
agent_state.sensor_states[
sensor
].position = agent_state.position + np.array([0, 1.5, 0])
self._sim.get_agent(0).set_state(agent_state)
# hard coded dimensions of maximum bounding box for all 3 default objects:
max_union_bb_dim = np.array([0.125, 0.19, 0.26])
# add some objects in a grid
object_lib_size = self._sim.get_physics_object_library_size()
object_init_grid_dim = (3, 1, 3)
object_init_grid = {}
assert (
object_lib_size > 0
), "!!!No objects loaded in library, aborting object instancing example!!!"
# clear the objects if we are re-running this initializer
for old_obj_id in self._sim.get_existing_object_ids():
self._sim.remove_object(old_obj_id)
for obj_id in range(num_objects):
# rand_obj_index = random.randint(0, object_lib_size - 1)
# rand_obj_index = 0 # overwrite for specific object only
rand_obj_index = self._sim_settings.get("test_object_index")
if rand_obj_index < 0: # get random object on -1
rand_obj_index = random.randint(0, object_lib_size - 1)
object_init_cell = (
random.randint(-object_init_grid_dim[0], object_init_grid_dim[0]),
random.randint(-object_init_grid_dim[1], object_init_grid_dim[1]),
random.randint(-object_init_grid_dim[2], object_init_grid_dim[2]),
)
while object_init_cell in object_init_grid:
object_init_cell = (
random.randint(-object_init_grid_dim[0], object_init_grid_dim[0]),
random.randint(-object_init_grid_dim[1], object_init_grid_dim[1]),
random.randint(-object_init_grid_dim[2], object_init_grid_dim[2]),
)
object_id = self._sim.add_object(rand_obj_index)
object_init_grid[object_init_cell] = object_id
object_offset = np.array(
[
max_union_bb_dim[0] * object_init_cell[0],
max_union_bb_dim[1] * object_init_cell[1],
max_union_bb_dim[2] * object_init_cell[2],
]
)
self._sim.set_translation(object_position + object_offset, object_id)
print(
"added object: "
+ str(object_id)
+ " of type "
+ str(rand_obj_index)
+ " at: "
+ str(object_position + object_offset)
+ " | "
+ str(object_init_cell)
)
def step(self, action, only_allowed=True):
"""
All angle calculations in this function is w.r.t habitat coordinate frame, on X-Z plane
where +Y is upward, -Z is forward and +X is rightward.
Angle 0 corresponds to +X, angle 90 corresponds to +y and 290 corresponds to 270.
:param action: action to be taken
:param only_allowed: if true, then can't step anywhere except allowed locations
:return:
Dict of observations
"""
assert self._is_episode_active, (
"episode is not active, environment not RESET or "
"STOP action called previously")
self._previous_step_collided = False
# STOP: 0, FORWARD: 1, LEFT: 2, RIGHT: 2
if action == HabitatSimActions.STOP:
self._is_episode_active = False
else:
prev_position_index = self._receiver_position_index
prev_rotation_angle = self._rotation_angle
if action == HabitatSimActions.MOVE_FORWARD:
# the agent initially faces -Z by default
self._previous_step_collided = True
for neighbor in self.graph[self._receiver_position_index]:
p1 = self.graph.nodes[
self._receiver_position_index]['point']
p2 = self.graph.nodes[neighbor]['point']
direction = int(
np.around(
np.rad2deg(np.arctan2(p2[2] - p1[2],
p2[0] - p1[0])))) % 360
if direction == self.get_orientation():
self._receiver_position_index = neighbor
self._previous_step_collided = False
break
elif action == HabitatSimActions.TURN_LEFT:
# agent rotates counterclockwise, so turning left means increasing rotation angle by 90
self._rotation_angle = (self._rotation_angle + 90) % 360
elif action == HabitatSimActions.TURN_RIGHT:
self._rotation_angle = (self._rotation_angle - 90) % 360
if self.config.CONTINUOUS_VIEW_CHANGE:
intermediate_observations = list()
fps = self.config.VIEW_CHANGE_FPS
if action == HabitatSimActions.MOVE_FORWARD:
prev_position = np.array(
self.graph.nodes[prev_position_index]['point'])
current_position = np.array(self.graph.nodes[
self._receiver_position_index]['point'])
for i in range(1, fps):
intermediate_position = prev_position + i / fps * (
current_position - prev_position)
self.set_agent_state(
intermediate_position.tolist(),
quat_from_angle_axis(
np.deg2rad(self._rotation_angle),
np.array([0, 1, 0])))
sim_obs = self._sim.get_sensor_observations()
observations = self._sensor_suite.get_observations(
sim_obs)
intermediate_observations.append(observations)
else:
for i in range(1, fps):
if action == HabitatSimActions.TURN_LEFT:
intermediate_rotation = prev_rotation_angle + i / fps * 90
elif action == HabitatSimActions.TURN_RIGHT:
intermediate_rotation = prev_rotation_angle - i / fps * 90
self.set_agent_state(
list(self.graph.nodes[
self._receiver_position_index]['point']),
quat_from_angle_axis(
np.deg2rad(intermediate_rotation),
np.array([0, 1, 0])))
sim_obs = self._sim.get_sensor_observations()
observations = self._sensor_suite.get_observations(
sim_obs)
intermediate_observations.append(observations)
if not self.config.USE_RENDERED_OBSERVATIONS:
self.set_agent_state(
list(self.graph.nodes[self._receiver_position_index]
['point']),
quat_from_angle_axis(np.deg2rad(self._rotation_angle),
np.array([0, 1, 0])))
else:
self._sim.set_agent_state(
list(self.graph.nodes[self._receiver_position_index]
['point']),
quat_from_angle_axis(np.deg2rad(self._rotation_angle),
np.array([0, 1, 0])))
self._episode_step_count += 1
# log debugging info
logging.debug(
'After taking action {}, s,r: {}, {}, orientation: {}, location: {}'
.format(action, self._source_position_index,
self._receiver_position_index, self.get_orientation(),
self.graph.nodes[self._receiver_position_index]['point']))
sim_obs = self._get_sim_observation()
if self.config.USE_RENDERED_OBSERVATIONS:
self._sim.set_sensor_observations(sim_obs)
self._prev_sim_obs = sim_obs
observations = self._sensor_suite.get_observations(sim_obs)
if self.config.CONTINUOUS_VIEW_CHANGE:
observations['intermediate'] = intermediate_observations
return observations
def run(self):
scene = self.sim.semantic_scene
objects = scene.objects
object_dict = {}
for obj in objects:
if obj == None or obj.category == None:
continue
if obj.category.name() not in object_dict:
object_dict[str(obj.category.name())] = 0
for obj in objects:
if obj == None or obj.category == None or obj.category.name(
) in self.ignore_classes: #not in self.include_classes:
continue
if object_dict[str(obj.category.name())] > 5:
continue
object_dict[str(obj.category.name())] += 1
print(object_dict[str(obj.category.name())])
# st()
#if self.verbose:
print(f"Object name is: {obj.category.name()}")
# Calculate distance to object center
obj_center = obj.obb.to_aabb().center
#print(obj_center)
obj_center = np.expand_dims(obj_center, axis=0)
#print(obj_center)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# if not valid_pts:
# continue
# plot valid points that we happen to select
# self.plot_navigable_points(valid_pts)
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
# pitch calculation
dxdz_norm = np.sqrt((dx * dx) + (dz * dz))
valid_pitch = np.degrees(np.arctan2(dy, dxdz_norm))
# binning yaw around object
nbins = 18
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
if num_valid_bins == 0:
continue
spawns_per_bin = int(self.num_views / num_valid_bins) + 2
if self.verbose:
print(f'spawns_per_bin: {spawns_per_bin}')
if self.visualize:
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt.show()
action = "do_nothing"
episodes = []
valid_pts_selected = []
cnt = 0
for b in range(nbins):
# get all angle indices in the current bin range
# st()
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
if inds_bin_cur[0].size == 0:
continue
for s in range(spawns_per_bin):
# st()
s_ind = np.random.choice(inds_bin_cur[0])
#s_ind = inds_bin_cur[0][0]
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
agent_state = habitat_sim.AgentState()
agent_state.position = pos_s + np.array([0, 1.5, 0])
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(
obj_center) - agent_state.position
agent_local_forward = np.array([0, 0, -1.0]) # y, z, x
flat_to_obj = np.array(
[agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
quat_yaw = quat_from_angle_axis(turn_angle,
np.array([0, 1.0, 0]))
# Set agent yaw rotation to look at object
agent_state.rotation = quat_yaw
# change sensor state to default
# need to move the sensors too
#print(self.agent.state.sensor_states)
for sensor in self.agent.state.sensor_states:
# st()
self.agent.state.sensor_states[
sensor].rotation = agent_state.rotation
self.agent.state.sensor_states[
sensor].position = agent_state.position # + np.array([0, 1.5, 0]) # ADDED IN UP TOP
# print("PRINT", self.agent.state.sensor_states[sensor].rotation)
# Calculate Pitch from head to object
turn_pitch = np.degrees(
math.atan2(agent_to_obj[1], flat_dist_to_obj))
num_turns = np.abs(
np.floor(turn_pitch / self.rot_interval)
).astype(
int
) # compute number of times to move head up or down by rot_interval
#print("MOVING HEAD ", num_turns, " TIMES")
movement = "look_up" if turn_pitch > 0 else "look_down"
# initiate agent
self.agent.set_state(agent_state)
self.sim.step(action)
# Rotate "head" of agent up or down based on calculated pitch angle to object to center it in view
if num_turns == 0:
pass
else:
for turn in range(num_turns):
# st()
self.sim.step(movement)
if self.verbose:
for sensor in self.agent.state.sensor_states:
print(self.agent.state.
sensor_states[sensor].rotation)
# get observations after centiering
observations = self.sim.step(action)
# Assuming all sensors have same rotation and position
observations["rotations"] = self.agent.state.sensor_states[
'color_sensor'].rotation #agent_state.rotation
observations["positions"] = self.agent.state.sensor_states[
'color_sensor'].position
mainobj_id = obj.id
main_id = int(mainobj_id[1:])
semantic = observations["semantic_sensor"]
mask = np.zeros_like(semantic)
mask[semantic == main_id] = 1
obj_is_visible = False if np.sum(mask) == 0 else True
if self.is_valid_datapoint(observations,
obj) and obj_is_visible:
if self.verbose:
print("episode is valid......")
episodes.append(observations)
if self.visualize:
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
depth = observations["depth_sensor"]
self.display_sample(rgb,
semantic,
depth,
mainobj=obj,
visualize=False)
# if self.visualize:
# rgb = observations["color_sensor"]
# semantic = observations["semantic_sensor"]
# depth = observations["depth_sensor"]
# self.display_sample(rgb, semantic, depth, mainobj=obj, visualize=False)
#print("agent_state: position", self.agent.state.position, "rotation", self.agent.state.rotation)
cnt += 1
if len(episodes) >= self.num_views:
print(f'num episodes: {len(episodes)}')
data_folder = obj.category.name() + '_' + obj.id
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
flat_obs = np.random.choice(episodes,
self.num_views,
replace=False)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(self.agent, obs, data_path, viewnum,
obj, True)
viewnum += 1
else:
print(f"Not enough episodes: f{len(episodes)}")
if self.visualize:
valid_pts_selected = np.vstack(valid_pts_selected)
self.plot_navigable_points(valid_pts_selected)
# Create observations array
observations = []
# @markdown Set how long the resutlant video should be, in seconds. The object will make 1 full revolution during this time.
video_length = 4.8 # @param {type:"slider", min:1.0, max:20.0, step:0.1}
# Sim time step
time_step = 1.0 / 60.0
# Amount to rotate per frame to make 1 full rotation
rot_amount = 2 * math.pi / (video_length / time_step)
# simulate with updated camera at each frame
start_time = sim.get_world_time()
while sim.get_world_time() - start_time < video_length:
sim.step_physics(time_step)
# rotate the agent to rotate the camera
agent_state.rotation *= ut.quat_from_angle_axis(
rot_amount, np.array([0, 1.0, 0]))
agent.set_state(agent_state)
observations.append(sim.get_sensor_observations())
# video rendering of carousel view
video_prefix = clip_short_name + "_scene_view"
if make_video:
vut.make_video(
observations,
"color_sensor_3rd_person",
"color",
output_path + video_prefix,
open_vid=show_video,
video_dims=[1280, 720],
)
def vis_panorama(env,
num,
model,
goals,
log=False,
forward_only=False,
classes=None):
device = get_device(model)
rotations = [
hutil.quat_from_angle_axis(a, np.array([0, 1, 0]))
for a in np.linspace(0, 2 * np.pi, endpoint=False, num=num)
]
pos, rot = env.agent_state()
ims = []
dists = min_dists(env, goals)
cols = 224
scale = (4.0 / num) - 0.05
rng = int(scale * cols / 2)
for r in rotations:
env.set_agent_state(pos, r * rot)
obs = env.get_observation()['rgb']
vals = model(numpy_to_imgnet(obs).unsqueeze(0).to(device))
if forward_only:
max_vals = vals[0, :, 0].detach().cpu().numpy()
else:
max_vals = vals.max(axis=2).values.squeeze().detach().cpu().numpy()
env.step(0)
dist_diff = -(min_dists(env, goals) - dists)
im = obs[0] if len(obs.shape) == 4 else obs
im = im[:, (cols // 2) - rng:(cols // 2) + rng, :]
# im[:,-1] = 0
if log:
max_vals = np.log(max_vals)
ims.append((im, max_vals, dist_diff))
im, max_vals, dist_diff = util.unzip(reversed(ims))
joined = np.concatenate(im, axis=1)
fig, axes = plt.subplots(6,
1,
gridspec_kw={
'hspace': 0,
"wspace": 0,
'height_ratios': [6, 0.5, 0.5, 0.5, 0.5, 0.5]
})
fig.subplots_adjust(hspace=0, wspace=0)
imax = axes[0]
# imax.axis('on')
pltaxes = axes[1:]
imax.set_xlim((0,joined.shape[1]))
imax.set_ylim((joined.shape[0],0))
# search for right hiehgt
low,high = 8,9
for _ in range(20):
mid = (high+low)/2
fig.set_figheight(mid)
fig.canvas.draw()
imwidth = imax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()).width
axwidth = pltaxes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted()).width
axwidth = pltaxes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted()).width
if imwidth == axwidth:
print('eq')
high = mid
else:
print('low')
low = mid
fig.set_figheight(high)
fig.savefig('vis/test.pdf',bbox_inches='tight',pad_inches=0.0)
imax.axis('on')
imax.tick_params(
axis='x', # changes apply to the x-axis
which='both', # both major and minor ticks are affected
bottom=False, # ticks along the bottom edge are off
top=False, # ticks along the top edge are off
labelbottom=False) # labels along the bottom edge are off
imax.tick_params(
axis='y', # changes apply to the x-axis
which='both', # both major and minor ticks are affected
left=False, # ticks along the bottom edge are off
labelleft=False) # labels along the bottom edge are off
vals = np.stack(max_vals).transpose()
for ax, va in zip(pltaxes, vals):
ax.imshow(va[None, :],
extent=[0, 12, 0, 1],
aspect='auto',
cmap='Wistia')
ax.set_xlim((0, 12))
ax.set_ylim((0,1))
for i, v in enumerate(va):
ax.text(i + 0.5,
0.45,
'%0.2f' % (v),
fontdict={'size': 16},
horizontalalignment='center',
verticalalignment='center')
ax.axis('on')
ax.tick_params(
axis='x', # changes apply to the x-axis
which='both', # both major and minor ticks are affected
bottom=False, # ticks along the bottom edge are off
top=False, # ticks along the top edge are off
labelbottom=False) # labels along the bottom edge are off
ax.tick_params(
axis='y', # changes apply to the x-axis
which='both', # both major and minor ticks are affected
left=False, # ticks along the bottom edge are off
labelleft=False) # labels along the bottom edge are off
ratio = joined.shape[0] / joined.shape[1]
imax.imshow(joined)
# fig.text(0.12, 0.41, 'Bed', horizontalalignment='right', fontdict={'size': 16})
# fig.text(0.12, 0.34, 'Chair', horizontalalignment='right', fontdict={'size': 16})
# fig.text(0.12, 0.271, 'Couch', horizontalalignment='right',
# fontdict={'size': 16})
# fig.text(0.12, 0.205, 'D. Table', horizontalalignment='right',
# fontdict={'size': 16})
# fig.text(0.12, 0.132, 'Toilet', horizontalalignment='right',
# fontdict={'size': 16})
# ratio = joined.shape[1] / (joined.shape[0] * 11 / 6)
ratio = joined.shape[1] / (joined.shape[0] * 8.5 / 6)
height = 8
# searched for right high to deal with text shift
# searched_height = 8.051948547363281
# fig.set_figheight(searched_height)
fig.set_figheight(height)
fig.set_figwidth(height * ratio)
fig.savefig('vis/test1.pdf',bbox_inches='tight',pad_inches=0.0)
import pdb; pdb.set_trace()
# set the agent back to where it was before
env.set_agent_state(pos, rot)
ims = list(reversed(ims))
ims = [(im, val, cor) for im, val, cor in ims]
annotated = []
dists = np.array([[v, d] for _, v, d in ims])
corrs = np.array(
[np.corrcoef(dists[:, 0, i], dists[:, 1, i])[0, 1] for i in range(5)])
return fig, corrs
def __init__(self, rot=None, trans=None):
self.rot = rot if rot is not None else quat_from_angle_axis(
0.0, geo.UP)
self.trans = trans if trans is not None else np.array([0.0, 0.0, 0.0])
def get_midpoint_obj_conf(self):
scene = self.sim.semantic_scene
objects = scene.objects
print(objects)
#objects = random.sample(list(objects), self.num_objects_per_episode)
xyz_obj_mids = []
print(self.num_objects_per_episode)
print(len(objects))
count = 0
while count < self.num_objects_per_episode:
print("GETTING OBJECT #", count)
obj_ind = np.random.randint(low=0, high=len(objects))
obj = objects[obj_ind]
if obj == None or obj.category == None or obj.category.name(
) not in self.include_classes:
continue
# st()
if self.verbose:
print(f"Object name is: {obj.category.name()}")
# Calculate distance to object center
obj_center = obj.obb.to_aabb().center
#print(obj_center)
obj_center = np.expand_dims(obj_center, axis=0)
#print(obj_center)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# if not valid_pts:
# continue
# plot valid points that we happen to select
# self.plot_navigable_points(valid_pts)
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
# pitch calculation
dxdz_norm = np.sqrt((dx * dx) + (dz * dz))
valid_pitch = np.degrees(np.arctan2(dy, dxdz_norm))
# binning yaw around object
nbins = 18
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
# if self.visualize:
# import matplotlib.cm as cm
# colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
# #plt.figure(2)
# #plt.clf()
# for bi in range(nbins):
# cur_bi = np.where(bin_yaw==(bi+1))
# points = valid_pts[cur_bi]
# x_sample = points[:,0]
# z_sample = points[:,2]
# plt.plot(z_sample, x_sample, 'o', color = next(colors))
# plt.plot(obj_center[:,2], obj_center[:,0], 'x', color = 'black')
# plt.show()
# plt.pause(0.5)
# plt.close()
action = "do_nothing"
episodes = []
valid_pts_selected = []
bin_inds = list(range(nbins))
bin_inds = random.sample(bin_inds, len(bin_inds))
# b_inds_notempty = []
# # get rid of empty bins
# for b in bin_inds:
# inds_bin_cur = np.where(bin_yaw==(b+1))
# if not inds_bin_cur[0].size == 0:
# b_inds_notempty.append(b)
for b in bin_inds: #b_inds_notempty:
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
if inds_bin_cur[0].size == 0:
continue
# get all angle indices in the current bin range
# st()
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
# st()
s_ind = np.random.choice(inds_bin_cur[0])
#s_ind = inds_bin_cur[0][0]
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
agent_state = habitat_sim.AgentState()
agent_state.position = pos_s + np.array([0, 1.5, 0])
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(obj_center) - agent_state.position
agent_local_forward = np.array([0, 0, -1.0]) # y, z, x
flat_to_obj = np.array([agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
quat_yaw = quat_from_angle_axis(turn_angle,
np.array([0, 1.0, 0]))
# Set agent yaw rotation to look at object
agent_state.rotation = quat_yaw
# change sensor state to default
# need to move the sensors too
print(self.agent.state.sensor_states)
for sensor in self.agent.state.sensor_states:
# st()
self.agent.state.sensor_states[
sensor].rotation = agent_state.rotation
self.agent.state.sensor_states[
sensor].position = agent_state.position # + np.array([0, 1.5, 0]) # ADDED IN UP TOP
# print("PRINT", self.agent.state.sensor_states[sensor].rotation)
# NOTE: for finding an object, i dont think we'd want to center it
# Calculate Pitch from head to object
turn_pitch = np.degrees(
math.atan2(agent_to_obj[1], flat_dist_to_obj))
num_turns = np.abs(
np.floor(turn_pitch / self.rot_interval)
).astype(
int
) # compute number of times to move head up or down by rot_interval
print("MOVING HEAD ", num_turns, " TIMES")
movement = "look_up" if turn_pitch > 0 else "look_down"
# initiate agent
self.agent.set_state(agent_state)
self.sim.step(action)
# Rotate "head" of agent up or down based on calculated pitch angle to object to center it in view
if num_turns == 0:
pass
else:
for turn in range(num_turns):
# st()
self.sim.step(movement)
if self.verbose:
for sensor in self.agent.state.sensor_states:
print(self.agent.state.sensor_states[sensor].
rotation)
# get observations after centiering
observations = self.sim.step(action)
####### %%%%%%%%%%%%%%%%%%%%%%% ######### MASK RCNN
im = observations["color_sensor"]
im = Image.fromarray(im, mode="RGBA")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
# plt.imshow(im)
# plt.show()
outputs = self.maskrcnn(im)
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(
self.cfg_det.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(
outputs['instances'].to("cpu"))
seg_im = out.get_image()
if self.visualize:
plt.imshow(seg_im)
plt.show()
pred_masks = outputs['instances'].pred_masks
pred_boxes = outputs['instances'].pred_boxes.tensor
pred_classes = outputs['instances'].pred_classes
pred_scores = outputs['instances'].scores
maskrcnn_to_catname = {
56: "chair",
59: "bed",
61: "toilet",
57: "couch",
58: "indoor-plant",
72: "refrigerator",
62: "tv",
60: "dining-table"
}
obj_ids = []
obj_catids = []
obj_scores = []
obj_masks = []
obj_all_catids = []
obj_all_scores = []
obj_all_boxes = []
for segs in range(len(pred_masks)):
if pred_classes[segs].item() in maskrcnn_to_catname:
if pred_scores[segs] >= 0.80:
obj_ids.append(segs)
obj_catids.append(pred_classes[segs].item())
obj_scores.append(pred_scores[segs].item())
obj_masks.append(pred_masks[segs])
obj_all_catids.append(pred_classes[segs].item())
obj_all_scores.append(pred_scores[segs].item())
y, x = torch.where(pred_masks[segs])
pred_box = torch.Tensor(
[min(y), min(x),
max(y), max(x)]) # ymin, xmin, ymax, xmax
obj_all_boxes.append(pred_box)
print("MASKS ", len(pred_masks))
print("VALID ", len(obj_scores))
print(obj_scores)
print(pred_scores.shape)
translation_ = self.agent.state.sensor_states[
'depth_sensor'].position
quaternion_ = self.agent.state.sensor_states[
'depth_sensor'].rotation
rotation_ = quaternion.as_rotation_matrix(quaternion_)
T_world_cam = np.eye(4)
T_world_cam[0:3, 0:3] = rotation_
T_world_cam[0:3, 3] = translation_
if not obj_masks:
continue
else:
# randomly choose a high confidence object
obj_mask_focus = random.choice(obj_masks)
depth = observations["depth_sensor"]
xs, ys = np.meshgrid(
np.linspace(-1 * 256 / 2., 1 * 256 / 2., 256),
np.linspace(1 * 256 / 2., -1 * 256 / 2., 256))
depth = depth.reshape(1, 256, 256)
xs = xs.reshape(1, 256, 256)
ys = ys.reshape(1, 256, 256)
xys = np.vstack(
(xs * depth, ys * depth, -depth, np.ones(depth.shape)))
xys = xys.reshape(4, -1)
xy_c0 = np.matmul(np.linalg.inv(self.K), xys)
xyz = xy_c0.T[:, :3].reshape(256, 256, 3)
#xyz = self.get_point_cloud_from_z(depth, self.camera_matrix, scale=1)
xyz_obj_masked = xyz[obj_mask_focus]
xyz_obj_masked = np.matmul(
rotation_, xyz_obj_masked.T) + translation_.reshape(
3, 1)
xyz_obj_mid = np.mean(xyz_obj_masked, axis=1)
print("MIDPOINT=", xyz_obj_mid)
xyz_obj_mids.append(xyz_obj_mid)
count += 1
break # got an object
#################################
# if self.visualize:
# rgb = observations["color_sensor"]
# semantic = observations["semantic_sensor"]
# depth = observations["depth_sensor"]
# self.display_sample(rgb, semantic, depth, mainobj=obj, visualize=False)
#print("agent_state: position", self.agent.state.position, "rotation", self.agent.state.rotation)
xyz_obj_mids = np.array(xyz_obj_mids)
return xyz_obj_mids
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
return DefaultPredictor(cfg)
def degree_to_rad(deg):
return deg / 180 * np.pi
map_resolution = 1500
success_distance = 1
MAX_STEPS = 500
NUM_ROTATIONS = 12
# angs = np.linspace(0, 2 * np.pi, num=NUM_ROTATIONS, endpoint=False)
angs = np.linspace(0, 2 * np.pi, num=NUM_ROTATIONS + 1)[1:]
rotations = [hutil.quat_from_angle_axis(a, np.array([0, 1, 0])) for a in angs]
# Quat is reference to pointing toward negative z direction
# i.e. 1,0,0,0 -> -z and rotation angle is counter clockwise around the y axis, towards negative x
def check_movement(config, env, start_ang, planner=None, rng=random):
points = []
for _ in range(100):
dist = rng.uniform(0.9, 2)
ang = rng.uniform(-degree_to_rad(7), degree_to_rad(7)) + start_ang
translation = np.array([-math.sin(ang), 0, -math.cos(ang)]) * dist
points.append(translation + env.pos)
if len(points) == 0: return None
point_index = planner.reachable_nearby(points)
if point_index is not None:
assert np.linalg.norm(s2.position - s1.position - expected.delta_pos) < 1e-5
assert (
angle_between_quats(s2.rotation * expected.delta_rot.inverse(), s1.rotation)
< 1e-5
)
default_body_control_testdata = [
("move_backward", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.BACK)),
("move_forward", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.FRONT)),
("move_right", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.RIGHT)),
("move_left", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.LEFT)),
(
"turn_right",
ExpectedDelta(
delta_rot=quat_from_angle_axis(np.deg2rad(10.0), habitat_sim.geo.GRAVITY)
),
),
(
"turn_left",
ExpectedDelta(
delta_rot=quat_from_angle_axis(np.deg2rad(10.0), habitat_sim.geo.UP)
),
),
]
@pytest.mark.parametrize("action,expected", default_body_control_testdata)
def test_default_body_contorls(action, expected):
scene_graph = habitat_sim.SceneGraph()
agent_config = habitat_sim.AgentConfiguration()
def main(dataset):
parser = argparse.ArgumentParser()
parser.add_argument(
"--config-path",
type=str,
default='baselines/config/{}/train_telephone/pointgoal_rgb.yaml'.format(dataset)
)
parser.add_argument(
"opts",
default=None,
nargs=argparse.REMAINDER,
help="Modify config options from command line",
)
args = parser.parse_args()
config = get_config(args.config_path, opts=args.opts)
config.defrost()
config.TASK_CONFIG.SIMULATOR.AGENT_0.SENSORS = ["RGB_SENSOR", "DEPTH_SENSOR"]
config.freeze()
simulator = None
scene_obs = defaultdict(dict)
num_obs = 0
scene_obs_dir = 'data/scene_observations/' + dataset
os.makedirs(scene_obs_dir, exist_ok=True)
metadata_dir = 'data/metadata/' + dataset
for scene in os.listdir(metadata_dir):
scene_obs = dict()
scene_metadata_dir = os.path.join(metadata_dir, scene)
points, graph = load_metadata(scene_metadata_dir)
if dataset == 'replica':
scene_mesh_dir = os.path.join('data/scene_datasets', dataset, scene, 'habitat/mesh_semantic.ply')
else:
scene_mesh_dir = os.path.join('data/scene_datasets', dataset, scene, scene + '.glb')
for node in graph.nodes():
agent_position = graph.nodes()[node]['point']
for angle in [0, 90, 180, 270]:
agent_rotation = quat_to_coeffs(quat_from_angle_axis(np.deg2rad(angle), np.array([0, 1, 0]))).tolist()
goal_radius = 0.00001
goal = NavigationGoal(
position=agent_position,
radius=goal_radius
)
episode = NavigationEpisode(
goals=[goal],
episode_id=str(0),
scene_id=scene_mesh_dir,
start_position=agent_position,
start_rotation=agent_rotation,
info={'sound': 'telephone'}
)
episode_sim_config = merge_sim_episode_config(config.TASK_CONFIG.SIMULATOR, episode)
if simulator is None:
simulator = SoundSpaces(episode_sim_config)
simulator.reconfigure(episode_sim_config)
obs, rotation_index = simulator.step(None)
scene_obs[(node, rotation_index)] = obs
num_obs += 1
print('Total number of observations: {}'.format(num_obs))
with open(os.path.join(scene_obs_dir, '{}.pkl'.format(scene)), 'wb') as fo:
pickle.dump(scene_obs, fo)
simulator.close()
del simulator
def run(self):
scene = self.sim.semantic_scene
objects = scene.objects
id = 0
for obj in objects:
if obj == None or obj.category == None or obj.category.name(
) not in self.include_classes:
continue
if self.verbose:
print(f"Object name is: {obj.category.name()}")
# Calculate distance to object center
obj_center = obj.obb.to_aabb().center
obj_center = np.expand_dims(obj_center, axis=0)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
nbins = 200
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
spawns_per_bin = 1 #int(self.num_views / num_valid_bins) + 2
print(f'spawns_per_bin: {spawns_per_bin}')
if self.visualize:
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt.show()
action = "do_nothing"
episodes = []
valid_pts_selected = []
cnt = 0
texts_images = []
texts_depths = []
rots_cam0 = []
camXs_T_camX0_4x4 = []
camX0_T_camXs_4x4 = []
origin_T_camXs = []
origin_T_camXs_t = []
rots_cam0.append([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
rots_origin = []
for b in range(nbins):
# get all angle indices in the current bin range
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
if inds_bin_cur[0].size == 0:
print("Continuing")
break
for s in range(spawns_per_bin):
if b == 0:
np.random.seed(1)
s_ind = np.random.choice(inds_bin_cur[0])
pos_s = valid_pts[s_ind]
pos_s_prev = pos_s
else:
# get closest valid position for smooth trajectory
pos_s_cur_all = valid_pts[inds_bin_cur[0]]
distances_to_prev = np.sqrt(
np.sum((pos_s_cur_all - pos_s_prev)**2, axis=1))
argmin_s_ind_cur = np.argmin(distances_to_prev)
s_ind = inds_bin_cur[0][argmin_s_ind_cur]
pos_s = valid_pts[s_ind]
pos_s_prev = pos_s
valid_pts_selected.append(pos_s)
agent_state = habitat_sim.AgentState()
agent_state.position = pos_s + np.array([0, 1.5, 0])
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(
obj_center) - agent_state.position
agent_local_forward = np.array([0, 0, -1.0]) # y, z, x
flat_to_obj = np.array(
[agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
quat_yaw = quat_from_angle_axis(turn_angle,
np.array([0, 1.0, 0]))
# Set agent yaw rotation to look at object
agent_state.rotation = quat_yaw
# change sensor state to default
# need to move the sensors too
if self.verbose:
print(self.agent.state.sensor_states)
for sensor in self.agent.state.sensor_states:
# st()
self.agent.state.sensor_states[
sensor].rotation = agent_state.rotation
self.agent.state.sensor_states[
sensor].position = agent_state.position # + np.array([0, 1.5, 0]) # ADDED IN UP TOP
# Calculate Pitch from head to object
turn_pitch = np.degrees(
math.atan2(agent_to_obj[1], flat_dist_to_obj))
num_turns = np.abs(
np.floor(turn_pitch / self.rot_interval)
).astype(
int
) # compute number of times to move head up or down by rot_interval
if self.verbose:
print("MOVING HEAD ", num_turns, " TIMES")
movement = "look_up" if turn_pitch > 0 else "look_down"
# initiate agent
self.agent.set_state(agent_state)
self.sim.step(action)
# Rotate "head" of agent up or down based on calculated pitch angle to object to center it in view
if num_turns == 0:
pass
else:
for turn in range(num_turns):
# st()
self.sim.step(movement)
if self.verbose:
for sensor in self.agent.state.sensor_states:
print(self.agent.state.
sensor_states[sensor].rotation)
# get observations after centiering
observations = self.sim.step(action)
# Assuming all sensors have same rotation and position
observations["rotations"] = self.agent.state.sensor_states[
'color_sensor'].rotation #agent_state.rotation
observations["positions"] = self.agent.state.sensor_states[
'color_sensor'].position
if self.is_valid_datapoint(observations, obj):
if self.verbose:
print("episode is valid......")
episodes.append(observations)
if self.visualize:
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
depth = observations["depth_sensor"]
self.display_sample(rgb,
semantic,
depth,
mainobj=obj,
visualize=False)
if cnt > 0:
origin_T_camX = quaternion.as_rotation_matrix(
episodes[cnt]["rotations"])
camX0_T_camX = np.matmul(camX0_T_origin, origin_T_camX)
# camX0_T_camX = np.matmul(camX0_T_origin, origin_T_camX)
origin_T_camXs.append(origin_T_camX)
origin_T_camXs_t.append(episodes[cnt]["positions"])
origin_T_camX_4x4 = np.eye(4)
origin_T_camX_4x4[0:3, 0:3] = origin_T_camX
origin_T_camX_4x4[:3, 3] = episodes[cnt]["positions"]
camX0_T_camX_4x4 = np.matmul(camX0_T_origin_4x4,
origin_T_camX_4x4)
camX_T_camX0_4x4 = self.safe_inverse_single(
camX0_T_camX_4x4)
camXs_T_camX0_4x4.append(camX_T_camX0_4x4)
camX0_T_camXs_4x4.append(camX0_T_camX_4x4)
# # trans_camX0_T_camX = episodes[cnt]["positions"] - episodes[0]["positions"]
# r_camX_T_camX0, t_camX_T_camX0, = self.split_rt_single(camX_T_camX0_4x4)
# r_camX_T_camX0_quat = quaternion.as_float_array(quaternion.from_rotation_matrix(r_camX_T_camX0))
# # print(t_camX_T_camX0)
# # print(trans_camX0_T_camX)
# # full_trans1 = np.hstack((cnt, trans, quat_diff1))
# full_trans_camX_T_camX0 = np.hstack((cnt, t_camX_T_camX0, r_camX_T_camX0_quat))
# # rots1.append(list(full_trans1))
# rots_cam0.append(list(full_trans_camX_T_camX0))
# # rots_origin.append(list(full_trans_origin))
camX0_T_camX_4x4 = self.safe_inverse_single(
camX_T_camX0_4x4)
origin_T_camX_4x4 = np.matmul(origin_T_camX0_4x4,
camX0_T_camX_4x4)
r_origin_T_camX, t_origin_T_camX, = self.split_rt_single(
origin_T_camX_4x4)
if self.verbose:
print(r_origin_T_camX)
print(origin_T_camX)
else:
origin_T_camX0_quat = episodes[0]["rotations"]
origin_T_camX0 = quaternion.as_rotation_matrix(
origin_T_camX0_quat)
camX0_T_origin = np.linalg.inv(origin_T_camX0)
# camX0_T_origin = self.safe_inverse_single(origin_T_camX0)
origin_T_camXs.append(origin_T_camX0)
origin_T_camXs_t.append(episodes[0]["positions"])
origin_T_camX0_4x4 = np.eye(4)
origin_T_camX0_4x4[0:3, 0:3] = origin_T_camX0
origin_T_camX0_4x4[:3, 3] = episodes[0]["positions"]
camX0_T_origin_4x4 = self.safe_inverse_single(
origin_T_camX0_4x4)
camXs_T_camX0_4x4.append(np.eye(4))
# r0 = quaternion.as_rotation_vector(episodes[0]["rotations"])
# quat_diff_origin = quaternion.as_float_array(episodes[0]["rotations"]) #quaternion.as_float_array(self.quaternion_from_two_vectors(self.origin_rot_vector,r0))
# full_trans_origin = np.hstack((cnt, episodes[0]["positions"], quat_diff_origin))
# rots_origin.append(list(full_trans_origin))
cnt += 1
id += 1
print("Generated ", len(episodes), " views.")
if len(episodes) < self.min_orbslam_views:
print("NOT ENOUGH VIEWS FOR ORBSLAM")
continue
if self.save_imgs:
with open(self.orbslam_path + "/" + 'rgb.txt', 'w') as f:
for item in texts_images:
f.write("%s\n" % item)
with open(self.orbslam_path + "/" + 'depth.txt', 'w') as f:
for item in texts_depths:
f.write("%s\n" % item)
# with open(self.orbslam_path + "/" + 'CamTraj1.txt', 'w') as f:
# for item in rots1:
# f.write("%s\n" % item)
# with open(self.orbslam_path + "/" + 'CamTraj.txt', 'w') as f:
# for item in rots_cam0:
# f.write("%s\n" % item)
############ ORBSLAM ################
if self.do_orbslam:
camXs_T_camX0_4x4 = np.array(camXs_T_camX0_4x4)
origin_T_camXs = np.array(origin_T_camXs)
origin_T_camXs_t = np.array(origin_T_camXs_t)
camX0_T_camXs_4x4 = np.array(camX0_T_camXs_4x4)
orbslam_dir = "/home/nel/ORB_SLAM2"
executable = os.path.join(orbslam_dir,
"Examples/RGB-D/rgbd_tum")
vocabulary_file = os.path.join(orbslam_dir,
"Vocabulary/ORBvoc.txt")
settings_file = os.path.join(orbslam_dir,
"Examples/RGB-D/REPLICA.yaml")
data_dir = os.path.join(orbslam_dir, "Examples/RGB-D/replica")
associations_file = os.path.join(
orbslam_dir, "Examples/RGB-D/associations/replica.txt")
associations_executable = os.path.join(
orbslam_dir, "Examples/RGB-D/associations/associate.py")
rgb_file = os.path.join(orbslam_dir,
"Examples/RGB-D/replica/rgb.txt")
depth_file = os.path.join(orbslam_dir,
"Examples/RGB-D/replica/depth.txt")
# associate rgb and depth files
print("ASSOCIATING FILES")
associate(rgb_file, depth_file, associations_file)
# run ORBSLAM2
print("RUNNING ORBSLAM2...")
try:
output = subprocess.run([
executable, vocabulary_file, settings_file, data_dir,
associations_file
],
capture_output=True,
text=True,
check=True,
timeout=800)
except:
print("PROCESS TIMED OUT")
continue
# load camera rotation and translation estimated by orbslam - these are wrt first frame
camXs_T_camX0_orb_output = np.loadtxt(self.ORBSLAM_Cam_Traj)
print(output)
assert len(episodes) == camXs_T_camX0_orb_output.shape[
0], f"{camXs_T_camX0_orb_output.shape[0]}, {len(episodes)}"
camXs_T_camX0_quant = []
camXs_T_camX0 = []
for i in range(camXs_T_camX0_orb_output.shape[0]):
cur = camXs_T_camX0_orb_output[i]
camX_T_camX0_quant = np.quaternion(cur[7], cur[4], cur[5],
cur[6])
camX_T_camX0 = quaternion.as_rotation_matrix(
camX_T_camX0_quant
) # Need to negative rotation because axes are flipped
camXs_T_camX0_quant.append(camX_T_camX0_quant)
camXs_T_camX0.append(camX_T_camX0)
camXs_T_camX0 = np.array(camXs_T_camX0)
camXs_T_camX0_quant = np.array(camXs_T_camX0_quant)
t = []
for i in range(camXs_T_camX0_orb_output.shape[0]):
cur = camXs_T_camX0_orb_output[i]
t_cur = np.array([-cur[1], -cur[2], -cur[3]
]) # i think x shouldn't be inverted
t.append(t_cur)
t = np.array(t)
camXs_T_camX0_4x4_orb = []
for i in range(camXs_T_camX0_orb_output.shape[0]):
# assert len(episodes) == camXs_T_camX0_orb_output.shape[0], f"{camXs_T_camX0_orb_output.shape[0]}, {len(episodes)}"
# get estimated 4x4
camX_T_camX0_4x4_orb = np.eye(4)
camX_T_camX0_4x4_orb[0:3, 0:3] = camXs_T_camX0[i]
camX_T_camX0_4x4_orb[:3, 3] = t[i]
# invert
camX0_T_camX_4x4_orb = self.safe_inverse_single(
camX_T_camX0_4x4_orb)
# convert to origin coordinates
origin_T_camX_4x4 = np.matmul(origin_T_camX0_4x4,
camX0_T_camX_4x4_orb)
r_origin_T_camX_orb, t_origin_T_camX_orb = self.split_rt_single(
origin_T_camX_4x4)
r_origin_T_camX_orb_quat = quaternion.from_rotation_matrix(
r_origin_T_camX_orb)
#save
episodes[i]["positions_orb"] = t_origin_T_camX_orb
episodes[i]["rotations_orb"] = r_origin_T_camX_orb_quat
camXs_T_camX0_4x4_orb.append(camX_T_camX0_4x4_orb)
camXs_T_camX0_4x4_orb = np.array(camXs_T_camX0_4x4_orb)
## PLOTTTING #########
# x_orbslam = []
# y_orbslam = []
# z_orbslam = []
# for i in range(camXs_T_camX0_orb_output.shape[0]):
# camX_T_camX0_4x4_orb_r, camX_T_camX0_4x4_orb_t = self.split_rt_single(camXs_T_camX0_4x4_orb[i])
# x_orbslam.append(camX_T_camX0_4x4_orb_t[0])
# y_orbslam.append(camX_T_camX0_4x4_orb_t[1])
# z_orbslam.append(camX_T_camX0_4x4_orb_t[2])
# x_orbslam = np.array(x_orbslam)
# y_orbslam = np.array(y_orbslam)
# z_orbslam = np.array(z_orbslam)
# x_gt = []
# y_gt = []
# z_gt = []
# for i in range(camXs_T_camX0_orb_output.shape[0]):
# camX_T_camX0_4x4_r, camX_T_camX0_4x4_t = self.split_rt_single(camXs_T_camX0_4x4[i])
# x_gt.append(camX_T_camX0_4x4_t[0])
# y_gt.append(camX_T_camX0_4x4_t[1])
# z_gt.append(camX_T_camX0_4x4_t[2])
# x_gt = np.array(x_gt)
# y_gt = np.array(y_gt)
# z_gt = np.array(z_gt)
print("PLOTTING")
x_orbslam = []
y_orbslam = []
z_orbslam = []
for i in range(camXs_T_camX0_orb_output.shape[0]):
x_orbslam.append(episodes[i]["positions_orb"][0])
y_orbslam.append(episodes[i]["positions_orb"][1])
z_orbslam.append(episodes[i]["positions_orb"][2])
x_orbslam = np.array(x_orbslam)
y_orbslam = np.array(y_orbslam)
z_orbslam = np.array(z_orbslam)
x_gt = []
y_gt = []
z_gt = []
for i in range(camXs_T_camX0_orb_output.shape[0]):
x_gt.append(episodes[i]["positions"][0])
y_gt.append(episodes[i]["positions"][1])
z_gt.append(episodes[i]["positions"][2])
x_gt = np.array(x_gt)
y_gt = np.array(y_gt)
z_gt = np.array(z_gt)
from mpl_toolkits.mplot3d import Axes3D
plt.figure()
ax = plt.axes(projection='3d')
ax.plot3D(x_orbslam, y_orbslam, z_orbslam, 'green')
ax.plot3D(x_gt, y_gt, z_gt, 'blue')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
# plt.plot(-np.array(x_orbslam), np.array(y_orbslam), label='ORB-SLAM', color='green', linestyle='dashed')
# plt.plot(x_gt, y_gt, label='GT', color='blue', linestyle='solid')
# plt.legend()
plt_name = 'maps/' + str(
self.ep_idx) + '_' + str(id) + '_' + 'map.png'
plt.savefig(plt_name)
if not self.do_orbslam:
if len(episodes) >= self.num_views:
print(f'num episodes: {len(episodes)}')
data_folder = obj.category.name() + '_' + obj.id
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
np.random.seed(1)
flat_obs = np.random.choice(episodes,
self.num_views,
replace=False)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(self.agent, obs, data_path,
viewnum, obj.id, True)
viewnum += 1
else:
print(f"Not enough episodes: f{len(episodes)}")
# the len of episodes is sometimes greater than camXs_T_camX0_orb_output.shape[0]
# so we need to sample from episode number less than camXs_T_camX0_orb_output.shape[0]
else:
num_valid_episodes = camXs_T_camX0_orb_output.shape[0]
if num_valid_episodes >= self.num_views:
print(f'num episodes: {num_valid_episodes}')
data_folder = obj.category.name() + '_' + obj.id
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
np.random.seed(1)
flat_obs = np.random.choice(episodes[:num_valid_episodes],
self.num_views,
replace=False)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(self.agent, obs, data_path,
viewnum, obj.id, True)
viewnum += 1
else:
print(f"Not enough episodes: f{len(episodes)}")
if self.visualize:
valid_pts_selected = np.vstack(valid_pts_selected)
self.plot_navigable_points(valid_pts_selected)
def test_kinematics():
cfg_settings = examples.settings.default_sim_settings.copy()
cfg_settings[
"scene"] = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
# enable the physics simulator: also clears available actions to no-op
cfg_settings["enable_physics"] = True
cfg_settings["depth_sensor"] = True
# test loading the physical scene
hab_cfg = examples.settings.make_cfg(cfg_settings)
with habitat_sim.Simulator(hab_cfg) as sim:
# get the rigid object attributes manager, which manages
# templates used to create objects
obj_template_mgr = sim.get_object_template_manager()
obj_template_mgr.load_configs("data/objects/example_objects/", True)
assert obj_template_mgr.get_num_templates() > 0
# get the rigid object manager, which provides direct
# access to objects
rigid_obj_mgr = sim.get_rigid_object_manager()
# test adding an object to the world
# get handle for object 0, used to test
obj_handle_list = obj_template_mgr.get_template_handles("cheezit")
cheezit_box = rigid_obj_mgr.add_object_by_template_handle(
obj_handle_list[0])
assert rigid_obj_mgr.get_num_objects() > 0
assert (len(rigid_obj_mgr.get_object_handles()) ==
rigid_obj_mgr.get_num_objects())
# test setting the motion type
cheezit_box.motion_type = habitat_sim.physics.MotionType.STATIC
assert cheezit_box.motion_type == habitat_sim.physics.MotionType.STATIC
cheezit_box.motion_type = habitat_sim.physics.MotionType.KINEMATIC
assert cheezit_box.motion_type == habitat_sim.physics.MotionType.KINEMATIC
# test kinematics
I = np.identity(4)
# test get and set translation
cheezit_box.translation = [0.0, 1.0, 0.0]
assert np.allclose(cheezit_box.translation, np.array([0.0, 1.0, 0.0]))
# test object SceneNode
assert np.allclose(cheezit_box.translation,
cheezit_box.root_scene_node.translation)
# test get and set transform
cheezit_box.transformation = I
assert np.allclose(cheezit_box.transformation, I)
# test get and set rotation
Q = quat_from_angle_axis(np.pi, np.array([0.0, 1.0, 0.0]))
expected = np.eye(4)
expected[0:3, 0:3] = quaternion.as_rotation_matrix(Q)
cheezit_box.rotation = quat_to_magnum(Q)
assert np.allclose(cheezit_box.transformation, expected)
assert np.allclose(quat_from_magnum(cheezit_box.rotation), Q)
# test object removal
rigid_obj_mgr.remove_object_by_id(cheezit_box.object_id)
assert rigid_obj_mgr.get_num_objects() == 0
obj_handle_list = obj_template_mgr.get_template_handles("cheezit")
cheezit_box = rigid_obj_mgr.add_object_by_template_handle(
obj_handle_list[0])
prev_time = 0.0
for _ in range(2):
# do some kinematics here (todo: translating or rotating instead of absolute)
cheezit_box.translation = np.random.rand(3)
T = cheezit_box.transformation # noqa : F841
# test getting observation
sim.step(random.choice(list(
hab_cfg.agents[0].action_space.keys())))
# check that time is increasing in the world
assert sim.get_world_time() > prev_time
prev_time = sim.get_world_time()
rigid_obj_mgr.remove_object_by_id(cheezit_box.object_id)
# test attaching/dettaching an Agent to/from physics simulation
agent_node = sim.agents[0].scene_node
obj_handle_list = obj_template_mgr.get_template_handles("cheezit")
cheezit_agent = rigid_obj_mgr.add_object_by_template_handle(
obj_handle_list[0], agent_node)
cheezit_agent.translation = np.random.rand(3)
assert np.allclose(agent_node.translation, cheezit_agent.translation)
rigid_obj_mgr.remove_object_by_id(
cheezit_agent.object_id,
delete_object_node=False) # don't delete the agent's node
assert agent_node.translation
# test get/set RigidState
cheezit_box = rigid_obj_mgr.add_object_by_template_handle(
obj_handle_list[0])
targetRigidState = habitat_sim.bindings.RigidState(
mn.Quaternion(), np.array([1.0, 2.0, 3.0]))
cheezit_box.rigid_state = targetRigidState
objectRigidState = cheezit_box.rigid_state
assert np.allclose(objectRigidState.translation,
targetRigidState.translation)
assert objectRigidState.rotation == targetRigidState.rotation
def run(self):
object_centers, cat_ids_objects = self.get_midpoint_obj_conf()
# check if any two 3d centers are too close to each other
print("OBJ CENTERS BEFORE=", object_centers.shape[0])
object_centers_unique = []
cat_ids_objects_unique = []
object_centers_unique.append(object_centers[-1, :])
cat_ids_objects_unique.append(cat_ids_objects[-1])
for i in range(object_centers.shape[0] - 1):
dists = np.sqrt(
np.sum((object_centers[i, :] - object_centers[i + 1:])**2,
axis=1))
far = dists > self.closeness_threshold
if np.all(far):
object_centers_unique.append(object_centers[i, :])
cat_ids_objects_unique.append(cat_ids_objects[i])
object_centers_unique = np.array(object_centers_unique)
print("OBJ CENTERS AFTER=", object_centers_unique.shape[0])
# Print all categories in the house that are relevant
print("OBJ CAT IDs UNIQUE=", cat_ids_objects_unique)
objs_in_house = []
scene = self.sim.semantic_scene
objects = scene.objects
for obj in objects:
if obj == None or obj.category == None or obj.category.name(
) not in self.include_classes:
continue
objs_in_house.append(obj.category.name())
print("ALL CATEGORIES IN HOUSE=", objs_in_house)
if self.visualize_obj:
plt.figure(1)
x_sample = self.nav_pts[:, 0]
z_sample = self.nav_pts[:, 2]
plt.plot(z_sample, x_sample, 'o', color='red')
plt.plot(object_centers_unique[:, 2],
object_centers_unique[:, 0],
'x',
color='green')
plt.figure(2)
x_sample = self.nav_pts[:, 0]
z_sample = self.nav_pts[:, 2]
plt.plot(z_sample, x_sample, 'o', color='red')
plt.plot(object_centers[:, 2],
object_centers[:, 0],
'x',
color='green')
plt.show()
for obj_idx in range(object_centers_unique.shape[0]):
# Calculate distance to object center
obj_center = object_centers_unique[obj_idx, :]
#print(obj_center)
#obj_center = np.expand_dims(object_centers_unique, axis=0)
obj_center = np.reshape(obj_center, (1, 3))
print("OBJECT CENTER: ", obj_center)
#print(self.nav_pts.shape)
#print(obj_center)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# if not valid_pts:
# continue
# plot valid points that we happen to select
# self.plot_navigable_points(valid_pts)
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
# pitch calculation
dxdz_norm = np.sqrt((dx * dx) + (dz * dz))
valid_pitch = np.degrees(np.arctan2(dy, dxdz_norm))
# binning yaw around object
nbins = 18
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
spawns_per_bin = int(self.num_views / num_valid_bins) + 2
print(f'spawns_per_bin: {spawns_per_bin}')
if self.visualize:
print("PLOTTING")
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt.show()
action = "do_nothing"
episodes = []
valid_pts_selected = []
cnt = 0
for b in range(nbins):
# get all angle indices in the current bin range
# st()
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
if inds_bin_cur[0].size == 0:
print("NO BINS")
continue
for s in range(spawns_per_bin):
# st()
s_ind = np.random.choice(inds_bin_cur[0])
#s_ind = inds_bin_cur[0][0]
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
agent_state = habitat_sim.AgentState()
agent_state.position = pos_s + np.array([0, 1.5, 0])
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(
obj_center) - agent_state.position
agent_local_forward = np.array([0, 0, -1.0]) # y, z, x
flat_to_obj = np.array(
[agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
quat_yaw = quat_from_angle_axis(turn_angle,
np.array([0, 1.0, 0]))
# Set agent yaw rotation to look at object
agent_state.rotation = quat_yaw
# change sensor state to default
# need to move the sensors too
for sensor in self.agent.state.sensor_states:
# st()
self.agent.state.sensor_states[
sensor].rotation = agent_state.rotation
self.agent.state.sensor_states[
sensor].position = agent_state.position # + np.array([0, 1.5, 0]) # ADDED IN UP TOP
# print("PRINT", self.agent.state.sensor_states[sensor].rotation)
# Calculate Pitch from head to object
turn_pitch = np.degrees(
math.atan2(agent_to_obj[1], flat_dist_to_obj))
num_turns = np.abs(
np.floor(turn_pitch / self.rot_interval)
).astype(
int
) # compute number of times to move head up or down by rot_interval
#print("MOVING HEAD ", num_turns, " TIMES")
movement = "look_up" if turn_pitch > 0 else "look_down"
# initiate agent
self.agent.set_state(agent_state)
self.sim.step(action)
# Rotate "head" of agent up or down based on calculated pitch angle to object to center it in view
if num_turns == 0:
pass
else:
for turn in range(num_turns):
# st()
self.sim.step(movement)
if self.verbose:
for sensor in self.agent.state.sensor_states:
print(self.agent.state.
sensor_states[sensor].rotation)
# get observations after centiering
observations = self.sim.step(action)
# Assuming all sensors have same rotation and position
observations["rotations"] = self.agent.state.sensor_states[
'color_sensor'].rotation #agent_state.rotation
observations["positions"] = self.agent.state.sensor_states[
'color_sensor'].position
if self.verbose:
print("episode is valid......")
episodes.append(observations)
if self.visualize:
im = observations["color_sensor"]
im = Image.fromarray(im, mode="RGBA")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
plt.imshow(im)
plt.show()
cnt += 1
print(f'num episodes: {len(episodes)}')
if len(episodes) >= self.num_views:
print(f'num episodes: {len(episodes)}')
data_folder = 'obj' + str(
obj_idx) #obj.category.name() + '_' + obj.id
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
flat_obs = np.random.choice(episodes,
self.num_views,
replace=False)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(self.agent, obs, data_path, viewnum,
obj_idx, True)
viewnum += 1
else:
print(f"Not enough episodes: f{len(episodes)}")
if self.visualize:
valid_pts_selected = np.vstack(valid_pts_selected)
self.plot_navigable_points(valid_pts_selected)
def _check_state_expected(s1, s2, expected: ExpectedDelta):
assert np.linalg.norm(s2.position - s1.position -
expected.delta_pos) < 1e-5
assert (angle_between_quats(s2.rotation * expected.delta_rot.inverse(),
s1.rotation) < 1e-5)
default_body_control_testdata = [
("move_backward", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.BACK)),
("move_forward", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.FRONT)),
("move_right", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.RIGHT)),
("move_left", ExpectedDelta(delta_pos=0.25 * habitat_sim.geo.LEFT)),
(
"turn_right",
ExpectedDelta(delta_rot=quat_from_angle_axis(np.deg2rad(10.0),
habitat_sim.geo.GRAVITY)),
),
(
"turn_left",
ExpectedDelta(delta_rot=quat_from_angle_axis(np.deg2rad(10.0),
habitat_sim.geo.UP)),
),
]
@pytest.mark.parametrize("action,expected", default_body_control_testdata)
def test_default_body_contorls(action, expected):
scene_graph = habitat_sim.SceneGraph()
agent_config = habitat_sim.AgentConfiguration()
agent_config.action_space = dict(
move_backward=habitat_sim.ActionSpec(
def run(self):
scene = self.sim.semantic_scene
objects = scene.objects
for obj in objects:
if obj == None or obj.category == None or obj.category.name(
) not in self.include_classes:
continue
# st()
if self.verbose:
print(f"Object name is: {obj.category.name()}")
# Calculate distance to object center
obj_center = obj.obb.to_aabb().center
#print(obj_center)
obj_center = np.expand_dims(obj_center, axis=0)
#print(obj_center)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# if not valid_pts:
# continue
# plot valid points that we happen to select
# self.plot_navigable_points(valid_pts)
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
valid_yaw = np.degrees(
np.arctan2(valid_pts_shift[:, 2], valid_pts_shift[:, 0]))
nbins = 18
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
spawns_per_bin = int(self.num_views / num_valid_bins) + 2
print(f'spawns_per_bin: {spawns_per_bin}')
if self.visualize:
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt.show()
# 1. angle calculation: https://math.stackexchange.com/questions/2521886/how-to-find-angle-between-2-points-in-3d-space
# 2. angle quantization (round each angle to the nearst bin, 163 --> 160, 47 --> 50, etc.) - STILL NEED
'''
# 3. loop around the sphere and get views
for vert_angle in vertical_angle:
for hori_angle in horizontal_angle:
# for i in range(num_view_per_angle)
valid_pts_bin[vert_angle, hori_angle][np.random.choice()]
# spawn agent
# rotate the agent with our calculated angle (so that the agent looks at the object)
# if is_valid_datapoint: save it to views
# self.display_sample(rgb, semantic, depth)
'''
action = "do_nothing"
# flat_views = [] #flatview corresponds to moveup=100
# any_views = []
episodes = []
# spawns_per_bin = 2
valid_pts_selected = []
for b in range(nbins):
# get all angle indices in the current bin range
# st()
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
if inds_bin_cur[0].size == 0:
continue
for s in range(spawns_per_bin):
# st()
s_ind = np.random.choice(inds_bin_cur[0])
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
agent_state = habitat_sim.AgentState()
agent_state.position = pos_s
agent_to_obj = np.squeeze(
obj_center) - agent_state.position
agent_local_forward = np.array([0, 0, -1.0]) # y, z, x
flat_to_obj = np.array(
[agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
quat_yaw = quat_from_angle_axis(turn_angle,
np.array([0, 1.0, 0]))
agent_state.rotation = quat_yaw
# change sensor state to default
# need to move the sensors too
for sensor in agent_state.sensor_states:
agent_state.sensor_states[
sensor].rotation = agent_state.rotation
agent_state.sensor_states[
sensor].position = agent_state.position + np.array[
0, 1.5, 0]
# agent_state.rotation = rot
self.agent.set_state(agent_state)
observations = self.sim.step(action)
observations["rotations"] = agent_state.rotation
observations["positions"] = agent_state.position
if self.is_valid_datapoint(observations, obj):
if self.verbose:
print("episode is valid......")
episodes.append(observations)
if self.visualize:
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
depth = observations["depth_sensor"]
self.display_sample(rgb,
semantic,
depth,
mainobj=obj,
visualize=False)
# if self.visualize:
# rgb = observations["color_sensor"]
# semantic = observations["semantic_sensor"]
# depth = observations["depth_sensor"]
# self.display_sample(rgb, semantic, depth, mainobj=obj, visualize=False)
#print("agent_state: position", self.agent.state.position, "rotation", self.agent.state.rotation)
if len(episodes) >= self.num_views:
print(f'num episodes: {len(episodes)}')
data_folder = obj.category.name() + '_' + obj.id
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
flat_obs = np.random.choice(episodes,
self.num_views,
replace=False)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(self.agent, obs, data_path, viewnum,
obj.id, True)
viewnum += 1
else:
print(f"Not enough episodes: f{len(episodes)}")
if self.visualize:
valid_pts_selected = np.vstack(valid_pts_selected)
self.plot_navigable_points(valid_pts_selected)