Initial commit for work to fix the multi-armature workflow

io_scene_psk_psa/psa/builder.py

@@ -1,10 +1,10 @@
 from bpy.types import Action, AnimData, Context, Object, PoseBone
 
 from .data import Psa
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Iterable
 from mathutils import Matrix, Quaternion, Vector
 
-from ..shared.helpers import create_psx_bones, get_coordinate_system_transform
+from ..shared.helpers import PsxBoneCollection, create_psx_bones, get_coordinate_system_transform
 
 
 class PsaBuildSequence:
@@ -14,8 +14,8 @@ class PsaBuildSequence:
             self.frame_start: int = 0
             self.frame_end: int = 0
 
-    def __init__(self, armature_object: Object, anim_data: AnimData):
-        self.armature_object = armature_object
+    def __init__(self, armature_objects: Iterable[Object], anim_data: AnimData):
+        self.armature_objects = list(armature_objects)
         self.anim_data = anim_data
         self.name: str = ''
         self.nla_state: PsaBuildSequence.NlaState = PsaBuildSequence.NlaState()
@@ -27,10 +27,9 @@ class PsaBuildSequence:
 class PsaBuildOptions:
     def __init__(self):
         self.armature_objects: List[Object] = []
-        self.animation_data: Optional[AnimData] = None
         self.sequences: List[PsaBuildSequence] = []
         self.bone_filter_mode: str = 'ALL'
-        self.bone_collection_indices: List[PsaBoneCollectionIndex] = []
+        self.bone_collection_indices: List[PsxBoneCollection] = []
         self.sequence_name_prefix: str = ''
         self.sequence_name_suffix: str = ''
         self.scale = 1.0
