From 526df424e3e7b5ef0d9014684531c7d30286b660 Mon Sep 17 00:00:00 2001 From: Colin Basnett Date: Fri, 29 Nov 2024 15:42:42 -0800 Subject: [PATCH] Added "Export Space" option to allow exporting in armature space Armature space exporting means that the user doesn't need to manually move the armature back to origin before exporting, since it was simply using world-space exporting before. Also moved all DFS functions into their own file & cleaned up the usage interface. Added `is_visible` and `is_selected` properties for better ergonomics. --- io_scene_psk_psa/psk/builder.py | 157 +++++++--------------- io_scene_psk_psa/psk/export/operators.py | 14 +- io_scene_psk_psa/psk/export/properties.py | 12 ++ io_scene_psk_psa/psk/import_/operators.py | 1 + io_scene_psk_psa/shared/dfs.py | 154 +++++++++++++++++++++ 5 files changed, 230 insertions(+), 108 deletions(-) create mode 100644 io_scene_psk_psa/shared/dfs.py diff --git a/io_scene_psk_psa/psk/builder.py b/io_scene_psk_psa/psk/builder.py index 65434d1..8782237 100644 --- a/io_scene_psk_psa/psk/builder.py +++ b/io_scene_psk_psa/psk/builder.py @@ -1,19 +1,20 @@ import typing -from typing import Optional, Set +from typing import Optional import bmesh import numpy as np -from bpy.types import Material, Collection, Context, LayerCollection, ViewLayer +from bpy.types import Material, Collection, Context, Mesh from mathutils import Matrix from .data import * from .properties import triangle_type_and_bit_flags_to_poly_flags +from ..shared.dfs import dfs_collection_objects, dfs_view_layer_objects, DfsObject from ..shared.helpers import * class PskInputObjects(object): def __init__(self): - self.mesh_objects: List[Tuple[Object, List[Object], Matrix]] = [] + self.mesh_objects: List[DfsObject] = [] self.armature_object: Optional[Object] = None @@ -25,26 +26,17 @@ class PskBuildOptions(object): self.materials: List[Material] = [] self.should_enforce_bone_name_restrictions = False self.scale = 1.0 + self.export_space = 'WORLD' -def get_mesh_objects_for_collection(collection: Collection, should_exclude_hidden_meshes: bool = True) -> Iterable[Tuple[Object, List[Object], Matrix]]: - for obj, instance_objects, matrix in dfs_collection_objects(collection): - if obj.type != 'MESH': - continue - if should_exclude_hidden_meshes: - if instance_objects: - if not instance_objects[-1].visible_get(): - continue - elif not obj.visible_get(): - continue - yield (obj, instance_objects, matrix) +def get_mesh_objects_for_collection(collection: Collection) -> Iterable[DfsObject]: + return filter(lambda x: x.obj.type == 'MESH', dfs_collection_objects(collection)) -def get_mesh_objects_for_context(context: Context) -> Iterable[Tuple[Object, List[Object], Matrix]]: - for (obj, instance_objects, matrix) in dfs_view_layer_objects(context.view_layer): - is_selected = obj.select_get() or any(x.select_get() for x in instance_objects) - if obj.type == 'MESH' and is_selected: - yield (obj, instance_objects, matrix) +def get_mesh_objects_for_context(context: Context) -> Iterable[DfsObject]: + for dfs_object in dfs_view_layer_objects(context.view_layer): + if dfs_object.obj.type == 'MESH' and dfs_object.is_selected: + yield dfs_object def get_armature_for_mesh_objects(mesh_objects: Iterable[Object]) -> Optional[Object]: @@ -69,17 +61,19 @@ def get_armature_for_mesh_objects(mesh_objects: Iterable[Object]) -> Optional[Ob return None -def _get_psk_input_objects(mesh_objects: List[Tuple[Object, List[Object], Matrix]]) -> PskInputObjects: +def _get_psk_input_objects(mesh_objects: Iterable[DfsObject]) -> PskInputObjects: + mesh_objects = list(mesh_objects) if len(mesh_objects) == 0: raise RuntimeError('At least one mesh must be selected') - for mesh_object, _, _ in mesh_objects: - if len(mesh_object.data.materials) == 0: - raise RuntimeError(f'Mesh "{mesh_object.name}" must have at least one material') + for dfs_object in mesh_objects: + mesh_data = cast(Mesh, dfs_object.obj.data) + if len(mesh_data.materials) == 0: + raise RuntimeError(f'Mesh "{dfs_object.obj.name}" must have at least one material') input_objects = PskInputObjects() input_objects.mesh_objects = mesh_objects - input_objects.armature_object = get_armature_for_mesh_objects([x[0] for x in mesh_objects]) + input_objects.armature_object = get_armature_for_mesh_objects([x.obj for x in mesh_objects]) return input_objects @@ -90,7 +84,9 @@ def get_psk_input_objects_for_context(context: Context) -> PskInputObjects: def get_psk_input_objects_for_collection(collection: Collection, should_exclude_hidden_meshes: bool = True) -> PskInputObjects: - mesh_objects = list(get_mesh_objects_for_collection(collection, should_exclude_hidden_meshes)) + mesh_objects = get_mesh_objects_for_collection(collection) + if should_exclude_hidden_meshes: + mesh_objects = filter(lambda x: x.is_visible, mesh_objects) return _get_psk_input_objects(mesh_objects) @@ -107,6 +103,16 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) psk = Psk() bones = [] + def get_export_space_matrix(): + match options.export_space: + case 'WORLD': + return Matrix.Identity(4) + case 'ARMATURE': + return armature_object.matrix_world.inverted() + case _: + raise ValueError(f'Invalid export space: {options.export_space}') + + export_space_matrix = get_export_space_matrix() # TODO: maybe neutralize the scale here? scale_matrix = Matrix.Scale(options.scale, 4) if armature_object is None or len(armature_object.data.bones) == 0: @@ -153,7 +159,16 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) parent_tail = inverse_parent_rotation @ bone.parent.tail location = (parent_tail - parent_head) + bone.head else: - armature_local_matrix = armature_object.matrix_local + def get_armature_local_matrix(): + match options.export_space: + case 'WORLD': + return armature_object.matrix_world + case 'ARMATURE': + return Matrix.Identity(4) + case _: + raise ValueError(f'Invalid export space: {options.export_space}') + + armature_local_matrix = get_armature_local_matrix() location = armature_local_matrix @ bone.head bone_rotation = bone.matrix.to_quaternion().conjugated() local_rotation = armature_local_matrix.to_3x3().to_quaternion().conjugated() @@ -162,6 +177,14 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) location = scale_matrix @ location + # If the armature object has been scaled, we need to scale the bone's location to match. + _, _, armature_object_scale = armature_object.matrix_world.decompose() + location.x *= armature_object_scale.x + location.y *= armature_object_scale.y + location.z *= armature_object_scale.z + + print(bone.name, location) + psk_bone.location.x = location.x psk_bone.location.y = location.y psk_bone.location.z = location.z @@ -171,12 +194,6 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) psk_bone.rotation.y = rotation.y psk_bone.rotation.z = rotation.z - # If the armature object has been scaled, we need to scale the bone's location to match. - _, _, armature_object_scale = armature_object.matrix_world.decompose() - psk_bone.location.x *= armature_object_scale.x - psk_bone.location.y *= armature_object_scale.y - psk_bone.location.z *= armature_object_scale.z - psk.bones.append(psk_bone) # MATERIALS @@ -197,7 +214,7 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) for object_index, input_mesh_object in enumerate(input_objects.mesh_objects): - obj, instance_objects, matrix_world = input_mesh_object + obj, instance_objects, matrix_world = input_mesh_object.obj, input_mesh_object.instance_objects, input_mesh_object.matrix_world should_flip_normals = False @@ -252,7 +269,7 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) raise ValueError(f'Invalid object evaluation state: {options.object_eval_state}') vertex_offset = len(psk.points) - matrix_world = scale_matrix @ mesh_object.matrix_world + matrix_world = scale_matrix @ export_space_matrix @ mesh_object.matrix_world # VERTICES for vertex in mesh_data.vertices: @@ -383,75 +400,3 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) result.psk = psk return result - - -def dfs_collection_objects_recursive(collection: Collection, visited: Optional[Set[Object]]=None) -> Iterable[Tuple[Object, List[Object], Matrix]]: - if visited is None: - visited = set() - yield from dfs_collection_objects(collection, visited=visited) - for child in collection.children: - yield from dfs_collection_objects_recursive(child, visited) - - -# Construct a list of objects in hierarchy order from `collection.objects`, only keeping those that are in the -# collection. -def dfs_object_children(obj: Object, collection: Collection): - yield obj - for child in obj.children: - if child in collection.objects: - yield from dfs_object_children(child, collection) - - -def dfs_objects_in_collection(collection: Collection): - # Return only the top-level objects in the collection. - objects_hierarchy = [] - for obj in collection.objects: - if obj.parent is None or obj.parent not in set(collection.objects): - objects_hierarchy.append(obj) - for obj in collection.objects: - yield from dfs_object_children(obj, collection) - - -def dfs_collection_objects( - collection: Collection, - instance_objects: Optional[List[Object]] = None, - matrix_world: Matrix = Matrix.Identity(4), - visited: Optional[Set[Object]]=None -) -> Iterable[Tuple[Object, List[Object], Matrix]]: - # We want to also yield the top-level instance object so that callers can inspect the selection status etc. - if visited is None: - visited = set() - - if instance_objects is None: - instance_objects = list() - - for child in collection.children: - yield from dfs_collection_objects(child, instance_objects, matrix_world.copy(), visited) - - for obj in dfs_objects_in_collection(collection): - visited_pair = (obj, instance_objects[-1] if instance_objects else None) - if visited_pair in visited: - continue - # If this an instance, we need to recurse into it. - if obj.instance_collection is not None: - # Calculate the instance transform. - instance_offset_matrix = Matrix.Translation(-obj.instance_collection.instance_offset) - # Recurse into the instance collection. - yield from dfs_collection_objects(obj.instance_collection, - instance_objects + [obj], - matrix_world @ (obj.matrix_world @ instance_offset_matrix), - visited) - else: - # Object is not an instance, yield it. - yield (obj, instance_objects, matrix_world @ obj.matrix_world) - visited.add(visited_pair) - - -def dfs_view_layer_objects(view_layer: ViewLayer) -> Iterable[Tuple[Object, List[Object], Matrix]]: - def layer_collection_objects_recursive(layer_collection: LayerCollection): - for child in layer_collection.children: - yield from layer_collection_objects_recursive(child) - # Iterate only the top-level objects in this collection first. - yield from dfs_collection_objects(layer_collection.collection) - - yield from layer_collection_objects_recursive(view_layer.layer_collection) diff --git a/io_scene_psk_psa/psk/export/operators.py b/io_scene_psk_psa/psk/export/operators.py index 99b3910..8f6dba4 100644 --- a/io_scene_psk_psa/psk/export/operators.py +++ b/io_scene_psk_psa/psk/export/operators.py @@ -5,7 +5,7 @@ from bpy.props import StringProperty, BoolProperty, EnumProperty, FloatProperty from bpy.types import Operator, Context, Object from bpy_extras.io_utils import ExportHelper -from .properties import object_eval_state_items +from .properties import object_eval_state_items, export_space_items from ..builder import build_psk, PskBuildOptions, get_psk_input_objects_for_context, \ get_psk_input_objects_for_collection from ..writer import write_psk @@ -115,6 +115,13 @@ class PSK_OT_export_collection(Operator, ExportHelper): min=0.0001, soft_max=100.0 ) + export_space: EnumProperty( + name='Export Space', + description='Space to export the mesh in', + items=export_space_items, + default='WORLD' + ) + def execute(self, context): collection = bpy.data.collections.get(self.collection) @@ -128,9 +135,10 @@ class PSK_OT_export_collection(Operator, ExportHelper): options = PskBuildOptions() options.bone_filter_mode = 'ALL' options.object_eval_state = self.object_eval_state - options.materials = get_materials_for_mesh_objects([x[0] for x in input_objects.mesh_objects]) + options.materials = get_materials_for_mesh_objects([x.obj for x in input_objects.mesh_objects]) options.should_enforce_bone_name_restrictions = self.should_enforce_bone_name_restrictions options.scale = self.scale + options.export_space = self.export_space try: result = build_psk(context, input_objects, options) @@ -155,6 +163,7 @@ class PSK_OT_export_collection(Operator, ExportHelper): flow.use_property_decorate = False flow.prop(self, 'scale') + flow.prop(self, 'export_space') # MESH mesh_header, mesh_panel = layout.panel('Mesh', default_closed=False) @@ -271,6 +280,7 @@ class PSK_OT_export(Operator, ExportHelper): options.materials = [m.material for m in pg.material_list] options.should_enforce_bone_name_restrictions = pg.should_enforce_bone_name_restrictions options.scale = pg.scale + options.export_space = pg.export_space try: result = build_psk(context, input_objects, options) diff --git a/io_scene_psk_psa/psk/export/properties.py b/io_scene_psk_psa/psk/export/properties.py index 551a27c..0506274 100644 --- a/io_scene_psk_psa/psk/export/properties.py +++ b/io_scene_psk_psa/psk/export/properties.py @@ -11,6 +11,11 @@ object_eval_state_items = ( ('ORIGINAL', 'Original', 'Use data from original object with no modifiers applied'), ) +export_space_items = [ + ('WORLD', 'World', 'Export in world space'), + ('ARMATURE', 'Armature', 'Export in armature space'), +] + class PSK_PG_material_list_item(PropertyGroup): material: PointerProperty(type=Material) index: IntProperty() @@ -49,6 +54,13 @@ class PSK_PG_export(PropertyGroup): min=0.0001, soft_max=100.0 ) + export_space: EnumProperty( + name='Export Space', + options=empty_set, + description='Space to export the mesh in', + items=export_space_items, + default='WORLD' + ) classes = ( diff --git a/io_scene_psk_psa/psk/import_/operators.py b/io_scene_psk_psa/psk/import_/operators.py index 4758ed9..9b1597c 100644 --- a/io_scene_psk_psa/psk/import_/operators.py +++ b/io_scene_psk_psa/psk/import_/operators.py @@ -152,6 +152,7 @@ class PSK_OT_import(Operator, ImportHelper): col.use_property_split = True col.use_property_decorate = False col.prop(self, 'scale') + col.prop(self, 'export_space') mesh_header, mesh_panel = layout.panel('mesh_panel_id', default_closed=False) mesh_header.prop(self, 'should_import_mesh') diff --git a/io_scene_psk_psa/shared/dfs.py b/io_scene_psk_psa/shared/dfs.py new file mode 100644 index 0000000..d3c51ff --- /dev/null +++ b/io_scene_psk_psa/shared/dfs.py @@ -0,0 +1,154 @@ +''' +Depth-first object iterator functions for Blender collections and view layers. + +These functions are used to iterate over objects in a collection or view layer in a depth-first manner, including +instances. This is useful for exporters that need to traverse the object hierarchy in a predictable order. +''' + +from typing import Optional, Set, Iterable, List + +from bpy.types import Collection, Object, ViewLayer, LayerCollection +from mathutils import Matrix + + +class DfsObject: + ''' + Represents an object in a depth-first search. + ''' + def __init__(self, obj: Object, instance_objects: List[Object], matrix_world: Matrix): + self.obj = obj + self.instance_objects = instance_objects + self.matrix_world = matrix_world + + @property + def is_visible(self) -> bool: + ''' + Check if the object is visible. + @return: True if the object is visible, False otherwise. + ''' + if self.instance_objects: + return self.instance_objects[-1].visible_get() + return self.obj.visible_get() + + @property + def is_selected(self) -> bool: + ''' + Check if the object is selected. + @return: True if the object is selected, False otherwise. + ''' + if self.instance_objects: + return self.instance_objects[-1].select_get() + return self.obj.select_get() + + + +def _dfs_object_children(obj: Object, collection: Collection) -> Iterable[Object]: + ''' + Construct a list of objects in hierarchy order from `collection.objects`, only keeping those that are in the + collection. + @param obj: The object to start the search from. + @param collection: The collection to search in. + @return: An iterable of objects in hierarchy order. + ''' + yield obj + for child in obj.children: + if child.name in collection.objects: + yield from _dfs_object_children(child, collection) + + +def dfs_objects_in_collection(collection: Collection) -> Iterable[Object]: + ''' + Returns a depth-first iterator over all objects in a collection, only keeping those that are directly in the + collection. + @param collection: The collection to search in. + @return: An iterable of objects in hierarchy order. + ''' + objects_hierarchy = [] + for obj in collection.objects: + if obj.parent is None or obj.parent not in set(collection.objects): + objects_hierarchy.append(obj) + for obj in objects_hierarchy: + yield from _dfs_object_children(obj, collection) + + +def dfs_collection_objects(collection: Collection, visible_only: bool = False) -> Iterable[DfsObject]: + ''' + Depth-first search of objects in a collection, including recursing into instances. + @param collection: The collection to search in. + @return: An iterable of tuples containing the object, the instance objects, and the world matrix. + ''' + yield from _dfs_collection_objects_recursive(collection) + + +def _dfs_collection_objects_recursive( + collection: Collection, + instance_objects: Optional[List[Object]] = None, + matrix_world: Matrix = Matrix.Identity(4), + visited: Optional[Set[Object]]=None +) -> Iterable[DfsObject]: + ''' + Depth-first search of objects in a collection, including recursing into instances. + This is a recursive function. + @param collection: The collection to search in. + @param instance_objects: The running hierarchy of instance objects. + @param matrix_world: The world matrix of the current object. + @param visited: A set of visited object-instance pairs. + @return: An iterable of tuples containing the object, the instance objects, and the world matrix. + ''' + + # We want to also yield the top-level instance object so that callers can inspect the selection status etc. + if visited is None: + visited = set() + + if instance_objects is None: + instance_objects = list() + + # First, yield all objects in child collections. + for child in collection.children: + yield from _dfs_collection_objects_recursive(child, instance_objects, matrix_world.copy(), visited) + + # Then, evaluate all objects in this collection. + for obj in dfs_objects_in_collection(collection): + visited_pair = (obj, instance_objects[-1] if instance_objects else None) + if visited_pair in visited: + continue + # If this an instance, we need to recurse into it. + if obj.instance_collection is not None: + # Calculate the instance transform. + instance_offset_matrix = Matrix.Translation(-obj.instance_collection.instance_offset) + # Recurse into the instance collection. + yield from _dfs_collection_objects_recursive(obj.instance_collection, + instance_objects + [obj], + matrix_world @ (obj.matrix_world @ instance_offset_matrix), + visited) + else: + # Object is not an instance, yield it. + yield DfsObject(obj, instance_objects, matrix_world @ obj.matrix_world) + visited.add(visited_pair) + + +def dfs_view_layer_objects(view_layer: ViewLayer) -> Iterable[DfsObject]: + ''' + Depth-first iterator over all objects in a view layer, including recursing into instances. + @param view_layer: The view layer to inspect. + @return: An iterable of tuples containing the object, the instance objects, and the world matrix. + ''' + def layer_collection_objects_recursive(layer_collection: LayerCollection): + for child in layer_collection.children: + yield from layer_collection_objects_recursive(child) + # Iterate only the top-level objects in this collection first. + yield from _dfs_collection_objects_recursive(layer_collection.collection) + + yield from layer_collection_objects_recursive(view_layer.layer_collection) + + +def _is_dfs_object_visible(obj: Object, instance_objects: List[Object]) -> bool: + ''' + Check if a DFS object is visible. + @param obj: The object. + @param instance_objects: The instance objects. + @return: True if the object is visible, False otherwise. + ''' + if instance_objects: + return instance_objects[-1].visible_get() + return obj.visible_get()