Renamed src folder to io_export_psk_psa

io_export_psk_psa/psa/builder.py

@@ -0,0 +1,130 @@
+from .data import *
+
+
+class PsaBuilderOptions(object):
+    def __init__(self):
+        self.actions = []
+
+
+# https://git.cth451.me/cth451/blender-addons/blob/master/io_export_unreal_psk_psa.py
+class PsaBuilder(object):
+    def __init__(self):
+        # TODO: add options in here (selected anims, eg.)
+        pass
+
+    def build(self, context, options) -> Psa:
+        object = context.view_layer.objects.active
+
+        if object.type != 'ARMATURE':
+            raise RuntimeError('Selected object must be an Armature')
+
+        armature = object
+
+        if armature.animation_data is None:
+            raise RuntimeError('No animation data for armature')
+
+        psa = Psa()
+
+        bones = list(armature.data.bones)
+
+        # The order of the armature bones and the pose bones is not guaranteed to be the same.
+        # As as a result, we need to reconstruct the list of pose bones in the same order as the
+        # armature bones.
+        bone_names = [x.name for x in bones]
+        pose_bones = [(bone_names.index(bone.name), bone) for bone in armature.pose.bones]
+        pose_bones.sort(key=lambda x: x[0])
+        pose_bones = [x[1] for x in pose_bones]
+
+        for bone in bones:
+            psa_bone = Psa.Bone()
+            psa_bone.name = bytes(bone.name, encoding='utf-8')
+            psa_bone.children_count = len(bone.children)
+
+            try:
+                psa_bone.parent_index = bones.index(bone.parent)
+            except ValueError:
+                psa_bone.parent_index = -1
+
+            if bone.parent is not None:
+                rotation = bone.matrix.to_quaternion()
+                rotation.x = -rotation.x
+                rotation.y = -rotation.y
+                rotation.z = -rotation.z
+                quat_parent = bone.parent.matrix.to_quaternion().inverted()
+                parent_head = quat_parent @ bone.parent.head
+                parent_tail = quat_parent @ bone.parent.tail
+                location = (parent_tail - parent_head) + bone.head
+            else:
+                location = armature.matrix_local @ bone.head
+                rot_matrix = bone.matrix @ armature.matrix_local.to_3x3()
+                rotation = rot_matrix.to_quaternion()
+
+            psa_bone.location.x = location.x
+            psa_bone.location.y = location.y
+            psa_bone.location.z = location.z
+
+            psa_bone.rotation.x = rotation.x
+            psa_bone.rotation.y = rotation.y
+            psa_bone.rotation.z = rotation.z
+            psa_bone.rotation.w = rotation.w
+
+            psa.bones.append(psa_bone)
+
+        frame_start_index = 0
+
+        for action in options.actions:
+            if len(action.fcurves) == 0:
+                continue
+
+            armature.animation_data.action = action
+            context.view_layer.update()
+
+            frame_min, frame_max = [int(x) for x in action.frame_range]
+
+            sequence = Psa.Sequence()
+            sequence.name = bytes(action.name, encoding='utf-8')
+            sequence.frame_count = frame_max - frame_min + 1
+            sequence.frame_start_index = frame_start_index
+            sequence.fps = 30  # TODO: fill in later with r
+
+            frame_count = frame_max - frame_min + 1
+
+            for frame in range(frame_count):
+                context.scene.frame_set(frame)
+
+                for bone in pose_bones:
+                    # TODO: is the cast-to-matrix necessary? (guessing no)
+                    key = Psa.Key()
+                    pose_bone_matrix = bone.matrix
+
+                    if bone.parent is not None:
+                        pose_bone_parent_matrix = bone.parent.matrix
+                        pose_bone_matrix = pose_bone_parent_matrix.inverted() @ pose_bone_matrix
+
+                    location = pose_bone_matrix.to_translation()
+                    rotation = pose_bone_matrix.to_quaternion().normalized()
+
+                    if bone.parent is not None:
+                        rotation.x = -rotation.x
+                        rotation.y = -rotation.y
+                        rotation.z = -rotation.z
+
+                    key.location.x = location.x
+                    key.location.y = location.y
+                    key.location.z = location.z
+                    key.rotation.x = rotation.x
+                    key.rotation.y = rotation.y
+                    key.rotation.z = rotation.z
+                    key.rotation.w = rotation.w
+                    key.time = 1.0 / sequence.fps
+
+                    psa.keys.append(key)
+
+                frame_start_index += 1
+
+                sequence.bone_count = len(pose_bones)
+                sequence.track_time = frame_count
+
+            psa.sequences.append(sequence)
+
+        return psa

io_export_psk_psa/psa/data.py

@@ -0,0 +1,44 @@
+from typing import List
+from ..data import *
+
+
+class Psa(object):
+
+    class Bone(Structure):
+        _fields_ = [
+            ('name', c_char * 64),
+            ('flags', c_int32),
+            ('children_count', c_int32),
+            ('parent_index', c_int32),
+            ('rotation', Quaternion),
+            ('location', Vector3),
+            ('padding', c_char * 16)
+        ]
+
+    class Sequence(Structure):
+        _fields_ = [
+            ('name', c_char * 64),
+            ('group', c_char * 64),
+            ('bone_count', c_int32),
+            ('root_include', c_int32),
+            ('compression_style', c_int32),
+            ('key_quotum', c_int32),  # what the fuck is a quotum
+            ('key_reduction', c_float),
+            ('track_time', c_float),
+            ('fps', c_float),
+            ('start_bone', c_int32),
+            ('frame_start_index', c_int32),
+            ('frame_count', c_int32)
+        ]
+
+    class Key(Structure):
+        _fields_ = [
+            ('location', Vector3),
+            ('rotation', Quaternion),
+            ('time', c_float)
+        ]
+
+    def __init__(self):
+        self.bones: List[Psa.Bone] = []
+        self.sequences: List[Psa.Sequence] = []
+        self.keys: List[Psa.Key] = []