@@ -58,7 +57,7 @@ def _get_pose_bone_location_and_rotation(
 
     if is_false_root_bone:
         pose_bone_matrix = coordinate_system_transform
-    elif pose_bone.parent is not None:
+    elif pose_bone is not None and pose_bone.parent is not None:
         pose_bone_matrix = pose_bone.matrix
         pose_bone_parent_matrix = pose_bone.parent.matrix
         pose_bone_matrix = pose_bone_parent_matrix.inverted() @ pose_bone_matrix
@@ -109,6 +108,7 @@ def build_psa(context: Context, options: PsaBuildOptions) -> Psa:
 
     psa = Psa()
 
+    # TODO: move this OUT!
     armature_objects_for_bones = options.armature_objects
     if options.sequence_source == 'ACTIVE_ACTION' and len(options.armature_objects) >= 2:
         # Make sure that the data-block for all the selected armature objects is the same.
@@ -143,7 +143,7 @@ def build_psa(context: Context, options: PsaBuildOptions) -> Psa:
         export_sequence.name = export_sequence.name.strip()
 
     # Save each armature object's current action and frame so that we can restore the state once we are done.
-    saved_armature_object_actions = {o: o.animation_data.action for o in options.armature_objects}
+    saved_armature_object_actions = {o: o.animation_data.action if o.animation_data else None for o in options.armature_objects}
     saved_frame_current = context.scene.frame_current
 
     # Now build the PSA sequences.
@@ -185,10 +185,9 @@ def build_psa(context: Context, options: PsaBuildOptions) -> Psa:
 
         frame = float(frame_start)
         
-        # Link the action to the animation data and update view layer.
-        for armature_object in options.armature_objects:
-            armature_object.animation_data.action = export_sequence.nla_state.action
+        export_sequence.anim_data.action = export_sequence.nla_state.action
 
+        assert context.view_layer
         context.view_layer.update()
 
         def add_key(location: Vector, rotation: Quaternion):
@@ -212,7 +211,7 @@ def build_psa(context: Context, options: PsaBuildOptions) -> Psa:
         armature_scales: Dict[Object, Vector] = {}
 
         # Extract the scale from the world matrix of the evaluated armature object.
-        for armature_object in options.armature_objects:
+        for armature_object in export_sequence.armature_objects:
             evaluated_armature_object = armature_object.evaluated_get(context.evaluated_depsgraph_get())
             _, _, scale = evaluated_armature_object.matrix_world.decompose()
             scale *= options.scale
@@ -223,6 +222,7 @@ def build_psa(context: Context, options: PsaBuildOptions) -> Psa:
         # locations.
         export_bones: List[PsaExportBone] = []
 
+        # TODO: we need different behavior here if it's ACTIVE_ACTION
         for psx_bone, armature_object in psx_bone_create_result.bones:
             if armature_object is None:
                 export_bones.append(PsaExportBone(None, None, Vector((1.0, 1.0, 1.0))))

io_scene_psk_psa/psa/export/operators.py

@@ -472,7 +472,7 @@ class PSA_OT_export(Operator, ExportHelper):
                 for action_item in filter(lambda x: x.is_selected, pg.action_list):
                     if len(action_item.action.fcurves) == 0:
                         continue
-                    export_sequence = PsaBuildSequence(context.active_object, animation_data)
+                    export_sequence = PsaBuildSequence(self.armature_objects, animation_data)
                     export_sequence.name = action_item.name
                     export_sequence.nla_state.action = action_item.action
                     export_sequence.nla_state.frame_start = action_item.frame_start
@@ -483,7 +483,7 @@ class PSA_OT_export(Operator, ExportHelper):
                     export_sequences.append(export_sequence)
             case 'TIMELINE_MARKERS':
                 for marker_item in filter(lambda x: x.is_selected, pg.marker_list):
-                    export_sequence = PsaBuildSequence(context.active_object, animation_data)
+                    export_sequence = PsaBuildSequence(self.armature_objects, animation_data)
                     export_sequence.name = marker_item.name
                     export_sequence.nla_state.frame_start = marker_item.frame_start
                     export_sequence.nla_state.frame_end = marker_item.frame_end
@@ -494,7 +494,7 @@ class PSA_OT_export(Operator, ExportHelper):
                     export_sequences.append(export_sequence)
             case 'NLA_TRACK_STRIPS':
                 for nla_strip_item in filter(lambda x: x.is_selected, pg.nla_strip_list):
-                    export_sequence = PsaBuildSequence(context.active_object, animation_data)
+                    export_sequence = PsaBuildSequence(self.armature_objects, animation_data)
                     export_sequence.name = nla_strip_item.name
                     export_sequence.nla_state.frame_start = nla_strip_item.frame_start
                     export_sequence.nla_state.frame_end = nla_strip_item.frame_end
@@ -504,7 +504,7 @@ class PSA_OT_export(Operator, ExportHelper):
                     export_sequences.append(export_sequence)
             case 'ACTIVE_ACTION':
                 for active_action_item in filter(lambda x: x.is_selected, pg.active_action_list):
-                    export_sequence = PsaBuildSequence(active_action_item.armature_object, active_action_item.armature_object.animation_data)
+                    export_sequence = PsaBuildSequence([active_action_item.armature_object], active_action_item.armature_object.animation_data)
                     action = active_action_item.action
                     export_sequence.name = action.name
                     export_sequence.nla_state.action = action
@@ -522,8 +522,6 @@ class PSA_OT_export(Operator, ExportHelper):
             return {'CANCELLED'}
 
         options = PsaBuildOptions()
-        options.armature_objects = self.armature_objects
-        options.animation_data = animation_data
         options.sequences = export_sequences
         options.bone_filter_mode = pg.bone_filter_mode
         options.bone_collection_indices = [PsxBoneCollection(x.armature_object_name, x.armature_data_name, x.index) for x in pg.bone_collection_list if x.is_selected]

io_scene_psk_psa/psa/export/ui.py

@@ -13,7 +13,7 @@ class PSA_UL_export_sequences(UIList):
         # Show the filtering options by default.
         self.use_filter_show = True
 
-    def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
+    def draw_item(self, context, layout, data, item, icon, active_data, active_property, index, flt_flag):
         item = typing_cast(PSA_PG_export_action_list_item, item)
 
         is_pose_marker = hasattr(item, 'is_pose_marker') and item.is_pose_marker
@@ -44,7 +44,7 @@ class PSA_UL_export_sequences(UIList):
             subrow.prop(pg, 'sequence_filter_pose_marker', icon_only=True, icon='PMARKER')
             subrow.prop(pg, 'sequence_filter_reversed', text='', icon='FRAME_PREV')
 
-    def filter_items(self, context, data, prop):
+    def filter_items(self, context, data, property):
         pg = getattr(context.scene, 'psa_export')
         actions = getattr(data, prop)
         flt_flags = filter_sequences(pg, actions)

io_scene_psk_psa/psa/importer.py

@@ -375,7 +375,7 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object,
         if animation_data is None:
             animation_data = armature_object.animation_data_create()
         for action in actions:
-            nla_track = armature_object.animation_data.nla_tracks.new()
+            nla_track = animation_data.nla_tracks.new()
             nla_track.name = action.name
             nla_track.mute = True
             nla_track.strips.new(name=action.name, start=0, action=action)

io_scene_psk_psa/psa/reader.py

@@ -17,7 +17,7 @@ def _try_fix_cue4parse_issue_103(sequences) -> bool:
         # Manually set the frame_start_index for each sequence. This assumes that the sequences are in order with
         # no shared frames between sequences (all exporters that I know of do this, so it's a safe assumption).
         frame_start_index = 0
-        for i, sequence in enumerate(sequences):
+        for sequence in sequences:
             sequence.frame_start_index = frame_start_index
             frame_start_index += sequence.frame_count
         return True

io_scene_psk_psa/psk/import_/operators.py

@@ -1,10 +1,14 @@
 import os
 from pathlib import Path
 
-from bpy.props import CollectionProperty, StringProperty
-from bpy.types import Context, FileHandler, Operator, OperatorFileListElement, UILayout
+from typing import cast as typing_cast
+from bpy.props import CollectionProperty, StringProperty, FloatProperty, EnumProperty
+from bpy.types import Armature, Context, FileHandler, Operator, OperatorFileListElement, UILayout
 from bpy_extras.io_utils import ImportHelper
 
+from ...shared.helpers import get_coordinate_system_transform
+from ...shared.types import AxisMixin
+
 from ..importer import PskImportOptions, import_psk
 from ..properties import PskImportMixin
 from ..reader import read_psk
@@ -162,6 +166,46 @@ class PSK_OT_import_drag_and_drop(Operator, PskImportMixin):
         return {'FINISHED'}
 
 
+class PSK_OT_create_bones_from_selected_objects(Operator, AxisMixin):
+    bl_idname = 'psk.create_bones_from_selected_objects'
+    bl_label = 'Create Bones from Selected Objects'
+    bl_options = {'UNDO'}
+
+    length: FloatProperty(name='Length', subtype='DISTANCE', default=0.01)
+    
+    @classmethod
+    def poll(cls, context: Context) -> bool:
+        return context.active_object is not None and context.active_object.type == 'ARMATURE'
+
+    def invoke(self, context, event):
+        assert context.window_manager
+        return context.window_manager.invoke_props_dialog(self)
+
+    def execute(self, context: Context):
+        armature_object = context.active_object
+        
+        assert armature_object
+
+        armature_data = typing_cast(Armature, armature_object.data)
+        axis_transform = get_coordinate_system_transform(self.forward_axis, self.up_axis)
+
+        import bpy
+        bpy.ops.object.mode_set(mode='EDIT')
+
+        for index, obj in enumerate(context.selected_objects):
+            if obj == armature_object:
+                continue
+            edit_bone_matrix = armature_object.matrix_world.inverted() @ obj.matrix_world
+            edit_bone = armature_data.edit_bones.new(f'{obj.name}_{index}')
+            # translation, rotation, _ = edit_bone_matrix.decompose()
+            edit_bone.length = self.length
+            edit_bone.matrix = edit_bone_matrix @ axis_transform
+        
+        bpy.ops.object.mode_set(mode='OBJECT')
+
+        return {'FINISHED'}
+
+
 # TODO: move to another file
 class PSK_FH_import(FileHandler):
     bl_idname = 'PSK_FH_import'

io_scene_psk_psa/psk/importer.py

@@ -53,7 +53,7 @@ class PskImportResult:
         self.mesh_object: Optional[Object] = None
 
     @property
-    def root_object(self) -> Object:
+    def root_object(self) -> Optional[Object]:
         return self.armature_object if self.armature_object is not None else self.mesh_object