From 09cc9e5d5115b1834adfdd7d70031d768e3ae13f Mon Sep 17 00:00:00 2001 From: Colin Basnett Date: Tue, 13 Feb 2024 14:03:04 -0800 Subject: [PATCH] Added PSA resampling Fixed PSA import resampling logic --- io_scene_psk_psa/psa/importer.py | 73 ++++++++++++++++++++++++++------ 1 file changed, 61 insertions(+), 12 deletions(-) diff --git a/io_scene_psk_psa/psa/importer.py b/io_scene_psk_psa/psa/importer.py index 630e3f6..3b1a185 100644 --- a/io_scene_psk_psa/psa/importer.py +++ b/io_scene_psk_psa/psa/importer.py @@ -2,7 +2,7 @@ import typing from typing import List, Optional import bpy -import numpy +import numpy as np from bpy.types import FCurve, Object, Context 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 +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: result = PsaImportResult() 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 _: raise ValueError(f'Unknown FPS source: {options.fps_source}') - keyframe_time_dilation = target_fps / sequence.fps - if options.should_write_keyframes: - # Remove existing f-curves (replace with action.fcurves.clear() in Blender 3.2) - while len(action.fcurves) > 0: - action.fcurves.remove(action.fcurves[-1]) + # 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(): @@ -226,19 +269,25 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, # Calculate the local-space key data for the bone. sequence_data_matrix[frame_index, bone_index] = _calculate_fcurve_data(import_bone, key_data) - # Write the keyframes out. - fcurve_data = numpy.zeros(2 * sequence.frame_count, dtype=float) + # 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, + 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): if import_bone is None: continue for fcurve_index, fcurve in enumerate(import_bone.fcurves): if fcurve is None: continue - fcurve_data[1::2] = sequence_data_matrix[:, bone_index, fcurve_index] - fcurve.keyframe_points.add(sequence.frame_count) + 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'