[Prototype] Mesh receptacle annotation automation

habitat-lab/habitat/datasets/rearrange/configs/all_receptacles_test.yaml

@@ -0,0 +1,52 @@
+---
+## All receptacles and objects are used.
+# Define your own dataset path,
+dataset_path: "data/replica_cad/replicaCAD.scene_dataset_config.json"
+additional_object_paths:
+  - "data/objects/ycb/configs/"
+correct_unstable_results: True
+scene_sets:
+  -
+    name: "all_scenes"
+    included_substrings:
+      - ""
+    excluded_substrings: ["NONE"]
+
+object_sets:
+  -
+    name: "simple_objects"
+    included_substrings:
+      - "002_master_chef_can"
+      - "003_cracker_box"
+      - "004_sugar_box"
+      - "005_tomato_soup_can"
+      - "007_tuna_fish_can"
+      - "008_pudding_box"
+      - "009_gelatin_box"
+      - "010_potted_meat_can"
+      - "024_bowl"
+    excluded_substrings: []
+receptacle_sets:
+  -
+    name: "all_receptacles"
+    included_object_substrings:
+      - ""
+    excluded_object_substrings: []
+    included_receptacle_substrings:
+      - ""
+    excluded_receptacle_substrings: []
+
+scene_sampler:
+  type: "subset"
+  params:
+    scene_sets: ["all_scenes"]
+
+object_samplers:
+  -
+    name: "simple_objects_sample"
+    type: "uniform"
+    params:
+      object_sets: ["simple_objects"]
+      receptacle_sets: ["all_receptacles"]
+      num_samples: [10, 20]
+      orientation_sampling: "up"

habitat-lab/habitat/datasets/rearrange/rearrange_generator.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import os.path as osp
+import time
 from collections import defaultdict
 
 try:
