├── .gitignore ├── README.md ├── batch_stl_to_obj.py ├── bone_rotation_exporter.py ├── joint_anim_exporter.py ├── joint_anim_importer.py ├── media ├── g1_run_mocap.gif └── mujoco_importer.png ├── mujoco_importer.py └── xml_import_test.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # blender_mujoco 2 | 3 | **This repository is in prototype mode** 4 | I've only tested it on one [xml file](https://github.com/danieldugas/DeepMimic_mujoco/blob/8ee8f10c4bc1c930d1cbbed2fe5b9a58a73b58bc/src/mujoco/humanoid_deepmimic/envs/asset/deepmimic_unitree_g1.xml) 5 | 6 | ### mujoco_importer 7 | 8 | Import a robot from a MuJoCo `.xml` file into blender. 9 | 10 | ![mujoco_importer](media/mujoco_importer.png) 11 | 12 | - **Blender** > **Scripting** (top tab) > **Open File** (folder icon) > mujoco_importer.py 13 | - **Run Script** (play button) 14 | - Make sure `.xml` file and all necessary meshes are in correct folder structure (in my case, `unitree_g1.xml`, and meshes in `assets` folder) 15 | - **File** > **Import** > **MuJoCo XML** 16 | 17 | ### joint_anim_importer 18 | 19 | Import animation of mujoco joints 20 | 21 | tested on [this mocap file](https://github.com/danieldugas/DeepMimic_mujoco/blob/967f4c255b41afcacb25d694c1ff2016129db2a4/src/mujoco/motions/unitree_g1_run.txt) 22 | 23 | ![joint_anim_importer](media/g1_run_mocap.gif) 24 | 25 | - Make sure the MuJoCo armature has been created 26 | - **Tools** (Right panel, screwdriver/wrench icon) > **Import Animation from JSON** 27 | 28 | ### joint_anim_exporter 29 | 30 | Export animation frames of joint rotations -------------------------------------------------------------------------------- /batch_stl_to_obj.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | 3 | names = [ 4 | "pelvis", 5 | "pelvis_contour_link", 6 | "left_hip_pitch_link", 7 | "left_hip_roll_link", 8 | "left_hip_yaw_link", 9 | "left_knee_link", 10 | "left_ankle_pitch_link", 11 | "left_ankle_roll_link", 12 | "right_hip_pitch_link", 13 | "right_hip_roll_link", 14 | "right_hip_yaw_link", 15 | "right_knee_link", 16 | "right_ankle_pitch_link", 17 | "right_ankle_roll_link", 18 | "torso_link", 19 | "head_link", 20 | "left_shoulder_pitch_link", 21 | "left_shoulder_roll_link", 22 | "left_shoulder_yaw_link", 23 | "left_elbow_pitch_link", 24 | "left_elbow_roll_link", 25 | "right_shoulder_pitch_link", 26 | "right_shoulder_roll_link", 27 | "right_shoulder_yaw_link", 28 | "right_elbow_pitch_link", 29 | "right_elbow_roll_link", 30 | "logo_link", 31 | "left_palm_link", 32 | "left_zero_link", 33 | "left_one_link", 34 | "left_two_link", 35 | "left_three_link", 36 | "left_four_link", 37 | "left_five_link", 38 | "left_six_link", 39 | "right_palm_link", 40 | "right_zero_link", 41 | "right_one_link", 42 | "right_two_link", 43 | "right_three_link", 44 | "right_four_link", 45 | "right_five_link", 46 | "right_six_link", 47 | ] 48 | 49 | # Filepath to your STL file 50 | for name in names: 51 | # log name 52 | print(name) 53 | 54 | stl_file_path = "/home/daniel/Code/DeepMimic_mujoco/src/mujoco/humanoid_deepmimic/envs/asset/assets/"+name+".STL" 55 | 56 | # Import the STL file 57 | #bpy.ops.import_mesh.stl(filepath=stl_file_path) 58 | bpy.ops.wm.stl_import(filepath=stl_file_path) 59 | 60 | # Set the export path for the OBJ file 61 | obj_file_path = "/home/daniel/Code/mujoco_wasm/examples/scenes/assets/"+name+".obj" 62 | 63 | # Select the imported object (ensure that it is the active object) 64 | imported_object = bpy.context.object 65 | bpy.ops.object.select_all(action='DESELECT') 66 | imported_object.select_set(True) 67 | bpy.context.view_layer.objects.active = imported_object 68 | 69 | # Export the selected object as an OBJ file 70 | bpy.ops.wm.obj_export(filepath=obj_file_path, up_axis="Z", forward_axis="Y") 71 | 72 | # delete the selected object 73 | bpy.ops.object.delete() -------------------------------------------------------------------------------- /bone_rotation_exporter.py: -------------------------------------------------------------------------------- 1 | bl_info = { 2 | "name": "Bone Rotation Exporter", 3 | "blender": (3, 5, 0), # Update to the version you are using 4 | "category": "Animation", 5 | } 6 | 7 | import bpy 8 | import mathutils 9 | 10 | 11 | class ExportBoneRotationsOperator(bpy.types.Operator): 12 | bl_idname = "export.bone_rotations" 13 | bl_label = "Export Bone Rotations" 14 | bl_options = {'REGISTER', 'UNDO'} 15 | 16 | def execute(self, context): 17 | return self.export_bone_rotations(context) 18 | 19 | def export_bone_rotations(self, context): 20 | obj = bpy.context.object 21 | if obj and obj.type == 'ARMATURE': 22 | armature = obj.data 23 | anim_data = obj.animation_data 24 | if not anim_data or not anim_data.action: 25 | self.report({'WARNING'}, "No animation found!") 26 | return {'CANCELLED'} 27 | 28 | action = anim_data.action 29 | frame_start = int(action.frame_range[0]) 30 | frame_end = int(action.frame_range[1]) 31 | 32 | # List to store all frames' bone rotations 33 | all_frame_data = [] 34 | 35 | for frame in range(frame_start, frame_end + 1): 36 | bpy.context.scene.frame_set(frame) 37 | frame_data = {} 38 | for bone_name, pose_bone in obj.pose.bones.items(): 39 | # Directly access the Euler rotation from the pose bone 40 | if pose_bone.rotation_mode == 'XYZ': 41 | rotation_euler = pose_bone.rotation_euler 42 | frame_data[bone_name] = rotation_euler.y 43 | self.report({'INFO'}, f"Bone {bone_name} rotation: {rotation_euler.y}") 44 | 45 | all_frame_data.append((frame, frame_data)) 46 | 47 | # Example of saving to a text file 48 | with open("/tmp/exported_data.txt", "w") as file: 49 | for frame, data in all_frame_data: 50 | file.write(f"Frame {frame}:\n") 51 | for bone, rotation in data.items(): 52 | file.write(f" Bone {bone}: {rotation}\n") 53 | 54 | self.report({'INFO'}, "Bone rotations exported successfully!") 55 | return {'FINISHED'} 56 | else: 57 | self.report({'ERROR'}, "Selected object is not an armature!") 58 | return {'CANCELLED'} 59 | 60 | class BoneRotationExporterPanel(bpy.types.Panel): 61 | bl_label = "Bone Rotation Exporter" 62 | bl_idname = "PANEL_PT_bone_rotation_exporter" 63 | bl_space_type = 'VIEW_3D' 64 | bl_region_type = 'UI' 65 | bl_category = 'Tool' 66 | 67 | def draw(self, context): 68 | layout = self.layout 69 | layout.operator("export.bone_rotations") 70 | 71 | def register(): 72 | bpy.utils.register_class(ExportBoneRotationsOperator) 73 | bpy.utils.register_class(BoneRotationExporterPanel) 74 | 75 | def unregister(): 76 | bpy.utils.unregister_class(ExportBoneRotationsOperator) 77 | bpy.utils.unregister_class(BoneRotationExporterPanel) 78 | 79 | if __name__ == "__main__": 80 | register() -------------------------------------------------------------------------------- /joint_anim_exporter.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | import json 3 | from mathutils import Quaternion, Euler, Vector 4 | 5 | bl_info = { 6 | "name": "Export JSON Animation", 7 | "blender": (3, 5, 0), # Adjust to your Blender version 8 | "category": "Animation", 9 | } 10 | 11 | class ExportJSONAnimationOperator(bpy.types.Operator): 12 | bl_idname = "export_animation.json" 13 | bl_label = "Export Animation to JSON" 14 | bl_options = {'REGISTER', 'UNDO'} 15 | 16 | filepath: bpy.props.StringProperty(subtype="FILE_PATH") 17 | 18 | def execute(self, context): 19 | return self.export_animation(context, self.filepath) 20 | 21 | def export_animation(self, context, filepath): 22 | obj = bpy.context.object 23 | if obj and obj.type == 'ARMATURE': 24 | armature = obj 25 | anim_data = obj.animation_data 26 | if not anim_data or not anim_data.action: 27 | self.report({'WARNING'}, "No animation found!") 28 | return {'CANCELLED'} 29 | 30 | action = anim_data.action 31 | frame_start = int(action.frame_range[0]) 32 | frame_end = int(action.frame_range[1]) 33 | 34 | labels = ["armature_pos_x", "armature_pos_y", "armature_pos_z", "armature_quat_w", "armature_quat_x", "armature_quat_y", "armature_quat_z"] 35 | 36 | for bone_name in armature.pose.bones.keys(): 37 | if "joint" in bone_name: 38 | labels.append(bone_name) 39 | 40 | frames = [] 41 | 42 | for frame in range(frame_start, frame_end + 1): 43 | bpy.context.scene.frame_set(frame) 44 | 45 | frame_data = [1.0 / bpy.context.scene.render.fps] # dt = 1/fps 46 | 47 | # Armature position and rotation 48 | pos = armature.location 49 | quat = armature.rotation_quaternion 50 | 51 | frame_data.extend([pos.x, pos.y, pos.z, quat.w, quat.x, quat.y, quat.z]) 52 | 53 | # Bone Y Euler rotation for each joint 54 | for bone_name in armature.pose.bones.keys(): 55 | if "joint" in bone_name: 56 | pose_bone = armature.pose.bones[bone_name] 57 | y_rotation = pose_bone.rotation_euler.y 58 | frame_data.append(y_rotation) 59 | 60 | frames.append(frame_data) 61 | 62 | # Write the JSON data 63 | animation_data = { 64 | "Labels": labels, 65 | "Frames": frames, 66 | "Format": "direct_qpos", 67 | } 68 | 69 | with open(filepath, 'w') as file: 70 | json.dump(animation_data, file, indent=4) 71 | 72 | self.report({'INFO'}, f"Animation exported to {filepath}") 73 | return {'FINISHED'} 74 | else: 75 | self.report({'ERROR'}, "No armature selected!") 76 | return {'CANCELLED'} 77 | 78 | def invoke(self, context, event): 79 | context.window_manager.fileselect_add(self) 80 | return {'RUNNING_MODAL'} 81 | 82 | class ExportJSONAnimationPanel(bpy.types.Panel): 83 | bl_label = "Export JSON Animation" 84 | bl_idname = "PANEL_PT_export_json_animation" 85 | bl_space_type = 'VIEW_3D' 86 | bl_region_type = 'UI' 87 | bl_category = 'Tool' 88 | 89 | def draw(self, context): 90 | layout = self.layout 91 | layout.operator("export_animation.json") 92 | 93 | def register(): 94 | bpy.utils.register_class(ExportJSONAnimationOperator) 95 | bpy.utils.register_class(ExportJSONAnimationPanel) 96 | 97 | def unregister(): 98 | bpy.utils.unregister_class(ExportJSONAnimationOperator) 99 | bpy.utils.unregister_class(ExportJSONAnimationPanel) 100 | 101 | if __name__ == "__main__": 102 | register() -------------------------------------------------------------------------------- /joint_anim_importer.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | import json 3 | from mathutils import Quaternion, Euler, Vector 4 | 5 | bl_info = { 6 | "name": "Import JSON Animation", 7 | "blender": (3, 5, 0), # Adjust to your Blender version 8 | "category": "Animation", 9 | } 10 | 11 | class ImportJSONAnimationOperator(bpy.types.Operator): 12 | bl_idname = "import_animation.json" 13 | bl_label = "Import Animation from JSON" 14 | bl_options = {'REGISTER', 'UNDO'} 15 | 16 | filepath: bpy.props.StringProperty(subtype="FILE_PATH") 17 | 18 | def execute(self, context): 19 | return self.import_animation(context, self.filepath) 20 | 21 | def import_animation(self, context, filepath): 22 | # Load the JSON file 23 | with open(filepath, 'r') as file: 24 | data = json.load(file) 25 | 26 | frames = data.get("Frames", []) 27 | labels = data.get("Labels", []) 28 | 29 | obj = bpy.context.object 30 | if obj and obj.type == 'ARMATURE': 31 | # Prepare the armature for animation 32 | armature = obj 33 | bpy.context.view_layer.objects.active = armature 34 | bpy.ops.object.mode_set(mode='POSE') 35 | 36 | for frame_index, frame_data in enumerate(frames): 37 | bpy.context.scene.frame_set(frame_index + 1) # Frames are 1-based in Blender 38 | 39 | dt = frame_data[0] 40 | frame_qpos = frame_data[1:] 41 | # Extract armature position and rotation (first 7 values) 42 | pos = Vector(frame_qpos[:3]) 43 | quat = Quaternion(frame_qpos[3:7]) 44 | 45 | armature.location = pos 46 | armature.rotation_mode = 'QUATERNION' 47 | armature.rotation_quaternion = quat 48 | 49 | armature.keyframe_insert(data_path="location", frame=frame_index + 1) 50 | armature.keyframe_insert(data_path="rotation_quaternion", frame=frame_index + 1) 51 | 52 | # Set each bone's Y rotation based on the labels 53 | for label_index, label in enumerate(labels): 54 | if "joint" not in label: 55 | continue 56 | bone_name = label 57 | if bone_name in armature.pose.bones: 58 | pose_bone = armature.pose.bones[bone_name] 59 | y_rotation = frame_data[label_index] 60 | 61 | pose_bone.rotation_mode = 'XYZ' 62 | pose_bone.rotation_euler = Euler((0, y_rotation, 0), 'XYZ') 63 | 64 | pose_bone.keyframe_insert(data_path="rotation_euler", frame=frame_index + 1) 65 | 66 | else: 67 | self.report({'ERROR'}, "No armature selected!") 68 | return {'CANCELLED'} 69 | 70 | self.report({'INFO'}, "Animation imported successfully!") 71 | return {'FINISHED'} 72 | 73 | def invoke(self, context, event): 74 | context.window_manager.fileselect_add(self) 75 | return {'RUNNING_MODAL'} 76 | 77 | class ImportJSONAnimationPanel(bpy.types.Panel): 78 | bl_label = "Import JSON Animation" 79 | bl_idname = "PANEL_PT_import_json_animation" 80 | bl_space_type = 'VIEW_3D' 81 | bl_region_type = 'UI' 82 | bl_category = 'Tool' 83 | 84 | def draw(self, context): 85 | layout = self.layout 86 | layout.operator("import_animation.json") 87 | 88 | def register(): 89 | bpy.utils.register_class(ImportJSONAnimationOperator) 90 | bpy.utils.register_class(ImportJSONAnimationPanel) 91 | 92 | def unregister(): 93 | bpy.utils.unregister_class(ImportJSONAnimationOperator) 94 | bpy.utils.unregister_class(ImportJSONAnimationPanel) 95 | 96 | if __name__ == "__main__": 97 | register() -------------------------------------------------------------------------------- /media/g1_run_mocap.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danieldugas/blender_mujoco/494a958a0cf0d5b51de0c78f1d13c3f220269d71/media/g1_run_mocap.gif -------------------------------------------------------------------------------- /media/mujoco_importer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danieldugas/blender_mujoco/494a958a0cf0d5b51de0c78f1d13c3f220269d71/media/mujoco_importer.png -------------------------------------------------------------------------------- /mujoco_importer.py: -------------------------------------------------------------------------------- 1 | bl_info = { 2 | "name": "MuJoCo XML Importer", 3 | "blender": (3, 0, 0), 4 | "category": "Import-Export", 5 | } 6 | 7 | BLENDER = True 8 | 9 | try: 10 | import bpy 11 | from bpy_extras.io_utils import ImportHelper 12 | except: 13 | BLENDER = False 14 | 15 | import xml.etree.ElementTree as ET 16 | import numpy as np 17 | import os 18 | 19 | if True: # transforms.py 20 | def normalized(vec): 21 | norm = np.linalg.norm(vec) 22 | if norm == 0: 23 | print("Warning: attempting to normalize zero vector") 24 | return vec 25 | return vec / norm 26 | 27 | class Quaternion(object): 28 | def __init__(self, x, y, z, w): 29 | self._x = x 30 | self._y = y 31 | self._z = z 32 | self._w = w 33 | 34 | def __repr__(self): 35 | return "Quaternion(x={:.9f}, y={:.9f}, z={:.9f}, w={:.9f})".format(*self.xyzw()) 36 | 37 | def as_quaternion(something): 38 | if isinstance(something, Quaternion): 39 | return something 40 | elif isinstance(something, tuple) and len(something) == 4: 41 | return Quaternion(*something) 42 | elif isinstance(something, list) and len(something) == 4: 43 | return Quaternion(*something) 44 | else: 45 | raise ValueError("Cannot convert {} to Quaternion".format(something)) 46 | 47 | def xyzw(self): 48 | return [self._x, self._y, self._z, self._w] 49 | 50 | def from_transform_matrix(self, matrix4x4): 51 | x, y, z, w = quaternion_from_transform_matrix(matrix4x4) 52 | return Quaternion(x, y, z, w) 53 | 54 | def to_transform_matrix(self): 55 | return transform_matrix_from_quaternion(*self.xyzw()) 56 | 57 | def to_rpy(self): 58 | """ returns roll, pitch, yaw in radians """ 59 | import math 60 | q0 = self._w 61 | q1 = self._x 62 | q2 = self._y 63 | q3 = self._z 64 | roll = math.atan2( 65 | 2 * ((q2 * q3) + (q0 * q1)), 66 | q0**2 - q1**2 - q2**2 + q3**2 67 | ) # radians 68 | pitch = math.asin(2 * ((q1 * q3) - (q0 * q2))) 69 | yaw = math.atan2( 70 | 2 * ((q1 * q2) + (q0 * q3)), 71 | q0**2 + q1**2 - q2**2 - q3**2 72 | ) 73 | return (roll, pitch, yaw) 74 | 75 | class Transform(object): 76 | def __init__(self, origin=None, x_axis=None, y_axis=None, quaternion=None): 77 | self._matrix = None 78 | self._quaternion = None 79 | if x_axis is not None or y_axis is not None: 80 | if x_axis is None or y_axis is None: 81 | raise ValueError("Underspecified transform: Must specify neither or both axes.") 82 | if origin is not None or x_axis is not None or y_axis is not None: 83 | self._matrix = np.eye(4) 84 | if x_axis is not None and y_axis is not None: 85 | if quaternion is not None: 86 | raise ValueError("Overspecified transform: Cannot specify both quaternion and axes") 87 | self._matrix[:3, 0] = x_axis 88 | self._matrix[:3, 1] = y_axis 89 | self._matrix[:3, 2] = np.cross(x_axis, y_axis) 90 | if origin is not None: 91 | self._matrix[:3, 3] = origin 92 | if quaternion is not None: 93 | self._quaternion = Quaternion.as_quaternion(quaternion) 94 | 95 | def to_json_dict(self): 96 | json_dict = { 97 | "type": "Transform", 98 | "_matrix": self._matrix.tolist() if self._matrix is not None else None, 99 | "_quaternion": self._quaternion.xyzw() if self._quaternion is not None else None, 100 | } 101 | return json_dict 102 | 103 | def from_json_dict(json_dict): 104 | new = Transform() 105 | new._matrix = np.array(json_dict["_matrix"]) if json_dict["_matrix"] is not None else None 106 | new._quaternion = Quaternion(*json_dict["_quaternion"]) if json_dict["_quaternion"] is not None else None 107 | return new 108 | 109 | def __repr__(self): 110 | return "Transform(origin={}, x_axis={}, y_axis={})".format(self.origin(), self.x_axis(), self.y_axis()) 111 | 112 | def __mul__(self, other): 113 | """ T_B_in_C * T_A_in_B = T_A_in_C """ 114 | if isinstance(other, Transform): 115 | return Transform.from_matrix(np.dot(self.matrix(), other.matrix())) 116 | else: 117 | raise NotImplementedError 118 | 119 | def is_right_handed(self): 120 | return np.linalg.det(self.matrix()[:3, :3]) > 0 121 | 122 | def from_matrix(matrix4x4): 123 | new = Transform() 124 | new._matrix = matrix4x4 125 | return new 126 | 127 | def from_quaternion(quaternion, origin=None): 128 | return Transform(origin=origin, quaternion=quaternion) 129 | 130 | def from_axis_angle(axis, angle_rad, translation=None): 131 | """ translation is not the axis origin! To implement rotation around an origin, use from_rotation_around_point """ 132 | new = Transform() 133 | new._matrix = transform_matrix_from_axis_angle(axis, angle_rad, translation) 134 | return new 135 | 136 | def from_rotation_around_point(axis, angle_rad, point): 137 | new = Transform() 138 | new._matrix = transform_matrix_from_axis_angle(axis, angle_rad) 139 | # to get translation, rotate point around axis at origin, compare to previous 140 | rotated_point = new.transform_points([point])[0] 141 | new._matrix[:3, 3] = point - rotated_point 142 | return new 143 | 144 | def inverse(self): 145 | return Transform.from_matrix(inverse(self.matrix())) 146 | 147 | def matrix(self): 148 | if self._quaternion is not None: 149 | rot_mat4x4 = self.quaternion().to_transform_matrix() 150 | if self._matrix is not None: 151 | if not np.allclose(self._matrix[:3, :3], np.eye(3)): 152 | raise ValueError("Overdefined transform: transform has a non-zero quaternion and non-zero rotation matrix.") 153 | rot_mat4x4[:3, 3] = self._matrix[:3, 3] 154 | return rot_mat4x4 155 | elif self._matrix is not None: 156 | return self._matrix 157 | else: 158 | return np.eye(4) 159 | 160 | def quaternion(self): 161 | if self._quaternion is not None: 162 | return self._quaternion 163 | elif self._matrix is not None: 164 | return Quaternion(*quaternion_from_transform_matrix(self.matrix())) 165 | else: 166 | return Quaternion(0, 0, 0, 1) 167 | 168 | def origin(self): 169 | return self.matrix()[:3, 3] 170 | 171 | def translation(self): 172 | return self.origin() 173 | 174 | def x_axis(self): 175 | return self.matrix()[:3, 0] 176 | 177 | def y_axis(self): 178 | return self.matrix()[:3, 1] 179 | 180 | def z_axis(self): 181 | return self.matrix()[:3, 2] 182 | 183 | def to_axis_angle(self): 184 | return axis_angle_from_transform_matrix(self.matrix()) 185 | 186 | def to_compas_frame(self): 187 | from compas.geometry import Frame 188 | return Frame(self.origin(), self.x_axis(), self.y_axis()) 189 | 190 | def from_compas_frame(frame): 191 | return Transform([frame.point.x, frame.point.y, frame.point.z], frame.xaxis, frame.yaxis) 192 | 193 | def to_pose_msg(self): 194 | from geometry_msgs.msg import Pose 195 | pose = Pose() 196 | x, y, z = self.origin() 197 | pose.position.x = x 198 | pose.position.y = y 199 | pose.position.z = z 200 | qx, qy, qz, qw = self.quaternion().xyzw() 201 | pose.orientation.x = qx 202 | pose.orientation.y = qy 203 | pose.orientation.z = qz 204 | pose.orientation.w = qw 205 | return pose 206 | 207 | def from_pose_msg(pose_msg): 208 | x = pose_msg.position.x 209 | y = pose_msg.position.y 210 | z = pose_msg.position.z 211 | qx = pose_msg.orientation.x 212 | qy = pose_msg.orientation.y 213 | qz = pose_msg.orientation.z 214 | qw = pose_msg.orientation.w 215 | return Transform.from_quaternion(Quaternion(qx, qy, qz, qw), [x, y, z]) 216 | 217 | def transform_vector(self, vector_in_A_frame): 218 | """ If this transform is the transform of A in B, then this returns the vector in B frame """ 219 | return transform_vector(vector_in_A_frame, self.matrix()) 220 | 221 | def transform_point(self, point): 222 | return transform_point(point, self.matrix()) 223 | 224 | def transform_points(self, points): 225 | return transform_points(points, self.matrix()) 226 | 227 | def plot_polyscope(self, name="A_in_B", axis_length=0.1): 228 | show_frame_in_polyscope(self.matrix(), name=name, axis_length=axis_length) 229 | 230 | def print_matrix(self): 231 | # 1 decimal 232 | for row in self.matrix(): 233 | for val in row: 234 | print("{:.1f}".format(val), end=" ") 235 | print() 236 | 237 | 238 | def inverse(transform_matrix_A_in_B): 239 | transform_matrix_B_in_A = np.linalg.inv(transform_matrix_A_in_B) 240 | return transform_matrix_B_in_A 241 | 242 | 243 | def show_frame_in_polyscope(frame_in_world_matrix, name="frame", axis_length=1.0): 244 | import polyscope as ps 245 | origin = frame_in_world_matrix[:3, 3] 246 | x_axis = frame_in_world_matrix[:3, 0] 247 | y_axis = frame_in_world_matrix[:3, 1] 248 | z_axis = frame_in_world_matrix[:3, 2] 249 | ps.register_curve_network( 250 | "{}_x_axis".format(name), 251 | np.array([origin, origin + x_axis * axis_length]), 252 | np.array([[0, 1]]), 253 | color=(1.0, 0.0, 0.0), 254 | ) 255 | ps.register_curve_network( 256 | "{}_y_axis".format(name), 257 | np.array([origin, origin + y_axis * axis_length]), 258 | np.array([[0, 1]]), 259 | color=(0.0, 1.0, 0.0), 260 | ) 261 | ps.register_curve_network( 262 | "{}_z_axis".format(name), 263 | np.array([origin, origin + z_axis * axis_length]), 264 | np.array([[0, 1]]), 265 | color=(0.0, 0.0, 1.0), 266 | ) 267 | 268 | 269 | def transform_matrix_from_origin_and_xy_axes(origin, x_axis, y_axis): 270 | """ 271 | returns the matrix for the transform A in B, where origin is the origin of A in B, and x_axis and y_axis are the x and y axes of A in B 272 | """ 273 | z_axis = np.cross(x_axis, y_axis) 274 | xx, xy, xz = x_axis 275 | yx, yy, yz = y_axis 276 | zx, zy, zz = z_axis 277 | ox, oy, oz = origin 278 | transform_matrix = np.array( 279 | [ 280 | [xx, yx, zx, ox], 281 | [xy, yy, zy, oy], 282 | [xz, yz, zz, oz], 283 | [0, 0, 0, 1], 284 | ] 285 | ) 286 | return transform_matrix 287 | 288 | def transform_matrix_from_translation(translation): 289 | transform_matrix = transform_matrix_from_origin_and_xy_axes(translation, [1, 0, 0], [0, 1, 0]) 290 | return transform_matrix 291 | 292 | def transform_matrix_from_axis_angle(axis, angle, translation=None): 293 | """ angle in radians """ 294 | ux, uy, uz = axis 295 | tx, ty, tz = translation if translation is not None else [0, 0, 0] 296 | transform_matrix = np.array( 297 | [ 298 | [ux * ux * (1 - np.cos(angle)) + np.cos(angle), ux * uy * (1 - np.cos(angle)) - uz * np.sin(angle), ux * uz * (1 - np.cos(angle)) + uy * np.sin(angle), tx], 299 | [uy * ux * (1 - np.cos(angle)) + uz * np.sin(angle), uy * uy * (1 - np.cos(angle)) + np.cos(angle), uy * uz * (1 - np.cos(angle)) - ux * np.sin(angle), ty], 300 | [uz * ux * (1 - np.cos(angle)) - uy * np.sin(angle), uz * uy * (1 - np.cos(angle)) + ux * np.sin(angle), uz * uz * (1 - np.cos(angle)) + np.cos(angle), tz], 301 | [0, 0, 0, 1], 302 | ] 303 | ) 304 | return transform_matrix 305 | 306 | def axis_angle_from_transform_matrix(transform_matrix): 307 | """ returns axis, angle """ 308 | R = transform_matrix[:3, :3] 309 | angle = np.arccos((np.trace(R) - 1) / 2) 310 | if angle == 0: 311 | return None, 0 312 | if np.allclose(angle, np.pi): 313 | pass 314 | axis = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]]) / (2 * np.sin(angle)) 315 | norm = np.linalg.norm(axis) 316 | if norm != 0: 317 | axis = axis / norm 318 | return axis, angle 319 | 320 | def axis_angle_from_transform_matrix(transform_matrix): 321 | """ returns axis, angle 322 | https://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToAngle/ 323 | """ 324 | R = transform_matrix[:3, :3] 325 | epsilon = 0.0001 # margin to allow for rounding errors 326 | epsilon2 = 0.1 # margin to distinguish between 0 and 180 degrees 327 | if ((np.abs(R[0, 1]-R[1, 0])< epsilon) and (np.abs(R[0, 2]-R[2, 0])< epsilon) and (np.abs(R[1, 2]-R[2, 1])< epsilon)) : 328 | # singularity found 329 | # first check for identity matrix which must have +1 for all terms 330 | # in leading diagonaland zero in other terms 331 | if ((np.abs(R[0, 1]+R[1, 0]) < epsilon2) and (np.abs(R[0, 2]+R[2, 0]) < epsilon2) and (np.abs(R[1, 2]+R[2, 1]) < epsilon2) and (np.abs(R[0, 0]+R[1, 1]+R[2, 2]-3) < epsilon2)) : 332 | # this singularity is identity matrix so angle = 0 333 | return None, 0 # zero angle, arbitrary axis 334 | 335 | # otherwise this singularity is angle = 180 336 | angle = np.pi 337 | xx = (R[0, 0]+1)/2 338 | yy = (R[1, 1]+1)/2 339 | zz = (R[2, 2]+1)/2 340 | xy = (R[0, 1]+R[1, 0])/4 341 | xz = (R[0, 2]+R[2, 0])/4 342 | yz = (R[1, 2]+R[2, 1])/4 343 | if ((xx > yy) and (xx > zz)) : # R[0, 0] is the largest diagonal term 344 | if (xx< epsilon) : 345 | x = 0 346 | y = 0.7071 347 | z = 0.7071 348 | else : 349 | x = np.sqrt(xx) 350 | y = xy/x 351 | z = xz/x 352 | 353 | elif (yy > zz) : # R[1, 1] is the largest diagonal term 354 | if (yy< epsilon) : 355 | x = 0.7071 356 | y = 0 357 | z = 0.7071 358 | else : 359 | y = np.sqrt(yy) 360 | x = xy/y 361 | z = yz/y 362 | 363 | else : # R[2, 2] is the largest diagonal term so base result on this 364 | if (zz< epsilon) : 365 | x = 0.7071 366 | y = 0.7071 367 | z = 0 368 | else : 369 | z = np.sqrt(zz) 370 | x = xz/z 371 | y = yz/z 372 | 373 | return np.array([x, y, z]), angle # return 180 deg rotation 374 | 375 | # as we have reached here there are no singularities so we can handle normally 376 | s = np.sqrt((R[2, 1] - R[1, 2])*(R[2, 1] - R[1, 2]) 377 | +(R[0, 2] - R[2, 0])*(R[0, 2] - R[2, 0]) 378 | +(R[1, 0] - R[0, 1])*(R[1, 0] - R[0, 1])) # used to normalise 379 | if (np.abs(s) < 0.001): 380 | s=1 381 | # prevent divide by zero, should not happen if matrix is orthogonal and should be 382 | # caught by singularity test above, but I've left it in just in case 383 | angle = np.arccos(( R[0, 0] + R[1, 1] + R[2, 2] - 1)/2) 384 | x = (R[2, 1] - R[1, 2])/s 385 | y = (R[0, 2] - R[2, 0])/s 386 | z = (R[1, 0] - R[0, 1])/s 387 | return np.array([x, y, z]), angle 388 | 389 | def transform_matrix_from_quaternion(x, y, z, w): 390 | transform_matrix = np.array([ 391 | [1 - 2 * y * y - 2 * z * z, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w, 0], 392 | [2 * x * y + 2 * z * w, 1 - 2 * x * x - 2 * z * z, 2 * y * z - 2 * x * w, 0], 393 | [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x * x - 2 * y * y, 0], 394 | [0, 0, 0, 1], 395 | ]) 396 | return transform_matrix 397 | 398 | def quaternion_from_transform_matrix(transform_matrix): 399 | """ Returns x, y, z, w components of the quaternion defined by the upper left part of the 4x4 transform_matrix """ 400 | q = np.empty((4, ), dtype=np.float64) 401 | M = np.array(transform_matrix, dtype=np.float64, copy=False)[:4, :4] 402 | t = np.trace(M) 403 | if t > M[3, 3]: 404 | q[3] = t 405 | q[2] = M[1, 0] - M[0, 1] 406 | q[1] = M[0, 2] - M[2, 0] 407 | q[0] = M[2, 1] - M[1, 2] 408 | else: 409 | i, j, k = 0, 1, 2 410 | if M[1, 1] > M[0, 0]: 411 | i, j, k = 1, 2, 0 412 | if M[2, 2] > M[i, i]: 413 | i, j, k = 2, 0, 1 414 | t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3] 415 | q[i] = t 416 | q[j] = M[i, j] + M[j, i] 417 | q[k] = M[k, i] + M[i, k] 418 | q[3] = M[k, j] - M[j, k] 419 | q *= 0.5 / np.sqrt(t * M[3, 3]) 420 | x, y, z, w = q 421 | return x, y, z, w 422 | 423 | def transform_points(points_in_A_frame, transform_matrix_A_in_B): 424 | points_in_A_frame = np.asanyarray(points_in_A_frame).reshape((-1, 3)) 425 | transform_matrix_A_in_B = np.asanyarray(transform_matrix_A_in_B) 426 | _N, _3 = points_in_A_frame.shape 427 | _4, _4 = transform_matrix_A_in_B.shape 428 | # add a dimension to points_in_A_frame 429 | points_in_A_frame = np.hstack([points_in_A_frame, np.ones((_N, 1))]) 430 | # transform_matrix_A_in_B 431 | points_in_B_frame = np.dot(points_in_A_frame, transform_matrix_A_in_B.T) 432 | # remove the dimension 433 | points_in_B_frame = points_in_B_frame[:, :-1] 434 | return points_in_B_frame 435 | 436 | def transform_point(point_in_A_frame, transform_matrix_A_in_B): 437 | return transform_points(np.asanyarray(point_in_A_frame).reshape((1, 3)), transform_matrix_A_in_B).reshape((3,)) 438 | 439 | def transform_vectors(vectors_in_A_frame, transform_matrix_A_in_B): 440 | rotation_matrix_A_in_B = np.zeros_like(transform_matrix_A_in_B) 441 | rotation_matrix_A_in_B[:3, :3] = transform_matrix_A_in_B[:3, :3] 442 | return transform_points(vectors_in_A_frame, rotation_matrix_A_in_B) 443 | 444 | def transform_vector(vector_in_A_frame, transform_matrix_A_in_B): 445 | return transform_vectors(vector_in_A_frame.reshape((1, 3)), transform_matrix_A_in_B).reshape((3,)) 446 | 447 | def transform_matrix_from_frame(frame): 448 | print("Warning: transform_matrix_from_frame is deprecated. Use transform_matrix_from_compas_frame instead.") 449 | return transform_matrix_from_compas_frame(frame) 450 | 451 | def transform_matrix_from_compas_frame(frame): 452 | from compas.geometry.transformations import Transformation 453 | return np.array(Transformation.from_frame(frame).matrix) 454 | 455 | def rotate_points_around_axis(points, axis_origin, axis, angle): 456 | translation_matrix = transform_matrix_from_translation(-np.array(axis_origin)) 457 | rotation_matrix = transform_matrix_from_axis_angle(axis, angle) 458 | reverse_translation_matrix = transform_matrix_from_translation(np.array(axis_origin)) 459 | return transform_points(points, reverse_translation_matrix @ rotation_matrix @ translation_matrix) 460 | 461 | def point_2d_to_3d(point_2d, z=0): 462 | return np.array([point_2d[0], point_2d[1], z]) 463 | 464 | def points_2d_to_3d(points_2d, z=0): 465 | _N, _2 = points_2d.shape 466 | if _2 != 2: 467 | raise ValueError("points_2d must be Nx2") 468 | return np.hstack([np.asanyarray(points_2d), np.ones((len(points_2d), 1)) * z]) 469 | 470 | class BoneInfo: 471 | def __init__(self, name, parent_bone_name, head, tail, is_joint): 472 | self._name = name 473 | self._parent_bone_name = parent_bone_name 474 | self._head_xyz = head 475 | self._tail_xyz = tail 476 | if np.allclose(head, tail): 477 | raise ValueError("Bone {} has zero length. Blender will refuse to create it.".format(name)) 478 | self._is_joint = is_joint 479 | 480 | def name(self): 481 | return self._name 482 | 483 | def parent_bone_name(self): 484 | return self._parent_bone_name 485 | 486 | def head(self): 487 | return self._head_xyz 488 | 489 | def tail(self): 490 | return self._tail_xyz 491 | 492 | def is_joint(self): 493 | return self._is_joint 494 | 495 | class JointInfo: 496 | JBL = 0.05 # joint bone length 497 | JBP = "POST" # joint bone is either PRE (before the body pose) or POST (after the body pose) 498 | ADD_FREEJOINT_BONE = False 499 | def __init__(self, name, body, body_parent, axis): 500 | self._name = name 501 | self._body = body 502 | self._body_parent = body_parent 503 | self._axis = axis 504 | 505 | def __repr__(self): 506 | return "JointInfo(name={})".format(self._name) 507 | 508 | def get_name(self): 509 | return self._name 510 | 511 | def get_initial_pos(self, in_world=True): 512 | if in_world: 513 | return self._body.get_initial_tf_in_world().origin() 514 | else: 515 | return self._body.get_initial_tf_in_armature().origin() 516 | 517 | def get_initial_axis(self, in_world=True): 518 | if in_world: 519 | return self._body.get_initial_tf_in_world().transform_vector(self._axis) 520 | else: 521 | return self._body.get_initial_tf_in_armature().transform_vector(self._axis) 522 | 523 | def get_initial_bonehead(self): 524 | in_world = JointInfo.ADD_FREEJOINT_BONE 525 | o = self.get_initial_pos(in_world) 526 | v = self.get_initial_axis(in_world) 527 | if JointInfo.JBP == "PRE": 528 | return o - JointInfo.JBL * v 529 | else: 530 | return o 531 | 532 | def get_initial_bonetail(self): 533 | in_world = JointInfo.ADD_FREEJOINT_BONE 534 | o = self.get_initial_pos(in_world) 535 | v = self.get_initial_axis(in_world) 536 | if JointInfo.JBP == "PRE": 537 | return o 538 | else: 539 | return o + JointInfo.JBL * v 540 | 541 | class MeshGeomInfo: 542 | def __init__(self, stl_path): 543 | self._stl_path = stl_path 544 | 545 | def get_stl_path(self): 546 | return self._stl_path 547 | 548 | class BodyInfo: 549 | def __init__(self, name, parent, body_in_parent_tf): 550 | self._name = name 551 | self._parent = parent 552 | self._body_in_parent_tf = body_in_parent_tf 553 | self._geoms = [] 554 | # infer initial world pose 555 | tf_parent_in_world = Transform() 556 | if self._parent is not None: 557 | tf_parent_in_world = self._parent.get_initial_tf_in_world() 558 | body_in_world = tf_parent_in_world * self._body_in_parent_tf 559 | self._initial_tf_in_world = body_in_world 560 | # infer initial pose in armature (root body) 561 | tf_parent_in_armature = Transform() 562 | if self._parent is not None: 563 | tf_parent_in_armature = self._parent.get_initial_tf_in_armature() 564 | body_in_armature = tf_parent_in_armature * self._body_in_parent_tf 565 | if self._parent is not None and self._parent.get_parent() is None: 566 | body_in_armature = Transform() # root bodies at armature origin 567 | self._initial_tf_in_armature = body_in_armature 568 | # infer depth 569 | self._depth = 0 570 | if self._parent is not None: 571 | self._depth = self._parent.get_depth() + 1 572 | # joint 573 | self._joint = None 574 | # end bone 575 | self._end_bone_name = None # not initialized 576 | self._end_bone_tail = None # not initialized 577 | 578 | def set_joint(self, joint_info): 579 | self._joint = joint_info 580 | 581 | def add_mesh_geom(self, mesh_geom_info): 582 | self._geoms.append(mesh_geom_info) 583 | 584 | def __repr__(self): 585 | return "BodyInfo(name={}, parent_name={})".format(self._name, self._parent.get_name() if self._parent is not None else "None") 586 | 587 | def get_name(self): 588 | return self._name 589 | 590 | def get_parent(self): 591 | return self._parent 592 | 593 | def get_initial_tf_in_world(self): 594 | return self._initial_tf_in_world 595 | 596 | def get_initial_tf_in_armature(self): 597 | return self._initial_tf_in_armature 598 | 599 | def get_depth(self): 600 | return self._depth 601 | 602 | def get_mesh_geoms(self): 603 | return self._geoms 604 | 605 | def create_bones(self): 606 | # bone going from parent joint-end to this body joint-start 607 | # bone going from this body joint-start to joint-end 608 | """ 609 | (PRE joint bone positioning) 610 | (no or free joint) 611 | body0 612 | +---------------x joint1 bone head 613 | / v 614 | / + body1 (axis joint) 615 | / 616 | + 617 | worldbody 618 | 619 | """ 620 | own_origin = self.get_initial_tf_in_world().origin() if JointInfo.ADD_FREEJOINT_BONE else self.get_initial_tf_in_armature().origin() 621 | if self.get_parent() is None: # no bones for worldbody 622 | self._end_bone_name = None 623 | self._end_bone_tail = own_origin 624 | return [] 625 | if not JointInfo.ADD_FREEJOINT_BONE: # parent is worldbody, create a small bone 626 | if self.get_parent().get_parent() is None: 627 | self._end_bone_name = "to_" + self.get_name() 628 | self._end_bone_tail = own_origin 629 | bone_head = own_origin - np.array([JointInfo.JBL, 0, 0]) 630 | parent_to_body_bone = BoneInfo(self._end_bone_name, None, bone_head, self._end_bone_tail, False) 631 | return [parent_to_body_bone] 632 | # check init order 633 | if self.get_parent()._end_bone_tail is None: 634 | raise ValueError("Child body's bones should be created after parent body's bones") 635 | # parent end bone 636 | parent_end_bone_name = self.get_parent()._end_bone_name 637 | parent_end_bone_tail = self.get_parent()._end_bone_tail 638 | # create bones 639 | if self._joint is None: # single bone, from parent to this body 640 | self._end_bone_name = "to_" + self.get_name() 641 | self._end_bone_tail = own_origin 642 | parent_to_body_bone = BoneInfo(self._end_bone_name, parent_end_bone_name, parent_end_bone_tail, self._end_bone_tail, False) 643 | return [parent_to_body_bone] 644 | else: # two bones, from parent to joint, from joint to this body 645 | bone1_name = "to_" + self.get_name() 646 | jbone_name = self._joint.get_name() 647 | jbone_head = self._joint.get_initial_bonehead() 648 | jbone_tail = self._joint.get_initial_bonetail() 649 | if np.allclose(parent_end_bone_tail, jbone_head): 650 | # This body has no translation/rot from parent. 651 | # Blender doesn't allow 0 length bones so we must skip it 652 | joint_bone = BoneInfo(jbone_name, parent_end_bone_name, parent_end_bone_tail, jbone_tail, True) 653 | self._end_bone_name = jbone_name 654 | self._end_bone_tail = jbone_tail 655 | return [joint_bone] 656 | parent_to_joint_bone = BoneInfo(bone1_name, parent_end_bone_name, parent_end_bone_tail, jbone_head, False) 657 | joint_bone = BoneInfo(jbone_name, bone1_name, jbone_head, jbone_tail, True) 658 | self._end_bone_name = jbone_name 659 | self._end_bone_tail = jbone_tail 660 | return [parent_to_joint_bone, joint_bone] 661 | 662 | class Scene: 663 | def __init__(self): 664 | pass 665 | 666 | def set_bodies(self, bodies): 667 | self.bodies = bodies 668 | 669 | def get_bodies_in_world_tfs(self): 670 | return [(b.get_name(), b.get_initial_tf_in_world()) for b in self.bodies] 671 | 672 | def get_all_bodies(self): 673 | return self.bodies 674 | 675 | def set_joints(self, joints): 676 | self.joints = joints 677 | 678 | def get_joints_in_world(self): 679 | return [(j.get_name(), j.get_initial_pos(), j.get_initial_axis()) for j in self.joints] 680 | 681 | def create_all_bones(self): 682 | # each body has one or more bones: 683 | # - a bone going from parent body's joint bone tail to this body's pos (a kind of "remainder bone") 684 | # - a bone for each joint, at the joint pos, with bone y == joint axis 685 | all_bones = [] 686 | for body in self.bodies: 687 | all_bones += body.create_bones() 688 | return all_bones 689 | 690 | def get_freejoint_pose(self): 691 | for body in self.bodies: 692 | if body.get_parent() is not None and body.get_parent().get_parent() is None: 693 | return body.get_initial_tf_in_world() 694 | 695 | def parse_mujoco_xml(filepath, blenderclass): 696 | tree = ET.parse(filepath) 697 | root = tree.getroot() 698 | 699 | # find all mesh files 700 | stl_dir = "" 701 | for c in root.findall("compiler"): 702 | stl_dir = c.get("meshdir", "") 703 | stl_dir_full = os.path.join(os.path.dirname(filepath), stl_dir) 704 | mesh_files = {} 705 | for a in root.findall("asset"): 706 | for m in a.findall("mesh"): 707 | filename = m.get("file", None) 708 | if filename is not None: 709 | mesh_files[os.path.splitext(os.path.split(filename)[-1])[0]] = os.path.join(stl_dir_full, filename) 710 | 711 | joints = [] 712 | bodies = [] 713 | bones = [] 714 | 715 | def process_body(body, parent_body): 716 | 717 | body_name = body.get("name", "unnamed_body") 718 | body_pos = np.array([float(x) for x in body.get("pos", "0 0 0").split()]) 719 | body_wxyz = np.array([float(x) for x in body.get("quat", "1 0 0 0").split()]) 720 | qw, qx, qy, qz = body_wxyz 721 | body_in_parent = Transform(origin=body_pos, quaternion=Quaternion(qx, qy, qz, qw)) 722 | body_info = BodyInfo(body_name, parent_body, body_in_parent) 723 | bodies.append(body_info) 724 | depth = body_info.get_depth() 725 | if BLENDER: 726 | # print in blender console 727 | blenderclass.report({'INFO'}, "-" * depth + body_name) 728 | else: 729 | print("-" * depth + body_name) 730 | 731 | n_joints = len(body.findall("joint")) 732 | if n_joints > 1: 733 | raise ValueError("Body {} has more than one joint. This is not supported as it doesn't translate well to bone hierarchies".format(body_name)) 734 | for joint in body.findall("joint"): 735 | joint_name = joint.get("name", "unnamed_joint") 736 | joint_axis = np.array([float(x) for x in joint.get("axis", "0 0 1").split()]) 737 | joint_info = JointInfo(joint_name, body_info, parent_body, joint_axis) 738 | body_info.set_joint(joint_info) 739 | joints.append(joint_info) 740 | 741 | # add geometries (meshes only for now) 742 | for geom in body.findall("geom"): 743 | geom_type = geom.get("type", "unknown") 744 | geom_class = geom.get("class", "unknown") 745 | if (geom_type == "mesh" or geom_type == "unknown") and (geom_class == "visual" or geom_class == "visual_light" or geom_class == "visual_dark"): 746 | geom_meshfile = geom.get("mesh", None) 747 | if geom_meshfile is not None: 748 | body_info.add_mesh_geom(MeshGeomInfo(mesh_files[os.path.splitext(geom_meshfile)[0]])) 749 | 750 | 751 | for child_body in body.findall("body"): 752 | process_body(child_body, body_info) 753 | 754 | for worldbody in root.findall("worldbody"): 755 | print("Worldbody:", worldbody.get("name", "worldbody")) 756 | process_body(worldbody, None) 757 | 758 | scene = Scene() 759 | scene.set_bodies(bodies) 760 | scene.set_joints(joints) 761 | 762 | return scene 763 | 764 | 765 | if BLENDER: 766 | class MUJOCO_OT_import(bpy.types.Operator, ImportHelper): 767 | """Import a MuJoCo XML File""" 768 | bl_idname = "import_scene.mujoco_xml" 769 | bl_label = "Import MuJoCo XML" 770 | filename_ext = ".xml" 771 | 772 | def execute(self, context): 773 | # Parse the XML file 774 | scene = parse_mujoco_xml(self.filepath, self) 775 | 776 | # Create armature, bone for each joint 777 | bpy.ops.object.armature_add(enter_editmode=True) 778 | armature = bpy.context.object 779 | armature.name = "MuJoCo_Armature" 780 | # move armature origin to the first freejoint position 781 | armature.location = scene.get_freejoint_pose().origin() 782 | 783 | bpy.ops.object.mode_set(mode='EDIT') 784 | edit_bones = armature.data.edit_bones 785 | default_bone = edit_bones[0] 786 | edit_bones.remove(default_bone) 787 | 788 | bone_dict = {} 789 | 790 | # Create a custom shape for joint bones (a simple disc) 791 | bpy.ops.mesh.primitive_circle_add(vertices=32, radius=2.0, fill_type='NOTHING', location=(0, 0, 0)) 792 | custom_shape_obj = bpy.context.object 793 | custom_shape_obj.name = "Joint_Bone_Shape" 794 | 795 | # Return to the armature object 796 | bpy.context.view_layer.objects.active = armature 797 | bpy.ops.object.mode_set(mode='EDIT') 798 | 799 | all_bones = scene.create_all_bones() 800 | for bone_info in all_bones: 801 | bone = edit_bones.new(bone_info.name()) 802 | bone.head = bone_info.head() 803 | bone.tail = bone_info.tail() 804 | if bone_info.parent_bone_name() is not None: 805 | bone.parent = bone_dict[bone_info.parent_bone_name()] 806 | self.report({'INFO'}, "Created Bone: {} (p: {})".format(bone_info.name(), bone_info.parent_bone_name())) 807 | 808 | # Store bone in the dictionary 809 | bone_dict[bone_info.name()] = bone 810 | 811 | # Exit Edit Mode to set bone custom shapes 812 | bpy.ops.object.mode_set(mode='POSE') 813 | 814 | # list all pose bones 815 | # for b in armature.pose.bones.keys(): 816 | # self.report({'INFO'}, b) 817 | 818 | pose_bones = armature.pose.bones 819 | for bone_info in all_bones: 820 | pose_bone = pose_bones[bone_info.name()] 821 | if bone_info.is_joint(): 822 | pose_bone.custom_shape = custom_shape_obj 823 | pose_bone.custom_shape_rotation_euler = (np.deg2rad(90), 0, 0) # Align along the Y-axis if needed 824 | pose_bone.rotation_mode = 'XYZ' 825 | 826 | # Attach STL meshes to corresponding bones 827 | bpy.ops.object.mode_set(mode='OBJECT') 828 | for body_info in scene.get_all_bodies(): 829 | for geom_info in body_info.get_mesh_geoms(): 830 | bone_name = body_info._end_bone_name 831 | if bone_name is None: 832 | self.report({'WARNING'}, "No bone for body {}".format(body_info.get_name())) 833 | continue 834 | stl_full_path = geom_info.get_stl_path() 835 | if stl_full_path is not None: 836 | if os.path.exists(stl_full_path): 837 | # bpy.ops.import_mesh.stl(filepath=stl_full_path) 838 | bpy.ops.wm.stl_import(filepath=stl_full_path) 839 | imported_mesh = bpy.context.object 840 | imported_mesh.name = body_info.get_name() + "_mesh" 841 | 842 | # Parent the mesh to the corresponding bone 843 | if bone_name in armature.data.bones: 844 | imported_mesh.parent = armature 845 | imported_mesh.parent_type = 'BONE' 846 | imported_mesh.parent_bone = bone_name 847 | imported_mesh.matrix_world = body_info.get_initial_tf_in_world().matrix().T 848 | 849 | return {'FINISHED'} 850 | 851 | def menu_func_import(self, context): 852 | self.layout.operator(MUJOCO_OT_import.bl_idname, text="MuJoCo XML (.xml)") 853 | 854 | def register(): 855 | bpy.utils.register_class(MUJOCO_OT_import) 856 | bpy.types.TOPBAR_MT_file_import.append(menu_func_import) 857 | 858 | def unregister(): 859 | bpy.utils.unregister_class(MUJOCO_OT_import) 860 | bpy.types.TOPBAR_MT_file_import.remove(menu_func_import) 861 | 862 | if __name__ == "__main__": 863 | if BLENDER: 864 | register() 865 | -------------------------------------------------------------------------------- /xml_import_test.py: -------------------------------------------------------------------------------- 1 | import xml.etree.ElementTree as ET 2 | import os 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | 6 | from mujoco_importer import parse_mujoco_xml 7 | 8 | def plot_joints(scene): 9 | fig, axs = plt.subplots(2, 2, figsize=(15, 5)) 10 | axdict = {"Side": axs[0][0], "Top": axs[1][0], "Front": axs[0][1]} 11 | axdims = {"Side": [0, 2], "Top": [0, 1], "Front": [1, 2]} 12 | for i, (view, ax) in enumerate(axdict.items()): 13 | ax.set_title(view + " View") 14 | ax.axis('equal') 15 | for body_name, body_tf in scene.get_bodies_in_world_tfs(): 16 | # Plot the body 17 | body_pos = body_tf.origin() 18 | body_x_axis = body_tf.x_axis() 19 | body_y_axis = body_tf.y_axis() 20 | body_z_axis = body_tf.z_axis() 21 | for view, ax in axdict.items(): 22 | S = 0.05 23 | ax.plot([body_pos[axdims[view][0]], body_pos[axdims[view][0]] + S * body_x_axis[axdims[view][0]]], [body_pos[axdims[view][1]], body_pos[axdims[view][1]] + S * body_x_axis[axdims[view][1]],], color='r') 24 | ax.plot([body_pos[axdims[view][0]], body_pos[axdims[view][0]] + S * body_y_axis[axdims[view][0]]], [body_pos[axdims[view][1]], body_pos[axdims[view][1]] + S * body_y_axis[axdims[view][1]],], color='g') 25 | ax.plot([body_pos[axdims[view][0]], body_pos[axdims[view][0]] + S * body_z_axis[axdims[view][0]]], [body_pos[axdims[view][1]], body_pos[axdims[view][1]] + S * body_z_axis[axdims[view][1]],], color='b') 26 | ax.text(body_pos[axdims[view][0]], body_pos[axdims[view][1]], body_name, fontsize=8, alpha=0.2) 27 | for joint_name, joint_pos, joint_axis in scene.get_joints_in_world(): 28 | for view, ax in axdict.items(): 29 | S = 0.1 30 | ax.plot([joint_pos[axdims[view][0]], joint_pos[axdims[view][0]] + S * joint_axis[axdims[view][0]]], [joint_pos[axdims[view][1]], joint_pos[axdims[view][1]] + S * joint_axis[axdims[view][1]],], color='y') 31 | ax.text(joint_pos[axdims[view][0]], joint_pos[axdims[view][1]], joint_name, fontsize=8, alpha=0.2) 32 | plt.tight_layout() 33 | plt.show() 34 | 35 | def plot_bones(scene): 36 | bones = scene.create_all_bones() 37 | fig, axs = plt.subplots(2, 2, figsize=(15, 5)) 38 | axdict = {"Side": axs[0][0], "Top": axs[1][0], "Front": axs[0][1]} 39 | axdims = {"Side": [0, 2], "Top": [0, 1], "Front": [1, 2]} 40 | for i, (view, ax) in enumerate(axdict.items()): 41 | ax.set_title(view + " View") 42 | ax.axis('equal') 43 | for bone in bones: 44 | print("{} ({})".format(bone.name(), bone.parent_bone_name())) 45 | for view, ax in axdict.items(): 46 | dim1, dim2 = axdims[view] 47 | head = bone.head() 48 | tail = bone.tail() 49 | delta = tail - head 50 | L = np.linalg.norm(delta) 51 | W = L * 0.2 52 | dd = np.array([delta[dim1], delta[dim2]]) 53 | pp = np.array([delta[dim2], -delta[dim1]]) / np.linalg.norm(dd) * W 54 | hh = np.array([head[dim1], head[dim2]]) 55 | tt = np.array([tail[dim1], tail[dim2]]) 56 | triangle = np.array([ 57 | hh + pp, 58 | tt, 59 | hh - pp, 60 | hh + pp 61 | ]) 62 | ax.plot(triangle[:, 0], triangle[:, 1], color=('r' if bone.is_joint() else 'k')) 63 | ax.text(*((hh+tt)/2.0), bone.name(), fontsize=8, alpha=0.2) 64 | plt.tight_layout() 65 | plt.show() 66 | 67 | 68 | 69 | 70 | 71 | 72 | if __name__ == "__main__": 73 | filepath = os.path.expanduser( 74 | "~/Code/DeepMimic_mujoco/src/mujoco/humanoid_deepmimic/envs/asset/deepmimic_unitree_g1.xml" 75 | ) 76 | scene = parse_mujoco_xml(filepath, None) 77 | # plot_joints(scene) 78 | plot_bones(scene) --------------------------------------------------------------------------------