def set_agent(self, sim_settings):
agent = self.sim.initialize_agent(sim_settings["default_agent"])
agent_state = habitat_sim.AgentState()
agent_state.position = np.array([1.5, 1.072447, 0.0])
#agent_state.position = np.array([1.0, 3.0, 1.0])
agent.set_state(agent_state)
agent_state = agent.get_state()
if self.verbose:
print("agent_state: position", agent_state.position, "rotation", agent_state.rotation)
self.agent = agent
def configure_agent(self, agent_settings):
# Robot init
agent = self.sim.initialize_agent(agent_settings)
# Set initial pose
agent_state = habitat_sim.AgentState()
agent_state.position = np.array([1.5, 1, 0.0]) # Global frame ref
agent.set_state(agent_state)
return agent
def place_agent(
sim: habitat_sim.Simulator,
agent_pos: list,
agent_orient: list,
) -> magnum.Matrix4:
"""Sets AgentState to some reasonable values and return a transformation matrix."""
# place our agent in the scene
agent_state = habitat_sim.AgentState()
agent_state.position = agent_pos
# agent_state.position = [-0.15, -1.6, 1.0]
agent_state.rotation = np.quaternion(*agent_orient)
agent = sim.initialize_agent(0, agent_state)
return agent.scene_node.transformation_matrix()
def test_sensors(scene, has_sem, sensor_type, sim, make_cfg_settings):
if not osp.exists(scene):
pytest.skip("Skipping {}".format(scene))
make_cfg_settings = {k: v for k, v in make_cfg_settings.items()}
make_cfg_settings["semantic_sensor"] = has_sem
make_cfg_settings["scene"] = scene
sim.reconfigure(make_cfg(make_cfg_settings))
with open(
osp.abspath(
osp.join(
osp.dirname(__file__),
"gt_data",
"{}-state.json".format(osp.basename(
osp.splitext(scene)[0])),
)),
"r",
) as f:
render_state = json.load(f)
state = habitat_sim.AgentState()
state.position = render_state["pos"]
state.rotation = habitat_sim.utils.quat_from_coeffs(
render_state["rot"])
sim.initialize_agent(0, state)
obs = sim.step("move_forward")
assert sensor_type in obs, f"{sensor_type} not in obs"
gt = np.load(
osp.abspath(
osp.join(
osp.dirname(__file__),
"gt_data",
"{}-{}.npy".format(osp.basename(osp.splitext(scene)[0]),
sensor_type),
)))
# Different GPUs and different driver version will produce slightly different images
assert np.linalg.norm(obs[sensor_type].astype(np.float) -
gt.astype(np.float)) < 1.5e-2 * np.linalg.norm(
gt.astype(
np.float)), f"Incorrect {sensor_type} output"
def _render_and_load_gt(sim, scene, sensor_type, gpu2gpu):
gt_data_pose_file = osp.abspath(
osp.join(
osp.dirname(__file__),
"gt_data",
"{}-state.json".format(osp.basename(osp.splitext(scene)[0])),
))
with open(gt_data_pose_file, "r") as f:
render_state = json.load(f)
state = habitat_sim.AgentState()
state.position = render_state["pos"]
state.rotation = quat_from_coeffs(render_state["rot"])
sim.initialize_agent(0, state)
obs = sim.step("move_forward")
assert sensor_type in obs, f"{sensor_type} not in obs"
# now that sensors are constructed, test some getter/setters
sim.get_agent(0)._sensors[sensor_type].fov = mn.Deg(80)
assert sim.get_agent(0)._sensors[sensor_type].fov == mn.Deg(
80), "fov not set correctly"
assert sim.get_agent(0)._sensors[sensor_type].hfov == mn.Deg(
80), "hfov not set correctly"
gt_obs_file = osp.abspath(
osp.join(
osp.dirname(__file__),
"gt_data",
"{}-{}.npy".format(osp.basename(osp.splitext(scene)[0]),
sensor_type),
))
gt = np.load(gt_obs_file)
if gpu2gpu:
torch = pytest.importorskip("torch")
for k, v in obs.items():
if torch.is_tensor(v):
obs[k] = v.cpu().numpy()
return obs, gt
def init_sim(sim_settings, start_pos, start_rot):
cfg = make_cfg(sim_settings)
sim = habitat_sim.Simulator(cfg)
# Actions should change the position of the agent as well.
action = habitat_sim.registry.get_move_fn("move_up")
action.body_action = True
action = habitat_sim.registry.get_move_fn("move_down")
action.body_action = True
random.seed(sim_settings["seed"])
sim.seed(sim_settings["seed"])
# Set agent state
agent = sim.initialize_agent(sim_settings["default_agent"])
agent_state = habitat_sim.AgentState()
agent_state.position = np.array(start_pos) # Agent start position set
agent_state.rotation = quaternion.quaternion(
*start_rot) # Agent start orientation set
agent.set_state(agent_state)
return sim, agent, cfg
def _render_and_load_gt(sim, scene, sensor_type, gpu2gpu):
gt_data_pose_file = osp.abspath(
osp.join(
osp.dirname(__file__),
"gt_data",
"{}-state.json".format(osp.basename(osp.splitext(scene)[0])),
)
)
with open(gt_data_pose_file, "r") as f:
render_state = json.load(f)
state = habitat_sim.AgentState()
state.position = render_state["pos"]
state.rotation = quat_from_coeffs(render_state["rot"])
sim.initialize_agent(0, state)
obs = sim.step("move_forward")
assert sensor_type in obs, f"{sensor_type} not in obs"
gt_obs_file = osp.abspath(
osp.join(
osp.dirname(__file__),
"gt_data",
"{}-{}.npy".format(osp.basename(osp.splitext(scene)[0]), sensor_type),
)
)
gt = np.load(gt_obs_file)
if gpu2gpu:
import torch
for k, v in obs.items():
if torch.is_tensor(v):
obs[k] = v.cpu().numpy()
return obs, gt
def get_midpoint_obj_conf(self):
#objects = random.sample(list(objects), self.num_objects_per_episode)
xyz_obj_mids = []
cat_ids_objects = []
count = 0
nav_points = np.array(self.nav_pts)
action = "do_nothing"
movement = "turn_right"
if self.visualize_maskrcnn:
self.plot_navigable_points(self.nav_pts)
# constraint if two fixation points are very close
# add in all confident masks
# 20 spawns
while count < self.num_spawns_per_episode:
print("GETTING OBJECT #", count)
rand_ind = np.random.randint(low=0, high=len(nav_points))
pos_rand = nav_points[rand_ind, :]
agent_state = habitat_sim.AgentState()
agent_state.position = pos_rand + np.array([0, 1.5, 0])
print("Random POS=", agent_state.position)
self.agent.set_state(agent_state)
# remove spawning points close to this spawn (so as to not spawn near there again)
distances = np.sqrt(np.sum((nav_points - pos_rand)**2, axis=1))
nav_points = nav_points[np.where(distances > self.radius_remove)]
# if self.visualize:
# x_sample = self.nav_pts[:,0]
# z_sample = self.nav_pts[:,2]
# plt.plot(z_sample, x_sample, 'o', color = 'red')
# plt.plot(agent_state.position[2], agent_state.position[0], 'x', 'blue')
# plt.show()
# rotate in place until get high confident object
# bin_angle_size = 60.0
# angles = np.arange(-180, 180, bin_angle_size)
angles = np.arange(0, 360, self.turn_angle)
for angle in angles: #b_inds_notempty:
# print("ANGLE=", angle)
# turn_angle = np.radians(angle)
# quat_yaw = quat_from_angle_axis(angle, np.array([0, 1.0, 0]))
# # Set agent yaw rotation to current angle
# 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)
# print(agent_state.rotation)
# # get observations after centiering
# self.agent.set_state(agent_state)
# rotate until find confident object
observations = self.sim.step(movement) #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)
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.90:
obj_ids.append(segs)
obj_catids.append(pred_classes[segs].item())
obj_scores.append(pred_scores[segs].item())
obj_masks.append(pred_masks[segs])
cat_ids_objects.append(maskrcnn_to_catname[int(
pred_classes[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
# instead of this I think we should iterate over ALL the high confident objects and fixate on them
# obj_mask_focus = random.choice(obj_masks)
for obj_mask in 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_obj_masked = xyz[obj_mask_focus]
xyz_obj_masked = xyz[obj_mask]
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
if self.visualize_maskrcnn:
plt.figure(1)
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()
plt.imshow(seg_im)
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(nav_points[:, 2],
nav_points[:, 0],
'o',
color='green')
plt.plot(agent_state.position[2],
agent_state.position[0], 'x', 'blue')
plt.show()
break # got a high confident view
xyz_obj_mids = np.array(xyz_obj_mids)
return xyz_obj_mids, cat_ids_objects
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 reconfigure_scene(env, scene_path):
"""This function takes a habitat scene and adds relevant changes to the
scene that make it useful for locobot assistant.
For example, it sets initial positions to be deterministic, and it
adds some humans to the environment a reasonable places
"""
sim = env._robot.base.sim
# these are useful to know to understand co-ordinate normalization
# and where to place objects to be within bounds. They change wildly
# from scene to scene
print("scene bounds: {}".format(sim.pathfinder.get_bounds()))
agent = sim.get_agent(0)
# keep the initial rotation consistent across runs
old_agent_state = agent.get_state()
new_agent_state = habitat_sim.AgentState()
new_agent_state.position = old_agent_state.position
new_agent_state.rotation = np.quaternion(1.0, 0.0, 0.0, 0.0)
agent.set_state(new_agent_state)
env._robot.base.init_state = agent.get_state()
env.initial_rotation = new_agent_state.rotation
###########################
# scene-specific additions
###########################
scene_name = os.path.basename(scene_path).split(".")[0]
if scene_name == "mesh_semantic":
# this must be Replica Dataset
scene_name = os.path.basename(
os.path.dirname(os.path.dirname(scene_path)))
supported_scenes = ["skokloster-castle", "van-gogh-room", "apartment_0"]
if scene_name not in supported_scenes:
print("Scene {} not in supported scenes, so skipping adding objects".
format(scene_name))
return
assets_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), "../test/test_assets"))
if hasattr(sim, 'get_object_template_manager'):
load_object_configs = sim.get_object_template_manager(
).load_object_configs
else:
load_object_configs = sim.load_object_configs
human_male_template_id = load_object_configs(
os.path.join(assets_path, "human_male"))[0]
human_female_template_id = load_object_configs(
os.path.join(assets_path, "human_female"))[0]
if scene_name == "apartment_0":
id_male = sim.add_object(human_male_template_id)
id_female = sim.add_object(human_female_template_id)
print("id_male, id_female: {} {}".format(id_male, id_female))
sim.set_translation([1.2, -0.81, 0.3],
id_female) # apartment_0, female
sim.set_translation([1.2, -0.75, -0.3], id_male) # apartment_0, male
rot = mn.Quaternion.rotation(mn.Deg(-75), mn.Vector3.y_axis())
sim.set_rotation(rot, id_male) # apartment_0
sim.set_rotation(rot, id_female) # apartment_0
elif scene_name == "skokloster-castle":
id_female = sim.add_object(human_female_template_id)
print("id_female: {}".format(id_female))
sim.set_translation([2.0, 3.0, 15.00], id_female) # skokloster castle
elif scene_name == "van-gogh-room":
id_female = sim.add_object(human_female_template_id)
print("id_female: {}".format(id_female))
sim.set_translation([1.0, 0.84, 0.00], id_female) # van-gogh-room
# make the objects STATIC so that collisions work
for obj in [id_male, id_female]:
sim.set_object_motion_type(habitat_sim.physics.MotionType.STATIC, obj)
navmesh_settings = habitat_sim.NavMeshSettings()
navmesh_settings.set_defaults()
sim.recompute_navmesh(sim.pathfinder,
navmesh_settings,
include_static_objects=True)
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(
) in self.ignore_classes:
continue
# st()
print(f"Object name is: {obj.category.name()}")
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))
# st()
closest_navpts = self.nav_pts[np.where(
distances < self.distance_thresh)]
action = "do_nothing"
flat_views = [] #flatview corresponds to moveup=100
any_views = []
for closest_navpt in closest_navpts[:15]:
print("Launch position is: ", closest_navpt)
agent_state = habitat_sim.AgentState()
agent_state.position = closest_navpt
self.agent.set_state(agent_state)
self.sim.step(action)
observations = self.sim.step("look_down_init")
for moveup in range(0, 160, int(self.rot_interval)):
actionup = "do_nothing" if moveup == 0 else "look_up"
print(
f"actionup is: {actionup}. Moveup value is: {moveup}")
observations = self.sim.step(actionup)
for moveleft in range(0, 360, int(self.rot_interval)):
actionleft = "do_nothing" if moveleft == 0 else "turn_left"
print(f"actionleft is {actionleft}")
observations = self.sim.step(actionleft)
if self.is_valid_datapoint(observations, obj):
if moveup == 100:
flat_views.append(observations)
else:
any_views.append(observations)
print("agent_state: position",
self.agent.state.position, "rotation",
self.agent.state.rotation)
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
depth = observations["depth_sensor"]
# self.display_sample(rgb, semantic, depth)
# cv2.waitKey(0)
# st()
if len(flat_views) >= self.num_flat_views and len(
any_views) >= self.num_any_views:
data_folder = obj.category.name() + '_' + obj.id
data_path = os.path.join(self.basepath, data_folder)
os.mkdir(data_path)
flat_obs = np.random.choice(flat_views,
self.num_flat_views,
replace=False)
any_obs = np.random.choice(any_views,
self.num_any_views,
replace=False)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(self.agent, obs, data_path, viewnum,
obj.id, True)
viewnum += 1
for obs in any_obs:
self.save_datapoint(self.agent, obs, data_path, viewnum,
obj.id, False)
viewnum += 1
def runVHACDSimulation(obj_path):
# data will store data from simulation
data = {
"observations": [],
}
cfg = make_configuration()
with habitat_sim.Simulator(cfg) as sim:
sim.initialize_agent(
0,
habitat_sim.AgentState(
[5.45, 2.4, 1.2],
np.quaternion(1, 0, 0, 0),
),
)
# get the physics object attributes manager
obj_templates_mgr = sim.get_object_template_manager()
# create a list that will store the object ids
obj_ids = []
# get object template & render asset handle
obj_id_1 = obj_templates_mgr.load_configs(
str(os.path.join(data_path, obj_path)))[0]
obj_template = obj_templates_mgr.get_template_by_ID(obj_id_1)
obj_handle = obj_template.render_asset_handle
obj_templates_mgr.register_template(obj_template,
force_registration=True)
obj_ids += [obj_id_1]
# == VHACD ==
# Now, create a new object template based on the recently created obj_template
# specify parameters for running CHD
# high resolution
params = habitat_sim._ext.habitat_sim_bindings.VHACDParameters()
params.resolution = 200000
params.max_convex_hulls = 32
# run VHACD, with params passed in and render
new_handle_1 = sim.apply_convex_hull_decomposition(
obj_handle, params, True, True)
new_obj_template_1 = obj_templates_mgr.get_template_by_handle(
new_handle_1)
obj_templates_mgr.register_template(new_obj_template_1,
force_registration=True)
obj_ids += [new_obj_template_1.ID]
# Medium resolution
params = habitat_sim.VHACDParameters()
params.resolution = 100000
params.max_convex_hulls = 4
# run VHACD, with params passed in and render
new_handle_2 = sim.apply_convex_hull_decomposition(
obj_handle, params, True, True)
new_obj_template_2 = obj_templates_mgr.get_template_by_handle(
new_handle_2)
obj_templates_mgr.register_template(new_obj_template_2,
force_registration=True)
obj_ids += [new_obj_template_2.ID]
# Low resolution
params = habitat_sim.VHACDParameters()
params.resolution = 10000
params.max_convex_hulls = 1
# run VHACD, with params passed in and render
new_handle_3 = sim.apply_convex_hull_decomposition(
obj_handle, params, True, True)
new_obj_template_3 = obj_templates_mgr.get_template_by_handle(
new_handle_3)
obj_templates_mgr.register_template(new_obj_template_3,
force_registration=True)
obj_ids += [new_obj_template_3.ID]
# now display objects
cur_ids = []
for i in range(len(obj_ids)):
cur_id = sim.add_object(obj_ids[i])
cur_ids.append(cur_id)
# get length
obj_node = sim.get_object_scene_node(cur_id)
obj_bb = obj_node.cumulative_bb
length = obj_bb.size().length() * 0.8
total_length = length * len(obj_ids)
dist = math.sqrt(3) * total_length / 2
set_object_state_from_agent(
sim,
cur_id,
offset=np.array([
-total_length / 2 + total_length * i / (len(obj_ids) - 1),
1.4,
-1 * dist,
]),
)
sim.set_object_motion_type(
habitat_sim.physics.MotionType.KINEMATIC, cur_id)
vel_control = sim.get_object_velocity_control(cur_id)
vel_control.controlling_ang_vel = True
vel_control.angular_velocity = np.array([0, -1.56, 0])
# simulate for 4 seconds
simulate(sim, dt=4, get_frames=True, data=data)
for cur_id in cur_ids:
vel_control = sim.get_object_velocity_control(cur_id)
vel_control.controlling_ang_vel = True
vel_control.angular_velocity = np.array([-1.56, 0, 0])
# simulate for 4 seconds
simulate(sim, dt=4, get_frames=True, data=data)
return data
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
def render_scene(test_scene):
os.makedirs(os.path.join(SAVE_FOLDER, 'rgb', test_scene), exist_ok=True)
os.makedirs(os.path.join(SAVE_FOLDER, 'observations', test_scene), exist_ok=True)
# house_idx = test_scene.split('/')[-2]
camera_path = os.path.join(DATASET_DIR, 'cameras', test_scene + '.pkl')
with open(camera_path, 'rb') as f:
cameras = pickle.load(f)
img_width = 640
img_height = 480
sim_settings = {
"width": img_width, # Spatial resolution of the observations
"height": img_height,
"scene": os.path.join(PLANE_FOLDER, test_scene, test_scene+'.glb'), # Scene path
"default_agent": 0,
"sensor_height": 0, # Height of sensors in meters
"color_sensor": True, # RGB sensor
"semantic_sensor": True, # Semantic sensor
"depth_sensor": True, # Depth sensor
"seed": 1,
}
plane_params = np.load(os.path.join(PLANE_FOLDER, test_scene, 'house_plane_params.npy'), allow_pickle=True)
cmap = random_colors(len(plane_params))
cmap = np.array(cmap)*255
cfg = utils.make_cfg(sim_settings, intrinsics=None)
sim = habitat_sim.Simulator(cfg)
# Print semantic annotation information (id, category, bounding box details)
# about levels, regions and objects in a hierarchical fashion
scene = sim.semantic_scene
# Set agent state
agent = sim.initialize_agent(sim_settings["default_agent"])
for camera in cameras:
agent_state = habitat_sim.AgentState()
# agent_state.position = grid + np.random.random(3)*np.array([GRID_SIZE, HEIGHT_RANGE, GRID_SIZE])
agent_state.position = camera['position']
agent_state.rotation = camera['rotation']
# print(agent_state)
agent.set_state(agent_state)
observations = sim.get_sensor_observations()
rgb = observations["color_sensor"]
plane_instance_id = observations["semantic_sensor"]
depth = observations["depth_sensor"]
# plane_instance_id = select_top_k_area(plane_instance_id)
plane_instance_id = erode_planes(plane_instance_id)
plane_instance_id = filter_by_depth(plane_instance_id, depth, plane_params, camera)
plane_instance_id = convex_hull_fill_holes(plane_instance_id, depth, plane_params, camera)
plane_instance_id = remove_small_plane(plane_instance_id)
plane_instance_id = convex_hull_fill_holes(plane_instance_id, depth, plane_params, camera)
quality = check_quality(plane_instance_id)
image = observations['color_sensor'][:,:,:3]
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plane_seg = cmap[plane_instance_id.flatten()%len(cmap)].reshape((img_height,img_width,3))
blend_pred = (plane_seg * 0.5 + image * 0.5).astype(np.uint8)
blend_pred[plane_instance_id == 0] = image[plane_instance_id==0]
"""
# DEBUG
save_path = os.path.join(SAVE_FOLDER, test_scene, 'vis_plane')
os.makedirs(save_path, exist_ok=True)
if quality:
cv2.imwrite(os.path.join(save_path, camera['img_name']), blend_pred)
else:
cv2.imwrite(os.path.join(save_path, camera['img_name'].split('.')[0]+'_bad.png'), blend_pred)
"""
semantic_info = {}
planes = []
planeSegmentation = np.zeros(plane_instance_id.shape)
semantic_id = 1
for instance_id in sorted(np.unique(plane_instance_id)):
if instance_id == 0:
continue
planeSegmentation[plane_instance_id == instance_id] = semantic_id
semantic_info[semantic_id] = {'plane_instance_id': instance_id}
semantic_id += 1
planes.append(plane_params[instance_id - 1])
planes = np.array(planes)
if len(planes) == 0:
quality = False
else:
planes = get_plane_params_in_local(planes, camera)
summary = {
'planes': planes,
'planeSegmentation': planeSegmentation,
'backgroundidx': 0,
'numPlanes': len(planes),
'semantic_info': semantic_info,
'good_quality': quality,
}
# with open(os.path.join(DATASET_DIR, 'planes_from_ply', test_scene, camera['img_name'].split('.')[0]+'.pkl'), 'wb') as f:
# Save Plane
save_path = os.path.join(SAVE_FOLDER, 'planes_ply_mp3dcoord_refined', test_scene, 'planes_from_ply')
os.makedirs(save_path, exist_ok=True)
with open(os.path.join(save_path, camera['img_name'].split('.')[0]+'.pkl'), 'wb') as f:
pickle.dump(summary, f)
# Save RGB
save_path = os.path.join(SAVE_FOLDER, 'rgb', test_scene, camera['img_name'].split('.')[0]+'.png')
rgb_img = Image.fromarray(observations["color_sensor"], mode="RGBA")
rgb_img.save(save_path)
# Save observations
obs_f = os.path.join(SAVE_FOLDER, 'observations', test_scene, camera['img_name'].split('.')[0]+'.pkl')
with open(obs_f, 'wb') as f:
pickle.dump(observations, f)
sim.close()
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)
def box_drop_test(
args,
objects,
num_objects=30,
box_size=2,
object_speed=2,
post_throw_settling_time=5,
useVHACD=False,
VHACDParams=def_params,
): # take in parameters here
"""Drops a specified number of objects into a box and returns metrics including the time to simulate each frame, render each frame, and the number of collisions each frame."""
data = {
"observations": [],
"physics_step_times": [],
"graphics_render_times": [],
"collisions": [],
}
cfg = make_configuration()
with habitat_sim.Simulator(cfg) as sim:
sim.initialize_agent(
0,
habitat_sim.AgentState(
[5.45 * box_size, 2.4 * box_size, 1.2 * box_size],
np.quaternion(-0.8000822, 0.1924500, -0.5345973, -0.1924500),
),
)
# get the physics object attributes manager
obj_templates_mgr = sim.get_object_template_manager()
# build box
cube_handle = obj_templates_mgr.get_template_handles("cube")[0]
box_parts = []
boxIDs = []
for _i in range(5):
box_parts += [
obj_templates_mgr.get_template_by_handle(cube_handle)
]
box_parts[0].scale = np.array([1.0, 0.1, 1.0]) * box_size
box_parts[1].scale = np.array([1.0, 0.6, 0.1]) * box_size
box_parts[2].scale = np.array([1.0, 0.6, 0.1]) * box_size
box_parts[3].scale = np.array([0.1, 0.6, 1.0]) * box_size
box_parts[4].scale = np.array([0.1, 0.6, 1.0]) * box_size
for i in range(5):
part_name = "box_part_" + str(i)
obj_templates_mgr.register_template(box_parts[i], part_name)
boxIDs += [sim.add_object_by_handle(part_name)]
sim.set_object_motion_type(
habitat_sim.physics.MotionType.KINEMATIC, boxIDs[i])
sim.set_translation(np.array([2.50, 0.05, 0.4]) * box_size, boxIDs[0])
sim.set_translation(np.array([2.50, 0.35, 1.30]) * box_size, boxIDs[1])
sim.set_translation(
np.array([2.50, 0.35, -0.50]) * box_size, boxIDs[2])
sim.set_translation(np.array([3.40, 0.35, 0.4]) * box_size, boxIDs[3])
sim.set_translation(np.array([1.60, 0.35, 0.4]) * box_size, boxIDs[4])
for i in range(5):
sim.set_object_motion_type(habitat_sim.physics.MotionType.STATIC,
boxIDs[i])
# load some object templates from configuration files
object_ids = []
for obj in objects:
obj_path = obj["path"]
object_ids += [
obj_templates_mgr.load_configs(
str(os.path.join(data_path, obj_path)))[0]
]
obj_template = obj_templates_mgr.get_template_by_ID(object_ids[-1])
if "scale" in obj:
obj_template.scale *= obj["scale"]
obj_handle = obj_template.render_asset_handle
if useVHACD:
new_handle = sim.apply_convex_hull_decomposition(
obj_handle, VHACDParams, True, False)
new_obj_template = obj_templates_mgr.get_template_by_handle(
new_handle)
if "scale" in obj:
new_obj_template.scale *= obj["scale"]
obj_templates_mgr.register_template(new_obj_template,
force_registration=True)
object_ids[-1] = new_obj_template.ID
print("Template Registered")
else:
obj_templates_mgr.register_template(obj_template,
force_registration=True)
# throw objects into box
for i in range(num_objects):
cur_id = sim.add_object(object_ids[i % len(object_ids)])
obj_node = sim.get_object_scene_node(cur_id)
obj_bb = obj_node.cumulative_bb
diagonal_length = obj_bb.size().length()
time_til_next_obj = diagonal_length / (object_speed * box_size) / 2
# set object position and velocity
sim.set_translation(
np.multiply(np.array([1.50, 2, 1.2]), box_size), cur_id)
sim.set_linear_velocity(
mn.Vector3(1, 0, -1).normalized() * object_speed * box_size,
cur_id)
# simulate a short amount of time, then add next object
simulate(sim,
dt=time_til_next_obj,
get_frames=args.make_video,
data=data)
simulate(sim,
dt=post_throw_settling_time,
get_frames=args.make_video,
data=data)
# [/basics]
# return total time to run, time to load, time to simulate physics, time for rendering
remove_all_objects(sim)
return data
def test_greedy_follower(test_navmesh, move_filter_fn, action_noise, pbar):
global num_fails
global num_tested
global total_spl
if not osp.exists(test_navmesh):
pytest.skip(f"{test_navmesh} not found")
pathfinder = habitat_sim.PathFinder()
pathfinder.load_nav_mesh(test_navmesh)
assert pathfinder.is_loaded
pathfinder.seed(0)
np.random.seed(seed=0)
scene_graph = habitat_sim.SceneGraph()
agent = habitat_sim.Agent(scene_graph.get_root_node().create_child())
agent.controls.move_filter_fn = getattr(pathfinder, move_filter_fn)
agent.agent_config.action_space["turn_left"].actuation.amount = TURN_DEGREE
agent.agent_config.action_space[
"turn_right"].actuation.amount = TURN_DEGREE
if action_noise:
# "_" prefix the perfect actions so that we can use noisy actions instead
agent.agent_config.action_space = {
"_" + k: v
for k, v in agent.agent_config.action_space.items()
}
agent.agent_config.action_space.update(**dict(
move_forward=habitat_sim.ActionSpec(
"pyrobot_noisy_move_forward",
habitat_sim.PyRobotNoisyActuationSpec(amount=0.25),
),
turn_left=habitat_sim.ActionSpec(
"pyrobot_noisy_turn_left",
habitat_sim.PyRobotNoisyActuationSpec(amount=TURN_DEGREE),
),
turn_right=habitat_sim.ActionSpec(
"pyrobot_noisy_turn_right",
habitat_sim.PyRobotNoisyActuationSpec(amount=TURN_DEGREE),
),
))
follower = habitat_sim.GreedyGeodesicFollower(
pathfinder,
agent,
forward_key="move_forward",
left_key="turn_left",
right_key="turn_right",
)
test_spl = 0.0
for _ in range(NUM_TESTS):
follower.reset()
state = habitat_sim.AgentState()
while True:
state.position = pathfinder.get_random_navigable_point()
goal_pos = pathfinder.get_random_navigable_point()
path = habitat_sim.ShortestPath()
path.requested_start = state.position
path.requested_end = goal_pos
if pathfinder.find_path(path) and path.geodesic_distance > 2.0:
break
agent.state = state
failed = False
gt_geo = path.geodesic_distance
agent_distance = 0.0
last_xyz = state.position
num_acts = 0
# If there is not action noise, then we can use find_path to get all the actions
if not action_noise:
try:
action_list = follower.find_path(goal_pos)
except habitat_sim.errors.GreedyFollowerError:
action_list = [None]
while True:
# If there is action noise, we need to plan a single action, actually take it, and repeat
if action_noise:
try:
next_action = follower.next_action_along(goal_pos)
except habitat_sim.errors.GreedyFollowerError:
break
else:
next_action = action_list[0]
action_list = action_list[1:]
if next_action is None:
break
agent.act(next_action)
agent_distance += np.linalg.norm(last_xyz - agent.state.position)
last_xyz = agent.state.position
num_acts += 1
if num_acts > 1e4:
break
end_state = agent.state
path.requested_start = end_state.position
pathfinder.find_path(path)
failed = path.geodesic_distance > follower.forward_spec.amount
spl = float(not failed) * gt_geo / max(gt_geo, agent_distance)
test_spl += spl
if test_all:
num_fails += float(failed)
num_tested += 1
total_spl += spl
pbar.set_postfix(
num_fails=num_fails,
failure_rate=num_fails / num_tested,
spl=total_spl / num_tested,
)
pbar.update()
if not test_all:
assert test_spl / NUM_TESTS >= ACCEPTABLE_SPLS[(move_filter_fn,
action_noise)]
def get_agent_state(self, agent_id: int = 0):
assert agent_id == 0, "No support of multi agent in {} yet.".format(
self.__class__.__name__)
state = habitat_sim.AgentState()
self._sim.get_agent(agent_id).get_state(state)
return state
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)
sim.close()
except NameError:
pass
sim = habitat_sim.Simulator(cfg)
# %% [markdown]
# ### Initialize the agent
#
# After we initialize the simulator, we could put the agent in the scene, set and query its state, such as position and orientation.
# %%
# initialize an agent
agent = sim.initialize_agent(sim_settings["default_agent"])
# Set agent state
agent_state = habitat_sim.AgentState()
agent_state.position = np.array([-0.6, 0.0, 0.0]) # in world space
agent.set_state(agent_state)
# Get agent state
agent_state = agent.get_state()
print("agent_state: position", agent_state.position, "rotation", agent_state.rotation)
# %% [markdown]
# ### Navigate and see
# %%
# obtain the default, discrete actions that an agent can perform
# default action space contains 3 actions: move_forward, turn_left, and turn_right
action_names = list(cfg.agents[sim_settings["default_agent"]].action_space.keys())
print("Discrete action space: ", action_names)
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)
