Fix bake node cleanup, add logging

Makefile

@@ -14,7 +14,7 @@ install:
 	cp -r quickbake/* $(INSTALL_DIR)/
 
 run: install
-	blender "example/Metal Box.blend"
+	blender --debug "example/Metal Box.blend"
 
 checks: lint test
 

quickbake/__init__.py

@@ -1,11 +1,22 @@
+import logging
+
 import bpy
 
 from .op import RENDER_OT_bake
 from .panel import RENDER_PT_main
 from .properties import QuickBakeToolPropertyGroup
+from .util import is_development
+
+
+_log = logging.getLogger(__name__)
+
+if is_development():
+    _log.setLevel(logging.DEBUG)
+    _log.addHandler(logging.StreamHandler())
 
 
 def register():
+    _log.debug("Register extension")
     bpy.utils.register_class(RENDER_OT_bake)
     bpy.utils.register_class(RENDER_PT_main)
     bpy.utils.register_class(QuickBakeToolPropertyGroup)
@@ -15,6 +26,7 @@ def register():
 
 
 def unregister():
+    _log.debug("Unregister extension")
     bpy.utils.unregister_class(RENDER_OT_bake)
     bpy.utils.unregister_class(RENDER_PT_main)
     bpy.utils.unregister_class(QuickBakeToolPropertyGroup)

quickbake/op.py

@@ -1,8 +1,12 @@
+import logging
+
 import bpy
 from bpy_extras.node_shader_utils import PrincipledBSDFWrapper
 
 from .properties import MaterialMode, QuickBakeToolPropertyGroup
 
+_log = logging.getLogger(__name__)
+
 
 class RENDER_OT_bake(bpy.types.Operator):
     """Do the bake."""
@@ -48,12 +52,15 @@ def poll(cls, context):
         return obj is not None and obj.type == "MESH"
 
     def execute(self, context: bpy.types.Context):
+        _log.info("Begin execution")
+
         # Keeping type hints happy, should not be possible
         scene = context.scene
         assert scene is not None, "Context must have a scene, got None"
 
         # Make sure cycles is the current render engine
         if scene.render.engine != "CYCLES":
+            _log.info("Setting render engine to cycles")
             scene.render.engine = "CYCLES"  # type: ignore
             self.report({"WARNING"}, "Changed render engine to Cycles")
 
@@ -64,11 +71,13 @@ def execute(self, context: bpy.types.Context):
 
         # This should be enforces by cls.poll() but is here to be sure
         if obj is None:
+            _log.error("No active object")
             self.report({"ERROR"}, "No active object")
             return {"CANCELLED"}  # canceled because nothing was altered / needs undo
 
         # This should be enforces by cls.poll() but is here to be sure
         if obj.type != "MESH":
+            _log.error("Expected active object to be mesh, got %s", obj.type)
             self.report({"ERROR"}, "Active object must be a mesh")
             return {"CANCELLED"}  # canceled because nothing was altered / needs undo
 
@@ -101,29 +110,40 @@ def execute(self, context: bpy.types.Context):
         if props.uv_enabled:
             passes.append(("UV", True))
 
+        _log.debug("Render passes will be %s", passes)
+
         # Keeping type hints happy
         assert isinstance(obj.data, bpy.types.Mesh), "Object is not a mesh"
         mesh = obj.data
 
         uv_layer = self.unwrap_object(mesh)
         bake_nodes = self.create_image_nodes(mesh)
-        images = {}
+        images: dict[str, bpy.types.Image] = {}
 
         for layer, is_data in passes:
+            _log.info("Starting layer %s", layer)
             self.report({"INFO"}, f"Starting layer {layer}")
 
             image_name = f"{props.bake_name}_{layer.lower()}"
 
             # Create image or use existing
             img = bpy.data.images.get(image_name)
             if img is None:
+                _log.info("Creating image %s", image_name)
                 img = bpy.data.images.new(
                     image_name, props.bake_size, props.bake_size, is_data=is_data
                 )
+            else:
+                _log.debug("Using existing image %s", image_name)
             images[layer] = img
 
             # Assign image to bake node in all materials
             for mat, texture_node in bake_nodes:
+                _log.debug(
+                    "Assigning image to texture node %s in material %s",
+                    texture_node.name,
+                    mat.name,
+                )
                 # TODO type ignore if it works
                 texture_node.image = img  # type: ignore
                 texture_node.select = True
@@ -133,9 +153,11 @@ def execute(self, context: bpy.types.Context):
             filepath = ""
             save_mode = "INTERNAL"
             if props.save_img:
-                filepath = f"{props.save_path}/{props.bake_name}_{layer}"
                 save_mode = "EXTERNAL"
+                filepath = f"{props.save_path}/{props.bake_name}_{layer}"
+                _log.debug("Images will be saved externally to %s", filepath)
 
+            _log.info("Starting bake for layer %s", layer)
             bpy.ops.object.bake(
                 type=layer,  # type: ignore
                 pass_filter={"COLOR"},  # TODO change this for other textures
@@ -144,6 +166,7 @@ def execute(self, context: bpy.types.Context):
                 save_mode=save_mode,
                 filepath=filepath,
             )
+            _log.info("Finished bake for layer %s", layer)
 
         self.cleanup_image_nodes(mesh)
 
@@ -155,47 +178,62 @@ def execute(self, context: bpy.types.Context):
 
         # Duplicate object and assign material to new
         if props.mat_mode == MaterialMode.DUPLICATE:
+            _log.info("Duplicating object before assigning material")
             bpy.ops.object.duplicate()
+            _log.debug("Hiding original object %s", obj.name)
             obj.hide_set(True)
             # Get new object
             obj = context.active_object
+            _log.debug(
+                "New object is named %s", obj.name if obj is not None else "None"
+            )
 
             # Keeping type hints happy
             assert obj is not None, "Object is None"
             assert isinstance(obj.data, bpy.types.Mesh), "Object is not a mesh"
 
         # Assign or Copy
         if props.mat_mode != MaterialMode.CREATE:
+            _log.info("Assigning material %s to object %s", mat.name, obj.name)
             obj.data.materials.clear()
             obj.active_material = mat
 
+        _log.info("Finished execution")
         return {"FINISHED"}
 
     def unwrap_object(self, mesh: bpy.types.Mesh) -> bpy.types.MeshUVLoopLayer:
         uv_name = "bake_uv"
+        _log.debug("Unwrapping mesh %s with uv layer %s", mesh.name, uv_name)
 
         # Use existing or create new uv layer for baking
         bake_uv = mesh.uv_layers.get(uv_name)
         if bake_uv is None:
+            _log.info("Creating new uv layer %s", uv_name)
             bake_uv = mesh.uv_layers.new(name=uv_name)
+        else:
+            _log.debug("Reusing existing uv layer %s", uv_name)
 
         # Store currently active layer
         active_layer = None
         for layer in mesh.uv_layers:
             if layer.active:
                 active_layer = layer
+                _log.debug("Currently active uv layer is %s", active_layer.name)
                 break
 
+        _log.debug("Start unwrapping mesh %s", mesh.name)
         # Unwrap the object
         bake_uv.active = True
         bpy.ops.object.mode_set(mode="EDIT")
         bpy.ops.mesh.select_all(action="SELECT")
         bpy.ops.uv.smart_project(island_margin=0.001)
         bpy.ops.object.mode_set(mode="OBJECT")
         bake_uv.active = False
+        _log.debug("Finished unwrapping mesh %s", mesh.name)
 
         # Restore active layer
         if active_layer is not None:
+            _log.debug("Restoring active uv layer %s", active_layer.name)
             active_layer.active = True
 
         return bake_uv
@@ -211,19 +249,27 @@ def create_image_nodes(
 
         for mat in mesh.materials:
             if mat is None or mat.node_tree is None:
+                _log.warning("Found null material in mesh %s", mesh.name)
                 null_count += 1
                 continue
 
+            _log.debug("Enabling nodes for material %s", mat.name)
             # Enable nodes if not already
             mat.use_nodes = True
 
             texture_node = mat.node_tree.get(node_name)
             if texture_node is None:
+                _log.info("Creating texture node for material %s", mat.name)
                 texture_node = mat.node_tree.nodes.new("ShaderNodeTexImage")
                 texture_node.name = node_name
 
+            else:
+                _log.debug("Using existing texture node %s", texture_node.name)
+
             image_nodes.append((mat, texture_node))
 
+        _log.info("Created %d nodes in mesh %s", len(image_nodes), mesh.name)
+
         # Notify user if any materials were unusable
         if null_count > 0:
             self.report({"WARNING"}, f"Mesh {mesh.name} has {null_count} null material")
@@ -232,25 +278,37 @@ def create_image_nodes(
 
     def cleanup_image_nodes(self, mesh: bpy.types.Mesh):
         node_name = "bake_image"
+        _log.info("Cleaning up bake texture node %s in msh %s", node_name, mesh.name)
 
         for mat in mesh.materials:
             if mat is None or mat.node_tree is None:
+                _log.warning("Found null material in mesh %s", mesh.name)
                 continue
 
-            node = mat.node_tree.get(node_name)
+            node = mat.node_tree.nodes.get(node_name)
             if node is not None:
+                _log.debug("Removing node %s from material %s", node.name, mat.name)
                 mat.node_tree.nodes.remove(node)
+            else:
+                _log.warning(
+                    "Failed to find node %s in material %s", node_name, mat.name
+                )
+                _log.debug(
+                    "Material %s has nodes %s", mat.name, list(mat.node_tree.nodes)
+                )
+                self.report({"WARNING"}, f"Failed to cleanup material {mat.name}")
 
     def create_material(
         self,
         props: QuickBakeToolPropertyGroup,
         uv_layer: bpy.types.MeshUVLoopLayer,
         passes: list[tuple[str, bool]],
-        images: dict,
+        images: dict[str, bpy.types.Image],
     ):
         # Create Material
         mat = bpy.data.materials.get(props.bake_name)
         if mat is None:
+            _log.debug("Creating material %s", props.bake_name)
             mat = bpy.data.materials.new(props.bake_name)
             mat.use_nodes = True
 
@@ -267,13 +325,15 @@ def create_material(
         # Texture coordinate node for uv map
         uv_node = shader_nodes.get("Texture Coordinate")
         if uv_node is None:
+            _log.debug("Creating uv node in material %s", mat.name)
             uv_node = shader_nodes.new(type="ShaderNodeUVMap")
             uv_node.location.x = -1100
         uv_node.uv_map = uv_layer.name  # type: ignore
 
         # Mapping node for position, scale, rotation
         mapping_node = shader_nodes.get("Texture Coordinate")
         if mapping_node is None:
+            _log.debug("Creating mapping node in material %s", mat.name)
             mapping_node = shader_nodes.new(type="ShaderNodeMapping")
             mapping_node.location.x = -900
 
@@ -287,31 +347,62 @@ def create_material(
 
             tex_node = mat.node_tree.get(layer)
             if tex_node is None:
+                _log.debug(
+                    "Creating image texture node for layer %s in material %s",
+                    layer,
+                    mat.name,
+                )
                 tex_node = shader_nodes.new(type="ShaderNodeTexImage")
                 tex_node.location.x = -700
                 tex_node.location.y = y
 
             tex_node.image = images[layer]  # type: ignore
+            _log.debug(
+                "Assigned image %s to image texture node %s",
+                images[layer].name,
+                tex_node.name,
+            )
             links.new(mapping_node.outputs["Vector"], tex_node.inputs["Vector"])
+            _log.debug(
+                "Linked mapping node %s output to image texture %s vector input",
+                mapping_node.name,
+                tex_node.name,
+            )
 
-            shader_input = self.layer_input_map.get(layer, "")
-            if shader_input:
+            if shader_input := self.layer_input_map.get(layer):
                 if layer == "NORMAL":
                     normal_map_node = shader_nodes.get("Normal Map")
                     if normal_map_node is None:
+                        _log.debug("Creating normal map node in material %s", mat.name)
                         normal_map_node = shader_nodes.new(type="ShaderNodeNormalMap")
                         normal_map_node.location.x = -400
                         normal_map_node.location.y = y
 
+                    _log.debug(
+                        "Linking image texture %s color output to normal map %s color input",
+                        tex_node.name,
+                        normal_map_node.name,
+                    )
                     links.new(
                         tex_node.outputs["Color"], normal_map_node.inputs["Color"]
                     )
+
+                    _log.debug(
+                        "Linking normal map %s color output to shader %s input",
+                        normal_map_node.name,
+                        shader_input,
+                    )
                     links.new(
                         normal_map_node.outputs["Normal"],
                         principled_node.inputs[shader_input],
                     )
 
                 else:
+                    _log.debug(
+                        "Connecting layer image texture for layer %s to shader input %s",
+                        layer,
+                        shader_input,
+                    )
                     links.new(
                         tex_node.outputs["Color"], principled_node.inputs[shader_input]
                     )

quickbake/util.py

@@ -0,0 +1,14 @@
+import os
+
+import tomllib
+
+
+def _get_version():
+    if os.path.exists("quickbake/blender_manifest.toml"):
+        with open("quickbake/blender_manifest.toml", "rb") as f:
+            data = tomllib.load(f)
+            return data.get("version")
+
+
+def is_development():
+    return _get_version() == "0.0.0"