Spaces:
Sleeping
Sleeping
| # 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) | |