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

Added PSA resampling

Fixed PSA import resampling logic
This commit is contained in:
Colin Basnett 2024-02-13 14:03:04 -08:00
parent d92f2d77d2
commit 09cc9e5d51

View File

@ -2,7 +2,7 @@ import typing
from typing import List, Optional from typing import List, Optional
import bpy import bpy
import numpy import numpy as np
from bpy.types import FCurve, Object, Context from bpy.types import FCurve, Object, Context
from mathutils import Vector, Quaternion from mathutils import Vector, Quaternion
@ -80,6 +80,52 @@ def _get_armature_bone_index_for_psa_bone(psa_bone_name: str, armature_bone_name
return None return None
def _resample_sequence_data_matrix(sequence_data_matrix: np.ndarray, time_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 target_frame_count: The number of frames to resample to.
@return: The resampled sequence data matrix, or sequence_data_matrix if no resampling is necessary.
'''
def get_sample_times(source_frame_count: int, time_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 += time_step
yield source_frame_count - 1
if time_step == 1.0:
# No resampling is necessary.
return sequence_data_matrix
source_frame_count, bone_count = sequence_data_matrix.shape[:2]
sample_times = list(get_sample_times(source_frame_count, time_step))
target_frame_count = len(sample_times)
resampled_sequence_data_matrix = np.zeros((target_frame_count, bone_count, 7), dtype=float)
for sample_index, sample_time in enumerate(sample_times):
frame_index = int(sample_time)
if sample_time % 1.0 == 0.0:
# Sample time has no fractional part, so just copy the frame.
resampled_sequence_data_matrix[sample_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_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_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: def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, options: PsaImportOptions) -> PsaImportResult:
result = PsaImportResult() result = PsaImportResult()
sequences = [psa_reader.sequences[x] for x in options.sequence_names] sequences = [psa_reader.sequences[x] for x in options.sequence_names]
@ -187,12 +233,9 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object,
case _: case _:
raise ValueError(f'Unknown FPS source: {options.fps_source}') raise ValueError(f'Unknown FPS source: {options.fps_source}')
keyframe_time_dilation = target_fps / sequence.fps
if options.should_write_keyframes: if options.should_write_keyframes:
# Remove existing f-curves (replace with action.fcurves.clear() in Blender 3.2) # Remove existing f-curves.
while len(action.fcurves) > 0: action.fcurves.clear()
action.fcurves.remove(action.fcurves[-1])
# Create f-curves for the rotation and location of each bone. # 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():
@ -226,19 +269,25 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object,
# Calculate the local-space key data for the bone. # Calculate the local-space key data for the bone.
sequence_data_matrix[frame_index, bone_index] = _calculate_fcurve_data(import_bone, key_data) sequence_data_matrix[frame_index, bone_index] = _calculate_fcurve_data(import_bone, key_data)
# Write the keyframes out. # Resample the sequence data to the target FPS.
fcurve_data = numpy.zeros(2 * sequence.frame_count, dtype=float) # 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,
time_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)
# Populate the keyframe time data.
fcurve_data[0::2] = [x * keyframe_time_dilation for x 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:
continue continue
for fcurve_index, fcurve in enumerate(import_bone.fcurves): for fcurve_index, fcurve in enumerate(import_bone.fcurves):
if fcurve is None: if fcurve is None:
continue continue
fcurve_data[1::2] = sequence_data_matrix[:, bone_index, fcurve_index] fcurve_data[1::2] = resampled_sequence_data_matrix[:, bone_index, fcurve_index]
fcurve.keyframe_points.add(sequence.frame_count) fcurve.keyframe_points.add(target_frame_count)
fcurve.keyframe_points.foreach_set('co', fcurve_data) fcurve.keyframe_points.foreach_set('co', fcurve_data)
for fcurve_keyframe in fcurve.keyframe_points: for fcurve_keyframe in fcurve.keyframe_points:
fcurve_keyframe.interpolation = 'LINEAR' fcurve_keyframe.interpolation = 'LINEAR'