1import bpy
2import math
3import os
4import glob
5import reader
6import mathutils
7from blend_ChParticle_utils import renderPsys
8import multiprocessing
9
10
11# resolution reduction for the rendered image, LOW speeds up things for development, HIGH is slow but produces HD images
12res_reduction = {
13    'HIGH' : 1,
14    'MID' : 4,
15    'LOW' : 16}
16
17# function to orient a camera to look at a point
18def look_at(obj_camera, point, up):
19    loc_camera = obj_camera.location
20    direction = point - loc_camera
21    # point the cameras '-Z' and use its 'Y' as up
22    rot_quat = direction.to_track_quat('-Z', up)
23
24    # assume we're using euler rotation
25    obj_camera.rotation_euler = rot_quat.to_euler()
26
27# converts creates Blender objects from .dat chrono output format
28def convertshape(shapedata, meshespaths):
29    # Triangular Mesh
30    if shapedata[0] == 5:
31        for dirpath in meshespaths:
32            file_path = dirpath + "/" + shapedata[-1] + '.obj'
33            if os.path.isfile(file_path ):
34                imported_object_3 = bpy.ops.import_scene.obj(filepath=file_path )
35                obj_object = bpy.context.object
36                olist =bpy.context.selected_objects
37                for o in olist:
38                    o.location.x = shapedata[3]
39                    o.location.y = shapedata[4]
40                    o.location.z = shapedata[5]
41                    o.rotation_mode = 'QUATERNION'
42                    q = (shapedata[6], shapedata[7], shapedata[8], shapedata[9])
43                    o.rotation_quaternion = q
44                return True
45
46        return False
47    # sphere, box cylinder
48    elif shapedata[0]==0 or shapedata[0]==2 or shapedata[0]==3:
49        obj = None
50        # Sphere
51        if shapedata[0] == 0:
52            bpy.ops.mesh.primitive_ico_sphere_add(size=shapedata[13], location=(shapedata[3], shapedata[4], shapedata[5]),
53                                                  rotation=mathutils.Quaternion(shapedata[6], shapedata[7], shapedata[8], shapedata[9]).to_euler())
54            obj = bpy.context.active_object
55
56
57        # Box
58        elif shapedata[0] == 2:
59            bpy.ops.mesh.primitive_cube_add(size=1, location=(shapedata[3], shapedata[4], shapedata[5]),
60                                                  rotation=mathutils.Quaternion((shapedata[6], shapedata[7], shapedata[8], shapedata[9])).to_euler())
61            obj = bpy.context.active_object
62            obj.scale= ( shapedata[13]*2, shapedata[14]*2, shapedata[15]*2)
63
64
65        # Cylinder
66        elif shapedata[0] == 3:
67            bodypos = mathutils.Vector((shapedata[3], shapedata[4], shapedata[5]))
68            bodyrot = mathutils.Quaternion((shapedata[6], shapedata[7], shapedata[8], shapedata[9]))
69            #p1 & p2
70            p1 = mathutils.Vector((shapedata[14], shapedata[15], shapedata[16]))
71            p2 = mathutils.Vector((shapedata[17], shapedata[18], shapedata[19]))
72            # p2 - p1 vector
73            axis = p2 - p1
74            #cyl center in loc ref then transform
75            glob_center = p1 + axis/2
76            glob_center.rotate(bodyrot)
77            # cyl center = body_ref +  (p1 + (p2-p1)/2) in abs ref
78            center =  bodypos +  glob_center
79            # direction -> Quat -> EulerAng
80            ang = (bodyrot @ axis.to_track_quat('Z', 'X')).to_euler()
81            #ang = axis.to_track_quat('Z', 'X').to_euler()
82            bpy.ops.mesh.primitive_cylinder_add(location=center, rotation=ang, radius=shapedata[13], depth=axis.length)
83            obj = bpy.context.active_object
84
85
86        mat = bpy.data.materials.new("Blue")
87        # Activate its nodes
88        mat.use_nodes = True
89        # Get the principled BSDF (created by default)
90        principled = mat.node_tree.nodes['Principled BSDF']
91        # Assign the color
92        principled.inputs['Base Color'].default_value = (shapedata[10], shapedata[11], shapedata[12],1)
93        # Assign the material to the object
94        obj.data.materials.append(mat)
95        return True
96    else:
97        print('Unrecognized shape type')
98        return False
99
100# go through all shapes, find specific one, return true + frame of the shape
101def find_camera_target(exported_shapes, body_id, type_id, shape_name = ''):
102    for shapedata in exported_shapes:
103        if shapedata[1] == body_id:
104            # if looking for mesh, is shape is mesh
105            if type_id == 5 and shapedata[0] == 5 and shape_name == shapedata[-1]:
106                return True, [shapedata[3],shapedata[4],shapedata[5]], [shapedata[6],shapedata[7],shapedata[8], shapedata[9]]
107            # if looking for generic shape on body:
108            elif type_id == shapedata[0] :
109                return True, [shapedata[3],shapedata[4],shapedata[5]], [shapedata[6],shapedata[7],shapedata[8], shapedata[9]]
110
111
112def render_image(dirpaths, filepath,file_index, out_dir, meshes_prefixes,res,  targ,
113                 camera_mode, camera_pos, use_sky, a_up, light_loc, light_energy):
114
115    # path to the chrono and chrono-particle systems outputs. Processed only if they exists
116    dat_path = dirpaths[0] + filepath + ".dat"
117    #particle_path = dirpaths[0] + filepath + ".chpf"
118
119
120    # delete cached objects (initial cube and shapes from previous frames
121    for oldobj in bpy.data.objects:
122        if oldobj != bpy.data.objects['Camera']:
123            bpy.data.objects.remove(oldobj)
124
125    # IF there is a chrono .dat output, it gets processed
126    expshapes = []
127    for dir in dirpaths:
128        if (os.path.exists(dir + filepath + ".dat")):
129            expshapes += reader.import_shapes(dir + filepath + ".dat")
130    for shapedata in expshapes:
131        isconv = convertshape(shapedata, meshes_prefixes)
132        if not isconv:
133            print('Error while importing meshes')
134
135    if (os.path.exists(dat_path)):
136        targetfound, targ_pos, targ_or = find_camera_target(expshapes, targ["bodyid"], targ["shapetypeid"], targ["name"])
137        if not targetfound:
138                print('Could not find target object')
139
140
141    scene = bpy.context.scene
142    scene.objects.keys()
143
144    if camera_mode == 'Lookat':
145        obj_camera = bpy.data.objects["Camera"]
146        obj_camera.location = camera_pos
147        look_at(obj_camera, mathutils.Vector(targ_pos), a_up)
148
149
150    elif camera_mode == 'Follow':
151        print('Follow camera mode')
152        obj_camera = bpy.data.objects["Camera"]
153        # get the distance in the global ref frame:
154        global_dist = mathutils.Vector(targ["distfrom"])
155        global_dist.rotate(mathutils.Quaternion(targ_or))
156        moving_camera_pos = mathutils.Vector(targ_pos) + global_dist
157        obj_camera.location = moving_camera_pos
158        look_at(obj_camera, mathutils.Vector(targ_pos), a_up)
159    else:
160        obj_camera = bpy.data.objects["Camera"]
161        obj_camera.location = camera_pos
162        look_at(obj_camera, mathutils.Vector(targ["position"]), a_up)
163
164
165    # IF there is a chrono .chpf output, it gets processed
166    for dirpath in dirpaths:
167        partpath = dirpath + filepath + ".chpf"
168        if (os.path.exists(partpath)):
169            renderPsys(partpath, bpy.data.objects["Camera"])
170
171    # distance above which objects are not rendered
172    #TODO: make this settable!!
173    bpy.context.scene.camera.data.clip_end = 1000
174    scene.cycles.device = 'GPU'
175    prefs = bpy.context.preferences
176    cprefs = prefs.addons['cycles'].preferences
177
178    # Attempt to set GPU device types if available
179    for compute_device_type in ('CUDA', 'OPENCL', 'NONE'):
180        try:
181            cprefs.compute_device_type = compute_device_type
182            break
183        except TypeError:
184            pass
185
186    # Enable all CPU and GPU devices
187    cprefs.get_devices()
188    for device in cprefs.devices:
189        device.use = True
190
191
192    if use_sky:
193        sky_texture = bpy.context.scene.world.node_tree.nodes.new("ShaderNodeTexSky")
194        bg = bpy.context.scene.world.node_tree.nodes["Background"]
195        bpy.context.scene.world.node_tree.links.new(bg.inputs["Color"], sky_texture.outputs["Color"])
196        sky_texture.sky_type = 'PREETHAM'
197        #sky_texture.turbidity = 2.0
198        #sky_texture.ground_albedo = 0.4
199        ##sky_texture.sun_direction = mathutils.Vector((1.0, 0.0, 1.0))  # add `import mathutils` at the beginning of the script
200
201    # create light datablock, set attributes
202    light_data = bpy.data.lights.new(name="light_2.80", type='POINT')
203    light_data.energy = light_energy
204    # create new object with our light datablock
205    light_object = bpy.data.objects.new(name="light_2.80", object_data=light_data)
206    # link light object
207    bpy.context.collection.objects.link(light_object)
208    # make it active
209    bpy.context.view_layer.objects.active = light_object
210    # change location
211    light_object.location = light_loc
212
213    bpy.context.scene.render.engine = 'CYCLES'
214    bpy.context.scene.cycles.device = 'GPU'
215    #bpy.context.scene.render.resolution_percentage = 200
216    bpy.context.scene.cycles.samples = 256
217    bpy.context.scene.render.resolution_x = 3840/res_reduction[res]
218    bpy.context.scene.render.resolution_y = 2160/res_reduction[res]
219    bpy.context.scene.render.filepath = out_dir + "/bl_render_out_" + str(file_index) + ".png"
220    #bpy.context.scene.render.image_settings.compression = 500
221    bpy.context.scene.render.image_settings.color_mode = 'RGBA'
222    bpy.context.scene.render.image_settings.file_format = 'PNG'
223    #bpy.ops.render.render(write_still=True)
224
225    bpy.ops.render.render(write_still=True)
226    print('Rendered frame')
227
228def bl_render(meshes_prefixes, out_dir, datadirs, res, camera_mode, use_sky, camera_pos, targ, up='Z', light_loc=(10, 10, 15), light_energy=13000 ):
229    if up == 'Z':
230        a_up = 'Y'
231    elif up == 'Y':
232        a_up = 'Z'
233    else:
234        print('Invalid up-axis choice')
235
236    # populates the list with all the chrono output files, truncating the suffix (since there might be both .dat and .chpf)
237    datafiles = []
238    #for file in os.listdir(datadir):
239    for root, dirs, files in os.walk(datadirs[0]):
240        for file in files:
241            files.sort()
242            if file.endswith(".dat") or file.endswith(".chpf"):
243                datafiles.append(file[:file.rfind(".")])
244
245
246    if os.name == 'nt':
247        for file_ind, datafile in enumerate(datafiles):
248            render_image(datadirs, datafile, file_ind, out_dir, meshes_prefixes, res, targ,
249                         camera_mode, camera_pos, use_sky, a_up, light_loc, light_energy)
250
251    elif os.name == 'posix':
252        print('MULTIPROCESSING ACTIVE')
253        pool = multiprocessing.Pool(processes=multiprocessing.cpu_count()-1)
254        pool.starmap(render_image, [(datadirs, datafile, file_ind, out_dir, meshes_prefixes, res, targ,
255                         camera_mode, camera_pos, use_sky, a_up, light_loc, light_energy)
256                                for file_ind, datafile in enumerate(datafiles)])
257        pool.terminate()
258