armory/blender/arm/lightmapper/utility/pack.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

309 lines
11 KiB
Python

import bpy, os, sys, math, mathutils, importlib
import numpy as np
from . rectpack import newPacker, PackingMode, PackingBin
def postpack():
cv_installed = False
cv2 = importlib.util.find_spec("cv2")
if cv2 is None:
print("CV2 not found - Ignoring postpacking")
return 0
else:
cv2 = importlib.__import__("cv2")
cv_installed = True
if cv_installed:
lightmap_directory = os.path.join(os.path.dirname(bpy.data.filepath), bpy.context.scene.TLM_EngineProperties.tlm_lightmap_savedir)
packedAtlas = {}
#TODO - TEST WITH ONLY 1 ATLAS AT FIRST (1 Atlas for each, but only 1 bin (no overflow))
#PackedAtlas = Packer
#Each atlas has bins
#Each bins has rects
#Each rect corresponds to a pack_object
scene = bpy.context.scene
sceneProperties = scene.TLM_SceneProperties
end = "_baked"
if sceneProperties.tlm_denoise_use:
end = "_denoised"
if sceneProperties.tlm_filtering_use:
end = "_filtered"
formatEnc = ".hdr"
image_channel_depth = cv2.IMREAD_ANYDEPTH
linear_straight = False
if sceneProperties.tlm_encoding_use and scene.TLM_EngineProperties.tlm_bake_mode != "Background":
if sceneProperties.tlm_encoding_device == "CPU":
if sceneProperties.tlm_encoding_mode_a == "HDR":
if sceneProperties.tlm_format == "EXR":
formatEnc = ".exr"
if sceneProperties.tlm_encoding_mode_a == "RGBM":
formatEnc = "_encoded.png"
image_channel_depth = cv2.IMREAD_UNCHANGED
else:
if sceneProperties.tlm_encoding_mode_b == "HDR":
if sceneProperties.tlm_format == "EXR":
formatEnc = ".exr"
if sceneProperties.tlm_encoding_mode_b == "LogLuv":
formatEnc = "_encoded.png"
image_channel_depth = cv2.IMREAD_UNCHANGED
linear_straight = True
if sceneProperties.tlm_encoding_mode_b == "RGBM":
formatEnc = "_encoded.png"
image_channel_depth = cv2.IMREAD_UNCHANGED
if sceneProperties.tlm_encoding_mode_b == "RGBD":
formatEnc = "_encoded.png"
image_channel_depth = cv2.IMREAD_UNCHANGED
packer = {}
for atlas in bpy.context.scene.TLM_PostAtlasList: #For each atlas
packer[atlas.name] = newPacker(PackingMode.Offline, PackingBin.BFF, rotation=False)
if scene.TLM_EngineProperties.tlm_setting_supersample == "2x":
supersampling_scale = 2
elif scene.TLM_EngineProperties.tlm_setting_supersample == "4x":
supersampling_scale = 4
else:
supersampling_scale = 1
atlas_resolution = int(int(atlas.tlm_atlas_lightmap_resolution) / int(scene.TLM_EngineProperties.tlm_resolution_scale) * int(supersampling_scale))
packer[atlas.name].add_bin(atlas_resolution, atlas_resolution, 1)
#AtlasList same name prevention?
rect = []
#For each object that targets the atlas
for obj in bpy.context.scene.objects:
if obj.TLM_ObjectProperties.tlm_mesh_lightmap_use:
if obj.TLM_ObjectProperties.tlm_postpack_object:
if obj.TLM_ObjectProperties.tlm_postatlas_pointer == atlas.name:
res = int(int(obj.TLM_ObjectProperties.tlm_mesh_lightmap_resolution) / int(scene.TLM_EngineProperties.tlm_resolution_scale) * int(supersampling_scale))
rect.append((res, res, obj.name))
for r in rect:
packer[atlas.name].add_rect(*r)
packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 3), dtype="float32")
#Continue here...overwrite value if using 8-bit encoding
if sceneProperties.tlm_encoding_use:
if sceneProperties.tlm_encoding_device == "CPU":
if sceneProperties.tlm_encoding_mode_a == "RGBM":
packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
if sceneProperties.tlm_encoding_mode_a == "RGBD":
packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
if sceneProperties.tlm_encoding_device == "GPU":
if sceneProperties.tlm_encoding_mode_b == "RGBM":
packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
if sceneProperties.tlm_encoding_mode_b == "RGBD":
packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
if sceneProperties.tlm_encoding_mode_b == "LogLuv":
packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
packer[atlas.name].pack()
for idy, rect in enumerate(packer[atlas.name].rect_list()):
aob = rect[5]
src = cv2.imread(os.path.join(lightmap_directory, aob + end + formatEnc), image_channel_depth) #"_baked.hdr"
print("Obj name is: " + aob)
x,y,w,h = rect[1],rect[2],rect[3],rect[4]
print(src.shape)
print(packedAtlas[atlas.name].shape)
packedAtlas[atlas.name][y:h+y, x:w+x] = src
obj = bpy.data.objects[aob]
for idx, layer in enumerate(obj.data.uv_layers):
if not obj.TLM_ObjectProperties.tlm_use_default_channel:
uv_channel = obj.TLM_ObjectProperties.tlm_uv_channel
else:
uv_channel = "UVMap_Lightmap"
if layer.name == uv_channel:
obj.data.uv_layers.active_index = idx
print("UVLayer set to: " + str(obj.data.uv_layers.active_index))
for uv_verts in obj.data.uv_layers.active.data:
#SET UV LAYER TO
atlasRes = atlas_resolution
texRes = rect[3]
x,y,w,z = x,y,texRes,texRes
ratio = atlasRes/texRes
if x == 0:
x_offset = 0
else:
x_offset = 1/(atlasRes/x)
if y == 0:
y_offset = 0
else:
y_offset = 1/(atlasRes/y)
vertex_x = (uv_verts.uv[0] * 1/(ratio)) + x_offset
vertex_y = (1 - ((uv_verts.uv[1] * 1/(ratio)) + y_offset))
uv_verts.uv[0] = vertex_x
uv_verts.uv[1] = vertex_y
scaleUV(obj.data.uv_layers.active, (1, -1), getBoundsCenter(obj.data.uv_layers.active))
print(getCenter(obj.data.uv_layers.active))
cv2.imwrite(os.path.join(lightmap_directory, atlas.name + end + formatEnc), packedAtlas[atlas.name])
print("Written: " + str(os.path.join(lightmap_directory, atlas.name + end + formatEnc)))
#Change the material for each material, slot
for obj in bpy.context.scene.objects:
if obj.TLM_ObjectProperties.tlm_mesh_lightmap_use:
if obj.TLM_ObjectProperties.tlm_postpack_object:
if obj.TLM_ObjectProperties.tlm_postatlas_pointer == atlas.name:
for slot in obj.material_slots:
nodetree = slot.material.node_tree
for node in nodetree.nodes:
if node.name == "TLM_Lightmap":
existing_image = node.image
atlasImage = bpy.data.images.load(os.path.join(lightmap_directory, atlas.name + end + formatEnc), check_existing=True)
if linear_straight:
if atlasImage.colorspace_settings.name != 'Linear':
atlasImage.colorspace_settings.name = 'Linear'
node.image = atlasImage
os.remove(os.path.join(lightmap_directory, obj.name + end + formatEnc))
existing_image.user_clear()
#Add dilation map here...
for obj in bpy.context.scene.objects:
if obj.TLM_ObjectProperties.tlm_mesh_lightmap_use:
if obj.TLM_ObjectProperties.tlm_postpack_object:
if obj.TLM_ObjectProperties.tlm_postatlas_pointer == atlas.name:
if atlas.tlm_atlas_dilation:
for slot in obj.material_slots:
nodetree = slot.material.node_tree
for node in nodetree.nodes:
if node.name == "TLM_Lightmap":
existing_image = node.image
atlasImage = bpy.data.images.load(os.path.join(lightmap_directory, atlas.name + end + formatEnc), check_existing=True)
img = cv2.imread(atlasImage.filepath_raw, image_channel_depth)
kernel = np.ones((5,5), dtype="float32")
img_dilation = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)
cv2.imshow('Dilation', img_dilation)
cv2.waitKey(0)
print("TODO: Adding dilation for: " + obj.name)
#TODO MASKING OPTION!
else:
print("OpenCV not installed. Skipping postpacking process.")
def getCenter(uv_layer):
total_x, total_y = 0,0
len = 0
for uv_verts in uv_layer.data:
total_x += uv_verts.uv[0]
total_y += uv_verts.uv[1]
len += 1
center_x = total_x / len
center_y = total_y / len
return (center_x, center_y)
def getBoundsCenter(uv_layer):
min_x = getCenter(uv_layer)[0]
max_x = getCenter(uv_layer)[0]
min_y = getCenter(uv_layer)[1]
max_y = getCenter(uv_layer)[1]
len = 0
for uv_verts in uv_layer.data:
if uv_verts.uv[0] < min_x:
min_x = uv_verts.uv[0]
if uv_verts.uv[0] > max_x:
max_x = uv_verts.uv[0]
if uv_verts.uv[1] < min_y:
min_y = uv_verts.uv[1]
if uv_verts.uv[1] > max_y:
max_y = uv_verts.uv[1]
center_x = (max_x - min_x) / 2 + min_x
center_y = (max_y - min_y) / 2 + min_y
return (center_x, center_y)
def scale2D(v, s, p):
return (p[0] + s[0]*(v[0] - p[0]), p[1] + s[1]*(v[1] - p[1]))
def scaleUV( uvMap, scale, pivot ):
for uvIndex in range( len(uvMap.data) ):
uvMap.data[uvIndex].uv = scale2D(uvMap.data[uvIndex].uv, scale, pivot)