diff --git a/core/import_pmx.py b/core/import_pmx.py index 19595ee..8c0ab3e 100644 --- a/core/import_pmx.py +++ b/core/import_pmx.py @@ -4,25 +4,69 @@ import bpy import struct import traceback import mathutils - from mathutils import Matrix, Vector -from bpy.types import Material, Operator, Context, Object, Image, Mesh, MeshUVLoopLayer, Float2AttributeValue, ShaderNodeTexImage, ShaderNodeBsdfPrincipled, ShaderNodeOutputMaterial +class PMXVertex: + def __init__(self, position, normal, uv, bone_indices, bone_weights, edge_scale, additional_uvs): + self.position = position + self.normal = normal + self.uv = uv + self.bone_indices = bone_indices + self.bone_weights = bone_weights + self.edge_scale = edge_scale + self.additional_uvs = additional_uvs -def replace_char(string, index, character): - temp = list(string) - temp[index] = character - return "".join(temp) +class PMXBone: + def __init__(self, name, english_name, position, parent_index, layer, flag, + tail_position, inherit_parent_index, inherit_influence, + fixed_axis, local_x, local_z, external_key, + ik_target_index, ik_loop_count, ik_limit_rad, ik_links): + self.name = name + self.english_name = english_name + self.position = position + self.parent_index = parent_index + self.layer = layer + self.flag = flag + self.tail_position = tail_position + self.inherit_parent_index = inherit_parent_index + self.inherit_influence = inherit_influence + self.fixed_axis = fixed_axis + self.local_x = local_x + self.local_z = local_z + self.external_key = external_key + self.ik_target_index = ik_target_index + self.ik_loop_count = ik_loop_count + self.ik_limit_rad = ik_limit_rad + self.ik_links = ik_links + +class PMXMaterial: + def __init__(self, name, english_name, diffuse, specular, specular_strength, + ambient, flag, edge_color, edge_size, texture_index, + sphere_texture_index, sphere_mode, toon_sharing_flag, + toon_texture_index, comment, surface_count): + self.name = name + self.english_name = english_name + self.diffuse = diffuse + self.specular = specular + self.specular_strength = specular_strength + self.ambient = ambient + self.flag = flag + self.edge_color = edge_color + self.edge_size = edge_size + self.texture_index = texture_index + self.sphere_texture_index = sphere_texture_index + self.sphere_mode = sphere_mode + self.toon_sharing_flag = toon_sharing_flag + self.toon_texture_index = toon_texture_index + self.comment = comment + self.surface_count = surface_count def read_pmx_header(file: BufferedReader): - # Read PMX header information magic = file.read(4) if magic != b'PMX ': raise ValueError("Invalid PMX file") version = struct.unpack(' bpy.types.Object: + armature = bpy.data.armatures.new(f"{model_name}_Armature") + armature_obj = bpy.data.objects.new(f"{model_name}_Armature", armature) + bpy.context.collection.objects.link(armature_obj) + bpy.context.view_layer.objects.active = armature_obj bpy.ops.object.mode_set(mode='EDIT') - edit_bones = [] - for i, bone_data in enumerate(bones_data): - try: - print(f"Creating bone {i}: {bone_data[0]}") - bone = armature_obj.data.edit_bones.new(bone_data[0]) - - # Convert and scale head position - head_pos = Vector(bone_data[2]).xzy * scale - bone.head = head_pos - - # Handle tail position - if bone_data[6][0] is not None: - tail_pos = Vector(bone_data[6]).xzy * scale - bone.tail = tail_pos - print(f"Using defined tail position for bone {bone_data[0]}") - else: - # Set a default tail position if not provided - bone.tail = head_pos + Vector((0, 0.1, 0)) * scale - print(f"Using default tail position for bone {bone_data[0]}") - - # Set parent if exists - if bone_data[3] != -1: - parent_bone = armature_obj.data.edit_bones[bones_data[bone_data[3]][0]] - bone.parent = parent_bone - print(f"Parented bone {bone_data[0]} to {parent_bone.name}") - - edit_bones.append(bone) - - except Exception as e: - print(f"Error creating bone {i}: {str(e)}") - continue + # First pass: Create bones with correct positions and sizes + edit_bones = [] # Using a list instead of dict for indexed access + for i, bone_data in enumerate(bones): + bone_name = f"bone_{i}" + edit_bone = armature.edit_bones.new(bone_name) + edit_bone.head = Vector(bone_data.position) - # Set bone hierarchy - for i, bone_data in enumerate(bones_data): - if bone_data[3] != -1: - edit_bones[i].parent = edit_bones[bone_data[3]] - - # Apply final transforms - bpy.ops.object.mode_set(mode='OBJECT') - armature_obj.rotation_euler[0] = 1.5708 - armature_obj.rotation_euler[2] = 3.14159 - armature_obj.select_set(True) - bpy.ops.object.transform_apply(location=True, rotation=True, scale=True) - - return edit_bones - -def read_morph(file: BufferedReader, morph_struct, morph_bytesize, vertex_struct, vertex_size, bone_struct, bone_size, material_struct, material_size, rigid_struct, rigid_size): - morph_name = str(file.read(struct.unpack('= 0: + parent_pos = Vector(bones[bone_data.parent_index].position) + if (Vector(bone_data.position) - parent_pos).length > 0.001: + direction = (Vector(bone_data.position) - parent_pos).normalized() * bone_length + edit_bone.tail = edit_bone.head + direction - # Assign vertex normals - custom_normals = [(Vector(i[1]).xzy).normalized() for i in vertices] - mesh.normals_split_custom_set_from_vertices(custom_normals) + edit_bones.append(edit_bone) + + # Second pass: Set up hierarchy and orientations + for i, bone_data in enumerate(bones): + edit_bone = edit_bones[i] - # Assign UV coordinates - uv_layer = mesh.uv_layers.new() - loop_indices_orig = tuple(i for f in faces for i in f) - uv_table = {vi:v for vi, v in enumerate([i[2] for i in vertices])} - uv_layer.data.foreach_set('uv', tuple(v for i in loop_indices_orig for v in uv_table[i])) - - cur_polygon_index: int = 0 - - # Assign materials - for material_data in materials: - material: bpy.types.Material - if material_data[0] in bpy.data.materials: - material = bpy.data.materials[material_data[0]] - else: - material = bpy.data.materials.new(material_data[0]) - - material.use_nodes = True - for node in [node for node in material.node_tree.nodes]: - material.node_tree.nodes.remove(node) + # Parent bones + if bone_data.parent_index >= 0: + parent_bone = edit_bones[bone_data.parent_index] + edit_bone.parent = parent_bone - # Create main nodes - principled_node = material.node_tree.nodes.new(type="ShaderNodeBsdfPrincipled") - principled_node.location = (0, 300) + # Connect bones only if they should be connected + if (Vector(bone_data.position) - Vector(parent_bone.tail)).length < 0.01: + edit_bone.use_connect = True + + # Handle bone orientation + if bone_data.fixed_axis != [0.0, 0.0, 0.0]: + edit_bone.align_roll(Vector(bone_data.fixed_axis)) + elif bone_data.local_x != [0.0, 0.0, 0.0]: + x_axis = Vector(bone_data.local_x).normalized() + z_axis = Vector(bone_data.local_z).normalized() + y_axis = z_axis.cross(x_axis) - output_node = material.node_tree.nodes.new(type="ShaderNodeOutputMaterial") - output_node.location = (300, 300) + # Create and apply orientation matrix + matrix = Matrix((x_axis, y_axis, z_axis)).to_3x3() + edit_bone.matrix = matrix + + bpy.ops.object.mode_set(mode='OBJECT') + return armature_obj - # Set up main texture - albedo_node = material.node_tree.nodes.new(type="ShaderNodeTexImage") - albedo_node.location = (-600, 400) - if textures[material_data[9]] in bpy.data.images: - albedo_node.image = bpy.data.images[textures[material_data[9]]] - else: - albedo_node.image = bpy.data.images.new(name=textures[material_data[9]], width=32, height=32) - albedo_node.image.filepath = os.path.join(os.path.dirname(filepath), textures[material_data[9]]) - albedo_node.image.source = 'FILE' - albedo_node.extension = 'REPEAT' - albedo_node.image.reload() - # Set material properties - principled_node.inputs["Base Color"].default_value = material_data[2] - principled_node.inputs["Specular IOR Level"].default_value = material_data[4] - principled_node.inputs["Roughness"].default_value = 0.5 - principled_node.inputs["Metallic"].default_value = 0.0 +def assign_vertex_weights(obj: bpy.types.Object, vertices: list[PMXVertex], bones: list[PMXBone]): + # Pre-create vertex groups + vertex_groups = {} + for bone in bones: + vertex_groups[bone.name] = obj.vertex_groups.new(name=bone.name) + + # Batch assign weights + for vertex_index, vertex in enumerate(vertices): + for bone_idx, weight in zip(vertex.bone_indices, vertex.bone_weights): + if bone_idx != -1 and weight > 0: + vertex_groups[bones[bone_idx].name].add([vertex_index], weight, 'REPLACE') - # Handle transparency - if material_data[2][3] < 0.99: - material.blend_method = 'HASHED' - material.use_backface_culling = False - material.alpha_threshold = 0.5 - material.show_transparent_back = True - - # Create mix shader for transparency - mix_shader = material.node_tree.nodes.new(type='ShaderNodeMixShader') - mix_shader.location = (100, 300) - transparent_node = material.node_tree.nodes.new(type='ShaderNodeBsdfTransparent') - transparent_node.location = (-200, 200) +def assign_materials(obj: bpy.types.Object, materials: list[PMXMaterial], textures: list[str], base_path: str): + current_face_index = 0 + + for material in materials: + # Create or get material + mat_name = material.name or f"Material_{len(obj.data.materials)}" + if mat_name in bpy.data.materials: + mat = bpy.data.materials[mat_name] + else: + mat = bpy.data.materials.new(name=mat_name) + + # Set up material nodes + texture_path = None + if material.texture_index >= 0 and material.texture_index < len(textures): + texture_path = os.path.join(base_path, textures[material.texture_index]) + + create_material_nodes(mat, texture_path, material.diffuse, material.specular, + material.specular_strength) + + # Assign material to mesh + if mat.name not in obj.data.materials: + obj.data.materials.append(mat) + + # Assign faces to material + mat_index = obj.data.materials.find(mat.name) + for face in obj.data.polygons[current_face_index:current_face_index + material.surface_count]: + face.material_index = mat_index + + current_face_index += material.surface_count - # Connect nodes for transparency - material.node_tree.links.new(mix_shader.inputs[0], albedo_node.outputs["Alpha"]) - material.node_tree.links.new(mix_shader.inputs[1], transparent_node.outputs[0]) - material.node_tree.links.new(mix_shader.inputs[2], principled_node.outputs[0]) - material.node_tree.links.new(output_node.inputs["Surface"], mix_shader.outputs[0]) - else: - material.blend_method = 'OPAQUE' - material.node_tree.links.new(output_node.inputs["Surface"], principled_node.outputs[0]) - - # Connect color - material.node_tree.links.new(principled_node.inputs["Base Color"], albedo_node.outputs["Color"]) - - if not (material.name in mesh.materials): - mesh.materials.append(material) - - end: int = cur_polygon_index + material_data[15] - 1 - for face in mesh.polygons.items()[cur_polygon_index:end]: - face[1].material_index = mesh.materials.find(material.name) +def import_pmx(filepath: str): + try: + with open(filepath, 'rb') as file: + # Read header + header_data = read_pmx_header(file) + version, encoding, additional_uvs, vertex_index_size, texture_index_size, \ + material_index_size, bone_index_size, morph_index_size, rigid_body_index_size, \ + model_name, model_english_name, model_comment, model_english_comment = header_data - cur_polygon_index = cur_polygon_index + material_data[15] - - # Create armature and assign bones - armature = bpy.data.armatures.new(model_name + "_Armature") - armature_obj = bpy.data.objects.new(model_name + "_Armature", armature) - bpy.context.collection.objects.link(armature_obj) - obj.parent = armature_obj - modifier = obj.modifiers.new("Armature", 'ARMATURE') - modifier.object = armature_obj - - bpy.context.view_layer.objects.active = armature_obj - bpy.ops.object.mode_set(mode='EDIT') - - print("Starting bone creation...") - print(f"Total bones to create: {len(bones)}") - - # Create the bones using our create_bones function - edit_bones = create_bones(armature_obj, bones, scale) - - # Now we can safely scale and position bones - for bone in armature.edit_bones: - bone_data = next(b for b in bones if b[0] == bone.name) - bone.head = Vector(bone_data[2]).xzy * scale - if bone_data[6][0] is not None: - bone.tail = Vector(bone_data[6]).xzy * scale - else: - bone.tail = bone.head + Vector((0, 0.1 * scale)) - - # Assign bone weights to the mesh - for i, vertex in enumerate(vertices): - for j in range(0, len(vertex[3])): - if vertex[3][j] != -1 and vertex[3][j] < len(bones): - bone_name = bones[vertex[3][j]][0] - weight = vertex[4][j] - - vertex_group = obj.vertex_groups.get(bone_name) - if not vertex_group: - vertex_group = obj.vertex_groups.new(name=bone_name) - - vertex_group.add([i], weight, 'REPLACE') - - # Assign morphs to the mesh - for morph_data in morphs: - morph_name = morph_data[0] - morph_type = morph_data[3] + # Set up index size formats + vertex_struct, vertex_size = read_index_size(vertex_index_size, 'BHi') + bone_struct, bone_size = read_index_size(bone_index_size, 'bhi') + texture_struct, texture_size = read_index_size(texture_index_size, 'bhi') + + # Read vertices + vertex_count = struct.unpack('