maxmustermann/armory
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity
armory/blender/arm/lightmapper/utility/utility.py
Alexander Kleemann ef8fb21536 Update lightmapper to Blender 2.9+ 
Finalized update to support Blender 2.9+ as well as new features, fixes and more stability
2021-03-18 18:49:30 +01:00

674 lines
23 KiB
Python
Raw Blame History

import bpy.ops as O
import bpy, os, re, sys, importlib, struct, platform, subprocess, threading, string, bmesh, shutil, glob, uuid
from io import StringIO
from threading import Thread
from queue import Queue, Empty
from dataclasses import dataclass
from dataclasses import field
from typing import List
###########################################################
###########################################################
# This set of utility functions are courtesy of LorenzWieseke
#
# Modified by Naxela
#
# https://github.com/Naxela/The_Lightmapper/tree/Lightmap-to-GLB
###########################################################
class Node_Types:
output_node = 'OUTPUT_MATERIAL'
ao_node = 'AMBIENT_OCCLUSION'
image_texture = 'TEX_IMAGE'
pbr_node = 'BSDF_PRINCIPLED'
diffuse = 'BSDF_DIFFUSE'
mapping = 'MAPPING'
normal_map = 'NORMAL_MAP'
bump_map = 'BUMP'
attr_node = 'ATTRIBUTE'
class Shader_Node_Types:
emission = "ShaderNodeEmission"
image_texture = "ShaderNodeTexImage"
mapping = "ShaderNodeMapping"
normal = "ShaderNodeNormalMap"
ao = "ShaderNodeAmbientOcclusion"
uv = "ShaderNodeUVMap"
mix = "ShaderNodeMixRGB"
def select_object(self,obj):
C = bpy.context
try:
O.object.select_all(action='DESELECT')
C.view_layer.objects.active = obj
obj.select_set(True)
except:
self.report({'INFO'},"Object not in View Layer")
def select_obj_by_mat(self,mat):
D = bpy.data
for obj in D.objects:
if obj.type == "MESH":
object_materials = [
slot.material for slot in obj.material_slots]
if mat in object_materials:
select_object(self,obj)
def save_image(image):
filePath = bpy.data.filepath
path = os.path.dirname(filePath)
try:
os.mkdir(path + "/tex")
except FileExistsError:
pass
try:
os.mkdir(path + "/tex/" + str(image.size[0]))
except FileExistsError:
pass
if image.file_format == "JPEG" :
file_ending = ".jpg"
elif image.file_format == "PNG" :
file_ending = ".png"
savepath = path + "/tex/" + \
str(image.size[0]) + "/" + image.name + file_ending
image.filepath_raw = savepath
image.save()
def get_file_size(filepath):
size = "Unpack Files"
try:
path = bpy.path.abspath(filepath)
size = os.path.getsize(path)
size /= 1024
except:
if bpy.context.scene.TLM_SceneProperties.tlm_verbose:
print("error getting file path for " + filepath)
return (size)
def scale_image(image, newSize):
if (image.org_filepath != ''):
image.filepath = image.org_filepath
image.org_filepath = image.filepath
image.scale(newSize[0], newSize[1])
save_image(image)
def check_only_one_pbr(self,material):
check_ok = True
# get pbr shader
nodes = material.node_tree.nodes
pbr_node_type = Node_Types.pbr_node
pbr_nodes = find_node_by_type(nodes,pbr_node_type)
# check only one pbr node
if len(pbr_nodes) == 0:
self.report({'INFO'}, 'No PBR Shader Found')
check_ok = False
if len(pbr_nodes) > 1:
self.report({'INFO'}, 'More than one PBR Node found ! Clean before Baking.')
check_ok = False
return check_ok
# is material already the baked one
def check_is_org_material(self,material):
check_ok = True
if "_Bake" in material.name:
self.report({'INFO'}, 'Change back to org. Material')
check_ok = False
return check_ok
def clean_empty_materials(self):
for obj in bpy.context.scene.objects:
for slot in obj.material_slots:
mat = slot.material
if mat is None:
if bpy.context.scene.TLM_SceneProperties.tlm_verbose:
print("Removed Empty Materials from " + obj.name)
bpy.ops.object.select_all(action='DESELECT')
obj.select_set(True)
bpy.ops.object.material_slot_remove()
def get_pbr_inputs(pbr_node):
base_color_input = pbr_node.inputs["Base Color"]
metallic_input = pbr_node.inputs["Metallic"]
specular_input = pbr_node.inputs["Specular"]
roughness_input = pbr_node.inputs["Roughness"]
normal_input = pbr_node.inputs["Normal"]
pbr_inputs = {"base_color_input":base_color_input, "metallic_input":metallic_input,"specular_input":specular_input,"roughness_input":roughness_input,"normal_input":normal_input}
return pbr_inputs
def find_node_by_type(nodes, node_type):
nodes_found = [n for n in nodes if n.type == node_type]
return nodes_found
def find_node_by_type_recusivly(material, note_to_start, node_type, del_nodes_inbetween=False):
nodes = material.node_tree.nodes
if note_to_start.type == node_type:
return note_to_start
for input in note_to_start.inputs:
for link in input.links:
current_node = link.from_node
if (del_nodes_inbetween and note_to_start.type != Node_Types.normal_map and note_to_start.type != Node_Types.bump_map):
nodes.remove(note_to_start)
return find_node_by_type_recusivly(material, current_node, node_type, del_nodes_inbetween)
def find_node_by_name_recusivly(node, idname):
if node.bl_idname == idname:
return node
for input in node.inputs:
for link in input.links:
current_node = link.from_node
return find_node_by_name_recusivly(current_node, idname)
def make_link(material, socket1, socket2):
links = material.node_tree.links
links.new(socket1, socket2)
def add_gamma_node(material, pbrInput):
nodeToPrincipledOutput = pbrInput.links[0].from_socket
gammaNode = material.node_tree.nodes.new("ShaderNodeGamma")
gammaNode.inputs[1].default_value = 2.2
gammaNode.name = "Gamma Bake"
# link in gamma
make_link(material, nodeToPrincipledOutput, gammaNode.inputs["Color"])
make_link(material, gammaNode.outputs["Color"], pbrInput)
def remove_gamma_node(material, pbrInput):
nodes = material.node_tree.nodes
gammaNode = nodes.get("Gamma Bake")
nodeToPrincipledOutput = gammaNode.inputs[0].links[0].from_socket
make_link(material, nodeToPrincipledOutput, pbrInput)
material.node_tree.nodes.remove(gammaNode)
def apply_ao_toggle(self,context):
all_materials = bpy.data.materials
ao_toggle = context.scene.toggle_ao
for mat in all_materials:
nodes = mat.node_tree.nodes
ao_node = nodes.get("AO Bake")
if ao_node is not None:
if ao_toggle:
emission_setup(mat,ao_node.outputs["Color"])
else:
pbr_node = find_node_by_type(nodes,Node_Types.pbr_node)[0]
remove_node(mat,"Emission Bake")
reconnect_PBR(mat, pbr_node)
def emission_setup(material, node_output):
nodes = material.node_tree.nodes
emission_node = add_node(material,Shader_Node_Types.emission,"Emission Bake")
# link emission to whatever goes into current pbrInput
emission_input = emission_node.inputs[0]
make_link(material, node_output, emission_input)
# link emission to materialOutput
surface_input = nodes.get("Material Output").inputs[0]
emission_output = emission_node.outputs[0]
make_link(material, emission_output, surface_input)
def link_pbr_to_output(material,pbr_node):
nodes = material.node_tree.nodes
surface_input = nodes.get("Material Output").inputs[0]
make_link(material,pbr_node.outputs[0],surface_input)
def reconnect_PBR(material, pbrNode):
nodes = material.node_tree.nodes
pbr_output = pbrNode.outputs[0]
surface_input = nodes.get("Material Output").inputs[0]
make_link(material, pbr_output, surface_input)
def mute_all_texture_mappings(material, do_mute):
nodes = material.node_tree.nodes
for node in nodes:
if node.bl_idname == "ShaderNodeMapping":
node.mute = do_mute
def add_node(material,shader_node_type,node_name):
nodes = material.node_tree.nodes
new_node = nodes.get(node_name)
if new_node is None:
new_node = nodes.new(shader_node_type)
new_node.name = node_name
new_node.label = node_name
return new_node
def remove_node(material,node_name):
nodes = material.node_tree.nodes
node = nodes.get(node_name)
if node is not None:
nodes.remove(node)
def lightmap_to_ao(material,lightmap_node):
nodes = material.node_tree.nodes
# -----------------------AO SETUP--------------------#
# create group data
gltf_settings = bpy.data.node_groups.get('glTF Settings')
if gltf_settings is None:
bpy.data.node_groups.new('glTF Settings', 'ShaderNodeTree')
# add group to node tree
ao_group = nodes.get('glTF Settings')
if ao_group is None:
ao_group = nodes.new('ShaderNodeGroup')
ao_group.name = 'glTF Settings'
ao_group.node_tree = bpy.data.node_groups['glTF Settings']
# create group inputs
if ao_group.inputs.get('Occlusion') is None:
ao_group.inputs.new('NodeSocketFloat','Occlusion')
# mulitply to control strength
mix_node = add_node(material,Shader_Node_Types.mix,"Adjust Lightmap")
mix_node.blend_type = "MULTIPLY"
mix_node.inputs["Fac"].default_value = 1
mix_node.inputs["Color2"].default_value = [3,3,3,1]
# position node
ao_group.location = (lightmap_node.location[0]+600,lightmap_node.location[1])
mix_node.location = (lightmap_node.location[0]+300,lightmap_node.location[1])
make_link(material,lightmap_node.outputs['Color'],mix_node.inputs['Color1'])
make_link(material,mix_node.outputs['Color'],ao_group.inputs['Occlusion'])
###########################################################
###########################################################
# This utility function is modified from blender_xatlas
# and calls the object without any explicit object context
# thus allowing blender_xatlas to pack from background.
###########################################################
# Code is courtesy of mattedicksoncom
# Modified by Naxela
#
# https://github.com/mattedicksoncom/blender-xatlas/
###########################################################
def gen_safe_name():
genId = uuid.uuid4().hex
# genId = "u_" + genId.replace("-","_")
return "u_" + genId
def Unwrap_Lightmap_Group_Xatlas_2_headless_call(obj):
blender_xatlas = importlib.util.find_spec("blender_xatlas")
if blender_xatlas is not None:
import blender_xatlas
else:
return 0
packOptions = bpy.context.scene.pack_tool
chartOptions = bpy.context.scene.chart_tool
sharedProperties = bpy.context.scene.shared_properties
#sharedProperties.unwrapSelection
context = bpy.context
#save whatever mode the user was in
startingMode = bpy.context.object.mode
selected_objects = bpy.context.selected_objects
#check something is actually selected
#external function/operator will select them
if len(selected_objects) == 0:
print("Nothing Selected")
self.report({"WARNING"}, "Nothing Selected, please select Something")
return {'FINISHED'}
#store the names of objects to be lightmapped
rename_dict = dict()
safe_dict = dict()
#make sure all the objects have ligthmap uvs
for obj in selected_objects:
if obj.type == 'MESH':
safe_name = gen_safe_name();
rename_dict[obj.name] = (obj.name,safe_name)
safe_dict[safe_name] = obj.name
context.view_layer.objects.active = obj
if obj.data.users > 1:
obj.data = obj.data.copy() #make single user copy
uv_layers = obj.data.uv_layers
#setup the lightmap uvs
uvName = "UVMap_Lightmap"
if sharedProperties.lightmapUVChoiceType == "NAME":
uvName = sharedProperties.lightmapUVName
elif sharedProperties.lightmapUVChoiceType == "INDEX":
if sharedProperties.lightmapUVIndex < len(uv_layers):
uvName = uv_layers[sharedProperties.lightmapUVIndex].name
if not uvName in uv_layers:
uvmap = uv_layers.new(name=uvName)
uv_layers.active_index = len(uv_layers) - 1
else:
for i in range(0, len(uv_layers)):
if uv_layers[i].name == uvName:
uv_layers.active_index = i
obj.select_set(True)
#save all the current edges
if sharedProperties.packOnly:
edgeDict = dict()
for obj in selected_objects:
if obj.type == 'MESH':
tempEdgeDict = dict()
tempEdgeDict['object'] = obj.name
tempEdgeDict['edges'] = []
print(len(obj.data.edges))
for i in range(0,len(obj.data.edges)):
setEdge = obj.data.edges[i]
tempEdgeDict['edges'].append(i)
edgeDict[obj.name] = tempEdgeDict
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.mesh.select_all(action='SELECT')
bpy.ops.mesh.quads_convert_to_tris(quad_method='FIXED', ngon_method='BEAUTY')
else:
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.mesh.select_all(action='SELECT')
bpy.ops.mesh.quads_convert_to_tris(quad_method='FIXED', ngon_method='BEAUTY')
bpy.ops.object.mode_set(mode='OBJECT')
#Create a fake obj export to a string
#Will strip this down further later
fakeFile = StringIO()
blender_xatlas.export_obj_simple.save(
rename_dict=rename_dict,
context=bpy.context,
filepath=fakeFile,
mainUVChoiceType=sharedProperties.mainUVChoiceType,
uvIndex=sharedProperties.mainUVIndex,
uvName=sharedProperties.mainUVName,
use_selection=True,
use_animation=False,
use_mesh_modifiers=True,
use_edges=True,
use_smooth_groups=False,
use_smooth_groups_bitflags=False,
use_normals=True,
use_uvs=True,
use_materials=False,
use_triangles=False,
use_nurbs=False,
use_vertex_groups=False,
use_blen_objects=True,
group_by_object=False,
group_by_material=False,
keep_vertex_order=False,
)
#print just for reference
# print(fakeFile.getvalue())
#get the path to xatlas
#file_path = os.path.dirname(os.path.abspath(__file__))
scriptsDir = bpy.utils.user_resource('SCRIPTS', "addons")
file_path = os.path.join(scriptsDir, "blender_xatlas")
if platform.system() == "Windows":
xatlas_path = os.path.join(file_path, "xatlas", "xatlas-blender.exe")
elif platform.system() == "Linux":
xatlas_path = os.path.join(file_path, "xatlas", "xatlas-blender")
#need to set permissions for the process on linux
subprocess.Popen(
'chmod u+x "' + xatlas_path + '"',
shell=True
)
#setup the arguments to be passed to xatlas-------------------
arguments_string = ""
for argumentKey in packOptions.__annotations__.keys():
key_string = str(argumentKey)
if argumentKey is not None:
print(getattr(packOptions,key_string))
attrib = getattr(packOptions,key_string)
if type(attrib) == bool:
if attrib == True:
arguments_string = arguments_string + " -" + str(argumentKey)
else:
arguments_string = arguments_string + " -" + str(argumentKey) + " " + str(attrib)
for argumentKey in chartOptions.__annotations__.keys():
if argumentKey is not None:
key_string = str(argumentKey)
print(getattr(chartOptions,key_string))
attrib = getattr(chartOptions,key_string)
if type(attrib) == bool:
if attrib == True:
arguments_string = arguments_string + " -" + str(argumentKey)
else:
arguments_string = arguments_string + " -" + str(argumentKey) + " " + str(attrib)
#add pack only option
if sharedProperties.packOnly:
arguments_string = arguments_string + " -packOnly"
arguments_string = arguments_string + " -atlasLayout" + " " + sharedProperties.atlasLayout
print(arguments_string)
#END setup the arguments to be passed to xatlas-------------------
#RUN xatlas process
xatlas_process = subprocess.Popen(
r'"{}"'.format(xatlas_path) + ' ' + arguments_string,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
shell=True
)
print(xatlas_path)
#shove the fake file in stdin
stdin = xatlas_process.stdin
value = bytes(fakeFile.getvalue() + "\n", 'UTF-8') #The \n is needed to end the input properly
stdin.write(value)
stdin.flush()
#Get the output from xatlas
outObj = ""
while True:
output = xatlas_process.stdout.readline()
if not output:
break
outObj = outObj + (output.decode().strip() + "\n")
#the objects after xatlas processing
# print(outObj)
#Setup for reading the output
@dataclass
class uvObject:
obName: string = ""
uvArray: List[float] = field(default_factory=list)
faceArray: List[int] = field(default_factory=list)
convertedObjects = []
uvArrayComplete = []
#search through the out put for STARTOBJ
#then start reading the objects
obTest = None
startRead = False
for line in outObj.splitlines():
line_split = line.split()
if not line_split:
continue
line_start = line_split[0] # we compare with this a _lot_
# print(line_start)
if line_start == "STARTOBJ":
print("Start reading the objects----------------------------------------")
startRead = True
# obTest = uvObject()
if startRead:
#if it's a new obj
if line_start == 'o':
#if there is already an object append it
if obTest is not None:
convertedObjects.append(obTest)
obTest = uvObject() #create new uv object
obTest.obName = line_split[1]
if obTest is not None:
#the uv coords
if line_start == 'vt':
newUv = [float(line_split[1]),float(line_split[2])]
obTest.uvArray.append(newUv)
uvArrayComplete.append(newUv)
#the face coords index
#faces are 1 indexed
if line_start == 'f':
#vert/uv/normal
#only need the uvs
newFace = [
int(line_split[1].split("/")[1]),
int(line_split[2].split("/")[1]),
int(line_split[3].split("/")[1])
]
obTest.faceArray.append(newFace)
#append the final object
convertedObjects.append(obTest)
print(convertedObjects)
#apply the output-------------------------------------------------------------
#copy the uvs to the original objects
# objIndex = 0
print("Applying the UVs----------------------------------------")
# print(convertedObjects)
for importObject in convertedObjects:
bpy.ops.object.select_all(action='DESELECT')
obTest = importObject
obTest.obName = safe_dict[obTest.obName] #probably shouldn't just replace it
bpy.context.scene.objects[obTest.obName].select_set(True)
context.view_layer.objects.active = bpy.context.scene.objects[obTest.obName]
bpy.ops.object.mode_set(mode = 'OBJECT')
obj = bpy.context.active_object
me = obj.data
#convert to bmesh to create the new uvs
bm = bmesh.new()
bm.from_mesh(me)
uv_layer = bm.loops.layers.uv.verify()
nFaces = len(bm.faces)
#need to ensure lookup table for some reason?
if hasattr(bm.faces, "ensure_lookup_table"):
bm.faces.ensure_lookup_table()
#loop through the faces
for faceIndex in range(nFaces):
faceGroup = obTest.faceArray[faceIndex]
bm.faces[faceIndex].loops[0][uv_layer].uv = (
uvArrayComplete[faceGroup[0] - 1][0],
uvArrayComplete[faceGroup[0] - 1][1])
bm.faces[faceIndex].loops[1][uv_layer].uv = (
uvArrayComplete[faceGroup[1] - 1][0],
uvArrayComplete[faceGroup[1] - 1][1])
bm.faces[faceIndex].loops[2][uv_layer].uv = (
uvArrayComplete[faceGroup[2] - 1][0],
uvArrayComplete[faceGroup[2] - 1][1])
# objIndex = objIndex + 3
# print(objIndex)
#assign the mesh back to the original mesh
bm.to_mesh(me)
#END apply the output-------------------------------------------------------------
#Start setting the quads back again-------------------------------------------------------------
if sharedProperties.packOnly:
bpy.ops.object.mode_set(mode='EDIT')
bpy.ops.mesh.select_all(action='DESELECT')
bpy.ops.object.mode_set(mode='OBJECT')
for edges in edgeDict:
edgeList = edgeDict[edges]
currentObject = bpy.context.scene.objects[edgeList['object']]
bm = bmesh.new()
bm.from_mesh(currentObject.data)
if hasattr(bm.edges, "ensure_lookup_table"):
bm.edges.ensure_lookup_table()
#assume that all the triangulated edges come after the original edges
newEdges = []
for edge in range(len(edgeList['edges']), len(bm.edges)):
newEdge = bm.edges[edge]
newEdge.select = True
newEdges.append(newEdge)
bmesh.ops.dissolve_edges(bm, edges=newEdges, use_verts=False, use_face_split=False)
bpy.ops.object.mode_set(mode='OBJECT')
bm.to_mesh(currentObject.data)
bm.free()
bpy.ops.object.mode_set(mode='EDIT')
#End setting the quads back again-------------------------------------------------------------
#select the original objects that were selected
for objectName in rename_dict:
if objectName[0] in bpy.context.scene.objects:
current_object = bpy.context.scene.objects[objectName[0]]
current_object.select_set(True)
context.view_layer.objects.active = current_object
bpy.ops.object.mode_set(mode=startingMode)
print("Finished Xatlas----------------------------------------")
return {'FINISHED'}
def transfer_assets(copy, source, destination):
for filename in glob.glob(os.path.join(source, '*.*')):
shutil.copy(filename, destination)
def transfer_load():
load_folder = bpy.path.abspath(os.path.join(os.path.dirname(bpy.data.filepath), bpy.context.scene.TLM_SceneProperties.tlm_load_folder))
lightmap_folder = os.path.join(os.path.dirname(bpy.data.filepath), bpy.context.scene.TLM_EngineProperties.tlm_lightmap_savedir)
print(load_folder)
print(lightmap_folder)
transfer_assets(True, load_folder, lightmap_folder)
Reference in a new issue View git blame Copy permalink