Moved things around for packaging on Blender extensions
This commit is contained in:
Colin Basnett
334
psa/importer.py
Normal file
334
psa/importer.py
Normal file
@@ -0,0 +1,334 @@
|
||||
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
|
||||
Reference in New Issue
Block a user