Utility for lightmap preparation

addons/io_hubs_addon/components/lightmaps.py

@@ -0,0 +1,188 @@
+import bpy
+import bmesh
+
+# Label for decoy texture
+HUBS_DECOY_IMAGE_TEXTURE = "HUBS_DECOY_IMAGE_TEXTURE"
+
+
+# Find and select the image texture associated with a MOZ_lightmap settings
+def findImageTexture(lightmapNode):
+    for inputs in lightmapNode.inputs:
+        for links in inputs.links:
+            return links.from_node
+    return None
+
+
+# Find the UV map associated with an image texture
+def findUvMap(imageTexture, material):
+    # Search for the parent UV Map
+    for inputs in imageTexture.inputs:
+        for links in inputs.links:
+            # Is this node a MOZ lightmap node?
+            if links.from_node.bl_idname == "ShaderNodeUVMap":
+                return links.from_node.uv_map
+            else:
+                raise ValueError(f"Unexpected node type '{links.from_node.bl_idname}' instead of 'ShaderNodeUVMap' on material '{material.name}'")
+    return None
+
+
+# Selects all the faces of the mesh that have been assigned the given material (important for UV packing for lightmaps)
+def selectMeshFacesFromMaterial(object, mesh, material):
+    materialSlotIndex = object.material_slots.find(material.name)
+    if materialSlotIndex < 0:
+        raise ValueError(f"Failed to find a slot with material '{material.name}' in '{mesh.name}' attached to object '{object.name}'")
+    bm = bmesh.new()
+    bm.from_mesh(object.data)
+    for f in bm.faces:
+        if f.material_index == materialSlotIndex:
+            f.select = True
+    bm.to_mesh(object.data)
+    bm.free()
+
+
+# Select the object that holds this mesh
+def selectObjectFromMesh(mesh, material):
+    for object in bpy.context.scene.objects:
+        if object.type == "MESH":
+            if object.data.name == mesh.name:
+                # Objects cannot be selected if they are hidden
+                object.hide_set(False)
+                object.select_set(True)
+                print(f" --- selected object '{object.name}' because it uses mesh '{mesh.name}'")
+                selectMeshFacesFromMaterial(object, mesh, material)
+
+
+# Select the UV input to the image texture for every mesh that uses the given material
+def selectUvMaps(imageTexture, material):
+    # Select the lightmap UVs on the associated mesh
+    uvMap = findUvMap(imageTexture, material)
+    if uvMap:
+        print(f" -- found UV Map Node '{uvMap}'")
+        # Search for meshes that use this material (can't find a parent property so this will have to do)
+        for mesh in bpy.data.meshes:
+            if mesh.materials.find(material.name) != -1:
+                print(f" -- found mesh '{mesh.name}' that uses this material")
+                selectObjectFromMesh(mesh, material)
+                if mesh.uv_layers.find(uvMap) != -1:
+                    uvLayer = mesh.uv_layers[uvMap]
+                    mesh.uv_layers.active = uvLayer
+                    print(f" --- UV layer '{uvMap}' is now active on '{mesh.name}'")
+                else:
+                    raise ValueError(f"Failed to find UV layer '{uvMap}' for mesh '{mesh.name}' using material '{material.name}'")
+    else:
+        raise ValueError(f"No UV map found for image texture '{imageTexture.name}' with image '{imageTexture.image.name}' in material '{material.name}'")
+
+
+# Selects all MOZ lightmap related components ready for baking
+def selectLightmapComponents(target):
+    # Force UI into OBJECT mode so scripts can manipulate meshes
+    try:
+        bpy.ops.object.mode_set(mode='OBJECT')
+    except Exception as e:
+        print(f"Failed to enter OBJECT mode (usually non-fatal): {str(e)}")
+    # Deslect all objects to start with (bake objects will then be selected)
+    for o in bpy.context.scene.objects:
+        o.select_set(False)
+        # Deselect and show all mesh faces (targetted faces will then be selected)
+        if o.type == "MESH":
+            bm = bmesh.new()
+            bm.from_mesh(o.data)
+            for f in bm.faces:
+                f.hide = False
+                f.select = False
+            bm.to_mesh(o.data)
+            bm.free()
+    # For every material
+    for material in bpy.data.materials:
+        if material.node_tree:
+            # Deactivate and unselect all nodes in the shader graph
+            #  - they can be active even if the UI doesn't show it and they will be baked
+            material.node_tree.nodes.active = None
+            for n in material.node_tree.nodes:
+                n.select = False
+            # For every node in the material graph
+            for shadernode in material.node_tree.nodes:
+                # Is this node a MOZ lightmap node?
+                if shadernode.bl_idname == "moz_lightmap.node":
+                    print(f"found '{shadernode.name}' ({shadernode.label}) on material '{material.name}'")
+                    imageTexture = findImageTexture(shadernode)
+                    if imageTexture:
+                        # Check image texture actually has an image
+                        if imageTexture.image is None:
+                            raise ValueError(f"No image found on image texture '{imageTexture.name}' ('{imageTexture.label}') in material '{material.name}'")
+                        # Is this lightmap texture image being targetted?
+                        if target == "" or target == imageTexture.image.name:
+                            # Select and activate the image texture node so it will be targetted by the bake
+                            imageTexture.select = True
+                            material.node_tree.nodes.active = imageTexture
+                            print(f" - selected image texture '{imageTexture.name}' ({imageTexture.label})")
+
+                            selectUvMaps(imageTexture, material)
+                        else:
+                            print(f" - ignoring image texture '{imageTexture.name}' because it uses image '{imageTexture.image.name}' and the target is '{target}'")
+                    else:
+                        raise ValueError(f"No image texture found on material '{material.name}'")
+                # Is it a decoy node?
+                elif shadernode.bl_idname == "ShaderNodeTexImage" and shadernode.label == "HUBS_DECOY_IMAGE_TEXTURE":
+                    # Select and activate the image texture node so it will be targetted by the bake
+                    shadernode.select = True
+                    material.node_tree.nodes.active = shadernode
+
+
+# List all the lightmap textures images
+def listLightmapImages():
+    result = set()
+    # For every material
+    for material in bpy.data.materials:
+        if material.node_tree:
+            # For every node in the material graph
+            for shadernode in material.node_tree.nodes:
+                # Is this node a MOZ lightmap node?
+                if shadernode.bl_idname == "moz_lightmap.node":
+                    imageTexture = findImageTexture(shadernode)
+                    if imageTexture:
+                        if imageTexture.image:
+                            result.add(imageTexture.image)
+    return result
+
+
+# Check for selected objects with non-lightmapped materials. They might get baked and have corrupted image textures
+def assertSelectedObjectsAreSafeToBake(addDecoyRatherThanThrow):
+    decoyCount = 0
+    for object in bpy.context.scene.objects:
+        if object.select_get():
+            for materialSlot in object.material_slots:
+                material = materialSlot.material
+                hasLightmapNode = False
+                hasAtRiskImageTexture = False
+                hasDecoyImageTexture = False
+                for shadernode in material.node_tree.nodes:
+                    if shadernode.bl_idname == "moz_lightmap.node":
+                        hasLightmapNode = True
+                    elif shadernode.bl_idname == "ShaderNodeTexImage":
+                        if shadernode.label == HUBS_DECOY_IMAGE_TEXTURE:
+                            hasDecoyImageTexture = True
+                        else:
+                            hasAtRiskImageTexture = True
+                if hasAtRiskImageTexture and not (hasLightmapNode or hasDecoyImageTexture):
+                    if addDecoyRatherThanThrow:
+                        decoy_image_texture = material.node_tree.nodes.new("ShaderNodeTexImage")
+                        decoy_image_texture.label = HUBS_DECOY_IMAGE_TEXTURE
+                        decoyCount += 1
+                    else:
+                        raise ValueError(f"Multi-material object '{object.name}' uses '{materialSlot.name}' with no lightmap or decoy texture. It will be corrupted by baking")
+    return decoyCount
+
+
+def removeAllDecoyImageTextures():
+    decoyCount = 0
+    # For every material
+    for material in bpy.data.materials:
+        if material.node_tree:
+            # For every node in the material graph
+            for shadernode in material.node_tree.nodes:
+                # Is this a decoy image?
+                if shadernode.bl_idname == "ShaderNodeTexImage" and shadernode.label == HUBS_DECOY_IMAGE_TEXTURE:
+                    material.node_tree.nodes.remove(shadernode)
+                    decoyCount += 1
+    return decoyCount

