diff --git a/.gitignore b/.gitignore index 08f6d91..2bd6553 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,5 @@ panda_drake tests/*.log tests/*.txt polygons.png +_isaac_sim_410 +InterDataEngine-docs \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/info.json b/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/info.json new file mode 100644 index 0000000..1d3cfda --- /dev/null +++ b/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/info.json @@ -0,0 +1,30 @@ +{ + "object_keypoints": { + "articulated_object_head": [ + -0.12221331991320702, + 0.004074180066694674, + 0.7207572775929765 + ], + "articulated_object_tail": [ + -0.573521257950959, + -0.0016857095412369238, + 0.49682923043048766 + ] + }, + "object_scale": [ + 0.2302160991881601, + 0.2302160991881601, + 0.2302160991881601, + 1.0 + ], + "object_name": "microwave7265", + "object_usd": "YOUT_PATH_TO_USD/instance.usd", + "object_link0_rot_axis": "y", + "object_link0_contact_axis": "-y", + "object_base_front_axis": "z", + "joint_index": 0, + "object_prim_path": "/microwave7265", + "object_link_path": "/microwave7265/instance/group_18", + "object_base_path": "/microwave7265/instance/group_0", + "object_revolute_joint_path": "/microwave7265/instance/group_18/RevoluteJoint" +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/keypoints.json b/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/keypoints.json new file mode 100644 index 0000000..10c183d --- /dev/null +++ b/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/keypoints.json @@ -0,0 +1,14 @@ +{ + "keypoints": { + "red": [ + -0.12221331991320702, + 0.004074180066694674, + 0.7207572775929765 + ], + "yellow": [ + -0.573521257950959, + -0.0016857095412369238, + 0.49682923043048766 + ] + } +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/keypoints_final.json b/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/keypoints_final.json new file mode 100644 index 0000000..ac206c4 --- /dev/null +++ b/workflows/simbox/tools/art/close_v/7265/usd/Kps/close_v/keypoints_final.json @@ -0,0 +1,20 @@ +{ + "keypoints": { + "red": [ + -0.12221331991320702, + 0.004074180066694674, + 0.7207572775929765 + ], + "yellow": [ + -0.573521257950959, + -0.0016857095412369238, + 0.49682923043048766 + ] + }, + "scale": [ + 0.2302160991881601, + 0.2302160991881601, + 0.2302160991881601, + 1.0 + ] +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/7265/usd/instance.usd b/workflows/simbox/tools/art/close_v/7265/usd/instance.usd new file mode 100644 index 0000000..1add067 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/7265/usd/instance.usd differ diff --git a/workflows/simbox/tools/art/close_v/7265/usd/keypoints_config.json b/workflows/simbox/tools/art/close_v/7265/usd/keypoints_config.json new file mode 100644 index 0000000..f36f7ba --- /dev/null +++ b/workflows/simbox/tools/art/close_v/7265/usd/keypoints_config.json @@ -0,0 +1,13 @@ +{ + "DIR": "/home/shixu/dev_shixu/DataEngine/workflows/simbox/tools/art/close_v/7265/usd", + "USD_NAME": "microwave_0.usd", + "INSTANCE_NAME": "microwave7265", + "link0_initial_prim_path": "/root/group_18", + "base_initial_prim_path": "/root/group_0", + "revolute_joint_initial_prim_path": "/root/group_18/RevoluteJoint", + "joint_index": 0, + "LINK0_ROT_AXIS": "y", + "BASE_FRONT_AXIS": "z", + "LINK0_CONTACT_AXIS": "-y", + "SCALED_VOLUME": 0.02 +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/7265/usd/microwave_0.usd b/workflows/simbox/tools/art/close_v/7265/usd/microwave_0.usd new file mode 100644 index 0000000..8a291f2 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/7265/usd/microwave_0.usd differ diff --git a/workflows/simbox/tools/art/close_v/7265/usd/textures/door_mesh_0_texture_0.jpg b/workflows/simbox/tools/art/close_v/7265/usd/textures/door_mesh_0_texture_0.jpg new file mode 100644 index 0000000..6afb854 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/7265/usd/textures/door_mesh_0_texture_0.jpg differ diff --git a/workflows/simbox/tools/art/close_v/7265/usd/textures/door_mesh_1_texture_0.jpg b/workflows/simbox/tools/art/close_v/7265/usd/textures/door_mesh_1_texture_0.jpg new file mode 100644 index 0000000..6afb854 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/7265/usd/textures/door_mesh_1_texture_0.jpg differ diff --git a/workflows/simbox/tools/art/close_v/BASE_FRONT_AXIS.jpg b/workflows/simbox/tools/art/close_v/BASE_FRONT_AXIS.jpg new file mode 100644 index 0000000..fabc434 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/BASE_FRONT_AXIS.jpg differ diff --git a/workflows/simbox/tools/art/close_v/LINK0_CONTACT_AXIS.jpg b/workflows/simbox/tools/art/close_v/LINK0_CONTACT_AXIS.jpg new file mode 100644 index 0000000..cbc9192 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/LINK0_CONTACT_AXIS.jpg differ diff --git a/workflows/simbox/tools/art/close_v/LINK0_ROT_AXIS.jpg b/workflows/simbox/tools/art/close_v/LINK0_ROT_AXIS.jpg new file mode 100644 index 0000000..979f65a Binary files /dev/null and b/workflows/simbox/tools/art/close_v/LINK0_ROT_AXIS.jpg differ diff --git a/workflows/simbox/tools/art/close_v/head.jpg b/workflows/simbox/tools/art/close_v/head.jpg new file mode 100644 index 0000000..130a046 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/head.jpg differ diff --git a/workflows/simbox/tools/art/close_v/readme.md b/workflows/simbox/tools/art/close_v/readme.md new file mode 100644 index 0000000..0d8ee95 --- /dev/null +++ b/workflows/simbox/tools/art/close_v/readme.md @@ -0,0 +1,45 @@ +# Annotation Documentation + +We provide an optimized and simplified annotation pipeline that removes many redundancies. No need to rename base_link, contact_link, etc. Keep the original hierarchy and naming as much as possible. + +## 🗂️ File Information + +| Configuration | Example | Description | +|---------------|---------|-------------| +| **DIR** | `/home/shixu/Downloads/peixun/7265/usd` | Directory where USD files are stored | +| **USD_NAME** | `microwave_0.usd` | Scene description file name | +| **INSTANCE_NAME** | `microwave7265` | Model identifier in the scene. You can name it yourself, preferably matching the generated file name | + +## 🔧 Model Structure Configuration + +| Component | Example | Description | +|-----------|---------|-------------| +| **link0_initial_prim_path** | `/root/group_18` | Absolute path in Isaac Sim for the "door" that interacts with the gripper. Check in the original USD | +| **base_initial_prim_path** | `/root/group_0` | Absolute path in Isaac Sim for the microwave base. Check in the original USD | +| **revolute_joint_initial_prim_path** | `/root/group_18/RevoluteJoint` | Absolute path in Isaac Sim for the revolute joint that opens/closes the microwave. Check in the original USD | +| **Joint Index** | `0` | Joint number, default is 0 | + +## 🧭 Axis Configuration + +| Axis Type | Example | Description | Visualization | +|-----------|---------|-------------|---------------| +| **LINK0_ROT_AXIS** | `y` | In the local coordinate system of the rotating joint, the axis direction pointing vertically upward | ![LINK0_ROT_AXIS Example](LINK0_ROT_AXIS.jpg) | +| **BASE_FRONT_AXIS** | `z` | In the local coordinate system of the microwave base link, the axis direction facing the door | ![BASE_FRONT_AXIS Example](BASE_FRONT_AXIS.jpg) | +| **LINK0_CONTACT_AXIS** | `-y` | In the local coordinate system of the contact link, the axis direction pointing vertically downward | ![LINK0_CONTACT_AXIS Example](LINK0_CONTACT_AXIS.jpg) | + +## 📏 Physical Parameters + +| Parameter | Example | Description | +|-----------|---------|-------------| +| **SCALED_VOLUME** | `0.02` | Default value 0.02 for microwave-like objects | + +--- + +# Point Annotation + +| Point Type | Description | Visualization | +|------------|-------------|---------------| +| First Point (articulated_object_head) | `Desired base position where the gripper contacts the microwave door` | ![First Point Diagram](head.jpg) | +| Second Point (articulated_object_tail) | `The line direction from the first point should be perpendicular to the microwave door's rotation axis` | ![Second Point Diagram](tail.jpg) | + +--- \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/tail.jpg b/workflows/simbox/tools/art/close_v/tail.jpg new file mode 100644 index 0000000..7933009 Binary files /dev/null and b/workflows/simbox/tools/art/close_v/tail.jpg differ diff --git a/workflows/simbox/tools/art/close_v/tool/keypoints_pipeline.sh b/workflows/simbox/tools/art/close_v/tool/keypoints_pipeline.sh new file mode 100755 index 0000000..9519cb9 --- /dev/null +++ b/workflows/simbox/tools/art/close_v/tool/keypoints_pipeline.sh @@ -0,0 +1,21 @@ +# pylint: skip-file +# flake8: noqa +# Replace the following keypoints_config path with your absolute path +CONFIG_PATH="YOUR_PATH_TO_7265/usd/keypoints_config.json" + +# Replace the following close_v_new path with your absolute path +cd workflows/simbox/tools/art/close_v/tool + +# Run the following scripts in sequence + +# 1. rehier - This should generate peixun/7265/usd/instance.usd file to indicate success +python rehier.py --config $CONFIG_PATH + +# 2. select points +python select_keypoint.py --config $CONFIG_PATH + +# 3. Transfer keypoints +python transfer_keypoints.py --config $CONFIG_PATH + +# 4. Overwrite keypoints +python overwrite_keypoints.py --config $CONFIG_PATH diff --git a/workflows/simbox/tools/art/close_v/tool/overwrite_keypoints.py b/workflows/simbox/tools/art/close_v/tool/overwrite_keypoints.py new file mode 100644 index 0000000..8de4ac7 --- /dev/null +++ b/workflows/simbox/tools/art/close_v/tool/overwrite_keypoints.py @@ -0,0 +1,57 @@ +# pylint: skip-file +# flake8: noqa +import os +import json +import numpy as np + +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--config", type=str, help="Path to config file") +args = parser.parse_args() + +with open(args.config, 'r') as f: + config = json.load(f) + +TASK = "close_v" +dir_name = config["DIR"] +dir_name_kps = os.path.join(config['DIR'], "Kps", TASK) +os.makedirs(dir_name_kps, exist_ok=True) + +usd_file = os.path.join(dir_name, "instance.usd") +keypoint_path = os.path.join(dir_name_kps, "keypoints_final.json") +target_keypoint_path = os.path.join(dir_name_kps, "info.json") + +if not os.path.exists(target_keypoint_path): + with open(target_keypoint_path,'w') as file: + data={"object_keypoints":{}} + json.dump(data,file,indent=4) + +if not os.path.exists(keypoint_path) or not os.path.exists(target_keypoint_path): + print(f"keypoint file {keypoint_path} or {target_keypoint_path} not found") + +kp = json.load(open(keypoint_path)) +tkp = json.load(open(target_keypoint_path)) + +tkp["object_keypoints"]["articulated_object_head"] = kp["keypoints"]["red"] +tkp["object_keypoints"]["articulated_object_tail"] = kp["keypoints"]["yellow"] + +tkp["object_scale"] = kp["scale"] +tkp["object_name"] = config["INSTANCE_NAME"] +tkp["object_usd"] = usd_file +tkp["object_link0_rot_axis"] = config["LINK0_ROT_AXIS"] +tkp["object_link0_contact_axis"] = config["LINK0_CONTACT_AXIS"] +tkp["object_base_front_axis"] = config["BASE_FRONT_AXIS"] +tkp["joint_index"] = config["joint_index"] + +tkp["object_prim_path"] = os.path.join("/", config["INSTANCE_NAME"]) +link0_initial_prim_path = (config["link0_initial_prim_path"]).replace("/root", "instance") +base_initial_prim_path = (config["base_initial_prim_path"]).replace("/root", "instance") +revolute_joint_initial_prim_path = (config["revolute_joint_initial_prim_path"]).replace("/root", "instance") + +tkp["object_link_path"] = os.path.join("/", config["INSTANCE_NAME"], link0_initial_prim_path) +tkp["object_base_path"] = os.path.join("/", config["INSTANCE_NAME"], base_initial_prim_path) +tkp["object_revolute_joint_path"] = os.path.join("/", config["INSTANCE_NAME"], revolute_joint_initial_prim_path) + +json.dump(tkp, open(target_keypoint_path, "w"), indent=4) +print("Saved keypoints to ", target_keypoint_path) diff --git a/workflows/simbox/tools/art/close_v/tool/rehier.py b/workflows/simbox/tools/art/close_v/tool/rehier.py new file mode 100755 index 0000000..2413bb3 --- /dev/null +++ b/workflows/simbox/tools/art/close_v/tool/rehier.py @@ -0,0 +1,302 @@ +# pylint: skip-file +# flake8: noqa +import os +import json +import argparse +from pathlib import Path +from pxr import Usd, UsdGeom, UsdPhysics, Gf, Sdf +from pdb import set_trace + +def remove_articulation_root(prim: Usd.Prim): + prim.RemoveAPI(UsdPhysics.ArticulationRootAPI) + allchildren = prim.GetAllChildren() + if len(allchildren) == 0: + return + else: + for child in allchildren: + remove_articulation_root(child) + +def remove_rigidbody(prim: Usd.Prim): + prim.RemoveAPI(UsdPhysics.RigidBodyAPI) + allchildren = prim.GetAllChildren() + if len(allchildren) == 0: + return + else: + for child in allchildren: + remove_rigidbody(child) + +def remove_mass(prim: Usd.Prim): + prim.RemoveAPI(UsdPhysics.MassAPI) + allchildren = prim.GetAllChildren() + if len(allchildren) == 0: + return + else: + for child in allchildren: + remove_mass(child) + +def add_rigidbody(prim: Usd.Prim): + UsdPhysics.RigidBodyAPI.Apply(prim) + +def add_mass(prim: Usd.Prim): + UsdPhysics.MassAPI.Apply(prim) + mass = prim.GetAttribute("physics:mass") + mass.Clear() + +def get_args(): + parser = argparse.ArgumentParser(description="USD Hierarchy and Physics Editor") + parser.add_argument("--config", type=str, required=True, help="Path to config file") + + return parser.parse_args() + +def load_config(config_path): + with open(config_path, 'r') as f: + return json.load(f) + +def safe_rename_prim(stage, old_path, new_path): + editor = Usd.NamespaceEditor(stage) + old_p = stage.GetPrimAtPath(old_path) + editor.RenamePrim(old_p, new_path.split('/')[-1]) + + if editor.CanApplyEdits(): + editor.ApplyEdits() + return True + else: + return False + +def modify_hierarchy(stage, config): + editor = Usd.NamespaceEditor(stage) + + """ Modify USD hierarchy structure """ + base_path = config["base_initial_prim_path"] + link0_path = config["link0_initial_prim_path"] + revolute_joint_path = config["revolute_joint_initial_prim_path"] + + # Get original root node + old_root_path = f"/{link0_path.split('/')[1]}" + instance_path = "/root/instance" + safe_rename_prim(stage, '/{}'.format(stage.GetDefaultPrim().GetName()), "/instance") + + return instance_path + +def modify_physics(stage, instance_root_path, config): + """ Modify physics properties and ensure colliders are uniformly set to Convex Hull """ + print(f"Applying physics modifications and setting colliders to ConvexHull...") + + # Traverse all nodes for processing + for prim in stage.Traverse(): + # 1. Clear instancing (if editing physics properties of individual instances) + if prim.IsInstanceable(): + prim.ClearInstanceable() + + # 2. Process Mesh colliders + if prim.IsA(UsdGeom.Mesh): + # Ensure base collision API is applied + if not prim.HasAPI(UsdPhysics.CollisionAPI): + UsdPhysics.CollisionAPI.Apply(prim) + + # Force apply MeshCollisionAPI to set collision approximation + mesh_collision_api = UsdPhysics.MeshCollisionAPI.Get(stage, prim.GetPath()) + if not mesh_collision_api: + mesh_collision_api = UsdPhysics.MeshCollisionAPI.Apply(prim) + + # Set collision approximation to 'convexHull' + # Optional values include: 'none', 'convexHull', 'convexDecomposition', 'meshSimplification', etc. + mesh_collision_api.CreateApproximationAttr().Set("convexHull") + + # Ensure physics collision is enabled + col_enabled_attr = prim.GetAttribute("physics:collisionEnabled") + if not col_enabled_attr.HasValue(): + col_enabled_attr.Set(True) + +def create_fixed_joint(stage, joint_path, body0_path, body1_path): + """ + Create a FixedJoint at the specified path and connect two rigid bodies. + """ + # 1. Define FixedJoint node + fixed_joint = UsdPhysics.FixedJoint.Define(stage, joint_path) + + # 2. Set Body0 and Body1 path references + # Note: Paths must be of Sdf.Path type + fixed_joint.GetBody0Rel().SetTargets([Sdf.Path(body0_path)]) + fixed_joint.GetBody1Rel().SetTargets([Sdf.Path(body1_path)]) + + # 3. (Optional) Set local offset (Local Pose) + # If not set, the joint defaults to the origin of both objects + # fixed_joint.GetLocalPos0Attr().Set(Gf.Vec3f(0, 0, 0)) + # fixed_joint.GetLocalRot0Attr().Set(Gf.Quatf(1, 0, 0, 0)) + + print(f"Successfully created FixedJoint: {joint_path}") + print(f" Connected: {body0_path} <---> {body1_path}") + return fixed_joint + +def process_joints(stage, revolute_joint_initial_prim_path): + # 1. Collect paths to process + paths_to_delete = [] + joints_to_convert = [] + + # Use TraverseAll to ensure no defined nodes are missed + for prim in stage.Traverse(): + # Check if it is a physics joint + if prim.IsA(UsdPhysics.Joint): + path = prim.GetPath() + + # Logic A: If FixedJoint -> delete + if prim.IsA(UsdPhysics.FixedJoint): + paths_to_delete.append(path) + + # Logic B: If not FixedJoint and path does not contain 'contact_link' -> convert to FixedJoint + # elif "contact_link" not in str(path).lower(): + # joints_to_convert.append(path) + elif str(path) != revolute_joint_initial_prim_path: + joints_to_convert.append(path) + + print(str(path)) + + # 2. Get current edit layer + layer = stage.GetEditTarget().GetLayer() + edit = Sdf.BatchNamespaceEdit() + + # Execute deletion logic + for path in paths_to_delete: + edit.Add(path, Sdf.Path.emptyPath) + print(f"[Delete] FixedJoint: {path}") + + # 3. Apply deletion edits + if paths_to_delete: + layer.Apply(edit) + + # 4. Execute type conversion logic + # In USD, changing type usually means re-Defining the new type at that path + for path in joints_to_convert: + # Record original Body0 and Body1 relationships to prevent loss after conversion + prim = stage.GetPrimAtPath(path) + joint = UsdPhysics.Joint(prim) + body0 = joint.GetBody0Rel().GetTargets() + body1 = joint.GetBody1Rel().GetTargets() + + # Redefine as FixedJoint + new_fixed_joint = UsdPhysics.FixedJoint.Define(stage, path) + + # Restore relationships + if body0: new_fixed_joint.GetBody0Rel().SetTargets(body0) + if body1: new_fixed_joint.GetBody1Rel().SetTargets(body1) + + safe_rename_prim(stage, str(new_fixed_joint.GetPath()), "/FixedJoint") + + print(f"[Convert] Regular joint -> FixedJoint: {path}") + + return stage + +def final_refine(stage, output_usd_path, revolute_joint_initial_prim_path): + root_prim = stage.GetPrimAtPath("/root") + instance_prim = stage.GetPrimAtPath("/root/instance") + ### remove articulation root ### + remove_articulation_root(root_prim) + + ### remove rigid body ### + remove_rigidbody(root_prim) + + ### remove mass ### + remove_mass(root_prim) + + ### add rigid body and mass ### + for child in instance_prim.GetAllChildren(): + + if child.GetTypeName() == "PhysicsRevoluteJoint" or child.GetTypeName() == "PhysicsPrismaticJoint": + continue + + if child.GetTypeName() == "Xform" : + print('name:', child.GetTypeName()) + add_rigidbody(child) + + stage = process_joints(stage, revolute_joint_initial_prim_path) + + ### add articulation root ### + UsdPhysics.ArticulationRootAPI.Apply(instance_prim) + stage.SetDefaultPrim(root_prim) + + for child in instance_prim.GetAllChildren(): + try: + attr = child.GetAttribute('physics:jointEnabled') + except: + continue + if attr.Get() is not None: + print(child) + attr.Set(True) + + modify_physics(stage, "/root/instance", 11) + stage.Export(output_usd_path) + + return stage + +def import_as_copy(source_usd_path, output_usd_path, root_name="root", sub_node_name="instance"): + """ + Create a new USD and copy the content from source_usd_path to /root/sub_node_name. + """ + # 1. Create target Stage and root node + stage = Usd.Stage.CreateNew(output_usd_path) + root_path = Sdf.Path(f"/{root_name}") + UsdGeom.Xform.Define(stage, root_path) + + # 2. Define copy destination path (e.g., /root/model_copy) + dest_path = root_path.AppendChild(sub_node_name) + + # 3. Open source file layer + source_layer = Sdf.Layer.FindOrOpen(source_usd_path) + + if not source_layer: + print(f"Error: Cannot find source file {source_usd_path}") + return + + # 4. Get source file's default prim (DefaultPrim) as copy target + # If source file has no default prim, use the first root prim + source_root_name = list(source_layer.rootPrims)[0].name + source_path = Sdf.Path(f"/{source_root_name}") + + # 5. Execute core copy operation (Sdf.CopySpec) + # This copies all attributes, topology, and properties from source file to new file + Sdf.CopySpec(source_layer, source_path, stage.GetRootLayer(), dest_path) + + # 6. Set default prim and save + stage.SetDefaultPrim(stage.GetPrimAtPath(root_path)) + stage.GetRootLayer().Save() + print(f"Success! Content copied to: {output_usd_path}") + + return stage, output_usd_path + +def main(): + args = get_args() + config = load_config(args.config) + + dir_name = config["DIR"] + usd_path = os.path.join(dir_name, config["USD_NAME"]) + output_path = os.path.join(dir_name, "instance.usd") + revolute_joint_initial_prim_path = (config["revolute_joint_initial_prim_path"]).replace("root", "root/instance") + + # --- Key improvement: Open Stage using Flatten --- + # This writes all reference data directly into the main layer, preventing reference loss after path changes + base_stage = Usd.Stage.Open(usd_path) + + stage = Usd.Stage.CreateInMemory() + stage.GetRootLayer().TransferContent(base_stage.Flatten()) + + # 1. Modify hierarchy + instance_path = modify_hierarchy(stage, config) + + # 3. Export + print(f"Exporting to: {output_path}") + stage.GetRootLayer().Export(output_path) + + stage, output_usd_path = import_as_copy(output_path, output_path.replace('.usd', '_refiened.usd')) + stage = final_refine(stage, output_usd_path.replace('.usd', '_final.usd'), revolute_joint_initial_prim_path) + + stage.Export(output_path) + + os.remove(output_path.replace('.usd', '_refiened.usd')) + os.remove(output_usd_path.replace('.usd', '_final.usd')) + + print("Done.") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/workflows/simbox/tools/art/close_v/tool/select_keypoint.py b/workflows/simbox/tools/art/close_v/tool/select_keypoint.py new file mode 100644 index 0000000..3385474 --- /dev/null +++ b/workflows/simbox/tools/art/close_v/tool/select_keypoint.py @@ -0,0 +1,202 @@ +# pylint: skip-file +# flake8: noqa +import os +import argparse +import numpy as np +import open3d as o3d +import json +from pxr import Usd, UsdGeom + +def mkdir_if_missing(dst_dir): + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + +colors = [ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0.5, 0, 0], + [0, 0.5, 0], + [0, 0, 0.5], +] # red, green, blue + +parser = argparse.ArgumentParser() +parser.add_argument("--config", type=str, help="Path to config file") +args = parser.parse_args() + +# Load configuration +with open(args.config, 'r') as f: + config = json.load(f) + +# Use configuration +TASK = "close_v" +dir_name = config['DIR'] + +dir_name_kps = os.path.join(config['DIR'], "Kps", TASK) +os.makedirs(dir_name_kps, exist_ok=True) +usd_file = os.path.join(dir_name, "instance.usd") +keypoint_path = os.path.join(dir_name_kps, "keypoints.json") +target_keypoint_path = os.path.join(dir_name_kps, "keypoints_final.json") + +def scale_pcd_to_unit(pcd): + # Get all points in the point cloud + points = np.asarray(pcd.points) + + # Find min and max values + min_bound = points.min(axis=0) + max_bound = points.max(axis=0) + + # Center and range + center = (min_bound + max_bound) / 2.0 + scale = (max_bound - min_bound).max() / 2.0 # Scale using the longest edge + + # Scale point cloud to [-1, 1] + scaled_points = (points - center) / scale + pcd.points = o3d.utility.Vector3dVector(scaled_points) + + return pcd + +def get_full_transformation(prim): + """ + Get the full transformation matrix of the object relative to the world coordinate system, + including transformations from all ancestor objects. + """ + transform_matrix = np.identity(4) # Initialize as identity matrix + + # Starting from the object, apply transformations from all ancestors level by level + current_prim = prim + while current_prim: + # Get the transformation matrix of the current object + xform = UsdGeom.Xform(current_prim) + local_transform = xform.GetLocalTransformation() + + # Apply the current object's transformation to the accumulated transformation + transform_matrix = np.dot(local_transform.GetTranspose(), transform_matrix) + + # Move to parent object + current_prim = current_prim.GetParent() + + return transform_matrix + +def convert_to_world_coordinates(prim, local_vertices): + """ + Convert local coordinates to world coordinates. + """ + # Get transformation matrix relative to world coordinate system + transform_matrix = get_full_transformation(prim) + + # Convert local vertices to world coordinate system + world_vertices = [] + for vertex in local_vertices: + # Convert to homogeneous coordinates + local_point = np.append(vertex, 1) # [x, y, z, 1] + + # Convert to world coordinate system + world_point = np.dot(transform_matrix, local_point)[:3] # Take only first 3 coordinates + world_vertices.append(world_point) + + return np.array(world_vertices) + +def extract_all_geometry_from_usd(usd_file): + # Load USD file + stage = Usd.Stage.Open(usd_file) + + # Store geometric information for all objects + all_meshes = [] + + # Traverse all Prims + for prim in stage.Traverse(): + if prim.IsA(UsdGeom.Mesh): + # Extract Mesh + mesh = UsdGeom.Mesh(prim) + + # Get vertices + points = mesh.GetPointsAttr().Get() # List of vertices + vertices = np.array([[p[0], p[1], p[2]] for p in points]) + + # If conversion to world coordinate system is needed, call conversion function + vertices = convert_to_world_coordinates(prim, vertices) + + # Get face vertex indices + face_indices = mesh.GetFaceVertexIndicesAttr().Get() # Indices of all face vertices + face_vertex_counts = mesh.GetFaceVertexCountsAttr().Get() # Number of vertices per face + + # Split indices and triangulate + faces = [] + index = 0 + for count in face_vertex_counts: + face = face_indices[index:index + count] + index += count + if len(face) == 3: + faces.append(face) # Already a triangle + elif len(face) > 3: + for i in range(1, len(face) - 1): + faces.append([face[0], face[i], face[i + 1]]) + + faces = np.array(faces) + + # Get normals (if available) + if mesh.GetNormalsAttr().IsAuthored() and mesh.GetNormalsAttr().Get() is not None: + normals = mesh.GetNormalsAttr().Get() + normals = np.array([[n[0], n[1], n[2]] for n in normals]) + else: + normals = None + + # Store geometric information for current object + all_meshes.append((vertices, faces, normals)) + + # If no geometry is found, raise exception + if not all_meshes: + raise ValueError("No geometry found in USD file.") + + return all_meshes + +def visualize_geometries(meshes): + # Build Open3D Mesh and visualize + all_meshes = o3d.geometry.TriangleMesh() + for idx , (vertice, face, _) in enumerate(meshes): # Ignore normals + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(vertice) + mesh.triangles = o3d.utility.Vector3iVector(face) + mesh.paint_uniform_color(np.random.random(size=3)) + all_meshes+=mesh + + pcd = all_meshes.sample_points_uniformly(number_of_points=10000000) + return pcd + + +# Extract geometric information +all_meshes = extract_all_geometry_from_usd(usd_file) +o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug) +viewer = o3d.visualization.VisualizerWithEditing() +viewer.create_window() + +# Visualize +pcd=visualize_geometries(all_meshes) +viewer.add_geometry(pcd) +opt = viewer.get_render_option() +opt.show_coordinate_frame = True + +viewer.run() +viewer.destroy_window() + +print("saving picked points") +picked_points = viewer.get_picked_points() + +if len(picked_points) == 0: + print("No points were picked") + exit() + +xyz = np.asarray(pcd.points) +print(picked_points) +picked_points = xyz[picked_points] +print(picked_points) +color_lists = ["red", "yellow", "blue", "green", "magenta", "purple", "orange"] + +keypoint_description_file = os.path.join(dir_name_kps, "keypoints.json") +keypoint_info = { + "keypoints": {c: p.tolist() for c, p in zip(color_lists, picked_points)}, +} +with open(keypoint_description_file, "w") as f: + json.dump(keypoint_info, f, indent=4, sort_keys=True) +print("keypoint_info saved to", keypoint_description_file) diff --git a/workflows/simbox/tools/art/close_v/tool/transfer_keypoints.py b/workflows/simbox/tools/art/close_v/tool/transfer_keypoints.py new file mode 100755 index 0000000..a71087a --- /dev/null +++ b/workflows/simbox/tools/art/close_v/tool/transfer_keypoints.py @@ -0,0 +1,82 @@ +# pylint: skip-file +# flake8: noqa +import os +import sys + +current_path = os.getcwd() +sys.path.append(f"{current_path}") + +import numpy as np +import argparse +import json + +from isaacsim import SimulationApp +simulation_app = SimulationApp({"headless": True}) +from omni.isaac.core.utils.prims import get_prim_at_path +from omni.isaac.core.utils.transformations import get_relative_transform +from omni.isaac.core.utils.stage import add_reference_to_stage +from omni.isaac.core import World +from omni.isaac.core.articulations.articulation import Articulation + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, help="Path to config file") + args = parser.parse_args() + + # Load configuration + with open(args.config, 'r') as f: + config = json.load(f) + + # Use configuration + dir_name = config['DIR'] + instance_prim_path = os.path.join("/", config["INSTANCE_NAME"]) + link0_initial_prim_path = (config["link0_initial_prim_path"]).replace("/root", "instance") + base_initial_prim_path = (config["base_initial_prim_path"]).replace("/root", "instance") + + TASK = "close_v" + dir_name_kps = os.path.join(config['DIR'], "Kps", TASK) + os.makedirs(dir_name_kps, exist_ok=True) + + usd_file = os.path.join(dir_name, "instance.usd") + keypoint_path = os.path.join(dir_name_kps, "keypoints.json") + target_keypoint_path = os.path.join(dir_name_kps, "keypoints_final.json") + + if not os.path.exists(keypoint_path): + print(f"keypoint file {keypoint_path} not found") + + my_world = World() + reference = add_reference_to_stage(usd_path=usd_file, prim_path=instance_prim_path) + prim_path = str(reference.GetPrimPath()) + prim = Articulation( + prim_path, + name=config['INSTANCE_NAME'] + ) + my_world.scene.add(prim) + instance2link_pose = get_relative_transform(get_prim_at_path(instance_prim_path),get_prim_at_path(os.path.join(instance_prim_path, link0_initial_prim_path))) + instance2base_pose = get_relative_transform(get_prim_at_path(instance_prim_path),get_prim_at_path(os.path.join(instance_prim_path, base_initial_prim_path))) + kploc2base = json.load(open(keypoint_path))["keypoints"] + kplocs = {} + for name, kploc in kploc2base.items(): + if name == "red" or name == "yellow": + kploc = np.append(kploc,1) + kplocs[name] = (instance2link_pose @ kploc).tolist()[:3] + elif name == "blue": + kploc = np.append(kploc,1) + kplocs[name] = (instance2base_pose @ kploc).tolist()[:3] + else: + kplocs[name] = kploc + + # compute scale + my_world.scene.enable_bounding_boxes_computations() + bbox = my_world.scene.compute_object_AABB(config['INSTANCE_NAME']) + volume = (bbox[1][0]-bbox[0][0])*(bbox[1][1]-bbox[0][1])*(bbox[1][2]-bbox[0][2]) + scaled_volume=config['SCALED_VOLUME'] + scale = (scaled_volume / volume) **(1/3) + + data = { + "keypoints": kplocs, + "scale" : [scale,scale,scale,1.0] + } + json.dump(data, open(target_keypoint_path, "w"), indent=4) + print("Saved keypoints to ", target_keypoint_path) diff --git a/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/info.json b/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/info.json new file mode 100644 index 0000000..b4ef3ef --- /dev/null +++ b/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/info.json @@ -0,0 +1,30 @@ +{ + "object_keypoints": { + "articulated_object_head": [ + -0.023042246798858412, + 0.3964970111846924, + 0.7334318733921993 + ], + "articulated_object_tail": [ + -0.46692867303273605, + 0.3964970111846924, + 0.4885498193849562 + ] + }, + "object_scale": [ + 0.2302160991881601, + 0.2302160991881601, + 0.2302160991881601, + 1.0 + ], + "object_name": "microwave7265", + "object_usd": "/home/shixu/dev_shixu/DataEngine/workflows/simbox/tools/art/open_v/7265/usd/instance.usd", + "object_link0_rot_axis": "y", + "object_link0_contact_axis": "-y", + "object_base_front_axis": "z", + "joint_index": 0, + "object_prim_path": "/microwave7265", + "object_link_path": "/microwave7265/instance/group_18", + "object_base_path": "/microwave7265/instance/group_0", + "object_revolute_joint_path": "/microwave7265/instance/group_18/RevoluteJoint" +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/keypoints.json b/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/keypoints.json new file mode 100644 index 0000000..50219cd --- /dev/null +++ b/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/keypoints.json @@ -0,0 +1,14 @@ +{ + "keypoints": { + "red": [ + -0.01690458027356498, + 0.39649701118469244, + 0.7550849855158882 + ], + "yellow": [ + -0.41020458227271933, + 0.3964970111846924, + 0.5357677460608468 + ] + } +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/keypoints_final.json b/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/keypoints_final.json new file mode 100644 index 0000000..842a7f5 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/7265/usd/Kps/open_v/keypoints_final.json @@ -0,0 +1,20 @@ +{ + "keypoints": { + "red": [ + -0.01690458027356498, + 0.39649701118469244, + 0.7550849855158882 + ], + "yellow": [ + -0.41020458227271933, + 0.3964970111846924, + 0.5357677460608468 + ] + }, + "scale": [ + 0.2302160991881601, + 0.2302160991881601, + 0.2302160991881601, + 1.0 + ] +} \ No newline at end of file diff --git a/workflows/simbox/tools/art/open_v/7265/usd/instance.usd b/workflows/simbox/tools/art/open_v/7265/usd/instance.usd new file mode 100644 index 0000000..d89950b Binary files /dev/null and b/workflows/simbox/tools/art/open_v/7265/usd/instance.usd differ diff --git a/workflows/simbox/tools/art/open_v/7265/usd/keypoints_config.json b/workflows/simbox/tools/art/open_v/7265/usd/keypoints_config.json new file mode 100644 index 0000000..5a1135e --- /dev/null +++ b/workflows/simbox/tools/art/open_v/7265/usd/keypoints_config.json @@ -0,0 +1,13 @@ +{ + "DIR": "/home/shixu/dev_shixu/DataEngine/workflows/simbox/tools/art/open_v/7265/usd", + "USD_NAME": "microwave_0.usd", + "INSTANCE_NAME": "microwave7265", + "link0_initial_prim_path": "/root/group_18", + "base_initial_prim_path": "/root/group_0", + "revolute_joint_initial_prim_path": "/root/group_18/RevoluteJoint", + "joint_index": 0, + "LINK0_ROT_AXIS": "y", + "BASE_FRONT_AXIS": "z", + "LINK0_CONTACT_AXIS": "-y", + "SCALED_VOLUME": 0.02 +} diff --git a/workflows/simbox/tools/art/open_v/7265/usd/microwave_0.usd b/workflows/simbox/tools/art/open_v/7265/usd/microwave_0.usd new file mode 100644 index 0000000..8a291f2 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/7265/usd/microwave_0.usd differ diff --git a/workflows/simbox/tools/art/open_v/7265/usd/textures/door_mesh_0_texture_0.jpg b/workflows/simbox/tools/art/open_v/7265/usd/textures/door_mesh_0_texture_0.jpg new file mode 100644 index 0000000..6afb854 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/7265/usd/textures/door_mesh_0_texture_0.jpg differ diff --git a/workflows/simbox/tools/art/open_v/7265/usd/textures/door_mesh_1_texture_0.jpg b/workflows/simbox/tools/art/open_v/7265/usd/textures/door_mesh_1_texture_0.jpg new file mode 100644 index 0000000..6afb854 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/7265/usd/textures/door_mesh_1_texture_0.jpg differ diff --git a/workflows/simbox/tools/art/open_v/BASE_FRONT_AXIS.jpg b/workflows/simbox/tools/art/open_v/BASE_FRONT_AXIS.jpg new file mode 100644 index 0000000..fabc434 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/BASE_FRONT_AXIS.jpg differ diff --git a/workflows/simbox/tools/art/open_v/LINK0_CONTACT_AXIS.jpg b/workflows/simbox/tools/art/open_v/LINK0_CONTACT_AXIS.jpg new file mode 100644 index 0000000..cbc9192 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/LINK0_CONTACT_AXIS.jpg differ diff --git a/workflows/simbox/tools/art/open_v/LINK0_ROT_AXIS.jpg b/workflows/simbox/tools/art/open_v/LINK0_ROT_AXIS.jpg new file mode 100644 index 0000000..979f65a Binary files /dev/null and b/workflows/simbox/tools/art/open_v/LINK0_ROT_AXIS.jpg differ diff --git a/workflows/simbox/tools/art/open_v/point_0.jpg b/workflows/simbox/tools/art/open_v/point_0.jpg new file mode 100644 index 0000000..7e4a850 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/point_0.jpg differ diff --git a/workflows/simbox/tools/art/open_v/point_1.jpg b/workflows/simbox/tools/art/open_v/point_1.jpg new file mode 100644 index 0000000..f69d449 Binary files /dev/null and b/workflows/simbox/tools/art/open_v/point_1.jpg differ diff --git a/workflows/simbox/tools/art/open_v/readme.md b/workflows/simbox/tools/art/open_v/readme.md new file mode 100644 index 0000000..79cac81 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/readme.md @@ -0,0 +1,45 @@ +# Annotation Documentation + +We provide an optimized and simplified annotation pipeline that removes many redundancies. No need to rename base_link, contact_link, etc. Keep the original hierarchy and naming as much as possible. + +## 🗂️ File Information + +| Configuration | Example | Description | +|---------------|---------|-------------| +| **DIR** | `YOUR_PATH_TO_DIR/usd` | Directory where USD files are stored | +| **USD_NAME** | `microwave_0.usd` | Scene description file name | +| **INSTANCE_NAME** | `microwave7265` | Model identifier in the scene. You can name it yourself, preferably matching the generated file name | + +## 🔧 Model Structure Configuration + +| Component | Example | Description | +|-----------|---------|-------------| +| **link0_initial_prim_path** | `/root/group_18` | Absolute path in Isaac Sim for the "door" that interacts with the gripper. Check in the original USD | +| **base_initial_prim_path** | `/root/group_0` | Absolute path in Isaac Sim for the microwave base. Check in the original USD | +| **revolute_joint_initial_prim_path** | `/root/group_18/RevoluteJoint` | Absolute path in Isaac Sim for the revolute joint that opens/closes the microwave. Check in the original USD | +| **Joint Index** | `0` | Joint number, default is 0 | + +## 🧭 Axis Configuration + +| Axis Type | Example | Description | Visualization | +|-----------|---------|-------------|---------------| +| **LINK0_ROT_AXIS** | `y` | In the local coordinate system of the rotating joint, the axis direction pointing vertically upward | ![LINK0_ROT_AXIS Example](LINK0_ROT_AXIS.jpg) | +| **BASE_FRONT_AXIS** | `z` | In the local coordinate system of the microwave base link, the axis direction facing the door | ![BASE_FRONT_AXIS Example](BASE_FRONT_AXIS.jpg) | +| **LINK0_CONTACT_AXIS** | `-y` | In the local coordinate system of the contact link, the axis direction pointing vertically downward | ![LINK0_CONTACT_AXIS Example](LINK0_CONTACT_AXIS.jpg) | + +## 📏 Physical Parameters + +| Parameter | Example | Description | +|-----------|---------|-------------| +| **SCALED_VOLUME** | `0.02` | Default value 0.02 for microwave-like objects | + +--- + +# Point Annotation + +| Point Type | Description | Visualization | +|------------|-------------|---------------| +| First Point (articulated_object_head) | `Desired base position where the gripper contacts the microwave door` | ![First Point Diagram](point_0.jpg) | +| Second Point (articulated_object_tail) | `The line direction from the first point should be perpendicular to the microwave door's rotation axis` | ![Second Point Diagram](point_1.jpg) | + +--- \ No newline at end of file diff --git a/workflows/simbox/tools/art/open_v/tools/keypoints_pipeline.sh b/workflows/simbox/tools/art/open_v/tools/keypoints_pipeline.sh new file mode 100644 index 0000000..8ffa090 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/tools/keypoints_pipeline.sh @@ -0,0 +1,16 @@ +#!/bin/bash +CONFIG_PATH="/home/shixu/dev_shixu/DataEngine/workflows/simbox/tools/art/open_v/7265/usd/keypoints_config.json" + +cd /home/shixu/dev_shixu/DataEngine/workflows/simbox/tools/art/open_v/tools + +# 1. rehier +python rehier.py --config $CONFIG_PATH + +# 2. select points +python select_keypoint.py --config $CONFIG_PATH + +# 3. Transfer keypoints +python transfer_keypoints.py --config $CONFIG_PATH + +# 4. Overwrite keypoints +python overwrite_keypoints.py --config $CONFIG_PATH \ No newline at end of file diff --git a/workflows/simbox/tools/art/open_v/tools/overwrite_keypoints.py b/workflows/simbox/tools/art/open_v/tools/overwrite_keypoints.py new file mode 100644 index 0000000..fe83038 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/tools/overwrite_keypoints.py @@ -0,0 +1,57 @@ +# pylint: skip-file +# flake8: noqa +import os +import json +import numpy as np + +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--config", type=str, help="Path to config file") +args = parser.parse_args() + +with open(args.config, 'r') as f: + config = json.load(f) + +TASK = "open_v" +dir_name = config["DIR"] +dir_name_kps = os.path.join(config['DIR'], "Kps", TASK) +os.makedirs(dir_name_kps, exist_ok=True) + +usd_file = os.path.join(dir_name, "instance.usd") +keypoint_path = os.path.join(dir_name_kps, "keypoints_final.json") +target_keypoint_path = os.path.join(dir_name_kps, "info.json") + +if not os.path.exists(target_keypoint_path): + with open(target_keypoint_path,'w') as file: + data={"object_keypoints":{}} + json.dump(data,file,indent=4) + +if not os.path.exists(keypoint_path) or not os.path.exists(target_keypoint_path): + print(f"keypoint file {keypoint_path} or {target_keypoint_path} not found") + +kp = json.load(open(keypoint_path)) +tkp = json.load(open(target_keypoint_path)) + +tkp["object_keypoints"]["articulated_object_head"] = kp["keypoints"]["red"] +tkp["object_keypoints"]["articulated_object_tail"] = kp["keypoints"]["yellow"] + +tkp["object_scale"] = kp["scale"] +tkp["object_name"] = config["INSTANCE_NAME"] +tkp["object_usd"] = usd_file +tkp["object_link0_rot_axis"] = config["LINK0_ROT_AXIS"] +tkp["object_link0_contact_axis"] = config["LINK0_CONTACT_AXIS"] +tkp["object_base_front_axis"] = config["BASE_FRONT_AXIS"] +tkp["joint_index"] = config["joint_index"] + +tkp["object_prim_path"] = os.path.join("/", config["INSTANCE_NAME"]) +link0_initial_prim_path = (config["link0_initial_prim_path"]).replace("/root", "instance") +base_initial_prim_path = (config["base_initial_prim_path"]).replace("/root", "instance") +revolute_joint_initial_prim_path = (config["revolute_joint_initial_prim_path"]).replace("/root", "instance") + +tkp["object_link_path"] = os.path.join("/", config["INSTANCE_NAME"], link0_initial_prim_path) +tkp["object_base_path"] = os.path.join("/", config["INSTANCE_NAME"], base_initial_prim_path) +tkp["object_revolute_joint_path"] = os.path.join("/", config["INSTANCE_NAME"], revolute_joint_initial_prim_path) + +json.dump(tkp, open(target_keypoint_path, "w"), indent=4) +print("Saved keypoints to ", target_keypoint_path) diff --git a/workflows/simbox/tools/art/open_v/tools/rehier.py b/workflows/simbox/tools/art/open_v/tools/rehier.py new file mode 100755 index 0000000..2413bb3 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/tools/rehier.py @@ -0,0 +1,302 @@ +# pylint: skip-file +# flake8: noqa +import os +import json +import argparse +from pathlib import Path +from pxr import Usd, UsdGeom, UsdPhysics, Gf, Sdf +from pdb import set_trace + +def remove_articulation_root(prim: Usd.Prim): + prim.RemoveAPI(UsdPhysics.ArticulationRootAPI) + allchildren = prim.GetAllChildren() + if len(allchildren) == 0: + return + else: + for child in allchildren: + remove_articulation_root(child) + +def remove_rigidbody(prim: Usd.Prim): + prim.RemoveAPI(UsdPhysics.RigidBodyAPI) + allchildren = prim.GetAllChildren() + if len(allchildren) == 0: + return + else: + for child in allchildren: + remove_rigidbody(child) + +def remove_mass(prim: Usd.Prim): + prim.RemoveAPI(UsdPhysics.MassAPI) + allchildren = prim.GetAllChildren() + if len(allchildren) == 0: + return + else: + for child in allchildren: + remove_mass(child) + +def add_rigidbody(prim: Usd.Prim): + UsdPhysics.RigidBodyAPI.Apply(prim) + +def add_mass(prim: Usd.Prim): + UsdPhysics.MassAPI.Apply(prim) + mass = prim.GetAttribute("physics:mass") + mass.Clear() + +def get_args(): + parser = argparse.ArgumentParser(description="USD Hierarchy and Physics Editor") + parser.add_argument("--config", type=str, required=True, help="Path to config file") + + return parser.parse_args() + +def load_config(config_path): + with open(config_path, 'r') as f: + return json.load(f) + +def safe_rename_prim(stage, old_path, new_path): + editor = Usd.NamespaceEditor(stage) + old_p = stage.GetPrimAtPath(old_path) + editor.RenamePrim(old_p, new_path.split('/')[-1]) + + if editor.CanApplyEdits(): + editor.ApplyEdits() + return True + else: + return False + +def modify_hierarchy(stage, config): + editor = Usd.NamespaceEditor(stage) + + """ Modify USD hierarchy structure """ + base_path = config["base_initial_prim_path"] + link0_path = config["link0_initial_prim_path"] + revolute_joint_path = config["revolute_joint_initial_prim_path"] + + # Get original root node + old_root_path = f"/{link0_path.split('/')[1]}" + instance_path = "/root/instance" + safe_rename_prim(stage, '/{}'.format(stage.GetDefaultPrim().GetName()), "/instance") + + return instance_path + +def modify_physics(stage, instance_root_path, config): + """ Modify physics properties and ensure colliders are uniformly set to Convex Hull """ + print(f"Applying physics modifications and setting colliders to ConvexHull...") + + # Traverse all nodes for processing + for prim in stage.Traverse(): + # 1. Clear instancing (if editing physics properties of individual instances) + if prim.IsInstanceable(): + prim.ClearInstanceable() + + # 2. Process Mesh colliders + if prim.IsA(UsdGeom.Mesh): + # Ensure base collision API is applied + if not prim.HasAPI(UsdPhysics.CollisionAPI): + UsdPhysics.CollisionAPI.Apply(prim) + + # Force apply MeshCollisionAPI to set collision approximation + mesh_collision_api = UsdPhysics.MeshCollisionAPI.Get(stage, prim.GetPath()) + if not mesh_collision_api: + mesh_collision_api = UsdPhysics.MeshCollisionAPI.Apply(prim) + + # Set collision approximation to 'convexHull' + # Optional values include: 'none', 'convexHull', 'convexDecomposition', 'meshSimplification', etc. + mesh_collision_api.CreateApproximationAttr().Set("convexHull") + + # Ensure physics collision is enabled + col_enabled_attr = prim.GetAttribute("physics:collisionEnabled") + if not col_enabled_attr.HasValue(): + col_enabled_attr.Set(True) + +def create_fixed_joint(stage, joint_path, body0_path, body1_path): + """ + Create a FixedJoint at the specified path and connect two rigid bodies. + """ + # 1. Define FixedJoint node + fixed_joint = UsdPhysics.FixedJoint.Define(stage, joint_path) + + # 2. Set Body0 and Body1 path references + # Note: Paths must be of Sdf.Path type + fixed_joint.GetBody0Rel().SetTargets([Sdf.Path(body0_path)]) + fixed_joint.GetBody1Rel().SetTargets([Sdf.Path(body1_path)]) + + # 3. (Optional) Set local offset (Local Pose) + # If not set, the joint defaults to the origin of both objects + # fixed_joint.GetLocalPos0Attr().Set(Gf.Vec3f(0, 0, 0)) + # fixed_joint.GetLocalRot0Attr().Set(Gf.Quatf(1, 0, 0, 0)) + + print(f"Successfully created FixedJoint: {joint_path}") + print(f" Connected: {body0_path} <---> {body1_path}") + return fixed_joint + +def process_joints(stage, revolute_joint_initial_prim_path): + # 1. Collect paths to process + paths_to_delete = [] + joints_to_convert = [] + + # Use TraverseAll to ensure no defined nodes are missed + for prim in stage.Traverse(): + # Check if it is a physics joint + if prim.IsA(UsdPhysics.Joint): + path = prim.GetPath() + + # Logic A: If FixedJoint -> delete + if prim.IsA(UsdPhysics.FixedJoint): + paths_to_delete.append(path) + + # Logic B: If not FixedJoint and path does not contain 'contact_link' -> convert to FixedJoint + # elif "contact_link" not in str(path).lower(): + # joints_to_convert.append(path) + elif str(path) != revolute_joint_initial_prim_path: + joints_to_convert.append(path) + + print(str(path)) + + # 2. Get current edit layer + layer = stage.GetEditTarget().GetLayer() + edit = Sdf.BatchNamespaceEdit() + + # Execute deletion logic + for path in paths_to_delete: + edit.Add(path, Sdf.Path.emptyPath) + print(f"[Delete] FixedJoint: {path}") + + # 3. Apply deletion edits + if paths_to_delete: + layer.Apply(edit) + + # 4. Execute type conversion logic + # In USD, changing type usually means re-Defining the new type at that path + for path in joints_to_convert: + # Record original Body0 and Body1 relationships to prevent loss after conversion + prim = stage.GetPrimAtPath(path) + joint = UsdPhysics.Joint(prim) + body0 = joint.GetBody0Rel().GetTargets() + body1 = joint.GetBody1Rel().GetTargets() + + # Redefine as FixedJoint + new_fixed_joint = UsdPhysics.FixedJoint.Define(stage, path) + + # Restore relationships + if body0: new_fixed_joint.GetBody0Rel().SetTargets(body0) + if body1: new_fixed_joint.GetBody1Rel().SetTargets(body1) + + safe_rename_prim(stage, str(new_fixed_joint.GetPath()), "/FixedJoint") + + print(f"[Convert] Regular joint -> FixedJoint: {path}") + + return stage + +def final_refine(stage, output_usd_path, revolute_joint_initial_prim_path): + root_prim = stage.GetPrimAtPath("/root") + instance_prim = stage.GetPrimAtPath("/root/instance") + ### remove articulation root ### + remove_articulation_root(root_prim) + + ### remove rigid body ### + remove_rigidbody(root_prim) + + ### remove mass ### + remove_mass(root_prim) + + ### add rigid body and mass ### + for child in instance_prim.GetAllChildren(): + + if child.GetTypeName() == "PhysicsRevoluteJoint" or child.GetTypeName() == "PhysicsPrismaticJoint": + continue + + if child.GetTypeName() == "Xform" : + print('name:', child.GetTypeName()) + add_rigidbody(child) + + stage = process_joints(stage, revolute_joint_initial_prim_path) + + ### add articulation root ### + UsdPhysics.ArticulationRootAPI.Apply(instance_prim) + stage.SetDefaultPrim(root_prim) + + for child in instance_prim.GetAllChildren(): + try: + attr = child.GetAttribute('physics:jointEnabled') + except: + continue + if attr.Get() is not None: + print(child) + attr.Set(True) + + modify_physics(stage, "/root/instance", 11) + stage.Export(output_usd_path) + + return stage + +def import_as_copy(source_usd_path, output_usd_path, root_name="root", sub_node_name="instance"): + """ + Create a new USD and copy the content from source_usd_path to /root/sub_node_name. + """ + # 1. Create target Stage and root node + stage = Usd.Stage.CreateNew(output_usd_path) + root_path = Sdf.Path(f"/{root_name}") + UsdGeom.Xform.Define(stage, root_path) + + # 2. Define copy destination path (e.g., /root/model_copy) + dest_path = root_path.AppendChild(sub_node_name) + + # 3. Open source file layer + source_layer = Sdf.Layer.FindOrOpen(source_usd_path) + + if not source_layer: + print(f"Error: Cannot find source file {source_usd_path}") + return + + # 4. Get source file's default prim (DefaultPrim) as copy target + # If source file has no default prim, use the first root prim + source_root_name = list(source_layer.rootPrims)[0].name + source_path = Sdf.Path(f"/{source_root_name}") + + # 5. Execute core copy operation (Sdf.CopySpec) + # This copies all attributes, topology, and properties from source file to new file + Sdf.CopySpec(source_layer, source_path, stage.GetRootLayer(), dest_path) + + # 6. Set default prim and save + stage.SetDefaultPrim(stage.GetPrimAtPath(root_path)) + stage.GetRootLayer().Save() + print(f"Success! Content copied to: {output_usd_path}") + + return stage, output_usd_path + +def main(): + args = get_args() + config = load_config(args.config) + + dir_name = config["DIR"] + usd_path = os.path.join(dir_name, config["USD_NAME"]) + output_path = os.path.join(dir_name, "instance.usd") + revolute_joint_initial_prim_path = (config["revolute_joint_initial_prim_path"]).replace("root", "root/instance") + + # --- Key improvement: Open Stage using Flatten --- + # This writes all reference data directly into the main layer, preventing reference loss after path changes + base_stage = Usd.Stage.Open(usd_path) + + stage = Usd.Stage.CreateInMemory() + stage.GetRootLayer().TransferContent(base_stage.Flatten()) + + # 1. Modify hierarchy + instance_path = modify_hierarchy(stage, config) + + # 3. Export + print(f"Exporting to: {output_path}") + stage.GetRootLayer().Export(output_path) + + stage, output_usd_path = import_as_copy(output_path, output_path.replace('.usd', '_refiened.usd')) + stage = final_refine(stage, output_usd_path.replace('.usd', '_final.usd'), revolute_joint_initial_prim_path) + + stage.Export(output_path) + + os.remove(output_path.replace('.usd', '_refiened.usd')) + os.remove(output_usd_path.replace('.usd', '_final.usd')) + + print("Done.") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/workflows/simbox/tools/art/open_v/tools/select_keypoint.py b/workflows/simbox/tools/art/open_v/tools/select_keypoint.py new file mode 100644 index 0000000..a329858 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/tools/select_keypoint.py @@ -0,0 +1,176 @@ +# pylint: skip-file +# flake8: noqa +import os +import argparse +import numpy as np +import open3d as o3d +import json +from pxr import Usd, UsdGeom + +def mkdir_if_missing(dst_dir): + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + +colors = [ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0.5, 0, 0], + [0, 0.5, 0], + [0, 0, 0.5], +] # red, green, blue + +parser = argparse.ArgumentParser() +parser.add_argument("--config", type=str, help="Path to config file") +args = parser.parse_args() + +# Load configuration +with open(args.config, 'r') as f: + config = json.load(f) + +# Use configuration +TASK = "open_v" +dir_name = config['DIR'] + +dir_name_kps = os.path.join(config['DIR'], "Kps", TASK) +os.makedirs(dir_name_kps, exist_ok=True) +usd_file = os.path.join(dir_name, "instance.usd") +keypoint_path = os.path.join(dir_name_kps, "keypoints.json") +target_keypoint_path = os.path.join(dir_name_kps, "keypoints_final.json") + +def scale_pcd_to_unit(pcd): + points = np.asarray(pcd.points) + + min_bound = points.min(axis=0) + max_bound = points.max(axis=0) + + center = (min_bound + max_bound) / 2.0 + scale = (max_bound - min_bound).max() / 2.0 + + scaled_points = (points - center) / scale + pcd.points = o3d.utility.Vector3dVector(scaled_points) + + return pcd + +def get_full_transformation(prim): + transform_matrix = np.identity(4) + + current_prim = prim + while current_prim: + xform = UsdGeom.Xform(current_prim) + local_transform = xform.GetLocalTransformation() + transform_matrix = np.dot(local_transform.GetTranspose(), transform_matrix) + current_prim = current_prim.GetParent() + return transform_matrix + +def convert_to_world_coordinates(prim, local_vertices): + transform_matrix = get_full_transformation(prim) + + world_vertices = [] + for vertex in local_vertices: + local_point = np.append(vertex, 1) # [x, y, z, 1] + + world_point = np.dot(transform_matrix, local_point)[:3] + world_vertices.append(world_point) + + return np.array(world_vertices) + +def extract_all_geometry_from_usd(usd_file, keyword, instance_name): + stage = Usd.Stage.Open(usd_file) + + all_meshes = [] + + for prim in stage.Traverse(): + if prim.IsA(UsdGeom.Mesh): + mesh = UsdGeom.Mesh(prim) + + path = prim.GetPrimPath() + curr_type = path.pathString + + if "microwave" in instance_name.lower(): + if keyword not in curr_type.lower(): + continue + + points = mesh.GetPointsAttr().Get() + vertices = np.array([[p[0], p[1], p[2]] for p in points]) + + vertices = convert_to_world_coordinates(prim, vertices) + + face_indices = mesh.GetFaceVertexIndicesAttr().Get() + face_vertex_counts = mesh.GetFaceVertexCountsAttr().Get() + + faces = [] + index = 0 + for count in face_vertex_counts: + face = face_indices[index:index + count] + index += count + if len(face) == 3: + faces.append(face) + elif len(face) > 3: + for i in range(1, len(face) - 1): + faces.append([face[0], face[i], face[i + 1]]) + + faces = np.array(faces) + + if mesh.GetNormalsAttr().IsAuthored() and mesh.GetNormalsAttr().Get() is not None: + normals = mesh.GetNormalsAttr().Get() + normals = np.array([[n[0], n[1], n[2]] for n in normals]) + else: + normals = None + + all_meshes.append((vertices, faces, normals)) + + if not all_meshes: + raise ValueError("No geometry found in USD file.") + + return all_meshes + +def visualize_geometries(meshes): + all_meshes = o3d.geometry.TriangleMesh() + for idx , (vertice, face, _) in enumerate(meshes): + mesh = o3d.geometry.TriangleMesh() + mesh.vertices = o3d.utility.Vector3dVector(vertice) + mesh.triangles = o3d.utility.Vector3iVector(face) + mesh.paint_uniform_color(np.random.random(size=3)) + all_meshes+=mesh + + pcd = all_meshes.sample_points_uniformly(number_of_points=10000000) + + return pcd + +instance_name = config['INSTANCE_NAME'] +keyword = os.path.basename(config['link0_initial_prim_path']) # should be "contact_link" +all_meshes = extract_all_geometry_from_usd(usd_file, keyword=keyword, instance_name=instance_name) +o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug) +viewer = o3d.visualization.VisualizerWithEditing() +viewer.create_window() + +# Visualize +pcd=visualize_geometries(all_meshes) +viewer.add_geometry(pcd) +opt = viewer.get_render_option() +opt.show_coordinate_frame = True + +viewer.run() +viewer.destroy_window() + +print("saving picked points") +picked_points = viewer.get_picked_points() + +if len(picked_points) == 0: + print("No points were picked") + exit() + +xyz = np.asarray(pcd.points) +print(picked_points) +picked_points = xyz[picked_points] +print(picked_points) +color_lists = ["red", "yellow", "blue", "green", "magenta", "purple", "orange"] + +keypoint_description_file = os.path.join(dir_name_kps, "keypoints.json") +keypoint_info = { + "keypoints": {c: p.tolist() for c, p in zip(color_lists, picked_points)}, +} +with open(keypoint_description_file, "w") as f: + json.dump(keypoint_info, f, indent=4, sort_keys=True) +print("keypoint_info saved to", keypoint_description_file) diff --git a/workflows/simbox/tools/art/open_v/tools/transfer_keypoints.py b/workflows/simbox/tools/art/open_v/tools/transfer_keypoints.py new file mode 100644 index 0000000..342c251 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/tools/transfer_keypoints.py @@ -0,0 +1,82 @@ +# pylint: skip-file +# flake8: noqa +import os +import sys + +current_path = os.getcwd() +sys.path.append(f"{current_path}") + +import numpy as np +import argparse +import json + +from isaacsim import SimulationApp +simulation_app = SimulationApp({"headless": True}) +from omni.isaac.core.utils.prims import get_prim_at_path +from omni.isaac.core.utils.transformations import get_relative_transform +from omni.isaac.core.utils.stage import add_reference_to_stage +from omni.isaac.core import World +from omni.isaac.core.articulations.articulation import Articulation + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, help="Path to config file") + args = parser.parse_args() + + # Load configuration + with open(args.config, 'r') as f: + config = json.load(f) + + # Use configuration + dir_name = config['DIR'] + instance_prim_path = os.path.join("/", config["INSTANCE_NAME"]) + link0_initial_prim_path = (config["link0_initial_prim_path"]).replace("/root", "instance") + base_initial_prim_path = (config["base_initial_prim_path"]).replace("/root", "instance") + + TASK = "open_v" + dir_name_kps = os.path.join(config['DIR'], "Kps", TASK) + os.makedirs(dir_name_kps, exist_ok=True) + + usd_file = os.path.join(dir_name, "instance.usd") + keypoint_path = os.path.join(dir_name_kps, "keypoints.json") + target_keypoint_path = os.path.join(dir_name_kps, "keypoints_final.json") + + if not os.path.exists(keypoint_path): + print(f"keypoint file {keypoint_path} not found") + + my_world = World() + reference = add_reference_to_stage(usd_path=usd_file, prim_path=instance_prim_path) + prim_path = str(reference.GetPrimPath()) + prim = Articulation( + prim_path, + name=config['INSTANCE_NAME'] + ) + my_world.scene.add(prim) + instance2link_pose = get_relative_transform(get_prim_at_path(instance_prim_path),get_prim_at_path(os.path.join(instance_prim_path, link0_initial_prim_path))) + instance2base_pose = get_relative_transform(get_prim_at_path(instance_prim_path),get_prim_at_path(os.path.join(instance_prim_path, base_initial_prim_path))) + kploc2base = json.load(open(keypoint_path))["keypoints"] + kplocs = {} + for name, kploc in kploc2base.items(): + if name == "red" or name == "yellow": + kploc = np.append(kploc,1) + kplocs[name] = (instance2link_pose @ kploc).tolist()[:3] + elif name == "blue": + kploc = np.append(kploc,1) + kplocs[name] = (instance2base_pose @ kploc).tolist()[:3] + else: + kplocs[name] = kploc + + # compute scale + my_world.scene.enable_bounding_boxes_computations() + bbox = my_world.scene.compute_object_AABB(config['INSTANCE_NAME']) + volume = (bbox[1][0]-bbox[0][0])*(bbox[1][1]-bbox[0][1])*(bbox[1][2]-bbox[0][2]) + scaled_volume=config['SCALED_VOLUME'] + scale = (scaled_volume / volume) **(1/3) + + data = { + "keypoints": kplocs, + "scale" : [scale,scale,scale,1.0] + } + json.dump(data, open(target_keypoint_path, "w"), indent=4) + print("Saved keypoints to ", target_keypoint_path) diff --git a/workflows/simbox/tools/art/open_v/tools/usd2mesh.py b/workflows/simbox/tools/art/open_v/tools/usd2mesh.py new file mode 100644 index 0000000..977d776 --- /dev/null +++ b/workflows/simbox/tools/art/open_v/tools/usd2mesh.py @@ -0,0 +1,183 @@ +from pxr import Usd, UsdGeom +import numpy as np + +import os +import os.path as osp +import trimesh +import json + +def get_full_transformation(prim): + """ + Get the full transformation matrix of the object relative to the world coordinate system, + including transformations from all ancestor objects. + """ + transform_matrix = np.identity(4) # Initialize as identity matrix + + # Starting from the object, apply transformations from all ancestors level by level + current_prim = prim + while current_prim: + # Get the transformation matrix of the current object + xform = UsdGeom.Xform(current_prim) + local_transform = xform.GetLocalTransformation() + + # Apply the current object's transformation to the accumulated transformation + transform_matrix = np.dot(local_transform.GetTranspose(), transform_matrix) + + # Move to parent object + current_prim = current_prim.GetParent() + + return transform_matrix + +def convert_to_world_coordinates(prim, local_vertices): + """ + Convert local coordinates to world coordinates. + """ + # Get transformation matrix relative to world coordinate system + transform_matrix = get_full_transformation(prim) + + # Convert local vertices to world coordinate system + world_vertices = [] + for vertex in local_vertices: + # Convert to homogeneous coordinates + local_point = np.append(vertex, 1) # [x, y, z, 1] + + # Convert to world coordinate system + world_point = np.dot(transform_matrix, local_point)[:3] # Take only first 3 coordinates + world_vertices.append(world_point) + + return np.array(world_vertices) + +def extract_all_geometry_from_usd(usd_file, local_frame=True): + # Load USD file + stage = Usd.Stage.Open(usd_file) + + # Store geometric information for all objects + all_meshes = { + "visuals": [], + "collisions": [], + } + + # Traverse all Prims + for prim in stage.Traverse(): + if prim.IsA(UsdGeom.Mesh): + + path = prim.GetPrimPath() + curr_type = path.pathString.split("/")[-2] + + if curr_type not in list(all_meshes.keys()): + curr_type = "visuals" + + # Extract Mesh + mesh = UsdGeom.Mesh(prim) + + # Get vertices + points = mesh.GetPointsAttr().Get() # List of vertices + vertices = np.array([[p[0], p[1], p[2]] for p in points]) + + if not local_frame: + # If conversion to world coordinate system is needed, call conversion function + vertices = convert_to_world_coordinates(prim, vertices) + + # Get face vertex indices + face_indices = mesh.GetFaceVertexIndicesAttr().Get() # Indices of all face vertices + face_vertex_counts = mesh.GetFaceVertexCountsAttr().Get() # Number of vertices per face + + # Split indices and triangulate + faces = [] + index = 0 + for count in face_vertex_counts: + face = face_indices[index:index + count] + index += count + if len(face) == 3: + faces.append(face) # Already a triangle + elif len(face) > 3: + for i in range(1, len(face) - 1): + faces.append([face[0], face[i], face[i + 1]]) + + faces = np.array(faces) + + # Get normals (if available) + if mesh.GetNormalsAttr().IsAuthored() and mesh.GetNormalsAttr().Get() is not None: + normals = mesh.GetNormalsAttr().Get() + normals = np.array([[n[0], n[1], n[2]] for n in normals]) + else: + normals = None + + # Store geometric information for current object + all_meshes[curr_type].append((vertices, faces, normals)) + + # If no geometry is found, raise exception + if not all_meshes: + raise ValueError("No geometry found in USD file.") + + return all_meshes + +if __name__ == "__main__": + + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, help="Path to config file") + args = parser.parse_args() + + # Load configuration + with open(args.config, 'r') as f: + config = json.load(f) + + asset_fn = osp.join(config["DIR"], "instance.usd") + + dir_name = osp.dirname(asset_fn) + + ######## Load annotated keypoints + keypoint_path = os.path.join(dir_name, "keypoints.json") # in global coordinate + keypoint_path_open = os.path.join(dir_name, "Kps/open_v", "keypoints.json") # in global coordinate + + ######## Global + + all_meshes = extract_all_geometry_from_usd(asset_fn, local_frame=False) + save_dir = osp.join(dir_name, "Meshes", "global_frame") + + os.makedirs(save_dir, exist_ok=True) + + for k, v in all_meshes.items(): + for i in range(len(v)): + m = trimesh.Trimesh(vertices=v[i][0], faces=v[i][1]) + os.makedirs(os.path.join(save_dir, k), exist_ok=True) + m.export(os.path.join(save_dir, k, "{}.obj".format(i))) + + ###### add keypoints -- Close Task + kps = json.load(open(keypoint_path))["keypoints"] + components = [] + radius = 0.02 + for k, v in kps.items(): + marker = trimesh.creation.icosphere(subdivisions=3, radius=radius) + marker.apply_translation(v) + if k == "blue": + marker.visual.vertex_colors[:, :3] = [0, 0, 255] + elif k == "red": + marker.visual.vertex_colors[:, :3] = [255, 0, 0] + elif k == "yellow": + marker.visual.vertex_colors[:, :3] = [255, 255, 0] + else: + raise NotImplementedError + components.append(marker) + merged_mesh = trimesh.util.concatenate(components) + merged_mesh.export(os.path.join(save_dir, "kps_close.obj")) + + ###### add keypoints -- Open Task + kps = json.load(open(keypoint_path_open))["keypoints"] + components = [] + for k, v in kps.items(): + marker = trimesh.creation.icosphere(subdivisions=3, radius=radius) + marker.apply_translation(v) + if k == "blue": + marker.visual.vertex_colors[:, :3] = [0, 0, 255] + elif k == "red": + marker.visual.vertex_colors[:, :3] = [255, 0, 0] + elif k == "yellow": + marker.visual.vertex_colors[:, :3] = [255, 255, 0] + else: + raise NotImplementedError + components.append(marker) + merged_mesh = trimesh.util.concatenate(components) + merged_mesh.export(os.path.join(save_dir, "kps_open.obj"))