blenderkit/overrides.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
  20 from blenderkit import utils
  21
  22 import bpy, mathutils
  23 from bpy.types import (
  24     Operator)
  25
  26
  27 def getNodes(nt, node_type='OUTPUT_MATERIAL'):
  28     chnodes = nt.nodes[:]
  29     nodes = []
  30     while len(chnodes) > 0:
  31         n = chnodes.pop()
  32         if n.type == node_type:
  33             nodes.append(n)
  34         if n.type == 'GROUP':
  35             chnodes.extend(n.node_tree.nodes)
  36     return nodes
  37
  38
  39 def getShadersCrawl(nt, chnodes):
  40     shaders = []
  41     done_nodes = chnodes[:]
  42
  43     while len(chnodes) > 0:
  44         check_node = chnodes.pop()
  45         is_shader = False
  46         for o in check_node.outputs:
  47             if o.type == 'SHADER':
  48                 is_shader = True
  49         for i in check_node.inputs:
  50             if i.type == 'SHADER':
  51                 is_shader = False  # this is for mix nodes and group inputs..
  52                 if len(i.links) > 0:
  53                     for l in i.links:
  54                         fn = l.from_node
  55                         if fn not in done_nodes:
  56                             done_nodes.append(fn)
  57                             chnodes.append(fn)
  58                             if fn.type == 'GROUP':
  59                                 group_outputs = getNodes(fn.node_tree, node_type='GROUP_OUTPUT')
  60                                 shaders.extend(getShadersCrawl(fn.node_tree, group_outputs))
  61
  62         if check_node.type == 'GROUP':
  63             is_shader = False
  64
  65         if is_shader:
  66             shaders.append((check_node, nt))
  67
  68     return (shaders)
  69
  70
  71 def addColorCorrectors(material):
  72     nt = material.node_tree
  73     output = getNodes(nt, 'OUTPUT_MATERIAL')[0]
  74     shaders = getShadersCrawl(nt, [output])
  75
  76     correctors = []
  77     for shader, nt in shaders:
  78
  79         if shader.type != 'BSDF_TRANSPARENT':  # exclude transparent for color tweaks
  80             for i in shader.inputs:
  81                 if i.type == 'RGBA':
  82                     if len(i.links) > 0:
  83                         l = i.links[0]
  84                         if not (l.from_node.type == 'GROUP' and l.from_node.node_tree.name == 'bkit_asset_tweaker'):
  85                             from_socket = l.from_socket
  86                             to_socket = l.to_socket
  87
  88                             g = nt.nodes.new(type='ShaderNodeGroup')
  89                             g.node_tree = bpy.data.node_groups['bkit_asset_tweaker']
  90                             g.location = shader.location
  91                             g.location.x -= 100
  92
  93                             nt.links.new(from_socket, g.inputs[0])
  94                             nt.links.new(g.outputs[0], to_socket)
  95                         else:
  96                             g = l.from_node
  97                         tweakers.append(g)
  98                     else:
  99                         g = nt.nodes.new(type='ShaderNodeGroup')
 100                         g.node_tree = bpy.data.node_groups['bkit_asset_tweaker']
 101                         g.location = shader.location
 102                         g.location.x -= 100
 103
 104                         nt.links.new(g.outputs[0], i)
 105                         correctors.append(g)
 106
 107
 108 def modelProxy():
 109     s = bpy.context.scene
 110     ao = bpy.context.active_object
 111     if utils.is_linked_asset(ao):
 112         utils.activate(ao)
 113
 114         g = ao.instance_collection
 115
 116         rigs = []
 117
 118         for ob in g.objects:
 119             if ob.type == 'ARMATURE':
 120                 rigs.append(ob)
 121
 122         if len(rigs) == 1:
 123
 124             ao.instance_collection = None
 125             bpy.ops.object.duplicate()
 126             new_ao = bpy.context.view_layer.objects.active
 127             new_ao.instance_collection = g
 128             new_ao.empty_display_type = 'SPHERE'
 129             new_ao.empty_display_size *= 0.1
 130
 131             bpy.ops.object.proxy_make(object=rigs[0].name)
 132             proxy = bpy.context.active_object
 133             bpy.context.view_layer.objects.active = ao
 134             ao.select_set(True)
 135             new_ao.select_set(True)
 136             new_ao.use_extra_recalc_object = True
 137             new_ao.use_extra_recalc_data = True
 138             bpy.ops.object.parent_set(type='OBJECT', keep_transform=True)
 139             return True
 140         else:  # TODO report this to ui
 141             utils.p('not sure what to proxify')
 142     return False
 143
 144
 145 eevee_transp_nodes = [
 146     'BSDF_GLASS',
 147     'BSDF_REFRACTION',
 148     'BSDF_TRANSPARENT',
 149     'PRINCIPLED_VOLUME',
 150     'VOLUME_ABSORPTION',
 151     'VOLUME_SCATTER'
 152 ]
 153
 154
 155 def ensure_eevee_transparency(m):
 156     ''' ensures alpha for transparent materials when the user didn't set it up correctly'''
 157     # if the blend mode is opaque, it means user probably ddidn't know or forgot to
 158     # set up material properly
 159     if m.blend_method == 'OPAQUE':
 160         alpha = False
 161         for n in m.node_tree.nodes:
 162             if n.type in eevee_transp_nodes:
 163                 alpha = True
 164             elif n.type == 'BSDF_PRINCIPLED':
 165                 i = n.inputs['Transmission']
 166                 if i.default_value > 0 or len(i.links) > 0:
 167                     alpha = True
 168         if alpha:
 169             m.blend_method = 'HASHED'
 170             m.shadow_method = 'HASHED'
 171
 172
 173 class BringToScene(Operator):
 174     """Bring linked object hierarchy to scene and make it editable"""
 175
 176     bl_idname = "object.blenderkit_bring_to_scene"
 177     bl_label = "BlenderKit bring objects to scene"
 178     bl_options = {'REGISTER', 'UNDO'}
 179
 180     @classmethod
 181     def poll(cls, context):
 182         return bpy.context.view_layer.objects.active is not None
 183
 184     def execute(self, context):
 185
 186         s = bpy.context.scene
 187         sobs = s.collection.all_objects
 188         aob = bpy.context.active_object
 189         dg = aob.instance_collection
 190         vlayer = bpy.context.view_layer
 191         instances_emptys = []
 192
 193         # first, find instances of this collection in the scene
 194         for ob in sobs:
 195             if ob.instance_collection == dg and ob not in instances_emptys:
 196                 instances_emptys.append(ob)
 197                 ob.instance_collection = None
 198                 ob.instance_type = 'NONE'
 199         # dg.make_local
 200         parent = None
 201         obs = []
 202         for ob in dg.objects:
 203             dg.objects.unlink(ob)
 204             try:
 205                 s.collection.objects.link(ob)
 206                 ob.select_set(True)
 207                 obs.append(ob)
 208                 if ob.parent == None:
 209                     parent = ob
 210                     bpy.context.view_layer.objects.active = parent
 211             except Exception as e:
 212                 print(e)
 213
 214         bpy.ops.object.make_local(type='ALL')
 215
 216         for i, ob in enumerate(obs):
 217             if ob.name in vlayer.objects:
 218                 obs[i] = vlayer.objects[ob.name]
 219                 try:
 220                     ob.select_set(True)
 221                 except Exception as e:
 222                     print('failed to select an object from the collection, getting a replacement.')
 223                     print(e)
 224
 225         related = []
 226
 227         for i, ob in enumerate(instances_emptys):
 228             if i > 0:
 229                 bpy.ops.object.duplicate(linked=True)
 230
 231             related.append([ob, bpy.context.active_object, mathutils.Vector(bpy.context.active_object.scale)])
 232
 233         for relation in related:
 234             bpy.ops.object.select_all(action='DESELECT')
 235             bpy.context.view_layer.objects.active = relation[0]
 236             relation[0].select_set(True)
 237             relation[1].select_set(True)
 238             relation[1].matrix_world = relation[0].matrix_world
 239             relation[1].scale.x = relation[2].x * relation[0].scale.x
 240             relation[1].scale.y = relation[2].y * relation[0].scale.y
 241             relation[1].scale.z = relation[2].z * relation[0].scale.z
 242             bpy.ops.object.parent_set(type='OBJECT', keep_transform=True)
 243
 244         return {'FINISHED'}
 245
 246
 247 class ModelProxy(Operator):
 248     """Attempt to create proxy armature from the asset"""
 249     bl_idname = "object.blenderkit_make_proxy"
 250     bl_label = "BlenderKit Make Proxy"
 251
 252     @classmethod
 253     def poll(cls, context):
 254         return bpy.context.view_layer.objects.active is not None
 255
 256     def execute(self, context):
 257         result = modelProxy()
 258         if not result:
 259             self.report({'INFO'}, 'No proxy made.There is no armature or more than one in the model.')
 260         return {'FINISHED'}
 261
 262
 263 class ColorCorrector(Operator):
 264     """Add color corector to the asset. """
 265     bl_idname = "object.blenderkit_color_corrector"
 266     bl_label = "Add color corrector"
 267
 268     @classmethod
 269     def poll(cls, context):
 270         return bpy.context.view_layer.objects.active is not None
 271
 272     def execute(self, context):
 273         ao = bpy.context.active_object
 274         g = ao.instance_collection
 275         ao['color correctors'] = []
 276         mats = []
 277
 278         for o in g.objects:
 279             for ms in o.material_slots:
 280                 if ms.material not in mats:
 281                     mats.append(ms.material)
 282         for mat in mats:
 283             correctors = addColorCorrectors(mat)
 284
 285         return 'FINISHED'
 286
 287
 288 def register_overrides():
 289     bpy.utils.register_class(BringToScene)
 290     bpy.utils.register_class(ModelProxy)
 291     bpy.utils.register_class(ColorCorrector)
 292
 293
 294 def unregister_overrides():
 295     bpy.utils.unregister_class(BringToScene)
 296     bpy.utils.unregister_class(ModelProxy)
 297     bpy.utils.unregister_class(ColorCorrector)