File size: 8,859 Bytes
ff66cf3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 |
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 <[email protected]>",
"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()
|