additioanl articulation tools
This commit is contained in:
Leon998
302
workflows/simbox/tools/art/open_h/tool/rehier.py
Executable file
302
workflows/simbox/tools/art/open_h/tool/rehier.py
Executable file
@@ -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()
|
||||
Reference in New Issue
Block a user