Python IO: Initial nodal shader support for import AND export.
[blender.git] / release / scripts / modules / bpy_extras / node_shader_utils.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 import bpy
22 from mathutils import Vector
23
24 __all__ = (
25     "PrincipledBSDFWrapper",
26     )
27
28
29 def _set_check(func):
30     from functools import wraps
31     @wraps(func)
32     def wrapper(self, *args, **kwargs):
33         if self.is_readonly:
34             assert(not "Trying to set value to read-only shader!")
35             return
36         return func(self, *args, **kwargs)
37     return wrapper
38
39
40 class ShaderWrapper():
41     """
42     Base class with minimal common ground for all types of shader interfaces we may want/need to implement.
43     """
44
45     # The two mandatory nodes any children class should support.
46     NODES_LIST = (
47         "node_out",
48
49         "_node_texcoords",
50         )
51
52     __slots__ = (
53         "is_readonly",
54         "material",
55         "_textures",
56         "_grid_locations",
57         ) + NODES_LIST
58
59     _col_size = 300
60     _row_size = 300
61
62     def _grid_to_location(self, x, y, dst_node=None, ref_node=None):
63         if ref_node is not None:  # x and y are relative to this node location.
64             nx = round(ref_node.location.x / self._col_size)
65             ny = round(ref_node.location.y / self._row_size)
66             x += nx
67             y += ny
68         loc = None
69         while True:
70             loc = (x * self._col_size, y * self._row_size)
71             if loc not in self._grid_locations:
72                 break
73             loc = (x * self._col_size, (y - 1) * self._row_size)
74             if loc not in self._grid_locations:
75                 break
76             loc = (x * self._col_size, (y - 2) * self._row_size)
77             if loc not in self._grid_locations:
78                 break
79             x -= 1
80         self._grid_locations.add(loc)
81         if dst_node is not None:
82             dst_node.location = loc
83             dst_node.width = min(dst_node.width, self._col_size - 20)
84         return loc
85
86     def __init__(self, material, is_readonly=True):
87         self.is_readonly = is_readonly
88         self.material = material
89         self.update()
90
91     def update(self):  # Should be re-implemented by children classes...
92         for node in self.NODES_LIST:
93             setattr(self, node, None)
94         self._textures = {}
95         self._grid_locations = set()
96
97     def use_nodes_get(self):
98         return self.material.use_nodes
99     @_set_check
100     def use_nodes_set(self, val):
101         self.material.use_nodes = val
102         self.update()
103     use_nodes = property(use_nodes_get, use_nodes_set)
104
105     def node_texcoords_get(self):
106         if not self.use_nodes:
107             return None
108         if self._node_texcoords is None:
109             for n in self.material.node_tree.nodes:
110                 if n.bl_idname == 'ShaderNodeTexCoord':
111                     self._node_texcoords = n
112                     self._grid_to_location(0, 0, ref_node=n)
113                     break
114         if self._node_texcoords is None and not self.is_readonly:
115             tree = self.material.node_tree
116             nodes = tree.nodes
117             links = tree.links
118             
119             node_texcoords = nodes.new(type='ShaderNodeTexCoord')
120             node_texcoords.label = "Texture Coords"
121             self._grid_to_location(-5, 1, dst_node=node_texcoords)
122             self._node_texcoords = node_texcoords
123         return self._node_texcoords
124     node_texcoords = property(node_texcoords_get)
125
126
127 class PrincipledBSDFWrapper(ShaderWrapper):
128     """
129     Hard coded shader setup, based in Principled BSDF.
130     Should cover most common cases on import, and gives a basic nodal shaders support for export.
131     Supports basic: diffuse/spec/reflect/transparency/normal, with texturing.
132     """
133     NODES_LIST = (
134         "node_out",
135         "node_principled_bsdf",
136
137         "_node_normalmap",
138         "_node_texcoords",
139         )
140
141     __slots__ = (
142         "is_readonly",
143         "material",
144         ) + NODES_LIST
145
146     NODES_LIST = ShaderWrapper.NODES_LIST + NODES_LIST
147
148     def __init__(self, material, is_readonly=True):
149         super(PrincipledBSDFWrapper, self).__init__(material, is_readonly)
150
151     def update(self):
152         super(PrincipledBSDFWrapper, self).update()
153
154         if not self.use_nodes:
155             return
156
157         tree = self.material.node_tree
158
159         nodes = tree.nodes
160         links = tree.links
161
162         # --------------------------------------------------------------------
163         # Main output and shader.
164         node_out = None
165         node_principled = None
166         for n in nodes:
167             if n.bl_idname == 'ShaderNodeOutputMaterial' and n.inputs[0].is_linked:
168                 node_out = n
169                 node_principled = n.inputs[0].links[0].from_node
170             elif n.bl_idname == 'ShaderNodeBsdfPrincipled' and n.outputs[0].is_linked:
171                 node_principled = n
172                 for lnk in n.outputs[0].links:
173                     node_out = lnk.to_node
174                     if node_out.bl_idname == 'ShaderNodeOutputMaterial':
175                         break
176             if (node_out is not None and node_principled is not None and
177                 node_out.bl_idname == 'ShaderNodeOutputMaterial' and
178                 node_principled.bl_idname == 'ShaderNodeBsdfPrincipled'):
179                 break
180             node_out = node_principled = None  # Could not find a valid pair, let's try again
181
182         if node_out is not None:
183             self._grid_to_location(0, 0, ref_node=node_out)
184         elif not self.is_readonly:
185             node_out = nodes.new(type='ShaderNodeOutputMaterial')
186             node_out.label = "Material Out"
187             node_out.target = 'ALL'
188             self._grid_to_location(1, 1, dst_node=node_out)
189         self.node_out = node_out
190
191         if node_principled is not None:
192             self._grid_to_location(0, 0, ref_node=node_principled)
193         elif not self.is_readonly:
194             node_principled = nodes.new(type='ShaderNodeBsdfPrincipled')
195             node_principled.label = "Principled BSDF"
196             self._grid_to_location(0, 1, dst_node=node_principled)
197             # Link
198             links.new(node_principled.outputs["BSDF"], self.node_out.inputs["Surface"])
199         self.node_principled_bsdf = node_principled
200
201         # --------------------------------------------------------------------
202         # Normal Map, lazy initialization...
203         self._node_normalmap = None
204
205         # --------------------------------------------------------------------
206         # Tex Coords, lazy initialization...
207         self._node_texcoords = None
208
209
210     def node_normalmap_get(self):
211         if not self.use_nodes:
212             return None
213         if self._node_normalmap is None and self.node_principled_bsdf is not None:
214             node_principled = self.node_principled_bsdf
215             if node_principled.inputs["Normal"].is_linked:
216                 node_normalmap = node_principled.inputs["Normal"].links[0].from_node
217                 if node_normalmap.bl_idname == 'ShaderNodeNormalMap':
218                     self._node_normalmap = node_normalmap
219                     self._grid_to_location(0, 0, ref_node=node_normalmap)
220             if self._node_normalmap is None and not self.is_readonly:
221                 node_normalmap = nodes.new(type='ShaderNodeNormalMap')
222                 node_normalmap.label = "Normal/Map"
223                 self._grid_to_location(-1, -2, dst_node=node_normalmap, ref_node=node_principled)
224                 # Link
225                 links.new(node_normalmap.outputs["Normal"], node_principled.inputs["Normal"])
226         return self._node_normalmap
227     node_normalmap = property(node_normalmap_get)
228
229
230     # --------------------------------------------------------------------
231     # Diffuse.
232
233     def diffuse_color_get(self):
234         if not self.use_nodes or self.node_principled_bsdf is None:
235             return self.material.diffuse_color
236         return self.node_principled_bsdf.inputs["Base Color"].default_value
237     @_set_check
238     def diffuse_color_set(self, color):
239         self.material.diffuse_color = color
240         if self.use_nodes and self.node_principled_bsdf is not None:
241             self.node_principled_bsdf.inputs["Base Color"].default_value = color
242     diffuse_color = property(diffuse_color_get, diffuse_color_set)
243
244     def diffuse_texture_get(self):
245         if not self.use_nodes or self.node_principled_bsdf is None:
246             return None
247         return ShaderImageTextureWrapper(self, self.node_principled_bsdf,
248                                          self.node_principled_bsdf.inputs["Base Color"],
249                                          grid_row_diff=1)
250     diffuse_texture = property(diffuse_texture_get)
251
252
253     # --------------------------------------------------------------------
254     # Specular.
255
256     def specular_get(self):
257         if not self.use_nodes or self.node_principled_bsdf is None:
258             return self.material.specular_intensity
259         return self.node_principled_bsdf.inputs["Specular"].default_value
260     @_set_check
261     def specular_set(self, value):
262         self.material.specular_intensity = value
263         if self.use_nodes and self.node_principled_bsdf is not None:
264             self.node_principled_bsdf.inputs["Specular"].default_value = value
265     specular = property(specular_get, specular_set)
266
267     def specular_tint_get(self):
268         if not self.use_nodes or self.node_principled_bsdf is None:
269             return 0.0
270         return self.node_principled_bsdf.inputs["Specular Tint"].default_value
271     @_set_check
272     def specular_tint_set(self, value):
273         if self.use_nodes and self.node_principled_bsdf is not None:
274             self.node_principled_bsdf.inputs["Specular Tint"].default_value = value
275     specular_tint = property(specular_tint_get, specular_tint_set)
276
277     # Will only be used as gray-scale one...
278     def specular_texture_get(self):
279         if not self.use_nodes or self.node_principled_bsdf is None:
280             return None
281         return ShaderImageTextureWrapper(self, self.node_principled_bsdf,
282                                          self.node_principled_bsdf.inputs["Specular"],
283                                          grid_row_diff=0)
284     specular_texture = property(specular_texture_get)
285
286
287     # --------------------------------------------------------------------
288     # Roughness (also sort of inverse of specular hardness...).
289
290     def roughness_get(self):
291         if not self.use_nodes or self.node_principled_bsdf is None:
292             return self.material.roughness
293         return self.node_principled_bsdf.inputs["Roughness"].default_value
294     @_set_check
295     def roughness_set(self, value):
296         self.material.roughness = value
297         if self.use_nodes and self.node_principled_bsdf is not None:
298             self.node_principled_bsdf.inputs["Roughness"].default_value = value
299     roughness = property(roughness_get, roughness_set)
300
301     # Will only be used as gray-scale one...
302     def roughness_texture_get(self):
303         if not self.use_nodes or self.node_principled_bsdf is None:
304             return None
305         return ShaderImageTextureWrapper(self, self.node_principled_bsdf,
306                                          self.node_principled_bsdf.inputs["Roughness"],
307                                          grid_row_diff=0)
308     roughness_texture = property(roughness_texture_get)
309
310
311     # --------------------------------------------------------------------
312     # Metallic (a.k.a reflection, mirror).
313
314     def metallic_get(self):
315         if not self.use_nodes or self.node_principled_bsdf is None:
316             return self.material.metallic
317         return self.node_principled_bsdf.inputs["Metallic"].default_value
318     @_set_check
319     def metallic_set(self, value):
320         self.material.metallic = value
321         if self.use_nodes and self.node_principled_bsdf is not None:
322             self.node_principled_bsdf.inputs["Metallic"].default_value = value
323     metallic = property(metallic_get, metallic_set)
324
325     # Will only be used as gray-scale one...
326     def metallic_texture_get(self):
327         if not self.use_nodes or self.node_principled_bsdf is None:
328             return None
329         return ShaderImageTextureWrapper(self, self.node_principled_bsdf,
330                                          self.node_principled_bsdf.inputs["Metallic"],
331                                          grid_row_diff=0)
332     metallic_texture = property(metallic_texture_get)
333
334
335     # --------------------------------------------------------------------
336     # Transparency settings.
337
338     def ior_get(self):
339         if not self.use_nodes or self.node_principled_bsdf is None:
340             return 1.0
341         return self.node_principled_bsdf.inputs["IOR"].default_value
342     @_set_check
343     def ior_set(self, value):
344         if self.use_nodes and self.node_principled_bsdf is not None:
345             self.node_principled_bsdf.inputs["IOR"].default_value = value
346     ior = property(ior_get, ior_set)
347
348     # Will only be used as gray-scale one...
349     def ior_texture_get(self):
350         if not self.use_nodes or self.node_principled_bsdf is None:
351             return None
352         return ShaderImageTextureWrapper(self, self.node_principled_bsdf,
353                                          self.node_principled_bsdf.inputs["IOR"],
354                                          grid_row_diff=-1)
355     ior_texture = property(ior_texture_get)
356
357     def transmission_get(self):
358         if not self.use_nodes or self.node_principled_bsdf is None:
359             return 0.0
360         return self.node_principled_bsdf.inputs["Transmission"].default_value
361     @_set_check
362     def transmission_set(self, value):
363         if self.use_nodes and self.node_principled_bsdf is not None:
364             self.node_principled_bsdf.inputs["Transmission"].default_value = value
365     transmission = property(transmission_get, transmission_set)
366
367     # Will only be used as gray-scale one...
368     def transmission_texture_get(self):
369         if not self.use_nodes or self.node_principled_bsdf is None:
370             return None
371         return ShaderImageTextureWrapper(self, self.node_principled_bsdf,
372                                          self.node_principled_bsdf.inputs["Transmission"],
373                                          grid_row_diff=-1)
374     transmission_texture = property(transmission_texture_get)
375
376     # TODO: Do we need more complex handling for alpha (allowing masking and such)?
377     #       Would need extra mixing nodes onto Base Color maybe, or even its own shading chain...
378
379
380     # --------------------------------------------------------------------
381     # Normal map.
382
383     def normalmap_strength_get(self):
384         if not self.use_nodes or self.node_normalmap is None:
385             return 0.0
386         return self.node_normalmap.inputs["Strength"].default_value
387     @_set_check
388     def normalmap_strength_set(self, value):
389         if self.use_nodes and self.node_normalmap is not None:
390             self.node_normalmap.inputs["Strength"].default_value = value
391     normalmap_strength = property(normalmap_strength_get, normalmap_strength_set)
392
393     def normalmap_texture_get(self):
394         if not self.use_nodes or self.node_normalmap is None:
395             return None
396         return ShaderImageTextureWrapper(self, self.node_normalmap,
397                                          self.node_normalmap.inputs["Color"],
398                                          grid_row_diff=-2)
399     normalmap_texture = property(normalmap_texture_get)
400
401
402 class ShaderImageTextureWrapper():
403     """
404     Generic 'image texture'-like wrapper, handling image node, some mapping (texture coordinates transformations),
405     and texture coordinates source.
406     """
407
408     # Note: this class assumes we are using nodes, otherwise it should never be used...
409
410     NODES_LIST = (
411         "node_dst",
412         "socket_dst",
413
414         "_node_image",
415         "_node_mapping",
416         )
417
418     __slots__ = (
419         "owner_shader",
420         "is_readonly",
421         "grid_row_diff",
422         "use_alpha",
423         ) + NODES_LIST
424
425     def __new__(cls, owner_shader: ShaderWrapper, node_dst, socket_dst, *args, **kwargs):
426         instance = owner_shader._textures.get((node_dst, socket_dst), None)
427         if instance is not None:
428             return instance
429         instance = super(ShaderImageTextureWrapper, cls).__new__(cls)
430         owner_shader._textures[(node_dst, socket_dst)] = instance
431         return instance
432
433     def __init__(self, owner_shader: ShaderWrapper, node_dst, socket_dst, grid_row_diff=0, use_alpha=False):
434         self.owner_shader = owner_shader
435         self.is_readonly = owner_shader.is_readonly
436         self.node_dst = node_dst
437         self.socket_dst = socket_dst
438         self.grid_row_diff = grid_row_diff
439         self.use_alpha = use_alpha
440
441         self._node_image = None
442         self._node_mapping = None
443
444         tree = node_dst.id_data
445         nodes = tree.nodes
446         links = tree.links
447
448         if socket_dst.is_linked:
449             self._node_image = socket_dst.links[0].from_node
450
451         if self.node_image is not None:
452             socket_dst = self.node_image.inputs["Vector"]
453             if socket_dst.is_linked:
454                 from_node = socket_dst.links[0].from_node
455                 if from_node.bl_idname == 'ShaderNodeMapping':
456                     self._node_mapping = from_node
457
458
459     # --------------------------------------------------------------------
460     # Image.
461
462     def node_image_get(self):
463         if self._node_image is None:
464             if self.socket_dst.is_linked:
465                 node_image = self.socket_dst.links[0].from_node
466                 if node_image.bl_idname == 'ShaderNodeTexImage':
467                     self._node_image = node_image
468                     self.owner_shader._grid_to_location(0, 0, ref_node=node_image)
469         if self._node_image is None and not self.is_readonly:
470             tree = self.owner_shader.material.node_tree
471
472             node_image = tree.nodes.new(type='ShaderNodeTexImage')
473             self.owner_shader._grid_to_location(-1, 0 + self.grid_row_diff, dst_node=node_image, ref_node=self.node_dst)
474
475             tree.links.new(node_image.outputs["Alpha" if self.use_alpha else "Color"], self.socket_dst)
476
477             self._node_image = node_image
478         return self._node_image
479     node_image = property(node_image_get)
480
481     def image_get(self):
482         return self.node_image.image if self.node_image is not None else None
483     @_set_check
484     def image_set(self, image):
485         self.node_image.image = image
486     image = property(image_get, image_set)
487
488     def projection_get(self):
489         return self.node_image.projection if self.node_image is not None else 'FLAT'
490     @_set_check
491     def projection_set(self, projection):
492         self.node_image.projection = projection
493     projection = property(projection_get, projection_set)
494
495     def texcoords_get(self):
496         if self.node_image is not None:
497             socket = (self.node_mapping if self._node_mapping is not None else self.node_image).inputs["Vector"]
498             if socket.is_linked:
499                 return socket.links[0].from_socket.name
500         return 'UV'
501     @_set_check
502     def texcoords_set(self, texcoords):
503         tree = self.node_image.id_data
504         links = tree.links
505         node_dst = self.node_mapping if self._node_mapping is not None else self.node_image
506         socket_src = self.owner_shader.node_texcoords.outputs[texcoords]
507         links.new(socket_src, node_dst.inputs["Vector"])
508     texcoords = property(texcoords_get, texcoords_set)
509
510
511     # --------------------------------------------------------------------
512     # Mapping.
513
514     def node_mapping_get(self):
515         if self._node_mapping is None:
516             if self.node_image is None:
517                 return None
518             if self.node_image.inputs["Vector"].is_linked:
519                 node_mapping = self.node_image.inputs["Vector"].links[0].from_node
520                 if node_mapping.bl_idname == 'ShaderNodeMapping':
521                     self._node_mapping = node_mapping
522                     self.owner_shader._grid_to_location(0, 0 + self.grid_row_diff, ref_node=node_mapping)
523         if self._node_mapping is None and not self.is_readonly:
524             # Find potential existing link into image's Vector input.
525             socket_dst = self.node_image.inputs["Vector"]
526             socket_src = socket_dst.links[0].from_socket if socket_dst.is_linked else None
527
528             tree = self.owner_shader.material.node_tree
529             node_mapping = tree.nodes.new(type='ShaderNodeMapping')
530             node_mapping.vector_type = 'TEXTURE'
531             self.owner_shader._grid_to_location(-1, 0, dst_node=node_mapping, ref_node=self.node_image)
532
533             # link mapping -> image node
534             tree.links.new(node_mapping.outputs["Vector"], socket_dst)
535             # And if already existing, re-link texcoords -> mapping
536             if socket_src is not None:
537                 tree.links.new(socket_src, node_mapping.inputs["Vector"])
538
539             self._node_mapping = node_mapping
540         return self._node_mapping
541     node_mapping = property(node_mapping_get)
542
543     def translation_get(self):
544         return self.node_mapping.translation if self.node_mapping is not None else Vector((0.0, 0.0, 0.0))
545     @_set_check
546     def translation_set(self, translation):
547         self.node_mapping.translation = translation
548     translation = property(translation_get, translation_set)
549
550     def rotation_get(self):
551         return self.node_mapping.rotation if self.node_mapping is not None else Vector((0.0, 0.0, 0.0))
552     @_set_check
553     def rotation_set(self, rotation):
554         self.node_mapping.rotation = rotation
555     rotation = property(rotation_get, rotation_set)
556
557     def scale_get(self):
558         return self.node_mapping.scale if self.node_mapping is not None else Vector((1.0, 1.0, 1.0))
559     @_set_check
560     def scale_set(self, scale):
561         self.node_mapping.scale = scale
562     scale = property(scale_get, scale_set)
563
564     def use_min_get(self):
565         return self.node_mapping.use_min if self_mapping.node is not None else False
566     @_set_check
567     def use_min_set(self, use_min):
568         self.node_mapping.use_min = use_min
569     use_min = property(use_min_get, use_min_set)
570
571     def use_max_get(self):
572         return self.node_mapping.use_max if self_mapping.node is not None else False
573     @_set_check
574     def use_max_set(self, use_max):
575         self.node_mapping.use_max = use_max
576     use_max = property(use_max_get, use_max_set)
577
578     def min_get(self):
579         return self.node_mapping.min if self.node_mapping is not None else Vector((0.0, 0.0, 0.0))
580     @_set_check
581     def min_set(self, min):
582         self.node_mapping.min = min
583     min = property(min_get, min_set)
584
585     def max_get(self):
586         return self.node_mapping.max if self.node_mapping is not None else Vector((0.0, 0.0, 0.0))
587     @_set_check
588     def max_set(self, max):
589         self.node_mapping.max = max
590     max = property(max_get, max_set)