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 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 = [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]: 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 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) 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:", exception=e) self.report({'ERROR'}, t("TextureAtlas.atlas_error")) raise e