import re from collections import Counter from typing import List, Iterable, Dict, Tuple import bpy from bpy.props import StringProperty from bpy.types import Context, Armature, Action, Object, AnimData, TimelineMarker from bpy_extras.io_utils import ExportHelper from bpy_types import Operator from .properties import PSA_PG_export, PSA_PG_export_action_list_item, filter_sequences from ..builder import build_psa, PsaBuildSequence, PsaBuildOptions from ..writer import write_psa from ...shared.helpers import populate_bone_collection_list, get_nla_strips_in_frame_range def is_action_for_armature(armature: Armature, action: Action): if len(action.fcurves) == 0: return False bone_names = set([x.name for x in armature.bones]) for fcurve in action.fcurves: match = re.match(r'pose\.bones\[\"([^\"]+)\"](\[\"([^\"]+)\"])?', fcurve.data_path) if not match: continue bone_name = match.group(1) if bone_name in bone_names: return True return False def update_actions_and_timeline_markers(context: Context, armature: Armature): pg = getattr(context.scene, 'psa_export') # Clear actions and markers. pg.action_list.clear() pg.marker_list.clear() # Get animation data. animation_data_object = get_animation_data_object(context) animation_data = animation_data_object.animation_data if animation_data_object else None if animation_data is None: return # Populate actions list. for action in bpy.data.actions: if not is_action_for_armature(armature, action): continue if action.name != '' and not action.name.startswith('#'): for (name, frame_start, frame_end) in get_sequences_from_action(action): item = pg.action_list.add() item.action = action item.name = name item.is_selected = False item.is_pose_marker = False item.frame_start = frame_start item.frame_end = frame_end # Pose markers are not guaranteed to be in frame-order, so make sure that they are. pose_markers = sorted(action.pose_markers, key=lambda x: x.frame) for pose_marker_index, pose_marker in enumerate(pose_markers): if pose_marker.name.strip() == '' or pose_marker.name.startswith('#'): continue for (name, frame_start, frame_end) in get_sequences_from_action_pose_markers(action, pose_markers, pose_marker, pose_marker_index): item = pg.action_list.add() item.action = action item.name = name item.is_selected = False item.is_pose_marker = True item.frame_start = frame_start item.frame_end = frame_end # Populate timeline markers list. marker_names = [x.name for x in context.scene.timeline_markers] sequence_frame_ranges = get_timeline_marker_sequence_frame_ranges(animation_data, context, marker_names) for marker_name in marker_names: if marker_name not in sequence_frame_ranges: continue if marker_name.strip() == '' or marker_name.startswith('#'): continue frame_start, frame_end = sequence_frame_ranges[marker_name] sequences = get_sequences_from_name_and_frame_range(marker_name, frame_start, frame_end) for (sequence_name, frame_start, frame_end) in sequences: item = pg.marker_list.add() item.name = sequence_name item.is_selected = False item.frame_start = frame_start item.frame_end = frame_end def get_sequence_fps(context: Context, fps_source: str, fps_custom: float, actions: Iterable[Action]) -> float: match fps_source: case 'SCENE': return context.scene.render.fps case 'CUSTOM': return fps_custom case 'ACTION_METADATA': # Get the minimum value of action metadata FPS values. return min([action.psa_export.fps for action in actions]) case _: raise RuntimeError(f'Invalid FPS source "{fps_source}"') def get_animation_data_object(context: Context) -> Object: pg: PSA_PG_export = getattr(context.scene, 'psa_export') active_object = context.view_layer.objects.active if active_object.type != 'ARMATURE': raise RuntimeError('Selected object must be an Armature') if pg.sequence_source != 'ACTIONS' and pg.should_override_animation_data: animation_data_object = pg.animation_data_override else: animation_data_object = active_object return animation_data_object def is_bone_filter_mode_item_available(context, identifier): if identifier == 'BONE_COLLECTIONS': armature = context.active_object.data if len(armature.collections) == 0: return False return True def get_timeline_marker_sequence_frame_ranges(animation_data: AnimData, context: Context, marker_names: List[str]) -> Dict: # Timeline markers need to be sorted so that we can determine the sequence start and end positions. sequence_frame_ranges = dict() sorted_timeline_markers = list(sorted(context.scene.timeline_markers, key=lambda x: x.frame)) sorted_timeline_marker_names = list(map(lambda x: x.name, sorted_timeline_markers)) for marker_name in marker_names: marker = context.scene.timeline_markers[marker_name] frame_start = marker.frame # Determine the final frame of the sequence based on the next marker. # If no subsequent marker exists, use the maximum frame_end from all NLA strips. marker_index = sorted_timeline_marker_names.index(marker_name) next_marker_index = marker_index + 1 frame_end = 0 if next_marker_index < len(sorted_timeline_markers): # There is a next marker. Use that next marker's frame position as the last frame of this sequence. frame_end = sorted_timeline_markers[next_marker_index].frame nla_strips = get_nla_strips_in_frame_range(animation_data, marker.frame, frame_end) if len(nla_strips) > 0: frame_end = min(frame_end, max(map(lambda nla_strip: nla_strip.frame_end, nla_strips))) frame_start = max(frame_start, min(map(lambda nla_strip: nla_strip.frame_start, nla_strips))) else: # No strips in between this marker and the next, just export this as a one-frame animation. frame_end = frame_start else: # There is no next marker. # Find the final frame of all the NLA strips and use that as the last frame of this sequence. for nla_track in animation_data.nla_tracks: if nla_track.mute: continue for strip in nla_track.strips: frame_end = max(frame_end, strip.frame_end) if frame_start > frame_end: continue sequence_frame_ranges[marker_name] = int(frame_start), int(frame_end) return sequence_frame_ranges def get_sequences_from_name_and_frame_range(name: str, frame_start: int, frame_end: int) -> List[Tuple[str, int, int]]: reversed_pattern = r'(.+)/(.+)' reversed_match = re.match(reversed_pattern, name) if reversed_match: forward_name = reversed_match.group(1) backwards_name = reversed_match.group(2) return [ (forward_name, frame_start, frame_end), (backwards_name, frame_end, frame_start) ] else: return [(name, frame_start, frame_end)] def get_sequences_from_action(action: Action) -> List[Tuple[str, int, int]]: frame_start = int(action.frame_range[0]) frame_end = int(action.frame_range[1]) return get_sequences_from_name_and_frame_range(action.name, frame_start, frame_end) def get_sequences_from_action_pose_markers(action: Action, pose_markers: List[TimelineMarker], pose_marker: TimelineMarker, pose_marker_index: int) -> List[Tuple[str, int, int]]: frame_start = pose_marker.frame sequence_name = pose_marker.name if pose_marker.name.startswith('!'): # If the pose marker name starts with an exclamation mark, only export the first frame. frame_end = frame_start sequence_name = sequence_name[1:] elif pose_marker_index + 1 < len(pose_markers): frame_end = pose_markers[pose_marker_index + 1].frame else: frame_end = int(action.frame_range[1]) return get_sequences_from_name_and_frame_range(sequence_name, frame_start, frame_end) def get_visible_sequences(pg: PSA_PG_export, sequences) -> List[PSA_PG_export_action_list_item]: visible_sequences = [] for i, flag in enumerate(filter_sequences(pg, sequences)): if bool(flag & (1 << 30)): visible_sequences.append(sequences[i]) return visible_sequences class PSA_OT_export(Operator, ExportHelper): bl_idname = 'psa_export.operator' bl_label = 'Export' bl_options = {'INTERNAL', 'UNDO'} __doc__ = 'Export actions to PSA' filename_ext = '.psa' filter_glob: StringProperty(default='*.psa', options={'HIDDEN'}) filepath: StringProperty( name='File Path', description='File path used for exporting the PSA file', maxlen=1024, default='') def __init__(self): self.armature_object = None @classmethod def poll(cls, context): try: cls._check_context(context) except RuntimeError as e: cls.poll_message_set(str(e)) return False return True def draw(self, context): layout = self.layout pg = getattr(context.scene, 'psa_export') # FPS layout.prop(pg, 'fps_source', text='FPS') if pg.fps_source == 'CUSTOM': layout.prop(pg, 'fps_custom', text='Custom') # SOURCE layout.prop(pg, 'sequence_source', text='Source') if pg.sequence_source in {'TIMELINE_MARKERS', 'NLA_TRACK_STRIPS'}: # ANIMDATA SOURCE layout.prop(pg, 'should_override_animation_data') if pg.should_override_animation_data: layout.prop(pg, 'animation_data_override', text='') if pg.sequence_source == 'NLA_TRACK_STRIPS': flow = layout.grid_flow() flow.use_property_split = True flow.use_property_decorate = False flow.prop(pg, 'nla_track') # SELECT ALL/NONE row = layout.row(align=True) row.label(text='Select') row.operator(PSA_OT_export_actions_select_all.bl_idname, text='All', icon='CHECKBOX_HLT') row.operator(PSA_OT_export_actions_deselect_all.bl_idname, text='None', icon='CHECKBOX_DEHLT') # ACTIONS if pg.sequence_source == 'ACTIONS': rows = max(3, min(len(pg.action_list), 10)) layout.template_list('PSA_UL_export_sequences', '', pg, 'action_list', pg, 'action_list_index', rows=rows) elif pg.sequence_source == 'TIMELINE_MARKERS': rows = max(3, min(len(pg.marker_list), 10)) layout.template_list('PSA_UL_export_sequences', '', pg, 'marker_list', pg, 'marker_list_index', rows=rows) elif pg.sequence_source == 'NLA_TRACK_STRIPS': rows = max(3, min(len(pg.nla_strip_list), 10)) layout.template_list('PSA_UL_export_sequences', '', pg, 'nla_strip_list', pg, 'nla_strip_list_index', rows=rows) col = layout.column() col.use_property_split = True col.use_property_decorate = False col.prop(pg, 'sequence_name_prefix') col.prop(pg, 'sequence_name_suffix') # Determine if there is going to be a naming conflict and display an error, if so. selected_items = [x for x in pg.action_list if x.is_selected] action_names = [x.name for x in selected_items] action_name_counts = Counter(action_names) for action_name, count in action_name_counts.items(): if count > 1: layout.label(text=f'Duplicate action: {action_name}', icon='ERROR') break layout.separator() # BONES row = layout.row(align=True) row.prop(pg, 'bone_filter_mode', text='Bones') if pg.bone_filter_mode == 'BONE_COLLECTIONS': row = layout.row(align=True) row.label(text='Select') row.operator(PSA_OT_export_bone_collections_select_all.bl_idname, text='All', icon='CHECKBOX_HLT') row.operator(PSA_OT_export_bone_collections_deselect_all.bl_idname, text='None', icon='CHECKBOX_DEHLT') rows = max(3, min(len(pg.bone_collection_list), 10)) layout.template_list('PSX_UL_bone_collection_list', '', pg, 'bone_collection_list', pg, 'bone_collection_list_index', rows=rows) layout.prop(pg, 'should_enforce_bone_name_restrictions') layout.separator() # ROOT MOTION layout.prop(pg, 'root_motion', text='Root Motion') @classmethod def _check_context(cls, context): if context.view_layer.objects.active is None: raise RuntimeError('An armature must be selected') if context.view_layer.objects.active.type != 'ARMATURE': raise RuntimeError('The selected object must be an armature') def invoke(self, context, _event): try: self._check_context(context) except RuntimeError as e: self.report({'ERROR_INVALID_CONTEXT'}, str(e)) pg: PSA_PG_export = getattr(context.scene, 'psa_export') self.armature_object = context.view_layer.objects.active if self.armature_object.animation_data is None: # This is required otherwise the action list will be empty if the armature has never had its animation # data created before (i.e. if no action was ever assigned to it). self.armature_object.animation_data_create() update_actions_and_timeline_markers(context, self.armature_object.data) populate_bone_collection_list(self.armature_object, pg.bone_collection_list) context.window_manager.fileselect_add(self) return {'RUNNING_MODAL'} def execute(self, context): pg = getattr(context.scene, 'psa_export') # Ensure that we actually have items that we are going to be exporting. if pg.sequence_source == 'ACTIONS' and len(pg.action_list) == 0: raise RuntimeError('No actions were selected for export') elif pg.sequence_source == 'TIMELINE_MARKERS' and len(pg.marker_list) == 0: raise RuntimeError('No timeline markers were selected for export') elif pg.sequence_source == 'NLA_TRACK_STRIPS' and len(pg.nla_strip_list) == 0: raise RuntimeError('No NLA track strips were selected for export') # Populate the export sequence list. animation_data_object = get_animation_data_object(context) animation_data = animation_data_object.animation_data if animation_data is None: raise RuntimeError(f'No animation data for object \'{animation_data_object.name}\'') export_sequences: List[PsaBuildSequence] = [] if pg.sequence_source == 'ACTIONS': for action_item in filter(lambda x: x.is_selected, pg.action_list): if len(action_item.action.fcurves) == 0: continue export_sequence = PsaBuildSequence() export_sequence.nla_state.action = action_item.action export_sequence.name = action_item.name export_sequence.nla_state.frame_start = action_item.frame_start export_sequence.nla_state.frame_end = action_item.frame_end export_sequence.fps = get_sequence_fps(context, pg.fps_source, pg.fps_custom, [action_item.action]) export_sequence.compression_ratio = action_item.action.psa_export.compression_ratio export_sequence.key_quota = action_item.action.psa_export.key_quota export_sequences.append(export_sequence) elif pg.sequence_source == 'TIMELINE_MARKERS': for marker_item in filter(lambda x: x.is_selected, pg.marker_list): export_sequence = PsaBuildSequence() export_sequence.name = marker_item.name export_sequence.nla_state.action = None export_sequence.nla_state.frame_start = marker_item.frame_start export_sequence.nla_state.frame_end = marker_item.frame_end nla_strips_actions = set( map(lambda x: x.action, get_nla_strips_in_frame_range(animation_data, marker_item.frame_start, marker_item.frame_end))) export_sequence.fps = get_sequence_fps(context, pg.fps_source, pg.fps_custom, nla_strips_actions) export_sequences.append(export_sequence) elif pg.sequence_source == 'NLA_TRACK_STRIPS': for nla_strip_item in filter(lambda x: x.is_selected, pg.nla_strip_list): export_sequence = PsaBuildSequence() export_sequence.name = nla_strip_item.name export_sequence.nla_state.action = None export_sequence.nla_state.frame_start = nla_strip_item.frame_start export_sequence.nla_state.frame_end = nla_strip_item.frame_end export_sequence.fps = get_sequence_fps(context, pg.fps_source, pg.fps_custom, [nla_strip_item.action]) export_sequence.compression_ratio = nla_strip_item.action.psa_export.compression_ratio export_sequence.key_quota = nla_strip_item.action.psa_export.key_quota export_sequences.append(export_sequence) else: raise ValueError(f'Unhandled sequence source: {pg.sequence_source}') options = PsaBuildOptions() options.animation_data = animation_data options.sequences = export_sequences options.bone_filter_mode = pg.bone_filter_mode options.bone_collection_indices = [x.index for x in pg.bone_collection_list if x.is_selected] options.should_ignore_bone_name_restrictions = pg.should_enforce_bone_name_restrictions options.sequence_name_prefix = pg.sequence_name_prefix options.sequence_name_suffix = pg.sequence_name_suffix options.root_motion = pg.root_motion try: psa = build_psa(context, options) self.report({'INFO'}, f'PSA export successful') except RuntimeError as e: self.report({'ERROR_INVALID_CONTEXT'}, str(e)) return {'CANCELLED'} write_psa(psa, self.filepath) return {'FINISHED'} class PSA_OT_export_actions_select_all(Operator): bl_idname = 'psa_export.sequences_select_all' bl_label = 'Select All' bl_description = 'Select all visible sequences' bl_options = {'INTERNAL'} @classmethod def get_item_list(cls, context): pg = context.scene.psa_export if pg.sequence_source == 'ACTIONS': return pg.action_list elif pg.sequence_source == 'TIMELINE_MARKERS': return pg.marker_list elif pg.sequence_source == 'NLA_TRACK_STRIPS': return pg.nla_strip_list return None @classmethod def poll(cls, context): pg = getattr(context.scene, 'psa_export') item_list = cls.get_item_list(context) visible_sequences = get_visible_sequences(pg, item_list) has_unselected_sequences = any(map(lambda item: not item.is_selected, visible_sequences)) return has_unselected_sequences def execute(self, context): pg = getattr(context.scene, 'psa_export') sequences = self.get_item_list(context) for sequence in get_visible_sequences(pg, sequences): sequence.is_selected = True return {'FINISHED'} class PSA_OT_export_actions_deselect_all(Operator): bl_idname = 'psa_export.sequences_deselect_all' bl_label = 'Deselect All' bl_description = 'Deselect all visible sequences' bl_options = {'INTERNAL'} @classmethod def get_item_list(cls, context): pg = context.scene.psa_export if pg.sequence_source == 'ACTIONS': return pg.action_list elif pg.sequence_source == 'TIMELINE_MARKERS': return pg.marker_list elif pg.sequence_source == 'NLA_TRACK_STRIPS': return pg.nla_strip_list return None @classmethod def poll(cls, context): item_list = cls.get_item_list(context) has_selected_items = any(map(lambda item: item.is_selected, item_list)) return len(item_list) > 0 and has_selected_items def execute(self, context): pg = getattr(context.scene, 'psa_export') item_list = self.get_item_list(context) for sequence in get_visible_sequences(pg, item_list): sequence.is_selected = False return {'FINISHED'} class PSA_OT_export_bone_collections_select_all(Operator): bl_idname = 'psa_export.bone_collections_select_all' bl_label = 'Select All' bl_description = 'Select all bone collections' bl_options = {'INTERNAL'} @classmethod def poll(cls, context): pg = getattr(context.scene, 'psa_export') item_list = pg.bone_collection_list has_unselected_items = any(map(lambda action: not action.is_selected, item_list)) return len(item_list) > 0 and has_unselected_items def execute(self, context): pg = getattr(context.scene, 'psa_export') for item in pg.bone_collection_list: item.is_selected = True return {'FINISHED'} class PSA_OT_export_bone_collections_deselect_all(Operator): bl_idname = 'psa_export.bone_collections_deselect_all' bl_label = 'Deselect All' bl_description = 'Deselect all bone collections' bl_options = {'INTERNAL'} @classmethod def poll(cls, context): pg = getattr(context.scene, 'psa_export') item_list = pg.bone_collection_list has_selected_actions = any(map(lambda action: action.is_selected, item_list)) return len(item_list) > 0 and has_selected_actions def execute(self, context): pg = getattr(context.scene, 'psa_export') for action in pg.bone_collection_list: action.is_selected = False return {'FINISHED'} classes = ( PSA_OT_export, PSA_OT_export_actions_select_all, PSA_OT_export_actions_deselect_all, PSA_OT_export_bone_collections_select_all, PSA_OT_export_bone_collections_deselect_all, )