update: update the init ori

changhaonan · Aug 29, 2023 · 4567f10 · 4567f10
1 parent 057fc81
commit 4567f10
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 7 deletions.
diff --git a/.DS_Store b/.DS_Store
diff --git a/lgmcts/components/encyclopedia/definitions.py b/lgmcts/components/encyclopedia/definitions.py
@@ -59,6 +59,7 @@ class ObjEntry(NamedTuple):
     template_file: Optional[str] = None
     symmetry: Optional[float] = None
     profile: Optional[ProfilePedia] = ProfilePedia.UNDETERMINED
+    init_ori: Optional[List[float]] = None
 
 
 # defines the data structure of each entry in the texture encyclopedia

diff --git a/lgmcts/env/base.py b/lgmcts/env/base.py
@@ -529,10 +529,7 @@ def get_random_pose(self, obj_size, prior=None, stack_prob=0.0):
         _, hmap, obj_mask = self.get_true_image()
 
         obj_stack_id = -1  # -1 is the base
-        if self.rng.random() < stack_prob and len(self.obj_support_tree.leaves) > 1:
-            # select an object to stack up
-            if len(self.obj_ids["rigid"]) == 0:
-                return [None, None], None
+        if self.rng.random() < stack_prob and len(self.obj_support_tree.leaves) > 1 and (len(self.obj_ids["rigid"]) > 0):
             leaf_nodes = self.obj_support_tree.leaves
             leaf_obj_ids = [leaf_node.name for leaf_node in leaf_nodes]
             obj_stack_id = self.rng.choice(leaf_obj_ids)

diff --git a/lgmcts/tasks/struct_rearrange.py b/lgmcts/tasks/struct_rearrange.py
@@ -72,7 +72,7 @@ def reset(self, env):
         self.distract_obj_ids = []
 
     def add_objects_to_pattern(
-        self, env, objs, colors, pattern_prior: np.ndarray, use_existing: bool=False, stack_prob: float=0.0):
+        self, env, objs, colors, pattern_prior: np.ndarray | None, use_existing: bool=False, stack_prob: float=0.0):
         """Set objects to a line, use_existing decides whether to add new object or not"""
         # Add object
         added_obj_ids = []
@@ -91,8 +91,9 @@ def add_objects_to_pattern(
         else:
             raise NotImplementedError("Not implemented yet")
         if len(added_obj_ids) == 0:
-            cv2.imshow("prior", pattern_prior)
-            cv2.waitKey(0)
+            if pattern_prior is not None:
+                cv2.imshow("prior", pattern_prior)
+                cv2.waitKey(0)
             assert False, "No object is added to the pattern"
         return added_obj_ids, obj_status