blender-tools/blender/material_id_baker/__init__.py

643 lines
22 KiB
Python

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=16,
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)