Yusarina
e4f3cdbf17
- Changed image replacement logic to reuse existing placeholder images instead of deleting and recreating them. This should prevents ReferenceError when multiple materials reference the same replacement image
327 lines
17 KiB
Python
327 lines
17 KiB
Python
from pathlib import Path
|
|
import numpy
|
|
import bpy
|
|
import os
|
|
from typing import List, Optional
|
|
from bpy.types import Material, Operator, Context, Object, Image, Mesh, MeshUVLoopLayer, Float2AttributeValue, ShaderNodeTexImage, ShaderNodeBsdfPrincipled, ShaderNodeNormalMap
|
|
from ..core.common import SceneMatClass, MaterialListBool, ProgressTracker
|
|
from ..core.packer.rectangle_packer import MaterialImageList, BinPacker
|
|
from ..core.translations import t
|
|
from ..core.logging_setup import logger
|
|
import traceback
|
|
|
|
class MaterialImageList:
|
|
def __init__(self):
|
|
self.albedo: Image = None
|
|
self.normal: Image = None
|
|
self.emission: Image = None
|
|
self.ambient_occlusion: Image = None
|
|
self.height: Image = None
|
|
self.roughness: Image = None
|
|
self.material: Material = None
|
|
self.parent_mesh: Object = None
|
|
self.w: int = 0
|
|
self.h: int = 0
|
|
self.fit = None
|
|
|
|
def scale_images_to_largest(images: List[Image]) -> tuple[int, int]:
|
|
x: int = 0
|
|
y: int = 0
|
|
|
|
valid_images = []
|
|
for img in images:
|
|
if img:
|
|
try:
|
|
if img.has_data:
|
|
valid_images.append(img)
|
|
except ReferenceError:
|
|
# Image has been removed from Blender's memory
|
|
pass
|
|
|
|
if not valid_images:
|
|
return 0, 0
|
|
|
|
for image in valid_images:
|
|
x = max(x, image.size[0])
|
|
y = max(y, image.size[1])
|
|
|
|
for image in valid_images:
|
|
image.scale(width=int(x), height=int(y))
|
|
|
|
return x, y
|
|
|
|
def MaterialImageList_to_Image_list(classitem: MaterialImageList) -> List[Image]:
|
|
return [
|
|
classitem.albedo,
|
|
classitem.normal,
|
|
classitem.emission,
|
|
classitem.ambient_occlusion,
|
|
classitem.height,
|
|
classitem.roughness
|
|
]
|
|
|
|
def get_material_images_from_scene(context: Context) -> list[MaterialImageList]:
|
|
material_image_list: list[MaterialImageList] = []
|
|
|
|
with ProgressTracker(context, len(context.scene.objects), "Processing Materials") as progress:
|
|
for obj in context.scene.objects:
|
|
if obj.type == 'MESH':
|
|
for mat_slot in obj.material_slots:
|
|
# Only process materials that are selected for atlas
|
|
if mat_slot.material and mat_slot.material.include_in_atlas is True:
|
|
new_mat_image_item = MaterialImageList()
|
|
try:
|
|
new_mat_image_item.albedo = bpy.data.images[mat_slot.material.texture_atlas_albedo]
|
|
except Exception:
|
|
name = mat_slot.material.name + "_albedo_replacement"
|
|
if name not in bpy.data.images:
|
|
new_mat_image_item.albedo = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
|
|
new_mat_image_item.albedo.pixels[:] = numpy.tile(numpy.array([0.0,0.0,0.0,1.0]), 32*32)
|
|
else:
|
|
new_mat_image_item.albedo = bpy.data.images[name]
|
|
try:
|
|
new_mat_image_item.normal = bpy.data.images[mat_slot.material.texture_atlas_normal]
|
|
except Exception:
|
|
name = mat_slot.material.name + "_normal_replacement"
|
|
if name not in bpy.data.images:
|
|
new_mat_image_item.normal = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
|
|
new_mat_image_item.normal.pixels[:] = numpy.tile(numpy.array([0.5,0.5,1.0,1.0]), 32*32)
|
|
else:
|
|
new_mat_image_item.normal = bpy.data.images[name]
|
|
try:
|
|
new_mat_image_item.emission = bpy.data.images[mat_slot.material.texture_atlas_emission]
|
|
except Exception:
|
|
name = mat_slot.material.name + "_emission_replacement"
|
|
if name not in bpy.data.images:
|
|
new_mat_image_item.emission = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
|
|
new_mat_image_item.emission.pixels[:] = numpy.tile(numpy.array([0.0,0.0,0.0,1.0]), 32*32)
|
|
else:
|
|
new_mat_image_item.emission = bpy.data.images[name]
|
|
try:
|
|
new_mat_image_item.ambient_occlusion = bpy.data.images[mat_slot.material.texture_atlas_ambient_occlusion]
|
|
except Exception:
|
|
name = mat_slot.material.name + "_ambient_occlusion_replacement"
|
|
if name not in bpy.data.images:
|
|
new_mat_image_item.ambient_occlusion = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
|
|
new_mat_image_item.ambient_occlusion.pixels[:] = numpy.tile(numpy.array([1.0,1.0,1.0,1.0]), 32*32)
|
|
else:
|
|
new_mat_image_item.ambient_occlusion = bpy.data.images[name]
|
|
try:
|
|
new_mat_image_item.height = bpy.data.images[mat_slot.material.texture_atlas_height]
|
|
except Exception:
|
|
name = mat_slot.material.name + "_height_replacement"
|
|
if name not in bpy.data.images:
|
|
new_mat_image_item.height = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
|
|
new_mat_image_item.height.pixels[:] = numpy.tile(numpy.array([0.5,0.5,0.5,1.0]), 32*32)
|
|
else:
|
|
new_mat_image_item.height = bpy.data.images[name]
|
|
try:
|
|
new_mat_image_item.roughness = bpy.data.images[mat_slot.material.texture_atlas_roughness]
|
|
except Exception:
|
|
name = mat_slot.material.name + "_roughness_replacement"
|
|
if name not in bpy.data.images:
|
|
new_mat_image_item.roughness = bpy.data.images.new(name=name, width=32, height=32, alpha=True)
|
|
new_mat_image_item.roughness.pixels[:] = numpy.tile(numpy.array([1.0,1.0,1.0,0.0]), 32*32)
|
|
else:
|
|
new_mat_image_item.roughness = bpy.data.images[name]
|
|
|
|
new_mat_image_item.material = mat_slot.material
|
|
new_mat_image_item.parent_mesh = obj
|
|
material_image_list.append(new_mat_image_item)
|
|
|
|
progress.step(f"Processed {obj.name}")
|
|
|
|
return material_image_list
|
|
|
|
def prep_images_in_scene(context: Context) -> List[MaterialImageList]:
|
|
preped_images = get_material_images_from_scene(context)
|
|
|
|
with ProgressTracker(context, len(preped_images), "Preparing Images") as progress:
|
|
for MaterialImageClass in preped_images:
|
|
ImageList = MaterialImageList_to_Image_list(MaterialImageClass)
|
|
MaterialImageClass.w, MaterialImageClass.h = scale_images_to_largest(ImageList)
|
|
progress.step(f"Scaled images for {MaterialImageClass.material.name}")
|
|
|
|
return preped_images
|
|
|
|
class AvatarToolKit_OT_AtlasMaterials(Operator):
|
|
bl_idname = "avatar_toolkit.atlas_materials"
|
|
bl_label = t("TextureAtlas.atlas_materials")
|
|
bl_description = t("TextureAtlas.atlas_materials_desc")
|
|
bl_options = {'REGISTER', 'UNDO'}
|
|
|
|
@classmethod
|
|
def poll(cls, context: Context) -> bool:
|
|
# Only allow operation if the file is saved and materials are selected.
|
|
if not bpy.data.filepath:
|
|
cls.poll_message_set(t("TextureAtlas.save_file_first"))
|
|
return False
|
|
return context.scene.avatar_toolkit.texture_atlas_Has_Mat_List_Shown
|
|
|
|
def execute(self, context: Context) -> set:
|
|
try:
|
|
selected_materials = [m for m in prep_images_in_scene(context)
|
|
if m.material and m.material.include_in_atlas]
|
|
|
|
if not selected_materials:
|
|
self.report({'WARNING'}, t("TextureAtlas.no_materials_selected"))
|
|
return {'CANCELLED'}
|
|
|
|
logger.info("Starting material atlas creation")
|
|
|
|
packer = BinPacker(selected_materials)
|
|
mat_images = packer.fit()
|
|
|
|
size = [
|
|
max([matimg.fit.w + matimg.albedo.size[0] for matimg in mat_images]),
|
|
max([matimg.fit.h + matimg.albedo.size[1] for matimg in mat_images])
|
|
]
|
|
|
|
atlased_mat = MaterialImageList()
|
|
|
|
# UV Remapping
|
|
with ProgressTracker(context, len(bpy.data.objects), "Remapping UVs") as progress:
|
|
for mat in mat_images:
|
|
x, y = int(mat.fit.x), int(mat.fit.y)
|
|
w, h = int(mat.albedo.size[0]), int(mat.albedo.size[1])
|
|
|
|
for obj in bpy.data.objects:
|
|
if obj.type == 'MESH':
|
|
mesh = obj.data
|
|
for layer in mesh.polygons:
|
|
if (obj.material_slots[layer.material_index].material and
|
|
obj.material_slots[layer.material_index].material == mat.material):
|
|
for loop_idx in layer.loop_indices:
|
|
for layer_loops in mesh.uv_layers:
|
|
uv_item = layer_loops.uv[loop_idx]
|
|
uv_item.vector.x = (uv_item.vector.x*(w/size[0]))+(x/size[0])
|
|
uv_item.vector.y = (uv_item.vector.y*(h/size[1]))+(y/size[1])
|
|
progress.step(f"Processed UVs for {obj.name}")
|
|
|
|
# Create atlas textures
|
|
texture_types = ["albedo", "normal", "emission", "ambient_occlusion", "height", "roughness"]
|
|
|
|
with ProgressTracker(context, len(texture_types), "Creating Atlas Textures") as progress:
|
|
for type_name in texture_types:
|
|
new_image_name = f"Atlas_{type_name}_{context.scene.name}_{Path(bpy.data.filepath).stem}"
|
|
logger.debug(f"Processing {type_name} atlas image")
|
|
|
|
if new_image_name in bpy.data.images:
|
|
bpy.data.images.remove(bpy.data.images[new_image_name])
|
|
|
|
canvas = bpy.data.images.new(name=new_image_name, width=int(size[0]),
|
|
height=int(size[1]), alpha=True)
|
|
c_w = canvas.size[0]
|
|
canvas_pixels = list(canvas.pixels[:])
|
|
|
|
for mat in mat_images:
|
|
x, y = int(mat.fit.x), int(mat.fit.y)
|
|
w, h = int(mat.albedo.size[0]), int(mat.albedo.size[1])
|
|
image_var = getattr(mat, type_name)
|
|
image_pixels = list(image_var.pixels[:])
|
|
|
|
for k in range(h):
|
|
for i in range(w):
|
|
for channel in range(4):
|
|
canvas_pixels[int((((k+y)*c_w)+(i+x))*4)+channel] = \
|
|
image_pixels[int(((k*w)+i)*4)+channel]
|
|
|
|
canvas.pixels[:] = canvas_pixels[:]
|
|
|
|
try:
|
|
save_dir = os.path.dirname(bpy.data.filepath)
|
|
canvas.save(filepath=os.path.join(save_dir, new_image_name+".png"))
|
|
except Exception as save_error:
|
|
logger.error(f"Failed to save atlas texture: {str(save_error)}")
|
|
self.report({'WARNING'}, f"Could not save texture to disk, This may be due to a lack of permissions.")
|
|
|
|
setattr(atlased_mat, type_name, canvas)
|
|
progress.step(f"Created {type_name} atlas")
|
|
|
|
# Create material nodes
|
|
atlased_mat.material = bpy.data.materials.new(
|
|
name=f"Atlas_Final_{context.scene.name}_{Path(bpy.data.filepath).stem}")
|
|
atlased_mat.material.use_nodes = True
|
|
atlased_mat.material.node_tree.nodes.clear()
|
|
|
|
principled_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeBsdfPrincipled")
|
|
principled_node.location.x = 7.29706335067749
|
|
principled_node.location.y = 298.918212890625
|
|
|
|
output_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeOutputMaterial")
|
|
output_node.location.x = 297.29705810546875
|
|
output_node.location.y = 298.918212890625
|
|
|
|
albedo_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
|
|
albedo_node.location.x = -588.6177978515625
|
|
albedo_node.location.y = 414.1948547363281
|
|
albedo_node.image = atlased_mat.albedo
|
|
|
|
emission_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
|
|
emission_node.location.x = -588.6177978515625
|
|
emission_node.location.y = -173.9259033203125
|
|
emission_node.image = atlased_mat.emission
|
|
|
|
normal_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
|
|
normal_node.location.x = -941.4189453125
|
|
normal_node.location.y = -20.8391780853271
|
|
normal_node.image = atlased_mat.normal
|
|
|
|
normal_map_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeNormalMap")
|
|
normal_map_node.location.x = -545.550537109375
|
|
normal_map_node.location.y = -0.7543716430664062
|
|
|
|
roughness_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
|
|
roughness_node.location.x = -592.1703491210938
|
|
roughness_node.location.y = 206.74075317382812
|
|
roughness_node.image = atlased_mat.roughness
|
|
|
|
ambient_occlusion_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
|
|
ambient_occlusion_node.location.x = -906.4371337890625
|
|
ambient_occlusion_node.location.y = -389.9602355957031
|
|
ambient_occlusion_node.image = atlased_mat.ambient_occlusion
|
|
|
|
height_node = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeTexImage")
|
|
height_node.location.x = -1222.383056640625
|
|
height_node.location.y = -375.48406982421875
|
|
height_node.image = atlased_mat.height
|
|
|
|
atlased_mat.material.node_tree.links.new(principled_node.inputs["Base Color"], albedo_node.outputs["Color"])
|
|
atlased_mat.material.node_tree.links.new(principled_node.inputs["Metallic"], roughness_node.outputs["Alpha"])
|
|
atlased_mat.material.node_tree.links.new(principled_node.inputs["Roughness"], roughness_node.outputs["Color"])
|
|
atlased_mat.material.node_tree.links.new(principled_node.inputs["Alpha"], albedo_node.outputs["Alpha"])
|
|
atlased_mat.material.node_tree.links.new(principled_node.inputs["Normal"], normal_map_node.outputs["Normal"])
|
|
atlased_mat.material.node_tree.links.new(principled_node.inputs["Emission Color"], emission_node.outputs["Color"])
|
|
atlased_mat.material.node_tree.links.new(output_node.inputs["Surface"], principled_node.outputs["BSDF"])
|
|
atlased_mat.material.node_tree.links.new(normal_map_node.inputs["Color"], normal_node.outputs["Color"])
|
|
|
|
# Update materials
|
|
with ProgressTracker(context, len(context.scene.objects), "Updating Materials") as progress:
|
|
for obj in context.scene.objects:
|
|
if obj.type == 'MESH':
|
|
mesh = obj.data
|
|
for i, mat_slot in enumerate(obj.material_slots):
|
|
if mat_slot.material and mat_slot.material.include_in_atlas:
|
|
mesh.materials[i] = atlased_mat.material
|
|
progress.step(f"Updated materials for {obj.name}")
|
|
|
|
MaterialListBool.old_list.pop(context.scene.name, None)
|
|
was_open = context.scene.avatar_toolkit.texture_atlas_Has_Mat_List_Shown
|
|
context.scene.avatar_toolkit.texture_atlas_Has_Mat_List_Shown = False
|
|
|
|
if was_open:
|
|
bpy.ops.avatar_toolkit.expand_section_materials()
|
|
|
|
for area in context.screen.areas:
|
|
if area.type == 'VIEW_3D':
|
|
area.tag_redraw()
|
|
|
|
logger.info("Material atlas creation completed successfully")
|
|
self.report({'INFO'}, t("TextureAtlas.atlas_completed"))
|
|
return {"FINISHED"}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating material atlas: {traceback.format_exc()}", exc_info=True)
|
|
self.report({'ERROR'}, t("TextureAtlas.atlas_error"))
|
|
raise e
|