From 84f168a9c0cd16ca43b049b84089b352a13efa76 Mon Sep 17 00:00:00 2001 From: Colin Basnett Date: Wed, 26 Jul 2023 16:25:52 -0700 Subject: [PATCH] Handle importing PSAs with duplicate bone names gracefully For Harry Potter 2 models, the rig and the PSA files have bones with identical names. This would cause the PSA importer to fail because it was expecting that each bone name would be unique. The new behavior is to ignore duplicate bones when importing the f-curve data, but print a warning to let the user know something is wacky with their PSA rig. For example, if there are two bones, A and B, with the same name, the first bone encountered (A) in the tree will be mapped to the armature bone, but bone B will be ignored. --- io_scene_psk_psa/psa/importer.py | 66 +++++++++++++++++++------------- 1 file changed, 39 insertions(+), 27 deletions(-) diff --git a/io_scene_psk_psa/psa/importer.py b/io_scene_psk_psa/psa/importer.py index f46b836..0a577fe 100644 --- a/io_scene_psk_psa/psa/importer.py +++ b/io_scene_psk_psa/psa/importer.py @@ -2,7 +2,6 @@ import fnmatch import os import re import typing -from collections import Counter from typing import List, Optional import bpy @@ -64,6 +63,23 @@ class PsaImportResult: 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 import_psa(psa_reader: PsaReader, armature_object: bpy.types.Object, options: PsaImportOptions) -> PsaImportResult: result = PsaImportResult() sequences = map(lambda x: psa_reader.sequences[x], options.sequence_names) @@ -71,37 +87,33 @@ def import_psa(psa_reader: PsaReader, armature_object: bpy.types.Object, options # 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') - try: - psa_to_armature_bone_indices[psa_bone_index] = armature_bone_names.index(psa_bone_name) - except ValueError: - # PSA bone could not be mapped directly to an armature bone by name. - # Attempt to create a bone mapping by ignoring the case of the names. - if options.bone_mapping_mode == 'CASE_INSENSITIVE': - for armature_bone_index, armature_bone_name in enumerate(armature_bone_names): - if armature_bone_name.upper() == psa_bone_name.upper(): - psa_to_armature_bone_indices[psa_bone_index] = armature_bone_index - psa_bone_name = armature_bone_name - break - psa_bone_names.append(psa_bone_name) + 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_names.index(psa_bone_name) + 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) - # Remove ambiguous bone mappings (where multiple PSA bones correspond to the same armature bone). - armature_bone_index_counts = Counter(psa_to_armature_bone_indices.values()) - for armature_bone_index, count in armature_bone_index_counts.items(): - if count > 1: - psa_bone_indices = [] - for psa_bone_index, mapped_bone_index in psa_to_armature_bone_indices: - if mapped_bone_index == armature_bone_index: - psa_bone_indices.append(psa_bone_index) - ambiguous_psa_bone_names = list(sorted([psa_bone_names[x] for x in psa_bone_indices])) - result.warnings.append( - f'Ambiguous mapping for bone {armature_bone_names[armature_bone_index]}!\n' - f'The following PSA bones all map to the same armature bone: {ambiguous_psa_bone_names}\n' - f'These bones will be ignored.' - ) + # 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 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))