customization for saving pose and image of current simulator view

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:

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):

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()

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()))