diff --git a/io_scene_psk_psa/psk/builder.py b/io_scene_psk_psa/psk/builder.py index 747915a..a784794 100644 --- a/io_scene_psk_psa/psk/builder.py +++ b/io_scene_psk_psa/psk/builder.py @@ -8,6 +8,8 @@ from psk_psa_py.shared.data import Vector3 from psk_psa_py.psk.data import Psk from .properties import triangle_type_and_bit_flags_to_poly_flags from ..shared.helpers import ( + ObjectNode, + ObjectTree, PskInputObjects, PsxBoneCollection, convert_string_to_cp1252_bytes, @@ -38,10 +40,12 @@ class PskBuildResult(object): self.warnings: List[str] = warnings -def _get_mesh_export_space_matrix(armature_object: Object | None, export_space: str, root_armature_object: Object | None) -> Matrix: - # TODO: this should be a bundle of armature objects. otherwise this creates a scenario where you can have - if armature_object is None or root_armature_object is None: +def _get_mesh_export_space_matrix(node: ObjectNode | None, export_space: str) -> Matrix: + if node is None: return Matrix.Identity(4) + + armature_object = node.object + root_armature_object = node.root.object def get_object_space_matrix(obj: Object) -> Matrix: translation, rotation, _ = obj.matrix_world.decompose() @@ -86,6 +90,7 @@ def build_psk(context: Context, input_objects: PskInputObjects, options: PskBuil assert context.window_manager armature_objects = list(input_objects.armature_objects) + armature_object_tree = ObjectTree(input_objects.armature_objects) warnings: list[str] = [] psk = Psk() @@ -138,18 +143,19 @@ def build_psk(context: Context, input_objects: PskInputObjects, options: PskBuil context.window_manager.progress_begin(0, len(input_objects.mesh_dfs_objects)) coordinate_system_matrix = get_coordinate_system_transform(options.forward_axis, options.up_axis) - root_armature_object = next(iter(input_objects.armature_objects), None) + root_armature_object = next(iter(armature_object_tree), None) # Calculate the export spaces for the armature objects. # This is used later to transform the mesh object geometry into the export space. armature_mesh_export_space_matrices: dict[Object | None, Matrix] = {None: Matrix.Identity(4)} + if options.export_space == 'ARMATURE': # For meshes without an armature modifier, we need to set the export space to the first armature object. - armature_mesh_export_space_matrices[None] = _get_mesh_export_space_matrix(root_armature_object, options.export_space, root_armature_object) + armature_mesh_export_space_matrices[None] = _get_mesh_export_space_matrix(root_armature_object, options.export_space) - for armature_object in armature_objects: - # TODO: also handle the case of multiple roots; dont' just assume we have one! - armature_mesh_export_space_matrices[armature_object] = _get_mesh_export_space_matrix(armature_object, options.export_space, root_armature_object) + # TODO: also handle the case of multiple roots; dont' just assume we have one! + for armature_node in iter(armature_object_tree): + armature_mesh_export_space_matrices[armature_node.object] = _get_mesh_export_space_matrix(armature_node, options.export_space) # Temporarily force the armature into the rest position. # The original pose position setting will be restored at the end. diff --git a/io_scene_psk_psa/shared/helpers.py b/io_scene_psk_psa/shared/helpers.py index d153925..0ca2d74 100644 --- a/io_scene_psk_psa/shared/helpers.py +++ b/io_scene_psk_psa/shared/helpers.py @@ -334,41 +334,47 @@ class PsxBoneCollection: class ObjectNode: def __init__(self, obj: Object): self.object = obj - self.children: List['ObjectNode'] = [] + self.parent: ObjectNode | None = None + self.children: List[ObjectNode] = [] + + @property + def root(self): + """ + Gets the root in the object hierarchy. This can return itself if this node has no parent. + """ + n = self + while n.parent is not None: + n = n.parent + return n class ObjectTree: - def __init__(self) -> None: + ''' + A tree of the armature objects based on their hierarchy. + ''' + def __init__(self, objects: Iterable[Object]): self.root_nodes: List[ObjectNode] = [] - - @staticmethod - def from_objects(objects: Iterable[Object]) -> 'ObjectTree': - ''' - Make a tree of the armature objects based on their hierarchy. - ''' - tree = ObjectTree() object_node_map: Dict[Object, ObjectNode] = {x: ObjectNode(x) for x in objects} for obj, object_node in object_node_map.items(): if obj.parent in object_node_map: parent_node = object_node_map[obj.parent] + object_node.parent = parent_node parent_node.children.append(object_node) else: - tree.root_nodes.append(object_node) - - return tree + self.root_nodes.append(object_node) def __iter__(self): """ An depth-first iterator over the armature tree. """ - node_stack = self.root_nodes + node_stack = [] + self.root_nodes while node_stack: node = node_stack.pop() yield node node_stack = node.children + node_stack - def objects_dfs(self): + def objects_iterator(self): for node in self: yield node.object @@ -401,7 +407,7 @@ def create_psx_bones( if bone_collection_indices is None: bone_collection_indices = [] - armature_tree = ObjectTree.from_objects(armature_objects) + armature_tree = ObjectTree(armature_objects) # Check that there is only one root bone. If there are multiple armature objects, the export space must be WORLD. if len(armature_tree.root_nodes) >= 2 and export_space != 'WORLD': @@ -655,7 +661,7 @@ def _get_psk_input_objects(mesh_dfs_objects: Iterable[DfsObject]) -> PskInputObj # Get the armature objects used on all the meshes being exported. armature_objects = get_armatures_for_mesh_objects(map(lambda x: x.obj, mesh_dfs_objects)) # Sort them in hierarchy order. - input_objects.armature_objects = list(ObjectTree.from_objects(armature_objects).objects_dfs()) + input_objects.armature_objects = list(ObjectTree(armature_objects).objects_iterator()) return input_objects