io_export_psk_psa/psa/exporter.py

@@ -0,0 +1,27 @@
+from typing import Type
+from .data import *
+
+
+class PsaExporter(object):
+    def __init__(self, psa: Psa):
+        self.psa: Psa = psa
+
+    # This method is shared by both PSA/K file formats, move this?
+    @staticmethod
+    def write_section(fp, name: bytes, data_type: Type[Structure] = None, data: list = None):
+        section = Section()
+        section.name = name
+        if data_type is not None and data is not None:
+            section.data_size = sizeof(data_type)
+            section.data_count = len(data)
+        fp.write(section)
+        if data is not None:
+            for datum in data:
+                fp.write(datum)
+
+    def export(self, path: str):
+        with open(path, 'wb') as fp:
+            self.write_section(fp, b'ANIMHEAD')
+            self.write_section(fp, b'BONENAMES', Psa.Bone, self.psa.bones)
+            self.write_section(fp, b'ANIMINFO', Psa.Sequence, self.psa.sequences)
+            self.write_section(fp, b'ANIMKEYS', Psa.Key, self.psa.keys)

io_export_psk_psa/psa/operator.py

@@ -0,0 +1,104 @@
+from bpy.types import Operator, Action, UIList, PropertyGroup
+from bpy_extras.io_utils import ExportHelper
+from bpy.props import StringProperty, BoolProperty, CollectionProperty, PointerProperty
+from .builder import PsaBuilder, PsaBuilderOptions
+from .exporter import PsaExporter
+import bpy
+import re
+
+
+class ActionListItem(PropertyGroup):
+    action: PointerProperty(type=Action)
+    is_selected: BoolProperty(default=False)
+
+    @property
+    def name(self):
+        return self.action.name
+
+
+class PSA_UL_ActionList(UIList):
+    def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
+        layout.alignment = 'LEFT'
+        layout.prop(item, 'is_selected', icon_only=True)
+        layout.label(text=item.action.name)
+
+    def filter_items(self, context, data, property):
+        # TODO: returns two lists, apparently
+        actions = getattr(data, property)
+        flt_flags = []
+        flt_neworder = []
+        if self.filter_name:
+            flt_flags = bpy.types.UI_UL_list.filter_items_by_name(self.filter_name, self.bitflag_filter_item, actions, 'name', reverse=self.use_filter_invert)
+        return flt_flags, flt_neworder
+
+
+class PsaExportOperator(Operator, ExportHelper):
+    bl_idname = 'export.psa'
+    bl_label = 'Export'
+    __doc__ = 'PSA Exporter (.psa)'
+    filename_ext = '.psa'
+    filter_glob : StringProperty(default='*.psa', options={'HIDDEN'})
+    filepath : StringProperty(
+        name='File Path',
+        description='File path used for exporting the PSA file',
+        maxlen=1024,
+        default='')
+
+    def __init__(self):
+        self.armature = None
+
+    def draw(self, context):
+        layout = self.layout
+        scene = context.scene
+        box = layout.box()
+        box.label(text='Actions', icon='ACTION')
+        row = box.row()
+        row.template_list('PSA_UL_ActionList', 'asd', scene, 'psa_action_list', scene, 'psa_action_list_index', rows=len(context.scene.psa_action_list))
+
+    def is_action_for_armature(self, action):
+        bone_names = [x.name for x in self.armature.data.bones]
+        print(bone_names)
+        for fcurve in action.fcurves:
+            match = re.match('pose\.bones\["(.+)"\].\w+', fcurve.data_path)
+            if not match:
+                continue
+            bone_name = match.group(1)
+            if bone_name not in bone_names:
+                return False
+        return True
+
+    def invoke(self, context, event):
+        if context.view_layer.objects.active.type != 'ARMATURE':
+            self.report({'ERROR_INVALID_CONTEXT'}, 'The selected object must be an armature.')
+            return {'CANCELLED'}
+
+        self.armature = context.view_layer.objects.active
+
+        context.scene.psa_action_list.clear()
+        for action in bpy.data.actions:
+            item = context.scene.psa_action_list.add()
+            item.action = action
+            if self.is_action_for_armature(action):
+                item.is_selected = True
+
+        if len(context.scene.psa_action_list) == 0:
+            self.report({'ERROR_INVALID_CONTEXT'}, 'There are no actions to export.')
+            return {'CANCELLED'}
+
+        context.window_manager.fileselect_add(self)
+        return {'RUNNING_MODAL'}
+
+    def execute(self, context):
+        actions = [x.action for x in context.scene.psa_action_list if x.is_selected]
+
+        if len(actions) == 0:
+            self.report({'ERROR_INVALID_CONTEXT'}, 'No actions were selected for export.')
+            return {'CANCELLED'}
+
+        options = PsaBuilderOptions()
+        options.actions = actions
+        builder = PsaBuilder()
+        psk = builder.build(context, options)
+        exporter = PsaExporter(psk)
+        exporter.export(self.filepath)
+        return {'FINISHED'}