forked from LeenkxTeam/LNXSDK
		
	
		
			
	
	
		
			354 lines
		
	
	
		
			13 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
		
		
			
		
	
	
			354 lines
		
	
	
		
			13 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) | ||
|  | 
 | ||
|  |             bpy.app.driver_namespace["logman"].append("Postpacking: " + str(atlas.name)) | ||
|  | 
 | ||
|  |             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)) | ||
|  | 
 | ||
|  |             #Add rect to bin | ||
|  |             for r in rect: | ||
|  |                 packer[atlas.name].add_rect(*r) | ||
|  | 
 | ||
|  |             print("Rects: " + str(rect)) | ||
|  |             print("Bins:" + str(packer[atlas.name])) | ||
|  | 
 | ||
|  |             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()): | ||
|  | 
 | ||
|  |                 print("Packing atlas at: " + str(rect)) | ||
|  | 
 | ||
|  |                 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("Obj Shape: " + str(src.shape)) | ||
|  |                 print("Atlas shape: " + str(packedAtlas[atlas.name].shape)) | ||
|  | 
 | ||
|  |                 print("Bin Pos: ",x,y,w,h) | ||
|  |                  | ||
|  | 
 | ||
|  |                 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)) | ||
|  | 
 | ||
|  |                 atlasRes = atlas_resolution | ||
|  |                 texRes = rect[3] #Any dimension w/h (square) | ||
|  |                 ratio = texRes/atlasRes | ||
|  |                 scaleUV(obj.data.uv_layers.active, (ratio, ratio), (0,1)) | ||
|  |                 print(rect) | ||
|  |                  | ||
|  |                 #Postpack error here... | ||
|  |                 for uv_verts in obj.data.uv_layers.active.data: | ||
|  |                     #For each vert | ||
|  | 
 | ||
|  |                     #NOTES! => X FUNKER | ||
|  |                     #TODO => Y | ||
|  | 
 | ||
|  |                     #[0] = bin index | ||
|  |                     #[1] = x | ||
|  |                     #[2] = y (? + 1) | ||
|  |                     #[3] = w | ||
|  |                     #[4] = h | ||
|  | 
 | ||
|  |                     vertex_x = uv_verts.uv[0] + (rect[1]/atlasRes) #WORKING! | ||
|  |                     vertex_y = uv_verts.uv[1] - (rect[2]/atlasRes) # + ((rect[2]-rect[4])/atlasRes) #            #         + (1-((rect[1]-rect[4])/atlasRes)) | ||
|  |                     #tr = "X: {0} + ({1}/{2})".format(uv_verts.uv[0],rect[1],atlasRes) | ||
|  |                     #print(tr) | ||
|  |                     #vertex_y = 1 - (uv_verts.uv[1]) uv_verts.uv[1] + (rect[1]/atlasRes) | ||
|  | 
 | ||
|  |                     #SET UV LAYER TO  | ||
|  | 
 | ||
|  |                     # atlasRes = atlas_resolution | ||
|  |                     # texRes = rect[3] #Any dimension w/h (square) | ||
|  |                     # print(texRes) | ||
|  |                     # #texRes = 0.0,0.0 | ||
|  |                     # #x,y,w,z = x,y,texRes,texRes | ||
|  |                     # x,y,w,z = x,y,0,0 | ||
|  |                      | ||
|  |                     # 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)) | ||
|  | 
 | ||
|  |                     #TO FIX: | ||
|  |                     #SELECT ALL | ||
|  |                     #Scale Y => -1 | ||
|  |                      | ||
|  |                     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 | ||
|  | 
 | ||
|  |                                         #print("Seeking for: " + atlasImage.filepath_raw) | ||
|  |                                         #print(x) | ||
|  | 
 | ||
|  |                                         if(os.path.exists(os.path.join(lightmap_directory, obj.name + end + formatEnc))): | ||
|  |                                             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) |