update: directly provide z

changhaonan · Sep 11, 2023 · d052fc8 · d052fc8
1 parent 589f960
commit d052fc8
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -10,8 +10,8 @@
             "request": "launch",
             "program": "${workspaceFolder}/lgmcts/scripts/data_generation/gen_lgmcts.py",
             "args": [
-                "--num_episodes=10",
-                "--debug",
+                "--num_episodes=100",
+                // "--debug",
             ],
             "console": "integratedTerminal",
             "justMyCode": false

diff --git a/lgmcts/env/base.py b/lgmcts/env/base.py
@@ -467,6 +467,7 @@ def get_obs(self):
         obs = {f"{modality}": {} for modality in self.modalities}  # sensing
         # obs["oracle"] = {}  # oracle information
         obs["point_cloud"] = {}  # point cloud
+        obs["scene_point_cloud"] = {}  # scene point cloud
         obs["poses"] = {}  # object poses
         # Sensing
         for view, config in self.agent_cams.items():
@@ -480,8 +481,7 @@ def get_obs(self):
             # Notice: convert depth back to mm
             real_depth = depth[0] * 20.0
             scene_pcd = misc_utils.get_pointcloud(real_depth, intrinsic_mat)
-            obs["point_cloud"][view] = {}
-            obs["poses"][view] = {}
+            obs["scene_point_cloud"][view] = scene_pcd
             #
             max_pcd_size = self.obs_img_size[0] * self.obs_img_size[1]
             obj_pcds = np.zeros([self.max_num_obj * max_pcd_size, 3])

diff --git a/lgmcts/scripts/data_generation/gen_strdiff.py b/lgmcts/scripts/data_generation/gen_strdiff.py
@@ -134,6 +134,11 @@ def _generate_data_for_one_task(
             f.create_dataset("proj_fov", data=proj_fov)
             f.create_dataset("proj_near", data=0.5)
             f.create_dataset("proj_far", data=5.0)
+            # point cloud
+            scene_point_cloud = obs.pop("scene_point_cloud")[view]
+            scene_point_cloud[:, :, :, 2] = scene_point_cloud[:, :, :, 2] - 1000.0  # because camera is at z=1000
+            f.create_dataset("scene_point_cloud", data=scene_point_cloud)
+
             # objs related
             poses = obs.pop("poses")
             poses = poses[view].transpose(1, 0, 2)