diff --git a/README.md b/README.md index 985fe0824b..17378c7d8f 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,14 @@ [![Supports Bullet](https://img.shields.io/static/v1?label=supports&message=Bullet%20Physics&color=informational&link=https://opensource.google/projects/bullet3)](https://opensource.google/projects/bullet3) [![Twitter Follow](https://img.shields.io/twitter/follow/ai_habitat?style=social)](https://twitter.com/ai_habitat) +## NOTE: +This code is slightly modified from original [Habitat-Sim](https://github.com/facebookresearch/habitat-sim) to obtain camera poses as transformation matrix and images as equirectangular projection. +Moreover, `skip` and `out_path` arguments are added to skip the frames and save outputs respectively. The modifications are made only for `examples/`. +You may run this version of code as follows: +``` +python3 examples/example.py --scene /path/to/data/scene_datasets/habitat-test-scenes/skokloster-castle.glb --save_png --max_frames 25 --height 2048 --width 4096 --skip 5 --out_path /path/to/data/scene_datasets/habitat-test-scenes/skokloster-castle/ +``` + # Habitat-Sim A high-performance physics-enabled 3D simulator with support for: diff --git a/examples/demo_runner.py b/examples/demo_runner.py index 42219c0cfe..5da46c8829 100644 --- a/examples/demo_runner.py +++ b/examples/demo_runner.py @@ -353,6 +353,7 @@ def init_common(self): def _bench_target(self, _idx=0): self.init_common() + print("inside _bench_target") best_perf = None for _ in range(3): diff --git a/examples/viewer.py b/examples/viewer.py index 8c42f70b33..0536302990 100644 --- a/examples/viewer.py +++ b/examples/viewer.py @@ -25,6 +25,8 @@ from habitat_sim.utils.common import quat_from_angle_axis from habitat_sim.utils.settings import default_sim_settings, make_cfg +from habitat_sim.utils.common import to_campose +from PIL import Image class HabitatSimInteractiveViewer(Application): # the maximum number of chars displayable in the app window @@ -444,6 +446,7 @@ def move_and_look(self, repetitions: int) -> None: for _ in range(int(repetitions)): [agent.act(x) for x in action_queue] + self.state = agent.get_state() #get current state of agent # update the grabber transform when our agent is moved if self.mouse_grabber is not None: # update location of grabbed object @@ -457,6 +460,33 @@ def invert_gravity(self) -> None: gravity: mn.Vector3 = self.sim.get_gravity() * -1 self.sim.set_gravity(gravity) + def save_pose_observation(self) -> None: + """ + Transforms Rotation quaternion and translation vector to 4x4 transformation matrix, + and saves the flattened matrix as txt + """ + T_mat = to_campose(self.state.rotation, self.state.position) + T_vec = T_mat.reshape(1, 16) + pose_path = os.path.join(self.sim_settings['out_path'], "transformations.txt") + + with open(pose_path, "ab") as f: + np.savetxt(f, T_vec, fmt='%5f', delimiter=' ') + + def save_color_observation(self, obs, sensor_type="color_sensor") -> None: + """ + Retrieves current view of simulator and saves it as image" + """ + color_obs = obs[sensor_type] + color_img = Image.fromarray(color_obs, mode="RGBA") + if self.sim_settings['out_path'] is None: + color_img.save("test.rgba.%05d.png" % time.time()) + else: + image_path = os.path.join(self.sim_settings['out_path'], "images") + if not os.path.exists(image_path): + print(f"{os.path.join(self.sim_settings['out_path'], 'images')} doesn't exist, so creating one.") + os.mkdir(image_path) + color_img.save(os.path.join(image_path,"test.rgba.%05d.png" % time.time())) + def key_press_event(self, event: Application.KeyEvent) -> None: """ Handles `Application.KeyEvent` on a key press by performing the corresponding functions. @@ -476,6 +506,16 @@ def key_press_event(self, event: Application.KeyEvent) -> None: self.exit_event(Application.ExitEvent) return + elif key == pressed.R: + # Press R to record data (image, pose) + observations = self.sim.get_sensor_observations(self.agent_id) + if sim_settings['save_png']: + self.save_pose_observation() + if sim_settings["equirect_rgba_sensor"]: + self.save_color_observation(observations, sensor_type="equirect_rgba_sensor") + else: + self.save_color_observation(observations) + elif key == pressed.H: self.print_help_text() @@ -1157,6 +1197,10 @@ def next_frame() -> None: type=int, help="Vertical resolution of the window.", ) + parser.add_argument("--silent", action="store_true") + parser.add_argument("--save_png", action="store_true") + parser.add_argument("--out_path", help="path to save images and transformations") + parser.add_argument("--disable_color_sensor", action="store_true") args = parser.parse_args() @@ -1180,6 +1224,11 @@ def next_frame() -> None: sim_settings["window_height"] = args.height sim_settings["pbr_image_based_lighting"] = args.ibl sim_settings["default_agent_navmesh"] = False + sim_settings["enable_hbao"] = args.hbao + sim_settings["silent"] = args.silent + sim_settings["save_png"] = args.save_png + sim_settings["out_path"] = args.out_path + sim_settings["equirect_rgba_sensor"] = args.disable_color_sensor # start the application HabitatSimInteractiveViewer(sim_settings).exec() diff --git a/src_python/habitat_sim/utils/common.py b/src_python/habitat_sim/utils/common.py index f6f309a0f5..e1ae9393d4 100755 --- a/src_python/habitat_sim/utils/common.py +++ b/src_python/habitat_sim/utils/common.py @@ -167,7 +167,13 @@ def random_quaternion(): ) return mn.Quaternion(qAxis, math.sqrt(1 - u[0]) * math.sin(2 * math.pi * u[1])) - +def to_campose(q: qt.quaternion, t: np.ndarray): + R = qt.as_rotation_matrix(q) + T = np.eye(4) + T[:3, :3] = R + T[:3, 3] = t + return T + def download_and_unzip(file_url, local_directory): response = urlopen(file_url) zipfile = ZipFile(BytesIO(response.read()))