Initial commit for handling the PSA config

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

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

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'}

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=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),  # Qx
+                    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),  # Qy
+                    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),  # Qz
+                    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),  # Lx
+                    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),  # Ly
+                    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),  # Lz
+                    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)