diff --git a/io_scene_psk_psa/__init__.py b/io_scene_psk_psa/__init__.py index f9d5968..7092d7e 100644 --- a/io_scene_psk_psa/__init__.py +++ b/io_scene_psk_psa/__init__.py @@ -30,6 +30,7 @@ if 'bpy' in locals(): importlib.reload(psk_import_operators) importlib.reload(psa_data) + importlib.reload(psa_config) importlib.reload(psa_reader) importlib.reload(psa_writer) importlib.reload(psa_builder) @@ -55,6 +56,7 @@ else: from .psk.import_ import operators as psk_import_operators from .psa import data as psa_data + from .psa import config as psa_config from .psa import reader as psa_reader from .psa import writer as psa_writer from .psa import builder as psa_builder diff --git a/io_scene_psk_psa/psa/config.py b/io_scene_psk_psa/psa/config.py new file mode 100644 index 0000000..f9e1515 --- /dev/null +++ b/io_scene_psk_psa/psa/config.py @@ -0,0 +1,48 @@ +import re +from configparser import ConfigParser +from typing import Dict + +from .reader import PsaReader + +REMOVE_TRACK_LOCATION = (1 << 0) +REMOVE_TRACK_ROTATION = (1 << 1) + + +class PsaConfig: + def __init__(self): + self.sequence_bone_flags: Dict[str, Dict[int, int]] = dict() + + +def read_psa_config(psa_reader: PsaReader, file_path: str) -> PsaConfig: + psa_config = PsaConfig() + + config = ConfigParser() + config.read(file_path) + + psa_sequence_names = list(psa_reader.sequences.keys()) + lowercase_sequence_names = [sequence_name.lower() for sequence_name in psa_sequence_names] + + if config.has_section('RemoveTracks'): + for key, value in config.items('RemoveTracks'): + match = re.match(f'^(.+)\.(\d+)$', key) + sequence_name = match.group(1) + bone_index = int(match.group(2)) + + # Map the sequence name onto the actual sequence name in the PSA file. + try: + sequence_name = psa_sequence_names[lowercase_sequence_names.index(sequence_name.lower())] + except ValueError: + pass + + if sequence_name not in psa_config.sequence_bone_flags: + psa_config.sequence_bone_flags[sequence_name] = dict() + + match value: + case 'all': + psa_config.sequence_bone_flags[sequence_name][bone_index] = (REMOVE_TRACK_LOCATION | REMOVE_TRACK_ROTATION) + case 'trans': + psa_config.sequence_bone_flags[sequence_name][bone_index] = REMOVE_TRACK_LOCATION + case 'rot': + psa_config.sequence_bone_flags[sequence_name][bone_index] = REMOVE_TRACK_ROTATION + + return psa_config diff --git a/io_scene_psk_psa/psa/import_/operators.py b/io_scene_psk_psa/psa/import_/operators.py index 08af962..afa964d 100644 --- a/io_scene_psk_psa/psa/import_/operators.py +++ b/io_scene_psk_psa/psa/import_/operators.py @@ -1,10 +1,12 @@ import os +from pathlib import Path from bpy.props import StringProperty from bpy.types import Operator, Event, Context from bpy_extras.io_utils import ImportHelper from .properties import get_visible_sequences +from ..config import read_psa_config from ..importer import import_psa, PsaImportOptions from ..reader import PsaReader @@ -169,6 +171,11 @@ class PSA_OT_import(Operator, ImportHelper): options.fps_source = pg.fps_source options.fps_custom = pg.fps_custom + # Read the PSA config file if it exists. + config_path = Path(self.filepath).with_suffix('.config') + if config_path.exists(): + options.psa_config = read_psa_config(psa_reader, str(config_path)) + if len(sequence_names) == 0: self.report({'ERROR_INVALID_CONTEXT'}, 'No sequences selected') return {'CANCELLED'} diff --git a/io_scene_psk_psa/psa/importer.py b/io_scene_psk_psa/psa/importer.py index c837e11..c4a66cc 100644 --- a/io_scene_psk_psa/psa/importer.py +++ b/io_scene_psk_psa/psa/importer.py @@ -6,6 +6,7 @@ import numpy 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 @@ -23,6 +24,7 @@ class PsaImportOptions(object): self.bone_mapping_mode = 'CASE_INSENSITIVE' self.fps_source = 'SEQUENCE' self.fps_custom: float = 30.0 + self.psa_config: PsaConfig = PsaConfig() class ImportBone(object): @@ -162,6 +164,11 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, 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: @@ -187,18 +194,21 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, # 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), # Qw - action.fcurves.new(rotation_data_path, index=1, action_group=pose_bone.name), # Qx - action.fcurves.new(rotation_data_path, index=2, action_group=pose_bone.name), # Qy - action.fcurves.new(rotation_data_path, index=3, action_group=pose_bone.name), # Qz - action.fcurves.new(location_data_path, index=0, action_group=pose_bone.name), # Lx - action.fcurves.new(location_data_path, index=1, action_group=pose_bone.name), # Ly - action.fcurves.new(location_data_path, index=2, action_group=pose_bone.name), # Lz + 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. @@ -216,11 +226,15 @@ def import_psa(context: Context, psa_reader: PsaReader, armature_object: Object, # Write the keyframes out. fcurve_data = numpy.zeros(2 * sequence.frame_count, dtype=float) + + # 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.keyframe_points.foreach_set('co', fcurve_data)