mirror of
https://github.com/DarklightGames/io_scene_psk_psa.git
synced 2025-02-13 00:24:26 +01:00
335 lines
16 KiB
Python
335 lines
16 KiB
Python
import typing
|
|
from typing import List, Optional
|
|
|
|
import bpy
|
|
import numpy as np
|
|
from bpy.types import FCurve, Object, Context
|
|
from mathutils import Vector, Quaternion
|
|
|
|
from .config import PsaConfig, REMOVE_TRACK_LOCATION, REMOVE_TRACK_ROTATION
|
|
from .data import Psa
|
|
from .reader import PsaReader
|
|
|
|
|
|
class PsaImportOptions(object):
|
|
def __init__(self):
|
|
self.should_use_fake_user = False
|
|
self.should_stash = False
|
|
self.sequence_names = []
|
|
self.should_overwrite = False
|
|
self.should_write_keyframes = True
|
|
self.should_write_metadata = True
|
|
self.action_name_prefix = ''
|
|
self.should_convert_to_samples = False
|
|
self.bone_mapping_mode = 'CASE_INSENSITIVE'
|
|
self.fps_source = 'SEQUENCE'
|
|
self.fps_custom: float = 30.0
|
|
self.should_use_config_file = True
|
|
self.psa_config: PsaConfig = PsaConfig()
|
|
|
|
|
|
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.original_location: Vector = Vector()
|
|
self.original_rotation: Quaternion = Quaternion()
|
|
self.post_rotation: Quaternion = Quaternion()
|
|
self.fcurves: List[FCurve] = []
|
|
|
|
|
|
def _calculate_fcurve_data(import_bone: ImportBone, key_data: typing.Iterable[float]):
|
|
# Convert world-space transforms to local-space transforms.
|
|
key_rotation = Quaternion(key_data[0:4])
|
|
key_location = Vector(key_data[4:])
|
|
q = import_bone.post_rotation.copy()
|
|
q.rotate(import_bone.original_rotation)
|
|
rotation = q
|
|
q = import_bone.post_rotation.copy()
|
|
if import_bone.parent is None:
|
|
q.rotate(key_rotation.conjugated())
|
|
else:
|
|
q.rotate(key_rotation)
|
|
rotation.rotate(q.conjugated())
|
|
location = key_location - import_bone.original_location
|
|
location.rotate(import_bone.post_rotation.conjugated())
|
|
return rotation.w, rotation.x, rotation.y, rotation.z, location.x, location.y, location.z
|
|
|
|
|
|
class PsaImportResult:
|
|
def __init__(self):
|
|
self.warnings: List[str] = []
|
|
|
|
|
|
def _get_armature_bone_index_for_psa_bone(psa_bone_name: str, armature_bone_names: List[str], bone_mapping_mode: str = 'EXACT') -> Optional[int]:
|
|
"""
|
|
@param psa_bone_name: The name of the PSA bone.
|
|
@param armature_bone_names: The names of the bones in the armature.
|
|
@param bone_mapping_mode: One of 'EXACT' or 'CASE_INSENSITIVE'.
|
|
@return: The index of the armature bone that corresponds to the given PSA bone, or None if no such bone exists.
|
|
"""
|
|
for armature_bone_index, armature_bone_name in enumerate(armature_bone_names):
|
|
if bone_mapping_mode == 'CASE_INSENSITIVE':
|
|
if armature_bone_name.lower() == psa_bone_name.lower():
|
|
return armature_bone_index
|
|
else:
|
|
if armature_bone_name == psa_bone_name:
|
|
return armature_bone_index
|
|
return None
|
|
|
|
def _get_sample_frame_times(source_frame_count: int, frame_step: float) -> typing.Iterable[float]:
|
|
# TODO: for correctness, we should also emit the target frame time as well (because the last frame can be a
|
|
# fractional frame).
|
|
time = 0.0
|
|
while time < source_frame_count - 1:
|
|
yield time
|
|
time += frame_step
|
|
yield source_frame_count - 1
|
|
|
|
def _resample_sequence_data_matrix(sequence_data_matrix: np.ndarray, frame_step: float = 1.0) -> np.ndarray:
|
|
"""
|
|
Resamples the sequence data matrix to the target frame count.
|
|
@param sequence_data_matrix: FxBx7 matrix where F is the number of frames, B is the number of bones, and X is the
|
|
number of data elements per bone.
|
|
@param frame_step: The step between frames in the resampled sequence.
|
|
@return: The resampled sequence data matrix, or sequence_data_matrix if no resampling is necessary.
|
|
"""
|
|
if frame_step == 1.0:
|
|
# No resampling is necessary.
|
|
return sequence_data_matrix
|
|
|
|
source_frame_count, bone_count = sequence_data_matrix.shape[:2]
|
|
sample_frame_times = list(_get_sample_frame_times(source_frame_count, frame_step))
|
|
target_frame_count = len(sample_frame_times)
|
|
resampled_sequence_data_matrix = np.zeros((target_frame_count, bone_count, 7), dtype=float)
|
|
|
|
for sample_frame_index, sample_frame_time in enumerate(sample_frame_times):
|
|
frame_index = int(sample_frame_time)
|
|
if sample_frame_time % 1.0 == 0.0:
|
|
# Sample time has no fractional part, so just copy the frame.
|
|
resampled_sequence_data_matrix[sample_frame_index, :, :] = sequence_data_matrix[frame_index, :, :]
|
|
else:
|
|
# Sample time has a fractional part, so interpolate between two frames.
|
|
next_frame_index = frame_index + 1
|
|
for bone_index in range(bone_count):
|
|
source_frame_1_data = sequence_data_matrix[frame_index, bone_index, :]
|
|
source_frame_2_data = sequence_data_matrix[next_frame_index, bone_index, :]
|
|
factor = sample_frame_time - frame_index
|
|
q = Quaternion((source_frame_1_data[:4])).slerp(Quaternion((source_frame_2_data[:4])), factor)
|
|
q.normalize()
|
|
l = Vector(source_frame_1_data[4:]).lerp(Vector(source_frame_2_data[4:]), factor)
|
|
resampled_sequence_data_matrix[sample_frame_index, bone_index, :] = q.w, q.x, q.y, q.z, l.x, l.y, l.z
|
|
|
|
return resampled_sequence_data_matrix
|
|
|
|
|
|
def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, options: PsaImportOptions) -> PsaImportResult:
|
|
result = PsaImportResult()
|
|
sequences = [psa_reader.sequences[x] for x in options.sequence_names]
|
|
armature_data = typing.cast(bpy.types.Armature, armature_object.data)
|
|
|
|
# Create an index mapping from bones in the PSA to bones in the target armature.
|
|
psa_to_armature_bone_indices = {}
|
|
armature_to_psa_bone_indices = {}
|
|
armature_bone_names = [x.name for x in armature_data.bones]
|
|
psa_bone_names = []
|
|
duplicate_mappings = []
|
|
|
|
for psa_bone_index, psa_bone in enumerate(psa_reader.bones):
|
|
psa_bone_name: str = psa_bone.name.decode('windows-1252')
|
|
armature_bone_index = _get_armature_bone_index_for_psa_bone(psa_bone_name, armature_bone_names, options.bone_mapping_mode)
|
|
if armature_bone_index is not None:
|
|
# Ensure that no other PSA bone has been mapped to this armature bone yet.
|
|
if armature_bone_index not in armature_to_psa_bone_indices:
|
|
psa_to_armature_bone_indices[psa_bone_index] = armature_bone_index
|
|
armature_to_psa_bone_indices[armature_bone_index] = psa_bone_index
|
|
else:
|
|
# This armature bone has already been mapped to a PSA bone.
|
|
duplicate_mappings.append((psa_bone_index, armature_bone_index, armature_to_psa_bone_indices[armature_bone_index]))
|
|
psa_bone_names.append(armature_bone_names[armature_bone_index])
|
|
else:
|
|
psa_bone_names.append(psa_bone_name)
|
|
|
|
# Warn about duplicate bone mappings.
|
|
if len(duplicate_mappings) > 0:
|
|
for (psa_bone_index, armature_bone_index, mapped_psa_bone_index) in duplicate_mappings:
|
|
psa_bone_name = psa_bone_names[psa_bone_index]
|
|
armature_bone_name = armature_bone_names[armature_bone_index]
|
|
mapped_psa_bone_name = psa_bone_names[mapped_psa_bone_index]
|
|
result.warnings.append(f'PSA bone {psa_bone_index} ({psa_bone_name}) could not be mapped to armature bone {armature_bone_index} ({armature_bone_name}) because the armature bone is already mapped to PSA bone {mapped_psa_bone_index} ({mapped_psa_bone_name})')
|
|
|
|
# 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:
|
|
result.warnings.append(
|
|
f'The armature \'{armature_object.name}\' is missing {len(missing_bone_names)} bones that exist in '
|
|
'the PSA:\n' +
|
|
str(list(sorted(missing_bone_names)))
|
|
)
|
|
del armature_bone_names
|
|
|
|
# Create intermediate bone data for import operations.
|
|
import_bones = []
|
|
psa_bone_names_to_import_bones = dict()
|
|
|
|
for (psa_bone_index, psa_bone), psa_bone_name in zip(enumerate(psa_reader.bones), psa_bone_names):
|
|
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)
|
|
import_bone.armature_bone = armature_data.bones[psa_bone_name]
|
|
import_bone.pose_bone = armature_object.pose.bones[psa_bone_name]
|
|
psa_bone_names_to_import_bones[psa_bone_name] = import_bone
|
|
import_bones.append(import_bone)
|
|
|
|
bones_with_missing_parents = []
|
|
|
|
for import_bone in filter(lambda x: x is not None, import_bones):
|
|
armature_bone = import_bone.armature_bone
|
|
has_parent = armature_bone.parent is not None
|
|
if has_parent:
|
|
if armature_bone.parent.name in psa_bone_names:
|
|
import_bone.parent = psa_bone_names_to_import_bones[armature_bone.parent.name]
|
|
else:
|
|
# Add a warning if the parent bone is not in the PSA.
|
|
bones_with_missing_parents.append(armature_bone)
|
|
# Calculate the original location & rotation of each bone (in world-space maybe?)
|
|
if has_parent:
|
|
import_bone.original_location = armature_bone.matrix_local.translation - armature_bone.parent.matrix_local.translation
|
|
import_bone.original_location.rotate(armature_bone.parent.matrix_local.to_quaternion().conjugated())
|
|
import_bone.original_rotation = armature_bone.matrix_local.to_quaternion()
|
|
import_bone.original_rotation.rotate(armature_bone.parent.matrix_local.to_quaternion().conjugated())
|
|
import_bone.original_rotation.conjugate()
|
|
else:
|
|
import_bone.original_location = armature_bone.matrix_local.translation.copy()
|
|
import_bone.original_rotation = armature_bone.matrix_local.to_quaternion().conjugated()
|
|
|
|
import_bone.post_rotation = import_bone.original_rotation.conjugated()
|
|
|
|
# Warn about bones with missing parents.
|
|
if len(bones_with_missing_parents) > 0:
|
|
count = len(bones_with_missing_parents)
|
|
message = f'{count} bone(s) have parents that are not present in the PSA:\n' + str([x.name for x in bones_with_missing_parents])
|
|
result.warnings.append(message)
|
|
|
|
context.window_manager.progress_begin(0, len(sequences))
|
|
|
|
# Create and populate the data for new sequences.
|
|
actions = []
|
|
for sequence_index, sequence in enumerate(sequences):
|
|
# Add the action.
|
|
sequence_name = sequence.name.decode('windows-1252')
|
|
action_name = options.action_name_prefix + sequence_name
|
|
|
|
# Get the bone track flags for this sequence, or an empty dictionary if none exist.
|
|
sequence_bone_track_flags = dict()
|
|
if sequence_name in options.psa_config.sequence_bone_flags.keys():
|
|
sequence_bone_track_flags = options.psa_config.sequence_bone_flags[sequence_name]
|
|
|
|
if options.should_overwrite and action_name in bpy.data.actions:
|
|
action = bpy.data.actions[action_name]
|
|
else:
|
|
action = bpy.data.actions.new(name=action_name)
|
|
|
|
# Calculate the target FPS.
|
|
match options.fps_source:
|
|
case 'CUSTOM':
|
|
target_fps = options.fps_custom
|
|
case 'SCENE':
|
|
target_fps = context.scene.render.fps
|
|
case 'SEQUENCE':
|
|
target_fps = sequence.fps
|
|
case _:
|
|
raise ValueError(f'Unknown FPS source: {options.fps_source}')
|
|
|
|
if options.should_write_keyframes:
|
|
# Remove existing f-curves.
|
|
action.fcurves.clear()
|
|
|
|
# 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():
|
|
bone_track_flags = sequence_bone_track_flags.get(psa_bone_index, 0)
|
|
import_bone = import_bones[psa_bone_index]
|
|
pose_bone = import_bone.pose_bone
|
|
rotation_data_path = pose_bone.path_from_id('rotation_quaternion')
|
|
location_data_path = pose_bone.path_from_id('location')
|
|
add_rotation_fcurves = (bone_track_flags & REMOVE_TRACK_ROTATION) == 0
|
|
add_location_fcurves = (bone_track_flags & REMOVE_TRACK_LOCATION) == 0
|
|
import_bone.fcurves = [
|
|
action.fcurves.new(rotation_data_path, index=0, action_group=pose_bone.name) if add_rotation_fcurves else None, # Qw
|
|
action.fcurves.new(rotation_data_path, index=1, action_group=pose_bone.name) if add_rotation_fcurves else None, # Qx
|
|
action.fcurves.new(rotation_data_path, index=2, action_group=pose_bone.name) if add_rotation_fcurves else None, # Qy
|
|
action.fcurves.new(rotation_data_path, index=3, action_group=pose_bone.name) if add_rotation_fcurves else None, # Qz
|
|
action.fcurves.new(location_data_path, index=0, action_group=pose_bone.name) if add_location_fcurves else None, # Lx
|
|
action.fcurves.new(location_data_path, index=1, action_group=pose_bone.name) if add_location_fcurves else None, # Ly
|
|
action.fcurves.new(location_data_path, index=2, action_group=pose_bone.name) if add_location_fcurves else None, # Lz
|
|
]
|
|
|
|
# Read the sequence data matrix from the PSA.
|
|
sequence_data_matrix = psa_reader.read_sequence_data_matrix(sequence_name)
|
|
|
|
# Convert the sequence's data from world-space to local-space.
|
|
for bone_index, import_bone in enumerate(import_bones):
|
|
if import_bone is None:
|
|
continue
|
|
for frame_index in range(sequence.frame_count):
|
|
# This bone has writeable keyframes for this frame.
|
|
key_data = sequence_data_matrix[frame_index, bone_index]
|
|
# Calculate the local-space key data for the bone.
|
|
sequence_data_matrix[frame_index, bone_index] = _calculate_fcurve_data(import_bone, key_data)
|
|
|
|
# Resample the sequence data to the target FPS.
|
|
# If the target frame count is the same as the source frame count, this will be a no-op.
|
|
resampled_sequence_data_matrix = _resample_sequence_data_matrix(sequence_data_matrix,
|
|
frame_step=sequence.fps / target_fps)
|
|
|
|
# Write the keyframes out.
|
|
# Note that the f-curve data consists of alternating time and value data.
|
|
target_frame_count = resampled_sequence_data_matrix.shape[0]
|
|
fcurve_data = np.zeros(2 * target_frame_count, dtype=float)
|
|
fcurve_data[0::2] = range(0, target_frame_count)
|
|
|
|
for bone_index, import_bone in enumerate(import_bones):
|
|
if import_bone is None:
|
|
continue
|
|
for fcurve_index, fcurve in enumerate(import_bone.fcurves):
|
|
if fcurve is None:
|
|
continue
|
|
fcurve_data[1::2] = resampled_sequence_data_matrix[:, bone_index, fcurve_index]
|
|
fcurve.keyframe_points.add(target_frame_count)
|
|
fcurve.keyframe_points.foreach_set('co', fcurve_data)
|
|
for fcurve_keyframe in fcurve.keyframe_points:
|
|
fcurve_keyframe.interpolation = 'LINEAR'
|
|
|
|
if options.should_convert_to_samples:
|
|
# Bake the curve to samples.
|
|
for fcurve in action.fcurves:
|
|
fcurve.convert_to_samples(start=0, end=sequence.frame_count)
|
|
|
|
# Write meta-data.
|
|
if options.should_write_metadata:
|
|
action.psa_export.fps = target_fps
|
|
|
|
action.use_fake_user = options.should_use_fake_user
|
|
|
|
actions.append(action)
|
|
|
|
context.window_manager.progress_update(sequence_index)
|
|
|
|
# If the user specifies, store the new animations as strips on a non-contributing NLA track.
|
|
if options.should_stash:
|
|
if armature_object.animation_data is None:
|
|
armature_object.animation_data_create()
|
|
for action in actions:
|
|
nla_track = armature_object.animation_data.nla_tracks.new()
|
|
nla_track.name = action.name
|
|
nla_track.mute = True
|
|
nla_track.strips.new(name=action.name, start=0, action=action)
|
|
|
|
context.window_manager.progress_end()
|
|
|
|
return result
|