Importing animations now mostly works, though the root bone is messed up.

This commit is contained in:
Colin Basnett
2022-01-14 21:04:18 -08:00
parent 9fa0780032
commit 5607788d1c
5 changed files with 199 additions and 69 deletions

View File

@@ -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] = []

View File

@@ -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'}

View File

@@ -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