# Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved. # NVIDIA CORPORATION and its licensors retain all intellectual property # and proprietary rights in and to this software, related documentation # and any modifications thereto. Any use, reproduction, disclosure or # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. import math import numpy as np import os import xml.etree.ElementTree as ET import warp as wp # SNU file format parser class MuscleUnit: def __init__(self): self.name = "" self.bones = [] self.points = [] class Skeleton: def __init__(self, root_xform, skeleton_file, muscle_file, builder, filter, armature=0.0): self.parse_skeleton(skeleton_file, builder, filter, root_xform, armature) self.parse_muscles(muscle_file, builder) def parse_skeleton(self, filename, builder, filter, root_xform, armature): file = ET.parse(filename) root = file.getroot() self.node_map = {} # map node names to link indices self.xform_map = {} # map node names to parent transforms self.mesh_map = {} # map mesh names to link indices objects self.coord_start = builder.joint_coord_count self.dof_start = builder.joint_dof_count type_map = { "Ball": wp.sim.JOINT_BALL, "Revolute": wp.sim.JOINT_REVOLUTE, "Prismatic": wp.sim.JOINT_PRISMATIC, "Free": wp.sim.JOINT_FREE, "Fixed": wp.sim.JOINT_FIXED, } builder.add_articulation() for child in root: if child.tag == "Node": body = child.find("Body") joint = child.find("Joint") name = child.attrib["name"] parent = child.attrib["parent"] parent_X_s = wp.transform_identity() if parent in self.node_map: parent_link = self.node_map[parent] parent_X_s = self.xform_map[parent] else: parent_link = -1 body_xform = body.find("Transformation") joint_xform = joint.find("Transformation") body_mesh = body.attrib["obj"] body_size = np.fromstring(body.attrib["size"], sep=" ") body_type = body.attrib["type"] body_mass = body.attrib["mass"] body_R_s = np.fromstring(body_xform.attrib["linear"], sep=" ").reshape((3, 3)) body_t_s = np.fromstring(body_xform.attrib["translation"], sep=" ") joint_R_s = np.fromstring(joint_xform.attrib["linear"], sep=" ").reshape((3, 3)) joint_t_s = np.fromstring(joint_xform.attrib["translation"], sep=" ") joint_type = type_map[joint.attrib["type"]] joint_lower = np.array([-1.0e3]) joint_upper = np.array([1.0e3]) try: joint_lower = np.fromstring(joint.attrib["lower"], sep=" ") joint_upper = np.fromstring(joint.attrib["upper"], sep=" ") except: pass if "axis" in joint.attrib: joint_axis = np.fromstring(joint.attrib["axis"], sep=" ") else: joint_axis = np.array((0.0, 0.0, 0.0)) body_X_s = wp.transform(body_t_s, wp.quat_from_matrix(body_R_s)) joint_X_s = wp.transform(joint_t_s, wp.quat_from_matrix(joint_R_s)) mesh_base = os.path.splitext(body_mesh)[0] mesh_file = mesh_base + ".usd" # ----------------------------------- # one time conversion, put meshes into local body space (and meter units) # stage = Usd.Stage.Open("./assets/snu/OBJ/" + mesh_file) # geom = UsdGeom.Mesh.Get(stage, "/" + mesh_base + "_obj/defaultobject/defaultobject") # body_X_bs = wp.transform_inverse(body_X_s) # joint_X_bs = wp.transform_inverse(joint_X_s) # points = geom.GetPointsAttr().Get() # for i in range(len(points)): # p = wp.transform_point(joint_X_bs, points[i]*0.01) # points[i] = Gf.Vec3f(p.tolist()) # cm -> meters # geom.GetPointsAttr().Set(points) # extent = UsdGeom.Boundable.ComputeExtentFromPlugins(geom, 0.0) # geom.GetExtentAttr().Set(extent) # stage.Save() # -------------------------------------- link = -1 if len(filter) == 0 or name in filter: joint_X_p = wp.transform_multiply(wp.transform_inverse(parent_X_s), joint_X_s) body_X_c = wp.transform_multiply(wp.transform_inverse(joint_X_s), body_X_s) if parent_link == -1: joint_X_p = wp.transform_identity() # add link link = builder.add_body( parent=parent_link, origin=wp.transform_multiply(root_xform, joint_X_s), joint_xform=joint_X_p, joint_axis=joint_axis, joint_type=joint_type, joint_target_ke=5.0, joint_target_kd=2.0, joint_limit_lower=joint_lower[0], joint_limit_upper=joint_upper[0], joint_limit_ke=1.0e3, joint_limit_kd=1.0e2, joint_armature=armature, ) # add shape shape = builder.add_shape_box( body=link, pos=body_X_c.p, rot=body_X_c.q, hx=body_size[0] * 0.5, hy=body_size[1] * 0.5, hz=body_size[2] * 0.5, ke=1.0e3 * 5.0, kd=1.0e2 * 2.0, kf=1.0e3, mu=0.5, ) # add lookup in name->link map # save parent transform self.xform_map[name] = joint_X_s self.node_map[name] = link self.mesh_map[mesh_base] = link def parse_muscles(self, filename, builder): # list of MuscleUnits muscles = [] file = ET.parse(filename) root = file.getroot() self.muscle_start = len(builder.muscle_activation) for child in root: if child.tag == "Unit": unit_name = child.attrib["name"] unit_f0 = float(child.attrib["f0"]) unit_lm = float(child.attrib["lm"]) unit_lt = float(child.attrib["lt"]) unit_lmax = float(child.attrib["lmax"]) unit_pen = float(child.attrib["pen_angle"]) m = MuscleUnit() m.name = unit_name incomplete = False for waypoint in child.iter("Waypoint"): way_bone = waypoint.attrib["body"] way_link = self.node_map[way_bone] way_loc = np.fromstring(waypoint.attrib["p"], sep=" ", dtype=np.float32) if way_link == -1: incomplete = True break # transform loc to joint local space joint_X_s = self.xform_map[way_bone] way_loc = wp.transform_point(wp.transform_inverse(joint_X_s), way_loc) m.bones.append(way_link) m.points.append(way_loc) if not incomplete: muscles.append(m) builder.add_muscle( m.bones, m.points, f0=unit_f0, lm=unit_lm, lt=unit_lt, lmax=unit_lmax, pen=unit_pen ) self.muscles = muscles def parse_snu(root_xform, skeleton_file, muscle_file, builder, filter, armature=0.0): return Skeleton(root_xform, skeleton_file, muscle_file, builder, filter, armature=0.0)