Files
Avatar-Toolkit/functions/atlas_materials.py
T
Yusarina fe8f5f69d5 Plugin Registration Changes
- Re-wrote how the plugin registers itself.
- No longer need @register_wrapper classes get auto detected and added.
- The new Auto loader is much better then the old way, no longer need "if "bpy" not in locals():" this was an old way of doing things and wasn't really efficient.

 using auto_load.py provides several advantages:

- It automatically discovers and loads all modules in the addon.
- It handles dependencies between classes correctly through topological sorting.
- It manages registration order automatically.
- It properly handles unregistration in the correct order.

This approach is much less error prone and I not had any issues so far. However it still needs testing fully.

I have also start to re-organise files into folders as well, this is going to be needed so we don't have a long list of files as Avatar Toolkit is getting larger then i originally planned.
2024-12-02 01:52:11 +00:00

297 lines
15 KiB
Python

from pathlib import Path
import numpy
import bpy
import os
from typing import List, Tuple, Optional
from bpy.types import Material, Operator, Context, Object, Image, Mesh, MeshUVLoopLayer, Float2AttributeValue, ShaderNodeTexImage, ShaderNodeBsdfPrincipled, ShaderNodeNormalMap
from ..core.common import SceneMatClass, MaterialListBool
from ..core.packer.rectangle_packer import MaterialImageList, BinPacker
from ..core.translations import t
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]) -> set:
x: int = 0
y: int = 0
# Filter out None or invalid images
valid_images = [img for img in images if img and img.has_data]
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]:
list_of_images: list[Image] = []
list_of_images.append(classitem.albedo)
list_of_images.append(classitem.normal)
list_of_images.append(classitem.emission)
list_of_images.append(classitem.ambient_occlusion)
list_of_images.append(classitem.height)
list_of_images.append(classitem.roughness)
return list_of_images
def get_material_images_from_scene(context: Context) -> list[MaterialImageList]:
material_image_list: list[MaterialImageList] = []
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 in bpy.data.images:
bpy.data.images.remove(image=bpy.data.images[name], do_unlink=True)
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)
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 in bpy.data.images:
bpy.data.images.remove(image=bpy.data.images[name], do_unlink=True)
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)
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 in bpy.data.images:
bpy.data.images.remove(image=bpy.data.images[name], do_unlink=True)
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)
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 in bpy.data.images:
bpy.data.images.remove(image=bpy.data.images[name], do_unlink=True)
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)
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 in bpy.data.images:
bpy.data.images.remove(image=bpy.data.images[name], do_unlink=True)
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)
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 in bpy.data.images:
bpy.data.images.remove(image=bpy.data.images[name], do_unlink=True)
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)
new_mat_image_item.material = mat_slot.material
new_mat_image_item.parent_mesh = obj
material_image_list.append(new_mat_image_item)
return material_image_list
def prep_images_in_scene(context: Context) -> list[MaterialImageList]:
preped_images: list[MaterialImageList] = get_material_images_from_scene(context)
for MaterialImageClass in preped_images:
ImageList: list[Image] = MaterialImageList_to_Image_list(MaterialImageClass)
MaterialImageClass.w, MaterialImageClass.h = scale_images_to_largest(ImageList)
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:
return context.scene.texture_atlas_Has_Mat_List_Shown
def execute(self, context: Context) -> set:
try:
# Get only materials that are explicitly marked for inclusion
selected_materials = [m for m in prep_images_in_scene(context) if m.material and m.material.include_in_atlas is True]
if not selected_materials:
self.report({'WARNING'}, t("TextureAtlas.no_materials_selected"))
return {'CANCELLED'}
packer: BinPacker = BinPacker(selected_materials)
mat_images = packer.fit()
size: list[int] = [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])]
print([matimg.fit.w + matimg.albedo.size[1] for matimg in mat_images])
atlased_mat: MaterialImageList = MaterialImageList()
for mat in mat_images:
x: int = int(mat.fit.x)
y: int = int(mat.fit.y)
w: int = int(mat.albedo.size[0])
h: int = int(mat.albedo.size[1])
for obj in bpy.data.objects:
if obj.type == 'MESH':
mesh: Mesh = obj.data
for layer in mesh.polygons:
if obj.material_slots[layer.material_index].material:
if obj.material_slots[layer.material_index].material == mat.material:
for loop_idx in layer.loop_indices:
layer_loops: MeshUVLoopLayer
for layer_loops in mesh.uv_layers:
uv_item: Float2AttributeValue = 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])
for type in ["albedo","normal", "emission","ambient_occlusion","height", "roughness"]:
new_image_name: str= "Atlas_"+type+"_"+context.scene.name+"_"+Path(bpy.data.filepath).stem
print("Processing "+type+" atlas image")
if new_image_name in bpy.data.images:
bpy.data.images.remove(bpy.data.images[new_image_name])
canvas: Image = 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[float] = list(canvas.pixels[:])
for mat in mat_images:
x: int = int(mat.fit.x)
y: int = int(mat.fit.y)
w: int = int(mat.albedo.size[0])
h: int = int(mat.albedo.size[1])
image_var: Image = eval("mat."+type)
image_pixels: list[float] = list(image_var.pixels[:])
print("writing image \""+image_var.name+"\" to canvas.")
print("x: \""+str(x)+"\" "+"y: \""+str(y)+"\" "+"w: \""+str(w)+"\" "+"h: \""+str(h)+"\" ")
for k in range(0,h):
for i in range(0, w):
for channel in range(0,4):
canvas_pixels[
int((((k+y)*c_w)
+
(i+x))*4)
+int(channel)
] = image_pixels[
int((
(k*w)
+i)*4)
+int(channel)]
canvas.pixels[:] = canvas_pixels[:]
canvas.save(filepath=os.path.join(os.path.dirname(bpy.data.filepath),new_image_name+".png"))
exec("atlased_mat."+type+" = canvas")
#I am sorry for the amount of nodes I'm instanciating here and their values.
#This is so that the nodes look pretty in the UI, which I think looks kinda nice. - @989onan
atlased_mat.material = bpy.data.materials.new(name="Atlas_Final_"+bpy.context.scene.name+"_"+Path(bpy.data.filepath).stem)
atlased_mat.material.use_nodes = True
atlased_mat.material.node_tree.nodes.clear()
principled_node: ShaderNodeBsdfPrincipled = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeBsdfPrincipled")
principled_node.location.x = 7.29706335067749
principled_node.location.y = 298.918212890625
output_node: ShaderNodeTexImage = atlased_mat.material.node_tree.nodes.new(type="ShaderNodeOutputMaterial")
output_node.location.x = 297.29705810546875
output_node.location.y = 298.918212890625
albedo_node: ShaderNodeTexImage = 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: ShaderNodeTexImage = 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: ShaderNodeTexImage = 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: ShaderNodeNormalMap = 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: ShaderNodeTexImage = 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: ShaderNodeTexImage = 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: ShaderNodeTexImage = 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"])
# Only update selected materials for meshes
for obj in context.scene.objects:
if obj.type == 'MESH':
mesh: Mesh = obj.data
for i, mat_slot in enumerate(obj.material_slots):
if mat_slot.material and mat_slot.material.include_in_atlas is True:
mesh.materials[i] = atlased_mat.material
self.report({'INFO'}, t("TextureAtlas.atlas_completed"))
return {"FINISHED"}
except Exception as e:
self.report({'ERROR'}, t("TextureAtlas.atlas_error"))
raise e
return {"FINISHED"}