diff --git a/io_export_psk_psa/data.py b/io_export_psk_psa/data.py index a9a9a8f..72bb979 100644 --- a/io_export_psk_psa/data.py +++ b/io_export_psk_psa/data.py @@ -13,6 +13,9 @@ class Vector3(Structure): yield self.y yield self.z + def __repr__(self): + return repr(tuple(self)) + class Quaternion(Structure): _fields_ = [ @@ -28,6 +31,9 @@ class Quaternion(Structure): yield self.y yield self.z + def __repr__(self): + return repr(tuple(self)) + class Section(Structure): _fields_ = [ diff --git a/io_export_psk_psa/psa/data.py b/io_export_psk_psa/psa/data.py index b43503c..819d81c 100644 --- a/io_export_psk_psa/psa/data.py +++ b/io_export_psk_psa/psa/data.py @@ -1,9 +1,13 @@ from typing import List, Dict from ..data import * +""" +Note that keys are not stored within the Psa object. +Use the PsaReader::get_sequence_keys to get a the keys for a sequence. +""" + class Psa(object): - class Bone(Structure): _fields_ = [ ('name', c_char * 64), @@ -38,7 +42,9 @@ class Psa(object): ('time', c_float) ] + def __repr__(self) -> str: + return repr((self.location, self.rotation, self.time)) + def __init__(self): self.bones: List[Psa.Bone] = [] self.sequences: Dict[Psa.Sequence] = {} - self.keys: List[Psa.Key] = [] diff --git a/io_export_psk_psa/psa/importer.py b/io_export_psk_psa/psa/importer.py index a3198f6..13be09e 100644 --- a/io_export_psk_psa/psa/importer.py +++ b/io_export_psk_psa/psa/importer.py @@ -1,7 +1,8 @@ import bpy import mathutils +from mathutils import Vector, Quaternion, Matrix from .data import Psa -from typing import List, AnyStr +from typing import List, AnyStr, Optional import bpy from bpy.types import Operator, Action, UIList, PropertyGroup, Panel, Armature from bpy_extras.io_utils import ExportHelper, ImportHelper @@ -13,64 +14,154 @@ class PsaImporter(object): def __init__(self): pass - def import_psa(self, psa: Psa, sequence_names: List[AnyStr], context): + def import_psa(self, psa_reader: PsaReader, sequence_names: List[AnyStr], context): + psa = psa_reader.psa properties = context.scene.psa_import sequences = map(lambda x: psa.sequences[x], sequence_names) - armature_object = properties.armature_object armature_data = armature_object.data + class ImportBone(object): + def __init__(self, psa_bone: Psa.Bone): + self.psa_bone: Psa.Bone = psa_bone + self.parent: Optional[ImportBone] = None + self.armature_bone = None + self.pose_bone = None + self.orig_loc: Vector = Vector() + self.orig_quat: Quaternion = Quaternion() + self.post_quat: Quaternion = Quaternion() + # TODO: this is UGLY, come up with a way to just map indices for these + self.fcurve_quat_w = None + self.fcurve_quat_x = None + self.fcurve_quat_y = None + self.fcurve_quat_z = None + self.fcurve_location_x = None + self.fcurve_location_y = None + self.fcurve_location_z = None + # create an index mapping from bones in the PSA to bones in the target armature. - bone_indices = {} - data_bone_names = [x.name for x in armature_data.bones] - for index, psa_bone in enumerate(psa.bones): - psa_bone_name = psa_bone.name.decode() + psa_to_armature_bone_indices = {} + armature_bone_names = [x.name for x in armature_data.bones] + psa_bone_names = [] + for psa_bone_index, psa_bone in enumerate(psa.bones): + psa_bone_name = psa_bone.name.decode('windows-1252') + psa_bone_names.append(psa_bone_name) try: - bone_indices[index] = data_bone_names.index(psa_bone_name) + psa_to_armature_bone_indices[psa_bone_index] = armature_bone_names.index(psa_bone_name) except ValueError: pass - del data_bone_names + # report if there are missing bones in the target armature + missing_bone_names = set(psa_bone_names).difference(set(armature_bone_names)) + if len(missing_bone_names) > 0: + print(f'The armature object \'{armature_object.name}\' is missing the following bones that exist in the PSA:') + print(list(sorted(missing_bone_names))) + del armature_bone_names + + # Create intermediate bone data for import operations. + import_bones = [] + for psa_bone_index, psa_bone in enumerate(psa.bones): + bone_name = psa_bone.name.decode('windows-1252') + if psa_bone_index not in psa_to_armature_bone_indices: + # PSA bone does not map to armature bone, skip it and leave an empty bone in its place. + import_bones.append(None) + continue + import_bone = ImportBone(psa_bone) + armature_bone = armature_data.bones[bone_name] + import_bone.pose_bone = armature_object.pose.bones[bone_name] + if psa_bone_index > 0: + import_bone.parent = import_bones[psa_bone.parent_index] + # Calculate the original location & rotation of each bone (in world space maybe?) + if import_bone.parent is not None: + import_bone.orig_loc = armature_bone.matrix_local.translation - armature_bone.parent.matrix_local.translation + import_bone.orig_loc.rotate(armature_bone.parent.matrix_local.to_quaternion().conjugated()) + import_bone.orig_quat = armature_bone.matrix_local.to_quaternion() + import_bone.orig_quat.rotate(armature_bone.parent.matrix_local.to_quaternion().conjugated()) + import_bone.orig_quat.conjugate() + else: + import_bone.orig_loc = armature_bone.matrix_local.translation.copy() + import_bone.orig_quat = armature_bone.matrix_local.to_quaternion() + import_bone.post_quat = import_bone.orig_quat.conjugated() + import_bones.append(import_bone) + + # Create and populate the data for new sequences. for sequence in sequences: action = bpy.data.actions.new(name=sequence.name.decode()) - for psa_bone_index, armature_bone_index in bone_indices.items(): - psa_bone = psa.bones[psa_bone_index] + # TODO: problem might be here (yea, we are confused about the ordering of these things!) + for psa_bone_index, armature_bone_index in psa_to_armature_bone_indices.items(): + import_bone = import_bones[psa_bone_index] pose_bone = armature_object.pose.bones[armature_bone_index] # rotation rotation_data_path = pose_bone.path_from_id('rotation_quaternion') - fcurve_quat_w = action.fcurves.new(rotation_data_path, index=0) - fcurve_quat_x = action.fcurves.new(rotation_data_path, index=0) - fcurve_quat_y = action.fcurves.new(rotation_data_path, index=0) - fcurve_quat_z = action.fcurves.new(rotation_data_path, index=0) + import_bone.fcurve_quat_w = action.fcurves.new(rotation_data_path, index=0) + import_bone.fcurve_quat_x = action.fcurves.new(rotation_data_path, index=1) + import_bone.fcurve_quat_y = action.fcurves.new(rotation_data_path, index=2) + import_bone.fcurve_quat_z = action.fcurves.new(rotation_data_path, index=3) # location location_data_path = pose_bone.path_from_id('location') - fcurve_location_x = action.fcurves.new(location_data_path, index=0) - fcurve_location_y = action.fcurves.new(location_data_path, index=1) - fcurve_location_z = action.fcurves.new(location_data_path, index=2) + import_bone.fcurve_location_x = action.fcurves.new(location_data_path, index=0) + import_bone.fcurve_location_y = action.fcurves.new(location_data_path, index=1) + import_bone.fcurve_location_z = action.fcurves.new(location_data_path, index=2) # add keyframes - fcurve_quat_w.keyframe_points.add(sequence.frame_count) - fcurve_quat_x.keyframe_points.add(sequence.frame_count) - fcurve_quat_y.keyframe_points.add(sequence.frame_count) - fcurve_quat_z.keyframe_points.add(sequence.frame_count) - fcurve_location_x.keyframe_points.add(sequence.frame_count) - fcurve_location_y.keyframe_points.add(sequence.frame_count) - fcurve_location_z.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_quat_w.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_quat_x.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_quat_y.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_quat_z.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_location_x.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_location_y.keyframe_points.add(sequence.frame_count) + import_bone.fcurve_location_z.keyframe_points.add(sequence.frame_count) + + fcurve_interpolation = 'LINEAR' + should_invert_root = False + + key_index = 0 + sequence_name = sequence.name.decode('windows-1252') + sequence_keys = psa_reader.get_sequence_keys(sequence_name) - raw_key_index = 0 # ? for frame_index in range(sequence.frame_count): - for psa_bone_index in range(len(psa.bones)): - if psa_bone_index not in bone_indices: + for import_bone in import_bones: + if import_bone is None: # bone does not exist in the armature, skip it - raw_key_index += 1 + key_index += 1 continue - psa_bone = psa.bones[psa_bone_index] - # ... + key_location = Vector(tuple(sequence_keys[key_index].location)) + key_rotation = Quaternion(tuple(sequence_keys[key_index].rotation)) - raw_key_index += 1 + # TODO: what is this doing exactly? + q = import_bone.post_quat.copy() + q.rotate(import_bone.orig_quat) + quat = q + q = import_bone.post_quat.copy() + if import_bone.parent is None and should_invert_root: + q.rotate(import_bone.orig_quat) + else: + q.rotate(key_rotation) + quat.rotate(q.conjugated()) + + loc = key_location - import_bone.orig_loc + loc.rotate(import_bone.post_quat.conjugated()) + + import_bone.fcurve_quat_w.keyframe_points[frame_index].co = frame_index, quat.w + import_bone.fcurve_quat_x.keyframe_points[frame_index].co = frame_index, quat.x + import_bone.fcurve_quat_y.keyframe_points[frame_index].co = frame_index, quat.y + import_bone.fcurve_quat_z.keyframe_points[frame_index].co = frame_index, quat.z + import_bone.fcurve_location_x.keyframe_points[frame_index].co = frame_index, loc.x + import_bone.fcurve_location_y.keyframe_points[frame_index].co = frame_index, loc.y + import_bone.fcurve_location_z.keyframe_points[frame_index].co = frame_index, loc.z + + import_bone.fcurve_quat_w.keyframe_points[frame_index].interpolation = fcurve_interpolation + import_bone.fcurve_quat_x.keyframe_points[frame_index].interpolation = fcurve_interpolation + import_bone.fcurve_quat_y.keyframe_points[frame_index].interpolation = fcurve_interpolation + import_bone.fcurve_quat_z.keyframe_points[frame_index].interpolation = fcurve_interpolation + import_bone.fcurve_location_x.keyframe_points[frame_index].interpolation = fcurve_interpolation + import_bone.fcurve_location_z.keyframe_points[frame_index].interpolation = fcurve_interpolation + import_bone.fcurve_location_z.keyframe_points[frame_index].interpolation = fcurve_interpolation + + key_index += 1 class PsaImportActionListItem(PropertyGroup): @@ -178,9 +269,9 @@ class PsaImportOperator(Operator): bl_label = 'Import' def execute(self, context): - psa = PsaReader().read(context.scene.psa_import.cool_filepath) + psa_reader = PsaReader(context.scene.psa_import.cool_filepath) sequence_names = [x.action_name for x in context.scene.psa_import.action_list if x.is_selected] - PsaImporter().import_psa(psa, sequence_names, context) + PsaImporter().import_psa(psa_reader, sequence_names, context) return {'FINISHED'} @@ -202,14 +293,14 @@ class PsaImportFileSelectOperator(Operator, ImportHelper): def execute(self, context): context.scene.psa_import.cool_filepath = self.filepath # Load the sequence names from the selected file - action_names = [] + sequence_names = [] try: - action_names = PsaReader().scan_sequence_names(self.filepath) + sequence_names = PsaReader.scan_sequence_names(self.filepath) except IOError: pass context.scene.psa_import.action_list.clear() - for action_name in action_names: + for sequence_name in sequence_names: item = context.scene.psa_import.action_list.add() - item.action_name = action_name.decode() + item.action_name = sequence_name.decode('windows-1252') item.is_selected = True return {'FINISHED'} diff --git a/io_export_psk_psa/psa/reader.py b/io_export_psk_psa/psa/reader.py index badd3c8..f127ccd 100644 --- a/io_export_psk_psa/psa/reader.py +++ b/io_export_psk_psa/psa/reader.py @@ -5,8 +5,10 @@ import ctypes class PsaReader(object): - def __init__(self): - pass + def __init__(self, path): + self.keys_data_offset = 0 + self.fp = open(path, 'rb') + self.psa = self._read(self.fp) @staticmethod def read_types(fp, data_class: ctypes.Structure, section: Section, data): @@ -17,7 +19,9 @@ class PsaReader(object): data.append(data_class.from_buffer_copy(buffer, offset)) offset += section.data_size - def scan_sequence_names(self, path) -> List[AnyStr]: + # TODO: this probably isn't actually needed anymore, we can just read it once. + @staticmethod + def scan_sequence_names(path) -> List[AnyStr]: sequences = [] with open(path, 'rb') as fp: while fp.read(1): @@ -30,26 +34,50 @@ class PsaReader(object): fp.seek(section.data_size * section.data_count, 1) return [] - def read(self, path) -> Psa: + def get_sequence_keys(self, sequence_name) -> List[Psa.Key]: + # Set the file reader to the beginning of the keys data + sequence = self.psa.sequences[sequence_name] + data_size = sizeof(Psa.Key) + bone_count = len(self.psa.bones) + buffer_length = data_size * bone_count * sequence.frame_count + print(f'data_size: {data_size}') + print(f'buffer_length: {buffer_length}') + print(f'bone_count: {bone_count}') + print(f'sequence.frame_count: {sequence.frame_count}') + print(f'self.keys_data_offset: {self.keys_data_offset}') + sequence_keys_offset = self.keys_data_offset + (sequence.frame_start_index * bone_count * data_size) + print(f'sequence_keys_offset: {sequence_keys_offset}') + self.fp.seek(sequence_keys_offset, 0) + buffer = self.fp.read(buffer_length) + offset = 0 + keys = [] + for _ in range(sequence.frame_count * bone_count): + key = Psa.Key.from_buffer_copy(buffer, offset) + keys.append(key) + offset += data_size + return keys + + def _read(self, fp) -> Psa: psa = Psa() - with open(path, 'rb') as fp: - while fp.read(1): - fp.seek(-1, 1) - section = Section.from_buffer_copy(fp.read(ctypes.sizeof(Section))) - if section.name == b'ANIMHEAD': - pass - elif section.name == b'BONENAMES': - PsaReader.read_types(fp, Psa.Bone, section, psa.bones) - elif section.name == b'ANIMINFO': - sequences = [] - PsaReader.read_types(fp, Psa.Sequence, section, sequences) - for sequence in sequences: - psa.sequences[sequence.name.decode()] = sequence - elif section.name == b'ANIMKEYS': - PsaReader.read_types(fp, Psa.Key, section, psa.keys) - elif section.name in [b'SCALEKEYS']: - fp.seek(section.data_size * section.data_count, 1) - else: - raise RuntimeError(f'Unrecognized section "{section.name}"') + while fp.read(1): + fp.seek(-1, 1) + section = Section.from_buffer_copy(fp.read(ctypes.sizeof(Section))) + if section.name == b'ANIMHEAD': + pass + elif section.name == b'BONENAMES': + PsaReader.read_types(fp, Psa.Bone, section, psa.bones) + elif section.name == b'ANIMINFO': + sequences = [] + PsaReader.read_types(fp, Psa.Sequence, section, sequences) + for sequence in sequences: + psa.sequences[sequence.name.decode()] = sequence + elif section.name == b'ANIMKEYS': + # Skip keys on this pass. We will keep this file open and read from it as needed. + self.keys_data_offset = fp.tell() + fp.seek(section.data_size * section.data_count, 1) + elif section.name in [b'SCALEKEYS']: + fp.seek(section.data_size * section.data_count, 1) + else: + raise RuntimeError(f'Unrecognized section "{section.name}"') return psa 1 \ No newline at end of file diff --git a/io_export_psk_psa/psk/importer.py b/io_export_psk_psa/psk/importer.py index 49ef02a..19af075 100644 --- a/io_export_psk_psa/psk/importer.py +++ b/io_export_psk_psa/psk/importer.py @@ -79,13 +79,12 @@ class PskImporter(object): edit_bone.parent = armature_data.edit_bones[bone.psk_bone.parent_index] elif not should_invert_root: bone.local_rotation.conjugate() - post_quat = bone.local_rotation.conjugated() edit_bone.tail = Vector((0.0, new_bone_size, 0.0)) - m = post_quat.copy() - m.rotate(bone.world_matrix) - m = m.to_matrix().to_4x4() - m.translation = bone.world_matrix.translation - edit_bone.matrix = m + edit_bone_matrix = bone.local_rotation.conjugated() + edit_bone_matrix.rotate(bone.world_matrix) + edit_bone_matrix = edit_bone_matrix.to_matrix().to_4x4() + edit_bone_matrix.translation = bone.world_matrix.translation + edit_bone.matrix = edit_bone_matrix # MESH mesh_data = bpy.data.meshes.new(name)