addons/io_hubs_addon/components/operators.py

@@ -11,6 +11,7 @@
 from .gizmos import update_gizmos
 from .utils import is_linked, redraw_component_ui
 import os
+from . import lightmaps
 
 
 class AddHubsComponent(Operator):
@@ -625,6 +626,51 @@ def invoke(self, context, event):
         return {'RUNNING_MODAL'}
 
 
+class PrepareHubsLightmaps(Operator):
+    bl_idname = "wm.prepare_hubs_lightmaps"
+    bl_label = "Select all lightmap elements for baking"
+    bl_description = "Select all MOZ_lightmap input textures, the matching mesh UV layer, and the objects ready for baking"
+
+    target: StringProperty(name="target")
+
+    def execute(self, context):
+        try:
+            lightmaps.selectLightmapComponents(self.target)
+            lightmaps.assertSelectedObjectsAreSafeToBake(False)
+            self.report({'INFO'}, "Lightmaps prepared and ready to bake")
+        except Exception as e:
+            self.report({'ERROR'}, str(e))
+        return {'FINISHED'}
+
+
+class AddDecoyImageTextures(Operator):
+    bl_idname = "wm.add_decoy_image_textures"
+    bl_label = "Add to Selected"
+    bl_description = "Adds decoy image textures to the materials of selected meshes if they are required"
+
+    def execute(self, context):
+        try:
+            decoyCount = lightmaps.assertSelectedObjectsAreSafeToBake(True)
+            self.report({'INFO'}, f"Added {decoyCount} decoy textures")
+        except Exception as e:
+            self.report({'ERROR'}, str(e))
+        return {'FINISHED'}
+
+
+class RemoveDecoyImageTextures(Operator):
+    bl_idname = "wm.remove_decoy_image_textures"
+    bl_label = "Remove All"
+    bl_description = "Remove all decoy image textures from materials regardless of selection"
+
+    def execute(self, context):
+        try:
+            decoyCount = lightmaps.removeAllDecoyImageTextures()
+            self.report({'INFO'}, f"Removed {decoyCount} decoy textures")
+        except Exception as e:
+            self.report({'ERROR'}, str(e))
+        return {'FINISHED'}
+
+
 def register():
     bpy.utils.register_class(AddHubsComponent)
     bpy.utils.register_class(RemoveHubsComponent)
