from __future__ import annotations import colorsys import json from pathlib import Path from typing import Any import bpy from bpy.props import ( BoolProperty, CollectionProperty, EnumProperty, FloatVectorProperty, IntProperty, PointerProperty, StringProperty, ) from bpy.types import Operator, Panel, PropertyGroup _TEMP_PREFIX = "__MID_BAKER_" def _distinct_color(index: int) -> tuple[float, float, float]: """Return a deterministic, vivid linear RGB color for a zero-based slot.""" hue = (0.03 + index * 0.618033988749895) % 1.0 saturation = 0.72 if index % 2 == 0 else 0.9 value = 1.0 if index % 3 else 0.82 return colorsys.hsv_to_rgb(hue, saturation, value) def _index_color(index: int) -> tuple[float, float, float]: """Encode a one-based material ID into 24 bits, reserving black for empty pixels.""" value = index + 1 return ( (value & 0xFF) / 255.0, ((value >> 8) & 0xFF) / 255.0, ((value >> 16) & 0xFF) / 255.0, ) def _slot_color( index: int, material: bpy.types.Material | None, palette_mode: str, ) -> tuple[float, float, float]: if palette_mode == "MATERIAL" and material is not None: return tuple(material.diffuse_color[:3]) if palette_mode == "INDEX": return _index_color(index) return _distinct_color(index) def _create_id_material( slot_index: int, color: tuple[float, float, float], image: bpy.types.Image, ) -> bpy.types.Material: material = bpy.data.materials.new(f"{_TEMP_PREFIX}Material_{slot_index + 1}") material.use_nodes = True nodes = material.node_tree.nodes nodes.clear() output = nodes.new("ShaderNodeOutputMaterial") output.location = (360.0, 0.0) emission = nodes.new("ShaderNodeEmission") emission.location = (80.0, 0.0) emission.inputs["Color"].default_value = (*color, 1.0) emission.inputs["Strength"].default_value = 1.0 image_node = nodes.new("ShaderNodeTexImage") image_node.location = (-260.0, -180.0) image_node.image = image image_node.select = True nodes.active = image_node material.node_tree.links.new(emission.outputs["Emission"], output.inputs["Surface"]) return material def _create_baked_ids_material(image: bpy.types.Image) -> bpy.types.Material: """Create an export-friendly, single material displaying the baked ID image.""" material = bpy.data.materials.new("baked ids") material.use_nodes = True nodes = material.node_tree.nodes nodes.clear() output = nodes.new("ShaderNodeOutputMaterial") output.location = (560.0, 0.0) principled = nodes.new("ShaderNodeBsdfPrincipled") principled.location = (260.0, 0.0) principled.inputs["Roughness"].default_value = 1.0 image_node = nodes.new("ShaderNodeTexImage") image_node.label = "Baked Material IDs" image_node.location = (-120.0, 0.0) image_node.image = image image_node.interpolation = "Closest" nodes.active = image_node material.node_tree.links.new(image_node.outputs["Color"], principled.inputs["Base Color"]) material.node_tree.links.new(principled.outputs["BSDF"], output.inputs["Surface"]) return material def _create_export_copy( source: bpy.types.Object, baked_mesh: bpy.types.Mesh, image: bpy.types.Image, apply_modifiers: bool, target_collection: bpy.types.Collection, ) -> tuple[bpy.types.Object, bpy.types.Mesh, bpy.types.Material]: """Create a persistent one-material copy suitable for later export.""" if apply_modifiers: export_mesh = baked_mesh.copy() export_object = bpy.data.objects.new( f"{source.name}_Baked_IDs", export_mesh, ) export_object.matrix_world = source.matrix_world.copy() else: export_object = source.copy() export_mesh = source.data.copy() export_object.data = export_mesh export_object.name = f"{source.name}_Baked_IDs" target_collection.objects.link(export_object) export_object.hide_set(False) export_object.hide_viewport = False export_object.hide_render = False material = _create_baked_ids_material(image) export_mesh.materials.clear() export_mesh.materials.append(material) for polygon in export_mesh.polygons: polygon.material_index = 0 return export_object, export_mesh, material def _normalise_png_path(filepath: str) -> Path: path = Path(bpy.path.abspath(filepath)).expanduser() if path.suffix.lower() != ".png": path = path.with_suffix(".png") return path def _normalise_glb_path(filepath: str) -> Path: path_text = filepath.strip() if not path_text: raise RuntimeError("Choose a GLB export path") path = Path(bpy.path.abspath(path_text)).expanduser() if path.suffix.lower() != ".glb": path = path.with_suffix(".glb") return path def _quick_export_glb( context: bpy.types.Context, export_object: bpy.types.Object, filepath: str, ) -> Path: output_path = _normalise_glb_path(filepath) output_path.parent.mkdir(parents=True, exist_ok=True) for obj in list(context.selected_objects): obj.select_set(False) export_object.select_set(True) context.view_layer.objects.active = export_object result = bpy.ops.export_scene.gltf( filepath=str(output_path), export_format="GLB", use_selection=True, export_materials="EXPORT", export_animations=False, ) if "FINISHED" not in result: raise RuntimeError("Blender did not finish the GLB export") return output_path def _snapshot_bake_settings(scene: bpy.types.Scene) -> dict[str, Any]: bake = scene.render.bake return { "engine": scene.render.engine, "margin": bake.margin, "target": bake.target, "use_clear": bake.use_clear, "use_selected_to_active": bake.use_selected_to_active, "cycles_samples": scene.cycles.samples, } def _restore_bake_settings(scene: bpy.types.Scene, state: dict[str, Any]) -> None: scene.render.engine = state["engine"] bake = scene.render.bake bake.margin = state["margin"] bake.target = state["target"] bake.use_clear = state["use_clear"] bake.use_selected_to_active = state["use_selected_to_active"] scene.cycles.samples = state["cycles_samples"] class MIDB_PaletteEntry(PropertyGroup): slot: IntProperty(name="Slot", default=1, min=1) material_name: StringProperty(name="Material") color: FloatVectorProperty( name="ID Color", subtype="COLOR", size=3, min=0.0, max=1.0, ) class MIDB_Settings(PropertyGroup): resolution: EnumProperty( name="Resolution", items=( ("256", "256", "256 x 256"), ("512", "512", "512 x 512"), ("1024", "1K", "1024 x 1024"), ("2048", "2K", "2048 x 2048"), ("4096", "4K", "4096 x 4096"), ("8192", "8K", "8192 x 8192"), ), default="2048", ) margin: IntProperty( name="Margin", description="Extend IDs beyond UV island edges by this many pixels", default=1, min=0, max=256, subtype="PIXEL", ) palette_mode: EnumProperty( name="Colors", items=( ( "DISTINCT", "Distinct", "Generate deterministic, easy-to-see colors for material slots", ), ( "MATERIAL", "Material Viewport Colors", "Use each material's viewport display color", ), ( "INDEX", "Exact Slot Index", "Encode the one-based slot number as a 24-bit RGB integer", ), ), default="DISTINCT", ) apply_modifiers: BoolProperty( name="Apply Modifiers", description="Bake evaluated modifier geometry instead of the base mesh", default=False, ) create_export_copy: BoolProperty( name="Create Export Copy", description=( "Create and select a duplicate with one 'baked ids' material using the baked image" ), default=False, ) quick_export_glb: BoolProperty( name="Quick Export GLB", description="Immediately export the generated one-material copy as a GLB", default=False, ) glb_filepath: StringProperty( name="GLB Path", description="Path for the Substance Painter-ready GLB export", subtype="FILE_PATH", default="//baked_ids.glb", ) image_name: StringProperty( name="Image Name", default="Material_ID", ) save_to_disk: BoolProperty( name="Save PNG", default=True, ) filepath: StringProperty( name="File Path", subtype="FILE_PATH", default="//material_id.png", ) write_legend: BoolProperty( name="Save JSON Legend", description="Write slot names and exact colors beside the PNG", default=True, ) last_image_name: StringProperty(name="Last Image") last_palette: CollectionProperty(type=MIDB_PaletteEntry) class MIDB_OT_Bake(Operator): bl_idname = "object.bake_material_id" bl_label = "Bake Material ID" bl_description = "Bake the active mesh's material assignments into an ID image" bl_options = {"REGISTER", "UNDO"} @classmethod def poll(cls, context: bpy.types.Context) -> bool: obj = context.active_object return bool( obj and obj.type == "MESH" and context.mode == "OBJECT" and obj.data.uv_layers ) def execute(self, context: bpy.types.Context) -> set[str]: settings = context.scene.material_id_baker source = context.active_object scene = context.scene view_layer = context.view_layer if source is None or source.type != "MESH": self.report({"ERROR"}, "Select an active mesh object") return {"CANCELLED"} if context.mode != "OBJECT": self.report({"ERROR"}, "Switch to Object Mode before baking") return {"CANCELLED"} if not source.data.uv_layers: self.report({"ERROR"}, "The active mesh needs a UV map") return {"CANCELLED"} selected_before = list(context.selected_objects) active_before = view_layer.objects.active bake_state = _snapshot_bake_settings(scene) temp_object: bpy.types.Object | None = None temp_mesh: bpy.types.Mesh | None = None temp_materials: list[bpy.types.Material] = [] image: bpy.types.Image | None = None export_object: bpy.types.Object | None = None export_mesh: bpy.types.Mesh | None = None export_material: bpy.types.Material | None = None succeeded = False try: if settings.apply_modifiers: depsgraph = context.evaluated_depsgraph_get() evaluated = source.evaluated_get(depsgraph) temp_mesh = bpy.data.meshes.new_from_object( evaluated, preserve_all_data_layers=True, depsgraph=depsgraph, ) else: temp_mesh = source.data.copy() if not temp_mesh.uv_layers: raise RuntimeError("The evaluated mesh has no UV map") source_uv = source.data.uv_layers.active baked_uv = temp_mesh.uv_layers.get(source_uv.name) if source_uv else None if baked_uv is None: baked_uv = temp_mesh.uv_layers.active temp_mesh.uv_layers.active = baked_uv for uv_layer in temp_mesh.uv_layers: uv_layer.active_render = uv_layer == baked_uv temp_object = bpy.data.objects.new( f"{_TEMP_PREFIX}{source.name}", temp_mesh, ) temp_object.matrix_world = source.matrix_world.copy() scene.collection.objects.link(temp_object) highest_slot = max( (polygon.material_index for polygon in temp_mesh.polygons), default=0, ) slot_count = max(1, len(temp_mesh.materials), highest_slot + 1) source_materials = [] for index in range(slot_count): material = None if index < len(source.material_slots): material = source.material_slots[index].material if material is None and index < len(temp_mesh.materials): material = temp_mesh.materials[index] source_materials.append(material) colors = [ _slot_color(index, material, settings.palette_mode) for index, material in enumerate(source_materials) ] resolution = int(settings.resolution) image = bpy.data.images.new( name=settings.image_name.strip() or "Material_ID", width=resolution, height=resolution, alpha=True, float_buffer=False, ) image.generated_color = (0.0, 0.0, 0.0, 0.0) image.alpha_mode = "STRAIGHT" try: image.colorspace_settings.name = "Non-Color" except TypeError: # Some color configurations use a differently named raw space. image.colorspace_settings.is_data = True for index, color in enumerate(colors): temp_materials.append(_create_id_material(index, color, image)) polygon_material_indices = [ polygon.material_index for polygon in temp_mesh.polygons ] temp_mesh.materials.clear() for material in temp_materials: temp_mesh.materials.append(material) for polygon, material_index in zip( temp_mesh.polygons, polygon_material_indices, strict=True, ): polygon.material_index = min(material_index, slot_count - 1) for obj in list(context.selected_objects): obj.select_set(False) temp_object.hide_set(False) temp_object.hide_render = False temp_object.select_set(True) view_layer.objects.active = temp_object # Cycles owns Blender's image-bake pipeline even though this flat # emission bake needs only one deterministic sample. scene.render.engine = "CYCLES" scene.cycles.samples = 1 scene.render.bake.margin = settings.margin scene.render.bake.target = "IMAGE_TEXTURES" scene.render.bake.use_clear = True scene.render.bake.use_selected_to_active = False context.window_manager.progress_begin(0, 1) try: result = bpy.ops.object.bake(type="EMIT") finally: context.window_manager.progress_end() if "FINISHED" not in result: raise RuntimeError("Blender did not finish the bake") palette = [ { "slot": index + 1, "material": material.name if material else f"Unassigned Slot {index + 1}", "rgb": [round(component, 8) for component in colors[index]], } for index, material in enumerate(source_materials) ] image["material_id_palette"] = json.dumps(palette, separators=(",", ":")) image["source_object"] = source.name image["uv_map"] = temp_mesh.uv_layers.active.name saved_path: Path | None = None if settings.save_to_disk: saved_path = _normalise_png_path(settings.filepath) saved_path.parent.mkdir(parents=True, exist_ok=True) image.filepath_raw = str(saved_path) image.file_format = "PNG" image.save() if settings.write_legend: legend_path = saved_path.with_suffix(".json") legend = { "image": saved_path.name, "source_object": source.name, "uv_map": temp_mesh.uv_layers.active.name, "palette_mode": settings.palette_mode, "materials": palette, } with legend_path.open("w", encoding="utf-8") as handle: json.dump(legend, handle, indent=2) handle.write("\n") settings.last_image_name = image.name settings.last_palette.clear() for item in palette: entry = settings.last_palette.add() entry.slot = item["slot"] entry.material_name = item["material"] entry.color = item["rgb"] if settings.create_export_copy: export_object, export_mesh, export_material = _create_export_copy( source, temp_mesh, image, settings.apply_modifiers, context.collection, ) exported_glb_path: Path | None = None if export_object is not None and settings.quick_export_glb: exported_glb_path = _quick_export_glb( context, export_object, settings.glb_filepath, ) succeeded = True if saved_path: message = f"Baked material IDs to {saved_path}" else: message = f"Baked material IDs to image '{image.name}'" if export_object is not None: message += f" and created '{export_object.name}'" if exported_glb_path is not None: message += f"; exported GLB to {exported_glb_path}" self.report({"INFO"}, message) return {"FINISHED"} except Exception as exc: self.report({"ERROR"}, f"Material ID bake failed: {exc}") return {"CANCELLED"} finally: if temp_object is not None: bpy.data.objects.remove(temp_object, do_unlink=True) if temp_mesh is not None and temp_mesh.users == 0: bpy.data.meshes.remove(temp_mesh) for material in temp_materials: if material.users == 0: bpy.data.materials.remove(material) if not succeeded and export_object is not None: bpy.data.objects.remove(export_object, do_unlink=True) if not succeeded and export_mesh is not None and export_mesh.users == 0: bpy.data.meshes.remove(export_mesh) if not succeeded and export_material is not None and export_material.users == 0: bpy.data.materials.remove(export_material) _restore_bake_settings(scene, bake_state) for obj in list(context.selected_objects): obj.select_set(False) for obj in selected_before: if obj.name in view_layer.objects: obj.select_set(True) if active_before and active_before.name in view_layer.objects: view_layer.objects.active = active_before if succeeded and export_object is not None: for obj in list(context.selected_objects): obj.select_set(False) export_object.select_set(True) view_layer.objects.active = export_object if not succeeded and image is not None and image.users == 0: bpy.data.images.remove(image) class MIDB_PT_Panel(Panel): bl_label = "Material ID Baker" bl_idname = "MIDB_PT_material_id_baker" bl_space_type = "VIEW_3D" bl_region_type = "UI" bl_category = "Material ID" def draw(self, context: bpy.types.Context) -> None: layout = self.layout settings = context.scene.material_id_baker obj = context.active_object if obj is None or obj.type != "MESH": layout.label(text="Select a mesh object", icon="INFO") return mesh = obj.data if mesh.uv_layers: uv_layer = mesh.uv_layers.active layout.label(text=f"UV Map: {uv_layer.name}", icon="GROUP_UVS") else: layout.label(text="A UV map is required", icon="ERROR") layout.label(text=f"Material Slots: {max(1, len(obj.material_slots))}") column = layout.column(align=True) column.prop(settings, "resolution") column.prop(settings, "margin") column.prop(settings, "palette_mode") column.prop(settings, "apply_modifiers") export_copy = layout.box() export_copy.prop(settings, "create_export_copy") quick_export = export_copy.column(align=True) quick_export.enabled = settings.create_export_copy quick_export.prop(settings, "quick_export_glb") if settings.quick_export_glb: quick_export.prop(settings, "glb_filepath") output = layout.box() output.label(text="Output") output.prop(settings, "image_name") output.prop(settings, "save_to_disk") if settings.save_to_disk: output.prop(settings, "filepath") output.prop(settings, "write_legend") operator_column = layout.column() operator_column.enabled = bool(mesh.uv_layers and context.mode == "OBJECT") operator_column.operator("object.bake_material_id", icon="RENDER_STILL") if context.mode != "OBJECT": layout.label(text="Switch to Object Mode to bake", icon="INFO") if settings.last_palette: palette_box = layout.box() palette_box.label(text=f"Last Bake: {settings.last_image_name}") for entry in settings.last_palette: row = palette_box.row(align=True) row.label(text=f"{entry.slot}: {entry.material_name}") row.prop(entry, "color", text="") _CLASSES = ( MIDB_PaletteEntry, MIDB_Settings, MIDB_OT_Bake, MIDB_PT_Panel, ) def register() -> None: for cls in _CLASSES: bpy.utils.register_class(cls) bpy.types.Scene.material_id_baker = PointerProperty(type=MIDB_Settings) def unregister() -> None: del bpy.types.Scene.material_id_baker for cls in reversed(_CLASSES): bpy.utils.unregister_class(cls)