from bpy.types import ( Operator, OperatorFileListElement, Panel ) from bpy.props import ( StringProperty, CollectionProperty ) from bpy_extras.io_utils import ImportHelper import bpy import pickle from os.path import splitext, join, basename bl_info = { "name": "PyBulletSimImporter", "author": "Huy Ha ", "version": (0, 0, 1), "blender": (2, 92, 0), "location": "3D View > Toolbox > Animation tab > PyBullet Simulation Importer", "description": "Imports PyBullet Simulation Results", "warning": "", "category": "Animation", } class ANIM_OT_import_pybullet_sim(Operator, ImportHelper): bl_label = "Import simulation" bl_idname = "pybulletsim.import" bl_description = "Imports a PyBullet Simulation" bl_options = {'REGISTER', 'UNDO'} files: CollectionProperty( name="Simulation files", type=OperatorFileListElement, ) directory: StringProperty(subtype='DIR_PATH') filename_ext = ".pkl" filter_glob: StringProperty( default='*.pkl', options={'HIDDEN'}) skip_frames: bpy.props.IntProperty( name="Skip Frames", default=10, min=1, max=100) max_frames: bpy.props.IntProperty( name="Max Frames", default=-1, min=-1, max=100000) def execute(self, context): for file in self.files: filepath = join(self.directory, file.name) print(f'Processing {filepath}') with open(filepath, 'rb') as pickle_file: data = pickle.load(pickle_file) collection_name = splitext(basename(filepath))[0] collection = bpy.data.collections.new(collection_name) bpy.context.scene.collection.children.link(collection) context.view_layer.active_layer_collection = \ context.view_layer.layer_collection.children[-1] for obj_key in data: pybullet_obj = data[obj_key] # Load mesh of each link # register material if pybullet_obj['type'] == 'mesh': extension = pybullet_obj['mesh_path'].split( ".")[-1].lower() # Handle different mesh formats if 'obj' in extension: bpy.ops.import_scene.obj( filepath=pybullet_obj['mesh_path'], axis_forward='Y', axis_up='Z') elif 'dae' in extension: bpy.ops.wm.collada_import( filepath=pybullet_obj['mesh_path']) elif 'stl' in extension: bpy.ops.import_mesh.stl( filepath=pybullet_obj['mesh_path']) else: print("Unsupported File Format:{}".format(extension)) pass elif pybullet_obj['type'] == 'cube': bpy.ops.mesh.primitive_cube_add() elif pybullet_obj['type'] == "sphere": bpy.ops.mesh.primitive_uv_sphere_add() # radius=pybullet_obj['mesh_scale'][0] elif pybullet_obj['type'] == "cylinder": bpy.ops.mesh.primitive_cylinder_add() # radius=pybullet_obj['mesh_scale'][0], length=pybullet_obj['mesh_scale'][-1] # Delete lights and camera parts = 0 final_objs = [] for import_obj in context.selected_objects: bpy.ops.object.select_all(action='DESELECT') import_obj.select_set(True) if 'Camera' in import_obj.name \ or 'Light' in import_obj.name\ or 'Lamp' in import_obj.name: bpy.ops.object.delete(use_global=True) else: scale = pybullet_obj['mesh_scale'] if scale is not None: # if type(scale) is list: import_obj.scale.x = scale[0] import_obj.scale.y = scale[1] import_obj.scale.z = scale[2] final_objs.append(import_obj) parts += 1 bpy.ops.object.select_all(action='DESELECT') for obj in final_objs: if obj.type == 'MESH': obj.select_set(True) if len(context.selected_objects): context.view_layer.objects.active =\ context.selected_objects[0] # join them bpy.ops.object.join() blender_obj = context.view_layer.objects.active blender_obj.name = obj_key if pybullet_obj['mesh_material_name'] is not None: # register material material_name = pybullet_obj['mesh_material_name'] material_color = pybullet_obj['mesh_material_color'] print("registering material:", material_name) mat = bpy.data.materials.new(name=material_name) mat.use_nodes = True nodes = mat.node_tree.nodes links = mat.node_tree.links for node in nodes: nodes.remove(node) mat_node = nodes.new(type='ShaderNodeBsdfPrincipled') mat_node.inputs["Base Color"].default_value = material_color mat_node_output = mat_node.outputs['BSDF'] output_node = nodes.new(type='ShaderNodeOutputMaterial') links.new(output_node.inputs['Surface'], mat_node_output) # attach material # obj.data.materials[0] = bpy.data.materials[material_name] if obj.data.materials: obj.data.materials[0] = bpy.data.materials[material_name] else: obj.data.materials.append(bpy.data.materials[material_name]) # Keyframe motion of imported object for frame_count, frame_data in enumerate( pybullet_obj['frames']): if frame_count % self.skip_frames != 0: continue if self.max_frames > 1 and frame_count > self.max_frames: print('Exceed max frame count') break percentage_done = frame_count / \ len(pybullet_obj['frames']) print(f'\r[{percentage_done*100:.01f}% | {obj_key}]', '#' * int(60*percentage_done), end='') pos = frame_data['position'] orn = frame_data['orientation'] context.scene.frame_set( frame_count // self.skip_frames) # Apply position and rotation blender_obj.location.x = pos[0] blender_obj.location.y = pos[1] blender_obj.location.z = pos[2] blender_obj.rotation_mode = 'QUATERNION' blender_obj.rotation_quaternion.x = orn[0] blender_obj.rotation_quaternion.y = orn[1] blender_obj.rotation_quaternion.z = orn[2] blender_obj.rotation_quaternion.w = orn[3] bpy.ops.anim.keyframe_insert_menu( type='Rotation') bpy.ops.anim.keyframe_insert_menu( type='Location') return {'FINISHED'} class VIEW3D_PT_pybullet_recorder(Panel): bl_space_type = 'VIEW_3D' bl_region_type = 'UI' bl_category = "Animation" bl_label = 'PyBulletSimImporter' def draw(self, context): layout = self.layout row = layout.row() row.operator("pybulletsim.import") def register(): bpy.utils.register_class(VIEW3D_PT_pybullet_recorder) bpy.utils.register_class(ANIM_OT_import_pybullet_sim) def unregister(): bpy.utils.unregister_class(VIEW3D_PT_pybullet_recorder) bpy.utils.unregister_class(ANIM_OT_import_pybullet_sim) if __name__ == "__main__": register()