@@ -636,6 +682,9 @@ def register():
     bpy.utils.register_class(ViewReportInInfoEditor)
     bpy.utils.register_class(CopyHubsComponent)
     bpy.utils.register_class(OpenImage)
+    bpy.utils.register_class(PrepareHubsLightmaps)
+    bpy.utils.register_class(AddDecoyImageTextures)
+    bpy.utils.register_class(RemoveDecoyImageTextures)
     bpy.types.WindowManager.hubs_report_scroll_index = IntProperty(default=0, min=0)
     bpy.types.WindowManager.hubs_report_scroll_percentage = IntProperty(
         name="Scroll Position", default=0, min=0, max=100, subtype='PERCENTAGE')
@@ -654,6 +703,9 @@ def unregister():
     bpy.utils.unregister_class(ViewReportInInfoEditor)
     bpy.utils.unregister_class(CopyHubsComponent)
     bpy.utils.unregister_class(OpenImage)
+    bpy.utils.unregister_class(PrepareHubsLightmaps)
+    bpy.utils.unregister_class(AddDecoyImageTextures)
+    bpy.utils.unregister_class(RemoveDecoyImageTextures)
     del bpy.types.WindowManager.hubs_report_scroll_index
     del bpy.types.WindowManager.hubs_report_scroll_percentage
     del bpy.types.WindowManager.hubs_report_last_title

addons/io_hubs_addon/components/ui.py

@@ -3,6 +3,8 @@
 from .types import PanelType
 from .components_registry import get_component_by_name, get_components_registry
 from .utils import get_object_source, dash_to_title, is_linked
+from . import operators
+from . import lightmaps
 
 
 def draw_component_global(panel, context):
@@ -179,6 +181,31 @@ def draw(self, context):
         draw_components_list(self, context)
 
 
+class HubsRenderPanel(bpy.types.Panel):
+    bl_label = 'Hubs'
+    bl_idname = "RENDER_PT_hubs"
+    bl_space_type = 'PROPERTIES'
+    bl_region_type = 'WINDOW'
+    bl_context = 'render'
+
+    def draw(self, context):
+        layout = self.layout
+        row = layout.row()
+        lightmapImages = sorted(lightmaps.listLightmapImages(), key=lambda it: it.name)
+        # Only show the panel if there are any lightmaps to prepare
+        if len(lightmapImages) > 0:
+            row.operator(operators.PrepareHubsLightmaps.bl_idname).target = ""
+        # Is their more than 1 lightmap texture?
+        if len(lightmapImages) > 1:
+            for lightmapImage in lightmapImages:
+                row = layout.row()
+                row.operator(operators.PrepareHubsLightmaps.bl_idname, text=f"Select '{lightmapImage.name}' elements for packing").target = lightmapImage.name
+        row = layout.row()
+        row.label(text="Decoy Textures")
+        row.operator(operators.AddDecoyImageTextures.bl_idname)
+        row.operator(operators.RemoveDecoyImageTextures.bl_idname)
+
+
 class TooltipLabel(bpy.types.Operator):
     bl_idname = "ui.hubs_tooltip_label"
     bl_label = "---"
@@ -213,6 +240,7 @@ def gizmo_display_popover_addition(self, context):
 
 
 def register():
+    bpy.utils.register_class(HubsRenderPanel)
     bpy.utils.register_class(HubsObjectPanel)
     bpy.utils.register_class(HubsScenePanel)
     bpy.utils.register_class(HubsMaterialPanel)
@@ -225,6 +253,7 @@ def register():
 
 
 def unregister():
+    bpy.utils.unregister_class(HubsRenderPanel)
     bpy.utils.unregister_class(HubsObjectPanel)
     bpy.utils.unregister_class(HubsScenePanel)
     bpy.utils.unregister_class(HubsMaterialPanel)