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)