From f8ac443bb149c49146bb19bf32affe2a50837753 Mon Sep 17 00:00:00 2001 From: Colin Basnett Date: Tue, 10 Dec 2019 02:52:09 -0800 Subject: [PATCH] Reorganized file structure and added WIP (non-working) PSA export. --- src/__init__.py | 38 ++++++++---- src/exporter.py | 77 ------------------------ src/psa/__init__.py | 0 src/psa/builder.py | 115 ++++++++++++++++++++++++++++++++++++ src/psa/data.py | 43 ++++++++++++++ src/psa/exporter.py | 27 +++++++++ src/psa/operator.py | 31 ++++++++++ src/psk/__init__.py | 0 src/{ => psk}/builder.py | 20 +++---- src/{psk.py => psk/data.py} | 29 +-------- src/psk/exporter.py | 38 ++++++++++++ src/psk/operator.py | 26 ++++++++ 12 files changed, 318 insertions(+), 126 deletions(-) delete mode 100644 src/exporter.py create mode 100644 src/psa/__init__.py create mode 100644 src/psa/builder.py create mode 100644 src/psa/data.py create mode 100644 src/psa/exporter.py create mode 100644 src/psa/operator.py create mode 100644 src/psk/__init__.py rename src/{ => psk}/builder.py (91%) rename src/{psk.py => psk/data.py} (75%) create mode 100644 src/psk/exporter.py create mode 100644 src/psk/operator.py diff --git a/src/__init__.py b/src/__init__.py index b8a8ec7..73cf11d 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -13,32 +13,48 @@ bl_info = { if 'bpy' in locals(): import importlib - importlib.reload(psk) - importlib.reload(exporter) - importlib.reload(builder) + importlib.reload(psk_data) + importlib.reload(psk_builder) + importlib.reload(psk_exporter) + importlib.reload(psk_operator) + importlib.reload(psa_data) + importlib.reload(psa_builder) + importlib.reload(psa_exporter) + importlib.reload(psa_operator) else: # if i remove this line, it can be enabled just fine - from . import psk - from . import exporter - from . import builder + from .psk import data as psk_data + from .psk import builder as psk_builder + from .psk import exporter as psk_exporter + from .psk import operator as psk_operator + from .psa import data as psa_data + from .psa import builder as psa_builder + from .psa import exporter as psa_exporter + from .psa import operator as psa_operator import bpy classes = [ - exporter.PskExportOperator + psk_operator.PskExportOperator, + psa_operator.PsaExportOperator ] -def menu_func(self, context): - self.layout.operator(exporter.PskExportOperator.bl_idname, text = "Unreal PSK (.psk)") +def psk_menu_func(self, context): + self.layout.operator(psk_operator.PskExportOperator.bl_idname, text ='Unreal PSK (.psk)') + +def psa_menu_func(self, context): + self.layout.operator(psa_operator.PsaExportOperator.bl_idname, text='Unreal PSA (.psa)') def register(): from bpy.utils import register_class for cls in classes: register_class(cls) - bpy.types.TOPBAR_MT_file_export.append(menu_func) + bpy.types.TOPBAR_MT_file_export.append(psk_menu_func) + bpy.types.TOPBAR_MT_file_export.append(psa_menu_func) def unregister(): - bpy.types.TOPBAR_MT_file_export.remove(menu_func) + bpy.types.TOPBAR_MT_file_export.remove(psa_menu_func) + bpy.types.TOPBAR_MT_file_export.remove(psk_menu_func) from bpy.utils import unregister_class for cls in reversed(classes): unregister_class(cls) diff --git a/src/exporter.py b/src/exporter.py deleted file mode 100644 index 7fc4c74..0000000 --- a/src/exporter.py +++ /dev/null @@ -1,77 +0,0 @@ -from bpy.types import Operator -from bpy_extras.io_utils import ExportHelper -from bpy.props import StringProperty, BoolProperty, FloatProperty -import ctypes -import struct -import io -from typing import Type -from .psk import Psk, Vector3, Quaternion -from .builder import PskBuilder - -# https://me3explorer.fandom.com/wiki/PSK_File_Format -# https://docs.unrealengine.com/udk/Two/rsrc/Two/BinaryFormatSpecifications/UnrealAnimDataStructs.h -class PskExportOperator(Operator, ExportHelper): - bl_idname = 'export.psk' - bl_label = 'Export' - __doc__ = 'PSK Exporter (.psk)' - filename_ext = '.psk' - # filter_glob : StringProperty(default='*.psk', options={'HIDDEN'}) - - filepath : StringProperty( - name='File Path', - description='File path used for exporting the PSK file', - maxlen=1024, - default='') - - def execute(self, context): - builder = PskBuilder() - psk = builder.build(context) - exporter = PskExporter(psk) - exporter.export(self.filepath) - return {'FINISHED'} - - -class PskExporter(object): - def __init__(self, psk: Psk): - self.psk: Psk = psk - - @staticmethod - def write_section(fp, name: bytes, data_type: Type[ctypes.Structure] = None, data: list = None): - section = Psk.Section() - section.name = name - if data_type is not None and data is not None: - section.data_size = ctypes.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): - # TODO: add logic somewhere to assert lengths of ctype structs (pack1) - with open(path, 'wb') as fp: - self.write_section(fp, b'ACTRHEAD') - - # POINTS - self.write_section(fp, b'PNTS0000', Vector3, self.psk.points) - - # WEDGES - # TODO: would be nice to have this implicit! - if len(self.psk.wedges) <= 65536: - wedge_type = Psk.Wedge16 - else: - wedge_type = Psk.Wedge32 - - self.write_section(fp, b'VTXW0000', wedge_type, self.psk.wedges) - - # FACES - self.write_section(fp, b'FACE0000', Psk.Face, self.psk.faces) - - # MATERIALS - self.write_section(fp, b'MATT0000', Psk.Material, self.psk.materials) - - # BONES - self.write_section(fp, b'REFSKELT', Psk.Bone, self.psk.bones) - - # WEIGHTS - self.write_section(fp, b'RAWWEIGHTS', Psk.Weight, self.psk.weights) diff --git a/src/psa/__init__.py b/src/psa/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/psa/builder.py b/src/psa/builder.py new file mode 100644 index 0000000..8909db9 --- /dev/null +++ b/src/psa/builder.py @@ -0,0 +1,115 @@ +import bpy +import mathutils +from .data import * + +class PsaBuilder(object): + def __init__(self): + # TODO: add options in here (selected anims, eg.) + pass + + def build(self, context) -> 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) + + 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) + + print('---- ACTIONS ----') + + frame_start_index = 0 + for action in bpy.data.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 + + for frame in range(frame_min, frame_max + 1): + context.scene.frame_set(frame) + + print(frame) + + for bone_index, bone in enumerate(armature.pose.bones): + # TODO: is the cast-to-matrix necesssary? (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 + + psa.sequences.append(sequence) + + return psa diff --git a/src/psa/data.py b/src/psa/data.py new file mode 100644 index 0000000..e9a4a6d --- /dev/null +++ b/src/psa/data.py @@ -0,0 +1,43 @@ +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] = [] diff --git a/src/psa/exporter.py b/src/psa/exporter.py new file mode 100644 index 0000000..d48cdbf --- /dev/null +++ b/src/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) diff --git a/src/psa/operator.py b/src/psa/operator.py new file mode 100644 index 0000000..26bd766 --- /dev/null +++ b/src/psa/operator.py @@ -0,0 +1,31 @@ +from bpy.types import Operator, Action +from bpy_extras.io_utils import ExportHelper +from bpy.props import StringProperty, BoolProperty, FloatProperty, CollectionProperty +from .builder import PsaBuilder +from .exporter import PsaExporter + + +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='') + + actions : CollectionProperty( + type=Action, + name='Sequences' + ) + + def execute(self, context): + builder = PsaBuilder() + psk = builder.build(context) + exporter = PsaExporter(psk) + exporter.export(self.filepath) + return {'FINISHED'} diff --git a/src/psk/__init__.py b/src/psk/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/builder.py b/src/psk/builder.py similarity index 91% rename from src/builder.py rename to src/psk/builder.py index cf95d37..4eaecda 100644 --- a/src/builder.py +++ b/src/psk/builder.py @@ -1,9 +1,9 @@ import bpy import bmesh -import mathutils -from .psk import Psk, Vector3, Quaternion +from .data import * +# TODO: move to another file def make_fquat(bquat): quat = Quaternion() # flip handedness for UT = set x,y,z to negative (rotate in other direction) @@ -22,30 +22,32 @@ def make_fquat_default(bquat): quat.w = bquat.w return quat + class PskBuilder(object): def __init__(self): + # TODO: add options in here pass def build(self, context) -> Psk: object = context.view_layer.objects.active if object.type != 'MESH': - raise RuntimeError('selected object must be a mesh') + raise RuntimeError('Selected object must be a mesh') if len(object.data.materials) == 0: - raise RuntimeError('the mesh must have at least one material') + raise RuntimeError('Mesh must have at least one material') # ensure that there is exactly one armature modifier modifiers = [x for x in object.modifiers if x.type == 'ARMATURE'] if len(modifiers) != 1: - raise RuntimeError('the mesh must have one armature modifier') + raise RuntimeError('Mesh must have one armature modifier') armature_modifier = modifiers[0] armature_object = armature_modifier.object if armature_object is None: - raise RuntimeError('the armature modifier has no linked object') + raise RuntimeError('Armature modifier has no linked object') # TODO: probably requires at least one bone? not sure mesh_data = object.data @@ -127,7 +129,6 @@ class PskBuilder(object): psk_bone.parent_index = 0 if bone.parent is not None: - # calc parented bone transform rotation = bone.matrix.to_quaternion() rotation.x = -rotation.x rotation.y = -rotation.y @@ -137,9 +138,8 @@ class PskBuilder(object): parent_tail = quat_parent @ bone.parent.tail location = (parent_tail - parent_head) + bone.head else: - # calc root bone transform - location = armature_object.matrix_local @ bone.head # ARMATURE OBJECT Location - rot_matrix = bone.matrix @ armature_object.matrix_local.to_3x3() # ARMATURE OBJECT Rotation + location = armature_object.matrix_local @ bone.head + rot_matrix = bone.matrix @ armature_object.matrix_local.to_3x3() rotation = rot_matrix.to_quaternion() psk_bone.location.x = location.x diff --git a/src/psk.py b/src/psk/data.py similarity index 75% rename from src/psk.py rename to src/psk/data.py index 17f03cc..5d4c8d1 100644 --- a/src/psk.py +++ b/src/psk/data.py @@ -1,36 +1,9 @@ from typing import List -from ctypes import * +from ..data import * -class Vector3(Structure): - _fields_ = [ - ('x', c_float), - ('y', c_float), - ('z', c_float), - ] - - -class Quaternion(Structure): - _fields_ = [ - ('x', c_float), - ('y', c_float), - ('z', c_float), - ('w', c_float), - ] class Psk(object): - class Section(Structure): - _fields_ = [ - ('name', c_char * 20), - ('type_flags', c_int32), - ('data_size', c_int32), - ('data_count', c_int32) - ] - - def __init__(self, *args, **kw): - super().__init__(*args, **kw) - self.type_flags = 1999801 - class Wedge16(Structure): _fields_ = [ ('point_index', c_int16), diff --git a/src/psk/exporter.py b/src/psk/exporter.py new file mode 100644 index 0000000..7dbb1b9 --- /dev/null +++ b/src/psk/exporter.py @@ -0,0 +1,38 @@ +from typing import Type +from .data import * + + +class PskExporter(object): + def __init__(self, psk: Psk): + self.psk: Psk = psk + + @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): + # TODO: add logic somewhere to assert lengths of ctype structs + with open(path, 'wb') as fp: + self.write_section(fp, b'ACTRHEAD') + self.write_section(fp, b'PNTS0000', Vector3, self.psk.points) + + # WEDGES + # TODO: this really should be on the level of the builder, not the exporter + if len(self.psk.wedges) <= 65536: + wedge_type = Psk.Wedge16 + else: + wedge_type = Psk.Wedge32 + + self.write_section(fp, b'VTXW0000', wedge_type, self.psk.wedges) + self.write_section(fp, b'FACE0000', Psk.Face, self.psk.faces) + self.write_section(fp, b'MATT0000', Psk.Material, self.psk.materials) + self.write_section(fp, b'REFSKELT', Psk.Bone, self.psk.bones) + self.write_section(fp, b'RAWWEIGHTS', Psk.Weight, self.psk.weights) diff --git a/src/psk/operator.py b/src/psk/operator.py new file mode 100644 index 0000000..7075ceb --- /dev/null +++ b/src/psk/operator.py @@ -0,0 +1,26 @@ +from bpy.types import Operator +from bpy_extras.io_utils import ExportHelper +from bpy.props import StringProperty, BoolProperty, FloatProperty +from .builder import PskBuilder +from .exporter import PskExporter + + +class PskExportOperator(Operator, ExportHelper): + bl_idname = 'export.psk' + bl_label = 'Export' + __doc__ = 'PSK Exporter (.psk)' + filename_ext = '.psk' + filter_glob : StringProperty(default='*.psk', options={'HIDDEN'}) + + filepath : StringProperty( + name='File Path', + description='File path used for exporting the PSK file', + maxlen=1024, + default='') + + def execute(self, context): + builder = PskBuilder() + psk = builder.build(context) + exporter = PskExporter(psk) + exporter.export(self.filepath) + return {'FINISHED'}