@@ -26,6 +27,7 @@
 from habitat.datasets.rearrange.rearrange_dataset import RearrangeEpisode
 from habitat.datasets.rearrange.samplers.receptacle import (
     OnTopOfReceptacle,
+    Receptacle,
     ReceptacleSet,
     ReceptacleTracker,
     find_receptacles,
@@ -391,29 +393,15 @@ def generate_scene(self) -> str:
 
     def visualize_scene_receptacles(self) -> None:
         """
-        Generate a wireframe bounding box for each receptacle in the scene, aim the camera at it and record 1 observation.
+        Generate a debug line representation for each receptacle in the scene, aim the camera at it and record 1 observation.
         """
         logger.info("visualize_scene_receptacles processing")
         receptacles = find_receptacles(self.sim)
         for receptacle in receptacles:
             logger.info("receptacle processing")
-            viz_objects = receptacle.add_receptacle_visualization(self.sim)
-
-            # sample points in the receptacles to display
-            # for sample in range(25):
-            #     sample_point = receptacle.sample_uniform_global(self.sim, 1.0)
-            #     sutils.add_viz_sphere(self.sim, 0.025, sample_point)
-
-            if viz_objects:
-                # point the camera at the 1st viz_object for the Receptacle
-                self.vdb.look_at(
-                    viz_objects[0].root_scene_node.absolute_translation
-                )
-                self.vdb.get_observation()
-            else:
-                logger.warning(
-                    f"visualize_scene_receptacles: no visualization object generated for Receptacle '{receptacle.name}'."
-                )
+            receptacle.debug_draw(self.sim)
+            self.vdb.look_at(receptacle.sample_uniform_global(self.sim, 1.0))
+            self.vdb.get_observation()
 
     def generate_episodes(
         self, num_episodes: int = 1, verbose: bool = False
@@ -545,7 +533,7 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
             self.vdb.make_debug_video(prefix="receptacles_")
 
         # sample object placements
-        object_to_containing_receptacle = {}
+        self.object_to_containing_receptacle = {}
         for sampler_name, obj_sampler in self._obj_samplers.items():
             object_sample_data = obj_sampler.sample(
                 self.sim,
@@ -558,7 +546,7 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
                 return None
             new_objects, receptacles = zip(*object_sample_data)
             for obj, rec in zip(new_objects, receptacles):
-                object_to_containing_receptacle[obj.handle] = rec
+                self.object_to_containing_receptacle[obj.handle] = rec
             if sampler_name not in self.episode_data["sampled_objects"]:
                 self.episode_data["sampled_objects"][
                     sampler_name
@@ -574,9 +562,13 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
             )
             # debug visualization showing each newly added object
             if self._render_debug_obs:
+                logger.info(
+                    f"Generating debug images for {len(new_objects)} objects..."
+                )
                 for new_object in new_objects:
                     self.vdb.look_at(new_object.translation)
                     self.vdb.get_observation()
+                logger.info("   ... done")
 
         # simulate the world for a few seconds to validate the placements
         if not self.settle_sim():
@@ -613,7 +605,7 @@ def generate_single_episode(self) -> Optional[RearrangeEpisode]:
                 vdb=self.vdb,
                 target_receptacles=target_receptacles[obj_sampler_name],
                 goal_receptacles=goal_receptacles[sampler_name],
-                object_to_containing_receptacle=object_to_containing_receptacle,
+                object_to_containing_receptacle=self.object_to_containing_receptacle,
             )
             if new_target_objects is None:
                 return None
@@ -694,7 +686,7 @@ def extract_recep_info(recep):
         ]
 
         name_to_receptacle = {
-            k: v.name for k, v in object_to_containing_receptacle.items()
+            k: v.name for k, v in self.object_to_containing_receptacle.items()
         }
 
         return RearrangeEpisode(
@@ -787,9 +779,12 @@ def initialize_sim(self, scene_name: str, dataset_path: str) -> None:
         self.sim.agents[0].scene_node.translation = scene_bb.center()
 
         # initialize the debug visualizer
-        self.vdb = DebugVisualizer(
-            self.sim, output_path="rearrange_ep_gen_output/"
+        output_path = (
+            "rearrange_ep_gen_output/"
+            if self.vdb is None
+            else self.vdb.output_path
         )
+        self.vdb = DebugVisualizer(self.sim, output_path=output_path)
 
     def settle_sim(
         self, duration: float = 5.0, make_video: bool = True
@@ -800,7 +795,9 @@ def settle_sim(
         """
         if len(self.ep_sampled_objects) == 0:
             return True
-        # assert len(self.ep_sampled_objects) > 0
+
+        settle_start_time = time.time()
+        logger.info("Running placement stability analysis...")
 
         scene_bb = (
             self.sim.get_active_scene_graph().get_root_node().cumulative_bb
@@ -824,11 +821,13 @@ def settle_sim(
             if self._render_debug_obs:
                 self.vdb.get_observation(obs_cache=settle_db_obs)
 
+        logger.info(f"   ...done in {time.time()-settle_start_time} seconds.")
         # check stability of placements
         logger.info("Computing placement stability report:")
+        logger.info("----------------------------------------")
         max_settle_displacement = 0
         error_eps = 0.1
-        unstable_placements = []
+        unstable_placements: List[str] = []  # list of unstable object handles
         for new_object in self.ep_sampled_objects:
             error = (
                 spawn_positions[new_object.handle] - new_object.translation
@@ -839,6 +838,21 @@ def settle_sim(
                 logger.info(
                     f"    Object '{new_object.handle}' unstable. Moved {error} units from placement."
                 )
+                if self._render_debug_obs:
+                    self.vdb.peek_rigid_object(
+                        obj=new_object,
+                        peek_all_axis=True,
+                        additional_savefile_prefix="unstable_",
+                        debug_lines=[
+                            (
+                                [
+                                    spawn_positions[new_object.handle],
+                                    new_object.translation,
+                                ],
+                                mn.Color4.red(),
+                            )
+                        ],
+                    )
         logger.info(
             f" : unstable={len(unstable_placements)}|{len(self.ep_sampled_objects)} ({len(unstable_placements)/len(self.ep_sampled_objects)*100}%) : {unstable_placements}."
         )
@@ -852,5 +866,72 @@ def settle_sim(
                 prefix="settle_", fps=30, obs_cache=settle_db_obs
             )
 
+        # detailed receptacle stability report
+        logger.info("  Detailed sampling stats:")
+
+        # receptacle: [num_objects, num_unstable_objects]
+        rec_num_obj_vs_unstable: Dict[Receptacle, List[int]] = {}
+        for obj_name, rec in self.object_to_containing_receptacle.items():
+            if rec not in rec_num_obj_vs_unstable:
+                rec_num_obj_vs_unstable[rec] = [0, 0]
+            rec_num_obj_vs_unstable[rec][0] += 1
+            if obj_name in unstable_placements:
+                rec_num_obj_vs_unstable[rec][1] += 1
+        for rec, details in rec_num_obj_vs_unstable.items():
+            logger.info(
+                f"      receptacle '{rec.name}': ({details[1]}/{details[0]}) (unstable/total) objects."
+            )
+
+        success = len(unstable_placements) == 0
+
+        # optionally salvage the episode by removing unstable objects
+        if self.cfg.correct_unstable_results and not success:
+            logger.info("  attempting to correct unstable placements...")
+            for sampler_name, objects in self.episode_data[
+                "sampled_objects"
+            ].items():
+                obj_names = [obj.handle for obj in objects]
+                sampler = self._obj_samplers[sampler_name]
+                unstable_subset = [
+                    obj_name
+                    for obj_name in unstable_placements
+                    if obj_name in obj_names
+                ]
+                # check that we have freedom to reject some objects
+                if (
+                    len(objects) - len(unstable_subset)
+                    >= sampler.num_objects[0]
+                ):
+                    # remove the unstable objects from datastructures
+                    self.episode_data["sampled_objects"][sampler_name] = [
+                        obj
+                        for obj in self.episode_data["sampled_objects"][
+                            sampler_name
+                        ]
+                        if obj.handle not in unstable_subset
+                    ]
+                    self.ep_sampled_objects = [
+                        obj
+                        for obj in self.ep_sampled_objects
+                        if obj.handle not in unstable_subset
+                    ]
+                else:
+                    logger.info(
+                        f"  ... could not remove all unstable placements without violating minimum object sampler requirements for {sampler_name}"
+                    )
+                    logger.info("----------------------------------------")
+                    return False
+            logger.info(
+                f"  ... corrected unstable placements successfully. Final object count = {len(self.ep_sampled_objects)}"
+            )
+            # we removed all unstable placements
+            success = True
+
+        logger.info("----------------------------------------")
+
+        if self._render_debug_obs and success:
+            for obj in self.ep_sampled_objects:
+                self.vdb.peek_rigid_object(obj, peek_all_axis=True)
+
         # return success or failure
-        return len(unstable_placements) == 0
+        return success

habitat-lab/habitat/datasets/rearrange/run_episode_generator.py

@@ -50,6 +50,8 @@ class RearrangeEpisodeGeneratorConfig:
     additional_object_paths: List[str] = field(
         default_factory=lambda: ["data/objects/ycb/"]
     )
+    # optionally correct unstable states by removing extra unstable objects (within minimum samples limitations)
+    correct_unstable_results: bool = False
     # ----- resource set definitions ------
     # Define the sets of scenes, objects, and receptacles which can be sampled from.
     # The SceneDataset will be searched for resources of each type with handles containing ANY "included" substrings and NO "excluded" substrings.

habitat-lab/habitat/datasets/rearrange/samplers/object_sampler.py

@@ -6,6 +6,7 @@
 
 import math
 import random
+import time
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple
 
@@ -388,6 +389,8 @@ def sample(
             f"    Trying to sample {self.target_objects_number} from range {self.num_objects}"
         )
 
+        sampling_start_time = time.time()
+        pairing_start_time = sampling_start_time
         while (
             len(new_objects) < self.target_objects_number
             and num_pairing_tries < self.max_sample_attempts
@@ -415,8 +418,18 @@ def sample(
                     self.receptacle_candidates = None
 
             if new_object is not None:
+                # when an object placement is successful, reset the try counter.
+                logger.info(
+                    f"    found obj|receptacle pairing ({len(new_objects)}/{self.target_objects_number}) in {num_pairing_tries} attempts ({time.time()-pairing_start_time}sec)."
+                )
+                num_pairing_tries = 0
+                pairing_start_time = time.time()
                 new_objects.append((new_object, receptacle))
 
+        logger.info(
+            f"    Sampling process completed in ({time.time()-sampling_start_time}sec)."
+        )
+
         if len(new_objects) >= self.num_objects[0]:
             return new_objects