2d0bc7fbae9f173bbaed9187354b60270742a9a4
[blender-addons-contrib.git] / render_cube_map.py
1 # ====================== BEGIN GPL LICENSE BLOCK ======================
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ======================= END GPL LICENSE BLOCK ========================
18
19 # <pep8 compliant>
20
21 # ########################################
22 # Render Cube Map
23 #
24 # Dalai Felinto
25 # --
26 # blendernetwork.org/dalai-felinto
27 # www.dalaifelinto.com
28 #
29 # Original code:
30 # Rio de Janeiro, September 2015
31 #
32 # Latest update:
33 # Rio de Janeiro, July 2016
34 # ########################################
35
36 import bpy
37
38 from bpy.app.handlers import persistent
39
40 from bpy.types import (
41         Operator,
42         Panel,
43         )
44
45 from bpy.props import (
46         BoolProperty,
47         )
48
49 bl_info = {
50     "name": "Cube Map",
51     "author": "Dalai Felinto",
52     "version": (1, 0),
53     "blender": (2, 77, 0),
54     "location": "Render Panel",
55     "description": "",
56     "warning": "",
57     "wiki_url": "https://github.com/dfelinto/render_cube_map",
58     "tracker_url": "",
59     "category": "Render"}
60
61
62 # ############################################################
63 # Global Check
64 # ############################################################
65
66 def do_run(cube_map, use_force):
67     if not (cube_map.use_cube_map or use_force):
68         return False
69
70     if cube_map.is_enabled and not use_force:
71         return False
72
73     return True
74
75
76 # ############################################################
77 # Callbacks
78 # ############################################################
79
80 class NodeTree:
81     def __init__(self, scene):
82         self._use_nodes = scene.use_nodes
83         self._use_compositing = scene.render.use_compositing
84
85         self._nodes_mute = {}
86         self._scene = scene
87
88         self._scene.render.use_compositing = True
89
90         if not self._use_nodes:
91             scene.use_nodes = True
92             self._muteNodes()
93         else:
94             self._storeNodes()
95             self._muteNodes()
96
97     def _storeNodes(self):
98         """
99         store the existent nodes and if they are muted
100         """
101         nodes = self._scene.node_tree.nodes
102         for node in nodes:
103             self._nodes_mute[hash(node)] = node.mute
104
105     def _muteNodes(self):
106         """
107         mute all the existent nodes
108         """
109         nodes = self._scene.node_tree.nodes
110         for node in nodes:
111             node.mute = True
112
113     def cleanupScene(self):
114         """
115         remove all the new nodes, and unmute original ones
116         """
117         scene = self._scene
118         scene.use_nodes = self._use_nodes
119         scene.render.use_compositing = self._use_compositing
120
121         self._cleanNodes()
122         self._unMuteNodes()
123
124     def _cleanNodes(self):
125         """
126         remove all the nodes created temporarily
127         """
128         nodes = self._scene.node_tree.nodes
129         to_del = []
130         keys = self._nodes_mute.keys()
131
132         for node in nodes:
133             if hash(node) not in keys:
134                 to_del.append(node)
135
136         for node in to_del:
137             nodes.remove(node)
138
139     def _unMuteNodes(self):
140         """
141         unmute all the existent nodes
142         """
143         nodes = self._scene.node_tree.nodes
144         for node in nodes:
145             node.mute = self._nodes_mute[hash(node)]
146
147
148 class View:
149     def __init__(self, name, euler_rotation):
150         self._name = name
151         self._scene = None
152         self._scene_camera = None
153         self._node = None
154         self._camera = None
155         self._euler_rotation = euler_rotation
156
157     def setScene(self, scene):
158         scene.name = self._name
159         self._scene = scene
160
161         scene.cube_map.use_cube_map = False
162         scene.render.use_compositing = False
163
164         self._setFilepath()
165
166     def _setFilepath(self):
167         import os
168
169         filepath = self._scene.render.filepath
170
171         dirname = os.path.dirname(filepath)
172         basename = os.path.basename(filepath)
173
174         path = os.path.join(dirname, "{0}{1}".format(self._name, basename))
175         self._scene.render.filepath = path
176
177     def setNode(self, node, links, node_output):
178         node.name = self._name
179         node.label = self._name
180         node.scene = self._scene
181         self._node = node
182
183         # TODO if there were nodetrees, duplicate them here
184
185         # connect to output
186         _input = node_output.layer_slots.new(self._name)
187         links.new(node.outputs[0], _input)
188
189     def setCamera(self, data, loc, zed):
190         self._scene_camera = self._scene.camera
191
192         self._camera = bpy.data.objects.new(self._name, data)
193         self._scene.objects.link(self._camera)
194
195         rotation = self._euler_rotation.copy()
196         rotation.z += zed
197
198         self._camera.rotation_euler = rotation
199         self._camera.location = loc
200
201         # change scene camera
202         self._scene.camera = self._camera
203
204     def resetCamera(self):
205         self._scene.objects.unlink(self._camera)
206         bpy.data.objects.remove(self._camera)
207         self._camera = None
208
209     @property
210     def scene(self):
211         return self._scene
212
213     @property
214     def name(self):
215         return self._name
216
217
218 @persistent
219 def cube_map_render_init(scene, use_force=False):
220     """
221     setup the cube map settings for all the render frames
222     """
223     from mathutils import Euler
224     from math import pi
225     half_pi = pi * 0.5
226
227     cube_map = scene.cube_map
228
229     if not do_run(cube_map, use_force):
230         return
231
232     main_scene = scene
233     hashes = [hash(scene) for scene in bpy.data.scenes]
234
235     views_raw = (
236             (
237                 'NORTH_',
238                 Euler((half_pi, 0.0,  0.0)),
239                 cube_map.use_view_north,
240                 ),
241             (
242                 'SOUTH_',
243                 Euler((half_pi, 0.0, pi)),
244                 cube_map.use_view_south,
245                 ),
246             (
247                 'WEST_',
248                 Euler((half_pi, 0.0, half_pi)),
249                 cube_map.use_view_west,
250                 ),
251             (
252                 'EAST_',
253                 Euler((half_pi, 0.0, -half_pi)),
254                 cube_map.use_view_east,
255                 ),
256             (
257                 'ZENITH_',
258                 Euler((pi, 0.0, 0.0)),
259                 cube_map.use_view_zenith,
260                 ),
261             (
262                 'NADIR_',
263                 Euler((0.0, 0.0, 0.0)),
264                 cube_map.use_view_nadir,
265                 ),
266             )
267
268     views = [
269             View(name, euler) for (name, euler, use) in views_raw
270             if use or not cube_map.is_advanced]
271
272     for view in views:
273         # create a scene per view
274         bpy.ops.scene.new(type='LINK_OBJECTS')
275         scene = [
276                 scene for scene in bpy.data.scenes if
277                 hash(scene) not in hashes][0]
278
279         # mark the scene to remove it afterwards
280         scene.cube_map.is_temporary = True
281
282         hashes.append(hash(scene))
283         view.setScene(scene)
284
285     # create a scene from scratch
286     node_tree_data = NodeTree(main_scene)
287
288     # created the necessary nodetrees there
289     node_tree = main_scene.node_tree
290
291     # output node
292     node_output = node_tree.nodes.new('CompositorNodeOutputFile')
293     node_output.inputs.clear()
294
295     for view in views:
296         node = node_tree.nodes.new('CompositorNodeRLayers')
297         view.setNode(node, node_tree.links, node_output)
298
299     # globals
300     bpy.cube_map_node_tree_data = node_tree_data
301     bpy.cube_map_views = views
302
303
304 # ############################################################
305 # Cameras Setup
306 # ############################################################
307
308 @persistent
309 def cube_map_render_pre(scene, use_force=False):
310
311     if not do_run(scene.cube_map, use_force):
312         return
313
314     from math import radians
315
316     camera = scene.camera
317     data = camera.data.copy()
318
319     data.lens_unit = 'FOV'
320     data.angle = radians(90)
321     data.type = 'PERSP'
322
323     mat = camera.matrix_world
324
325     loc = mat.to_translation()
326     rot = mat.to_euler()
327     zed = rot.z
328
329     views = bpy.cube_map_views
330
331     for view in views:
332         view.setCamera(data, loc, zed)
333
334
335 @persistent
336 def cube_map_render_post(scene, use_force=False):
337
338     if not do_run(scene.cube_map, use_force):
339         return
340
341     views = bpy.cube_map_views
342
343     for view in views:
344         view.resetCamera()
345
346
347 # ############################################################
348 # Clean-Up
349 # ############################################################
350
351 @persistent
352 def cube_map_render_cancel(scene):
353     cube_map_cleanup(scene)
354
355
356 @persistent
357 def cube_map_render_complete(scene):
358     cube_map_cleanup(scene)
359
360
361 def cube_map_cleanup(scene, use_force=False):
362     """
363     remove all the temporary data created for the cube map
364     """
365
366     if not do_run(scene.cube_map, use_force):
367         return
368
369     bpy.cube_map_node_tree_data.cleanupScene()
370     del bpy.cube_map_node_tree_data
371     del bpy.cube_map_views
372
373     bpy.app.handlers.scene_update_post.append(cube_map_post_update_cleanup)
374
375
376 def cube_map_post_update_cleanup(scene):
377     """
378     delay removal of scenes (otherwise we get a crash)
379     """
380     scenes_temp = [
381             scene for scene in bpy.data.scenes if
382             scene.cube_map.is_temporary]
383
384     if not scenes_temp:
385         bpy.app.handlers.scene_update_post.remove(cube_map_post_update_cleanup)
386
387     else:
388         scenes_temp[0].user_clear()
389         try:
390             bpy.data.scenes.remove(scenes_temp[0], do_unlink=False)
391         except TypeError:
392             bpy.data.scenes.remove(scenes_temp[0])
393
394
395 # ############################################################
396 # Setup Operator
397 # ############################################################
398
399 class CubeMapSetup(Operator):
400     """"""
401     bl_idname = "render.cube_map_setup"
402     bl_label = "Cube Map Render Setup"
403     bl_description = ""
404
405     action: bpy.props.EnumProperty(
406         description="",
407         items=(("SETUP", "Setup", "Created linked scenes and setup cube map"),
408                ("RESET", "Reset", "Delete added scenes"),
409                ),
410         default="SETUP",
411         options={'SKIP_SAVE'},
412         )
413
414     @classmethod
415     def poll(cls, context):
416         return True
417
418     def setup(self, window, scene):
419         cube_map = scene.cube_map
420         cube_map.is_enabled = True
421
422         cube_map_render_init(scene, use_force=True)
423         cube_map_render_pre(scene, use_force=True)
424
425         # set initial scene back as the main scene
426         window.screen.scene = scene
427
428     def reset(self, scene):
429         cube_map = scene.cube_map
430         cube_map.is_enabled = False
431
432         cube_map_render_post(scene, use_force=True)
433         cube_map_cleanup(scene, use_force=True)
434
435     def invoke(self, context, event):
436         scene = context.scene
437         cube_map = scene.cube_map
438
439         is_enabled = cube_map.is_enabled
440
441         if self.action == 'RESET':
442
443             if is_enabled:
444                 if cube_map.is_temporary:
445                     self.report(
446                             {'ERROR'},
447                             "Cannot reset cube map from one of "
448                             "the created scenes")
449
450                     return {'CANCELLED'}
451                 else:
452                     self.reset(scene)
453                     return {'FINISHED'}
454             else:
455                 self.report({'ERROR'}, "Cube Map render is not setup")
456                 return {'CANCELLED'}
457
458         else:  # SETUP
459             if is_enabled:
460                 self.report({'ERROR'}, "Cube Map render is already setup")
461                 return {'CANCELLED'}
462             else:
463                 self.setup(context.window, scene)
464                 return {'FINISHED'}
465
466
467 # ############################################################
468 # User Interface
469 # ############################################################
470
471 class RENDER_PT_cube_map(Panel):
472     bl_space_type = 'PROPERTIES'
473     bl_region_type = 'WINDOW'
474     bl_context = "render"
475     bl_label = "Cube Map"
476
477     @classmethod
478     def poll(cls, context):
479         scene = context.scene
480         return scene and (scene.render.engine != 'BLENDER_GAME')
481
482     def draw_header(self, context):
483         self.layout.prop(context.scene.cube_map, "use_cube_map", text="")
484
485     def draw(self, context):
486         layout = self.layout
487         col = layout.column()
488
489         scene = context.scene
490         cube_map = scene.cube_map
491
492         if not cube_map.is_enabled:
493             col.operator(
494                     "render.cube_map_setup",
495                     text="Scene Setup").action = 'SETUP'
496         else:
497             col.operator(
498                     "render.cube_map_setup",
499                     text="Scene Reset", icon="X").action = 'RESET'
500
501         col = layout.column()
502         col.active = cube_map.use_cube_map
503         col.prop(cube_map, "is_advanced")
504         if cube_map.is_advanced:
505             box = col.box()
506             box.active = cube_map.use_cube_map and cube_map.is_advanced
507             row = box.row()
508             row.prop(cube_map, "use_view_north")
509             row.prop(cube_map, "use_view_west")
510             row.prop(cube_map, "use_view_zenith")
511
512             row = box.row()
513             row.prop(cube_map, "use_view_south")
514             row.prop(cube_map, "use_view_east")
515             row.prop(cube_map, "use_view_nadir")
516
517
518 # ############################################################
519 # Scene Properties
520 # ############################################################
521
522 class CubeMapInfo(bpy.types.PropertyGroup):
523     use_cube_map: BoolProperty(
524             name="Cube Map",
525             default=False,
526             )
527
528     is_temporary: BoolProperty(
529             name="Temporary",
530             default=False,
531             )
532
533     is_enabled: BoolProperty(
534             name="Enabled",
535             default=False,
536             )
537
538     # per view settings
539     is_advanced: BoolProperty(
540             name="Advanced",
541             default=False,
542             description="Decide which views to render",
543             )
544
545     use_view_north: BoolProperty(
546             name="North",
547             default=True,
548             )
549
550     use_view_south: BoolProperty(
551             name="South",
552             default=True,
553             )
554
555     use_view_west: BoolProperty(
556             name="West",
557             default=True,
558             )
559
560     use_view_east: BoolProperty(
561             name="East",
562             default=True,
563             )
564
565     use_view_zenith: BoolProperty(
566             name="Zenith",
567             default=True,
568             )
569
570     use_view_nadir: BoolProperty(
571             name="Nadir",
572             default=True,
573             )
574
575
576 # ############################################################
577 # Un/Registration
578 # ############################################################
579
580 def register():
581     bpy.utils.register_class(CubeMapInfo)
582     bpy.utils.register_class(CubeMapSetup)
583     bpy.types.Scene.cube_map = bpy.props.PointerProperty(
584             name="cube_map",
585             type=CubeMapInfo,
586             options={'HIDDEN'},
587             )
588
589     bpy.utils.register_class(RENDER_PT_cube_map)
590
591     bpy.app.handlers.render_init.append(cube_map_render_init)
592     bpy.app.handlers.render_pre.append(cube_map_render_pre)
593     bpy.app.handlers.render_post.append(cube_map_render_post)
594     bpy.app.handlers.render_cancel.append(cube_map_render_cancel)
595     bpy.app.handlers.render_complete.append(cube_map_render_complete)
596
597
598 def unregister():
599     bpy.utils.unregister_class(CubeMapInfo)
600     bpy.utils.unregister_class(CubeMapSetup)
601     bpy.utils.unregister_class(RENDER_PT_cube_map)
602
603     bpy.app.handlers.render_init.remove(cube_map_render_init)
604     bpy.app.handlers.render_pre.remove(cube_map_render_pre)
605     bpy.app.handlers.render_post.remove(cube_map_render_post)
606     bpy.app.handlers.render_cancel.remove(cube_map_render_cancel)
607     bpy.app.handlers.render_complete.remove(cube_map_render_complete)
608
609     del bpy.types.Scene.cube_map
610
611
612 if __name__ == '__main__':
613     register()