addons-contrib: objects.link/unlink syntax update
[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._collection = None
152         self._scene = None
153         self._scene_camera = None
154         self._node = None
155         self._camera = None
156         self._euler_rotation = euler_rotation
157
158     def setScene(self, scene):
159         scene.name = self._name
160         self._scene = scene
161
162         scene.cube_map.use_cube_map = False
163         scene.render.use_compositing = False
164
165         self._setFilepath()
166
167     def _setFilepath(self):
168         import os
169
170         filepath = self._scene.render.filepath
171
172         dirname = os.path.dirname(filepath)
173         basename = os.path.basename(filepath)
174
175         path = os.path.join(dirname, "{0}{1}".format(self._name, basename))
176         self._scene.render.filepath = path
177
178     def setNode(self, node, links, node_output):
179         node.name = self._name
180         node.label = self._name
181         node.scene = self._scene
182         self._node = node
183
184         # TODO if there were nodetrees, duplicate them here
185
186         # connect to output
187         _input = node_output.layer_slots.new(self._name)
188         links.new(node.outputs[0], _input)
189
190     def setCamera(self, data, loc, zed):
191         self._scene_camera = self._scene.camera
192
193         self._camera = bpy.data.objects.new(self._name, data)
194         self._collection.objects.link(self._camera)
195
196         rotation = self._euler_rotation.copy()
197         rotation.z += zed
198
199         self._camera.rotation_euler = rotation
200         self._camera.location = loc
201
202         # change scene camera
203         self._scene.camera = self._camera
204
205     def resetCamera(self):
206         self._collection.objects.unlink(self._camera)
207         bpy.data.objects.remove(self._camera)
208         self._camera = None
209
210     @property
211     def scene(self):
212         return self._scene
213
214     @property
215     def name(self):
216         return self._name
217
218
219 @persistent
220 def cube_map_render_init(scene, use_force=False):
221     """
222     setup the cube map settings for all the render frames
223     """
224     from mathutils import Euler
225     from math import pi
226     half_pi = pi * 0.5
227
228     cube_map = scene.cube_map
229
230     if not do_run(cube_map, use_force):
231         return
232
233     main_scene = scene
234     hashes = [hash(scene) for scene in bpy.data.scenes]
235
236     views_raw = (
237             (
238                 'NORTH_',
239                 Euler((half_pi, 0.0,  0.0)),
240                 cube_map.use_view_north,
241                 ),
242             (
243                 'SOUTH_',
244                 Euler((half_pi, 0.0, pi)),
245                 cube_map.use_view_south,
246                 ),
247             (
248                 'WEST_',
249                 Euler((half_pi, 0.0, half_pi)),
250                 cube_map.use_view_west,
251                 ),
252             (
253                 'EAST_',
254                 Euler((half_pi, 0.0, -half_pi)),
255                 cube_map.use_view_east,
256                 ),
257             (
258                 'ZENITH_',
259                 Euler((pi, 0.0, 0.0)),
260                 cube_map.use_view_zenith,
261                 ),
262             (
263                 'NADIR_',
264                 Euler((0.0, 0.0, 0.0)),
265                 cube_map.use_view_nadir,
266                 ),
267             )
268
269     views = [
270             View(name, euler) for (name, euler, use) in views_raw
271             if use or not cube_map.is_advanced]
272
273     for view in views:
274         # create a scene per view
275         bpy.ops.scene.new(type='LINK_OBJECTS')
276         scene = [
277                 scene for scene in bpy.data.scenes if
278                 hash(scene) not in hashes][0]
279
280         # mark the scene to remove it afterwards
281         scene.cube_map.is_temporary = True
282
283         hashes.append(hash(scene))
284         view.setScene(scene)
285         # have Dalai to look at this?
286         view._collection = bpy.context.collection  # XXX TODO better fix
287
288     # create a scene from scratch
289     node_tree_data = NodeTree(main_scene)
290
291     # created the necessary nodetrees there
292     node_tree = main_scene.node_tree
293
294     # output node
295     node_output = node_tree.nodes.new('CompositorNodeOutputFile')
296     node_output.inputs.clear()
297
298     for view in views:
299         node = node_tree.nodes.new('CompositorNodeRLayers')
300         view.setNode(node, node_tree.links, node_output)
301
302     # globals
303     bpy.cube_map_node_tree_data = node_tree_data
304     bpy.cube_map_views = views
305
306
307 # ############################################################
308 # Cameras Setup
309 # ############################################################
310
311 @persistent
312 def cube_map_render_pre(scene, use_force=False):
313
314     if not do_run(scene.cube_map, use_force):
315         return
316
317     from math import radians
318
319     camera = scene.camera
320     data = camera.data.copy()
321
322     data.lens_unit = 'FOV'
323     data.angle = radians(90)
324     data.type = 'PERSP'
325
326     mat = camera.matrix_world
327
328     loc = mat.to_translation()
329     rot = mat.to_euler()
330     zed = rot.z
331
332     views = bpy.cube_map_views
333
334     for view in views:
335         view.setCamera(data, loc, zed)
336
337
338 @persistent
339 def cube_map_render_post(scene, use_force=False):
340
341     if not do_run(scene.cube_map, use_force):
342         return
343
344     views = bpy.cube_map_views
345
346     for view in views:
347         view.resetCamera()
348
349
350 # ############################################################
351 # Clean-Up
352 # ############################################################
353
354 @persistent
355 def cube_map_render_cancel(scene):
356     cube_map_cleanup(scene)
357
358
359 @persistent
360 def cube_map_render_complete(scene):
361     cube_map_cleanup(scene)
362
363
364 def cube_map_cleanup(scene, use_force=False):
365     """
366     remove all the temporary data created for the cube map
367     """
368
369     if not do_run(scene.cube_map, use_force):
370         return
371
372     bpy.cube_map_node_tree_data.cleanupScene()
373     del bpy.cube_map_node_tree_data
374     del bpy.cube_map_views
375
376     bpy.app.handlers.scene_update_post.append(cube_map_post_update_cleanup)
377
378
379 def cube_map_post_update_cleanup(scene):
380     """
381     delay removal of scenes (otherwise we get a crash)
382     """
383     scenes_temp = [
384             scene for scene in bpy.data.scenes if
385             scene.cube_map.is_temporary]
386
387     if not scenes_temp:
388         bpy.app.handlers.scene_update_post.remove(cube_map_post_update_cleanup)
389
390     else:
391         scenes_temp[0].user_clear()
392         try:
393             bpy.data.scenes.remove(scenes_temp[0], do_unlink=False)
394         except TypeError:
395             bpy.data.scenes.remove(scenes_temp[0])
396
397
398 # ############################################################
399 # Setup Operator
400 # ############################################################
401
402 class CubeMapSetup(Operator):
403     """"""
404     bl_idname = "render.cube_map_setup"
405     bl_label = "Cube Map Render Setup"
406     bl_description = ""
407
408     action: bpy.props.EnumProperty(
409         description="",
410         items=(("SETUP", "Setup", "Created linked scenes and setup cube map"),
411                ("RESET", "Reset", "Delete added scenes"),
412                ),
413         default="SETUP",
414         options={'SKIP_SAVE'},
415         )
416
417     @classmethod
418     def poll(cls, context):
419         return True
420
421     def setup(self, window, scene):
422         cube_map = scene.cube_map
423         cube_map.is_enabled = True
424
425         cube_map_render_init(scene, use_force=True)
426         cube_map_render_pre(scene, use_force=True)
427
428         # set initial scene back as the main scene
429         window.screen.scene = scene
430
431     def reset(self, scene):
432         cube_map = scene.cube_map
433         cube_map.is_enabled = False
434
435         cube_map_render_post(scene, use_force=True)
436         cube_map_cleanup(scene, use_force=True)
437
438     def invoke(self, context, event):
439         scene = context.scene
440         cube_map = scene.cube_map
441
442         is_enabled = cube_map.is_enabled
443
444         if self.action == 'RESET':
445
446             if is_enabled:
447                 if cube_map.is_temporary:
448                     self.report(
449                             {'ERROR'},
450                             "Cannot reset cube map from one of "
451                             "the created scenes")
452
453                     return {'CANCELLED'}
454                 else:
455                     self.reset(scene)
456                     return {'FINISHED'}
457             else:
458                 self.report({'ERROR'}, "Cube Map render is not setup")
459                 return {'CANCELLED'}
460
461         else:  # SETUP
462             if is_enabled:
463                 self.report({'ERROR'}, "Cube Map render is already setup")
464                 return {'CANCELLED'}
465             else:
466                 self.setup(context.window, scene)
467                 return {'FINISHED'}
468
469
470 # ############################################################
471 # User Interface
472 # ############################################################
473
474 class RENDER_PT_cube_map(Panel):
475     bl_space_type = 'PROPERTIES'
476     bl_region_type = 'WINDOW'
477     bl_context = "render"
478     bl_label = "Cube Map"
479
480     @classmethod
481     def poll(cls, context):
482         scene = context.scene
483         return scene and (scene.render.engine != 'BLENDER_GAME')
484
485     def draw_header(self, context):
486         self.layout.prop(context.scene.cube_map, "use_cube_map", text="")
487
488     def draw(self, context):
489         layout = self.layout
490         col = layout.column()
491
492         scene = context.scene
493         cube_map = scene.cube_map
494
495         if not cube_map.is_enabled:
496             col.operator(
497                     "render.cube_map_setup",
498                     text="Scene Setup").action = 'SETUP'
499         else:
500             col.operator(
501                     "render.cube_map_setup",
502                     text="Scene Reset", icon="X").action = 'RESET'
503
504         col = layout.column()
505         col.active = cube_map.use_cube_map
506         col.prop(cube_map, "is_advanced")
507         if cube_map.is_advanced:
508             box = col.box()
509             box.active = cube_map.use_cube_map and cube_map.is_advanced
510             row = box.row()
511             row.prop(cube_map, "use_view_north")
512             row.prop(cube_map, "use_view_west")
513             row.prop(cube_map, "use_view_zenith")
514
515             row = box.row()
516             row.prop(cube_map, "use_view_south")
517             row.prop(cube_map, "use_view_east")
518             row.prop(cube_map, "use_view_nadir")
519
520
521 # ############################################################
522 # Scene Properties
523 # ############################################################
524
525 class CubeMapInfo(bpy.types.PropertyGroup):
526     use_cube_map: BoolProperty(
527             name="Cube Map",
528             default=False,
529             )
530
531     is_temporary: BoolProperty(
532             name="Temporary",
533             default=False,
534             )
535
536     is_enabled: BoolProperty(
537             name="Enabled",
538             default=False,
539             )
540
541     # per view settings
542     is_advanced: BoolProperty(
543             name="Advanced",
544             default=False,
545             description="Decide which views to render",
546             )
547
548     use_view_north: BoolProperty(
549             name="North",
550             default=True,
551             )
552
553     use_view_south: BoolProperty(
554             name="South",
555             default=True,
556             )
557
558     use_view_west: BoolProperty(
559             name="West",
560             default=True,
561             )
562
563     use_view_east: BoolProperty(
564             name="East",
565             default=True,
566             )
567
568     use_view_zenith: BoolProperty(
569             name="Zenith",
570             default=True,
571             )
572
573     use_view_nadir: BoolProperty(
574             name="Nadir",
575             default=True,
576             )
577
578
579 # ############################################################
580 # Un/Registration
581 # ############################################################
582
583 def register():
584     bpy.utils.register_class(CubeMapInfo)
585     bpy.utils.register_class(CubeMapSetup)
586     bpy.types.Scene.cube_map = bpy.props.PointerProperty(
587             name="cube_map",
588             type=CubeMapInfo,
589             options={'HIDDEN'},
590             )
591
592     bpy.utils.register_class(RENDER_PT_cube_map)
593
594     bpy.app.handlers.render_init.append(cube_map_render_init)
595     bpy.app.handlers.render_pre.append(cube_map_render_pre)
596     bpy.app.handlers.render_post.append(cube_map_render_post)
597     bpy.app.handlers.render_cancel.append(cube_map_render_cancel)
598     bpy.app.handlers.render_complete.append(cube_map_render_complete)
599
600
601 def unregister():
602     bpy.utils.unregister_class(CubeMapInfo)
603     bpy.utils.unregister_class(CubeMapSetup)
604     bpy.utils.unregister_class(RENDER_PT_cube_map)
605
606     bpy.app.handlers.render_init.remove(cube_map_render_init)
607     bpy.app.handlers.render_pre.remove(cube_map_render_pre)
608     bpy.app.handlers.render_post.remove(cube_map_render_post)
609     bpy.app.handlers.render_cancel.remove(cube_map_render_cancel)
610     bpy.app.handlers.render_complete.remove(cube_map_render_complete)
611
612     del bpy.types.Scene.cube_map
613
614
615 if __name__ == '__main__':
616     register()