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 <hqh2101@columbia.edu>",
"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()