articulation tasks commit

This commit is contained in:
Leon998
2026-03-18 00:03:11 +08:00
parent 7411e0e004
commit cf8b043bc8
4 changed files with 623 additions and 27 deletions

View File

@@ -19,15 +19,16 @@ class ArticulatedObject(Articulation):
def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
self.asset_root = asset_root
self.object_name = cfg["name"]
self.asset_relative_path = cfg["path"]
self.object_dir = os.path.join(asset_root, cfg["path"])
self.usd_path = os.path.join(asset_root, cfg["path"])
self._root_prim_path = root_prim_path
object_info_path = os.path.join(asset_root, cfg["obj_info_path"])
info_name = cfg["info_name"]
object_info_path = self.usd_path.replace("instance.usd", f"Kps/{info_name}/info.json")
with open(object_info_path, "r", encoding="utf-8") as f:
object_info = json.load(f)
self.category = cfg["category"]
self.cfg = cfg
self.get_articulated_info(object_info)
self.cfg["scale"] = self.object_scale[:3]
super().__init__(prim_path=self.object_prim_path, name=cfg["name"], *args, **kwargs)
def update_articulated_info(self, obj_info_path):
@@ -81,7 +82,6 @@ class ArticulatedObject(Articulation):
for key, item in self.object_keypoints.items():
self.object_keypoints[key] = np.append(item, [1.0], axis=0)
self.object_scale = np.array(object_info["object_scale"])
self.object_usd = os.path.join(self.object_dir, "instance.usd")
for key, item in object_info.items():
if key in prim_path_list:
setattr(self, key, self._root_prim_path + object_info[key])
@@ -96,7 +96,7 @@ class ArticulatedObject(Articulation):
self.object_joint_number = 0
# Contact plane normal
self.contact_plane_normal = None
add_reference_to_stage(usd_path=self.object_usd, prim_path=self.object_prim_path)
add_reference_to_stage(usd_path=self.usd_path, prim_path=self.object_prim_path)
def get_joint_position(self, stage):
joint_parent_prim = stage.GetPrimAtPath(self.object_joint_path.rsplit("/", 1)[0])

View File

@@ -108,17 +108,11 @@ def update_articulated_objs(cfg):
for obj_cfg in cfg["objects"]:
apply_randomization = obj_cfg.get("apply_randomization", False)
if apply_randomization and obj_cfg["target_class"] == "ArticulatedObject":
dirs = os.path.join(cfg["asset_root"], os.path.dirname(obj_cfg["path"]))
dirs = os.path.join(cfg["asset_root"], os.path.dirname(os.path.dirname(obj_cfg["path"])))
paths = glob.glob(os.path.join(dirs, "*"))
paths.sort()
path = random.choice(paths)
# left hearth 0.5: [1, 2, 5, 6, 13, ] ;
# left hearth 0.785 [3, 4, 7, 8, 9, 11, 12, 14, 15, 16, 17]
# left hearth no planning [0, 10, 18, 19]
# right hearth 0.5: [0, 1, 4, 10, 11]
# right hearth 0.785: [2, 3, 5, 6, 7, 8, 9, ]
info_name = obj_cfg["info_name"]
info_path = f"{path}/Kps/{info_name}/info.json"
with open(info_path, "r", encoding="utf-8") as file:
@@ -126,7 +120,7 @@ def update_articulated_objs(cfg):
scale = info["object_scale"][:3]
asset_root = cfg["asset_root"]
obj_cfg["path"] = path.replace(f"{asset_root}/", "", 1)
obj_cfg["path"] = path.replace(f"{asset_root}/", "", 1) + "/instance.usd"
obj_cfg["category"] = path.split("/")[-2]
obj_cfg["obj_info_path"] = info_path.replace(f"{asset_root}/", "", 1)
obj_cfg["scale"] = scale

View File

@@ -87,21 +87,22 @@ class SimBoxDualWorkFlow(NimbusWorkFlow):
for obj_cfg in self.task_cfg["objects"]:
if obj_cfg["target_class"] == "ArticulatedObject":
asset_root = self.task_cfg["asset_root"]
art_paths = glob.glob(os.path.join(asset_root, obj_cfg["art_cat"], "*"))
art_paths.sort()
path = random.choice(art_paths)
info_name = obj_cfg["info_name"]
info_path = f"{path}/Kps/{info_name}/info.json"
with open(info_path, "r", encoding="utf-8") as f:
info = json.load(f)
scale = info["object_scale"][:3]
if obj_cfg.get("apply_randomization", False):
asset_root = self.task_cfg["asset_root"]
art_paths = glob.glob(os.path.join(asset_root, obj_cfg["art_cat"], "*"))
art_paths.sort()
path = random.choice(art_paths)
info_name = obj_cfg["info_name"]
info_path = f"{path}/Kps/{info_name}/info.json"
with open(info_path, "r", encoding="utf-8") as f:
info = json.load(f)
scale = info["object_scale"][:3]
obj_cfg["path"] = path.replace(f"{asset_root}/", "", 1)
obj_cfg["category"] = path.split("/")[-2]
obj_cfg["obj_info_path"] = info_path.replace(f"{asset_root}/", "", 1)
obj_cfg["scale"] = scale
self.task_cfg["data"]["collect_info"] = obj_cfg["category"]
obj_cfg["path"] = path.replace(f"{asset_root}/", "", 1) + "/instance.usd"
obj_cfg["category"] = path.split("/")[-2]
obj_cfg["obj_info_path"] = info_path.replace(f"{asset_root}/", "", 1)
obj_cfg["scale"] = scale
self.task_cfg["data"]["collect_info"] = obj_cfg["category"]
self.task_cfg.pop("arena_file", None)
self.task_cfg.pop("camera_file", None)