Cool_Tools/io_scene_psk_psa
1
0
Fork 0
mirror of https://github.com/DarklightGames/io_scene_psk_psa.git synced 2024-11-28 00:20:48 +01:00
Code

Fixed a bug where animations would be incorrectly imported.

Browse Source
This commit is contained in:
Colin Basnett 2022-01-18 16:06:54 -08:00
parent 7fd0c6de81
commit 728f70a356
2 changed files with 54 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
io_export_psk_psa/psa/importer.py
View File

@ -1,14 +1,12 @@
import bpy import bpy
import mathutils import os
from mathutils import Vector, Quaternion, Matrix from mathutils import Vector, Quaternion, Matrix
from .data import Psa from .data import Psa
from typing import List, AnyStr, Optional from typing import List, AnyStr, Optional
import bpy
from bpy.types import Operator, Action, UIList, PropertyGroup, Panel, Armature, FileSelectParams from bpy.types import Operator, Action, UIList, PropertyGroup, Panel, Armature, FileSelectParams
from bpy_extras.io_utils import ExportHelper, ImportHelper from bpy_extras.io_utils import ExportHelper, ImportHelper
from bpy.props import StringProperty, BoolProperty, CollectionProperty, PointerProperty, IntProperty from bpy.props import StringProperty, BoolProperty, CollectionProperty, PointerProperty, IntProperty
from .reader import PsaReader from .reader import PsaReader
import numpy as np
class PsaImporter(object): class PsaImporter(object):
@ -16,9 +14,8 @@ class PsaImporter(object):
pass pass
def import_psa(self, psa_reader: PsaReader, 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 properties = context.scene.psa_import
sequences = map(lambda x: psa.sequences[x], sequence_names) sequences = map(lambda x: psa_reader.sequences[x], sequence_names)
armature_object = properties.armature_object armature_object = properties.armature_object
armature_data = armature_object.data armature_data = armature_object.data
@ -33,11 +30,11 @@ class PsaImporter(object):
self.post_quat: Quaternion = Quaternion() self.post_quat: Quaternion = Quaternion()
self.fcurves = [] self.fcurves = []
# create an index mapping from bones in the PSA to bones in the target armature. # Create an index mapping from bones in the PSA to bones in the target armature.
psa_to_armature_bone_indices = {} psa_to_armature_bone_indices = {}
armature_bone_names = [x.name for x in armature_data.bones] armature_bone_names = [x.name for x in armature_data.bones]
psa_bone_names = [] psa_bone_names = []
for psa_bone_index, psa_bone in enumerate(psa.bones): for psa_bone_index, psa_bone in enumerate(psa_reader.bones):
psa_bone_name = psa_bone.name.decode('windows-1252') psa_bone_name = psa_bone.name.decode('windows-1252')
psa_bone_names.append(psa_bone_name) psa_bone_names.append(psa_bone_name)
try: try:
@ -45,7 +42,7 @@ class PsaImporter(object):
except ValueError: except ValueError:
pass pass
# report if there are missing bones in the target armature # Report if there are missing bones in the target armature.
missing_bone_names = set(psa_bone_names).difference(set(armature_bone_names)) missing_bone_names = set(psa_bone_names).difference(set(armature_bone_names))
if len(missing_bone_names) > 0: 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(f'The armature object \'{armature_object.name}\' is missing the following bones that exist in the PSA:')
@ -56,7 +53,7 @@ class PsaImporter(object):
import_bones = [] import_bones = []
import_bones_dict = dict() import_bones_dict = dict()
for psa_bone_index, psa_bone in enumerate(psa.bones): for psa_bone_index, psa_bone in enumerate(psa_reader.bones):
bone_name = psa_bone.name.decode('windows-1252') bone_name = psa_bone.name.decode('windows-1252')
if psa_bone_index not in psa_to_armature_bone_indices: # TODO: replace with bone_name in armature_data.bones if psa_bone_index not in psa_to_armature_bone_indices: # TODO: replace with bone_name in armature_data.bones
# PSA bone does not map to armature bone, skip it and leave an empty bone in its place. # PSA bone does not map to armature bone, skip it and leave an empty bone in its place.
@ -93,14 +90,14 @@ class PsaImporter(object):
# Create and populate the data for new sequences. # Create and populate the data for new sequences.
for sequence in sequences: for sequence in sequences:
action = bpy.data.actions.new(name=sequence.name.decode()) action = bpy.data.actions.new(name=sequence.name.decode())
# Create f-curves for the rotation and location of each bone.
for psa_bone_index, armature_bone_index in psa_to_armature_bone_indices.items(): for psa_bone_index, armature_bone_index in psa_to_armature_bone_indices.items():
import_bone = import_bones[psa_bone_index] import_bone = import_bones[psa_bone_index]
pose_bone = import_bone.pose_bone pose_bone = import_bone.pose_bone
# create fcurves from rotation and location data
rotation_data_path = pose_bone.path_from_id('rotation_quaternion') rotation_data_path = pose_bone.path_from_id('rotation_quaternion')
location_data_path = pose_bone.path_from_id('location') location_data_path = pose_bone.path_from_id('location')
import_bone.fcurves.extend([ import_bone.fcurves = [
action.fcurves.new(rotation_data_path, index=0), # Qw action.fcurves.new(rotation_data_path, index=0), # Qw
action.fcurves.new(rotation_data_path, index=1), # Qx action.fcurves.new(rotation_data_path, index=1), # Qx
action.fcurves.new(rotation_data_path, index=2), # Qy action.fcurves.new(rotation_data_path, index=2), # Qy
@ -108,14 +105,14 @@ class PsaImporter(object):
action.fcurves.new(location_data_path, index=0), # Lx action.fcurves.new(location_data_path, index=0), # Lx
action.fcurves.new(location_data_path, index=1), # Ly action.fcurves.new(location_data_path, index=1), # Ly
action.fcurves.new(location_data_path, index=2), # Lz action.fcurves.new(location_data_path, index=2), # Lz
]) ]
key_index = 0
# Read the sequence keys from the PSA file. # Read the sequence keys from the PSA file.
sequence_name = sequence.name.decode('windows-1252') sequence_name = sequence.name.decode('windows-1252')
sequence_keys = psa_reader.read_sequence_keys(sequence_name) sequence_keys = psa_reader.read_sequence_keys(sequence_name)
# Add keyframes for each frame of the sequence.
key_index = 0
for frame_index in range(sequence.frame_count): for frame_index in range(sequence.frame_count):
for bone_index, import_bone in enumerate(import_bones): for bone_index, import_bone in enumerate(import_bones):
if import_bone is None: if import_bone is None:
@ -134,17 +131,22 @@ class PsaImporter(object):
else: else:
q.rotate(key_rotation) q.rotate(key_rotation)
quat.rotate(q.conjugated()) quat.rotate(q.conjugated())
key_location = Vector(tuple(sequence_keys[key_index].location)) key_location = Vector(tuple(sequence_keys[key_index].location))
loc = key_location - import_bone.orig_loc loc = key_location - import_bone.orig_loc
loc.rotate(import_bone.post_quat.conjugated()) loc.rotate(import_bone.post_quat.conjugated())
# Add keyframe data for each of the associated f-curves.
bone_fcurve_data = quat.w, quat.x, quat.y, quat.z, loc.x, loc.y, loc.z bone_fcurve_data = quat.w, quat.x, quat.y, quat.z, loc.x, loc.y, loc.z
for fcurve, datum in zip(import_bone.fcurves, bone_fcurve_data): for fcurve, datum in zip(import_bone.fcurves, bone_fcurve_data):
fcurve.keyframe_points.insert(frame_index, datum) fcurve.keyframe_points.insert(frame_index, datum, options={'FAST'})
key_index += 1 key_index += 1
# Explicitly update the f-curves.
for import_bone in filter(lambda x: x is not None, import_bones):
for fcurve in import_bone.fcurves:
fcurve.update()
class PsaImportActionListItem(PropertyGroup): class PsaImportActionListItem(PropertyGroup):
action_name: StringProperty() action_name: StringProperty()
@ -156,23 +158,27 @@ class PsaImportActionListItem(PropertyGroup):
return self.action_name return self.action_name
def on_psa_filepath_updated(property, context): def on_psa_file_path_updated(property, context):
context.scene.psa_import.action_list.clear() context.scene.psa_import.action_list.clear()
try: try:
# Read the file and populate the action list. # Read the file and populate the action list.
psa = PsaReader(context.scene.psa_import.psa_filepath).psa p = os.path.abspath(context.scene.psa_import.psa_file_path)
for sequence in psa.sequences.values(): print(p)
psa_reader = PsaReader(p)
for sequence in psa_reader.sequences.values():
item = context.scene.psa_import.action_list.add() item = context.scene.psa_import.action_list.add()
item.action_name = sequence.name.decode('windows-1252') item.action_name = sequence.name.decode('windows-1252')
item.frame_count = sequence.frame_count item.frame_count = sequence.frame_count
item.is_selected = True item.is_selected = True
except IOError: except IOError as e:
print('ERROR READING FILE')
print(e)
# TODO: set an error somewhere so the user knows the PSA could not be read. # TODO: set an error somewhere so the user knows the PSA could not be read.
pass pass
class PsaImportPropertyGroup(bpy.types.PropertyGroup): class PsaImportPropertyGroup(bpy.types.PropertyGroup):
psa_filepath: StringProperty(default='', subtype='FILE_PATH', update=on_psa_filepath_updated) psa_file_path: StringProperty(default='', subtype='FILE_PATH', update=on_psa_file_path_updated)
armature_object: PointerProperty(name='Armature', type=bpy.types.Object) armature_object: PointerProperty(name='Armature', type=bpy.types.Object)
action_list: CollectionProperty(type=PsaImportActionListItem) action_list: CollectionProperty(type=PsaImportActionListItem)
action_list_index: IntProperty(name='', default=0) action_list_index: IntProperty(name='', default=0)
@ -255,7 +261,7 @@ class PSA_PT_ImportPanel(Panel):
layout = self.layout layout = self.layout
scene = context.scene scene = context.scene
row = layout.row() row = layout.row()
row.prop(scene.psa_import, 'psa_filepath', text='PSA File') row.prop(scene.psa_import, 'psa_file_path', text='PSA File')
row = layout.row() row = layout.row()
row.prop_search(scene.psa_import, 'armature_object', bpy.data, 'objects') row.prop_search(scene.psa_import, 'armature_object', bpy.data, 'objects')
box = layout.box() box = layout.box()
@ -282,7 +288,7 @@ class PsaImportOperator(Operator):
return has_selected_actions and armature_object is not None return has_selected_actions and armature_object is not None
def execute(self, context): def execute(self, context):
psa_reader = PsaReader(context.scene.psa_import.psa_filepath) psa_reader = PsaReader(context.scene.psa_import.psa_file_path)
sequence_names = [x.action_name for x in context.scene.psa_import.action_list if x.is_selected] sequence_names = [x.action_name for x in context.scene.psa_import.action_list if x.is_selected]
PsaImporter().import_psa(psa_reader, sequence_names, context) PsaImporter().import_psa(psa_reader, sequence_names, context)
return {'FINISHED'} return {'FINISHED'}
@ -304,6 +310,6 @@ class PsaImportFileSelectOperator(Operator, ImportHelper):
return {'RUNNING_MODAL'} return {'RUNNING_MODAL'}
def execute(self, context): def execute(self, context):
context.scene.psa_import.psa_filepath = self.filepath context.scene.psa_import.psa_file_path = self.filepath
# Load the sequence names from the selected file # Load the sequence names from the selected file
return {'FINISHED'} return {'FINISHED'}

45
io_export_psk_psa/psa/reader.py
View File

@ -1,17 +1,28 @@
from .data import * from .data import *
from typing import AnyStr
import ctypes import ctypes
class PsaReader(object): class PsaReader(object):
"""
This class will read the sequences and bone information immediately upon instantiation and hold onto a file handle.
The key data is not read into memory upon instantiation due to it's potentially very large size.
To read the key data for a particular sequence, call `read_sequence_keys`.
"""
def __init__(self, path): def __init__(self, path):
self.keys_data_offset = 0 self.keys_data_offset: int = 0
self.fp = open(path, 'rb') self.fp = open(path, 'rb')
self.psa = self._read(self.fp) self.psa: Psa = self._read(self.fp)
@property
def bones(self):
return self.psa.bones
@property
def sequences(self):
return self.psa.sequences
@staticmethod @staticmethod
def read_types(fp, data_class: ctypes.Structure, section: Section, data): def _read_types(fp, data_class: ctypes.Structure, section: Section, data):
buffer_length = section.data_size * section.data_count buffer_length = section.data_size * section.data_count
buffer = fp.read(buffer_length) buffer = fp.read(buffer_length)
offset = 0 offset = 0
@ -19,22 +30,12 @@ class PsaReader(object):
data.append(data_class.from_buffer_copy(buffer, offset)) data.append(data_class.from_buffer_copy(buffer, offset))
offset += section.data_size offset += section.data_size
# 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):
fp.seek(-1, 1)
section = Section.from_buffer_copy(fp.read(ctypes.sizeof(Section)))
if section.name == b'ANIMINFO':
PsaReader.read_types(fp, Psa.Sequence, section, sequences)
return [sequence.name for sequence in sequences]
else:
fp.seek(section.data_size * section.data_count, 1)
return []
def read_sequence_keys(self, sequence_name) -> List[Psa.Key]: def read_sequence_keys(self, sequence_name) -> List[Psa.Key]:
""" Reads and returns the key data for a sequence.
:param sequence_name: The name of the sequence.
:return: A list of Psa.Keys.
"""
# Set the file reader to the beginning of the keys data # Set the file reader to the beginning of the keys data
sequence = self.psa.sequences[sequence_name] sequence = self.psa.sequences[sequence_name]
data_size = sizeof(Psa.Key) data_size = sizeof(Psa.Key)
@ -59,10 +60,10 @@ class PsaReader(object):
if section.name == b'ANIMHEAD': if section.name == b'ANIMHEAD':
pass pass
elif section.name == b'BONENAMES': elif section.name == b'BONENAMES':
PsaReader.read_types(fp, Psa.Bone, section, psa.bones) PsaReader._read_types(fp, Psa.Bone, section, psa.bones)
elif section.name == b'ANIMINFO': elif section.name == b'ANIMINFO':
sequences = [] sequences = []
PsaReader.read_types(fp, Psa.Sequence, section, sequences) PsaReader._read_types(fp, Psa.Sequence, section, sequences)
for sequence in sequences: for sequence in sequences:
psa.sequences[sequence.name.decode()] = sequence psa.sequences[sequence.name.decode()] = sequence
elif section.name == b'ANIMKEYS': elif section.name == b'ANIMKEYS':