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