DarklightGames/io_scene_psk_psa
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added "Export Space" option to allow exporting in armature space

Browse Source 
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.
This commit is contained in:
Colin Basnett
2024-11-29 15:42:42 -08:00
parent ed42b2e227
commit 526df424e3
5 changed files with 230 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
io_scene_psk_psa/psk/builder.py
View File

@@ -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)

14
io_scene_psk_psa/psk/export/operators.py
View File

@@ -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)

12
io_scene_psk_psa/psk/export/properties.py
View File

@@ -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 = (

1
io_scene_psk_psa/psk/import_/operators.py
View File

@@ -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')

154
io_scene_psk_psa/shared/dfs.py Normal file
View File

@@ -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()