From 10a25dc0365f61b69da61fa6d7dda9d62720df97 Mon Sep 17 00:00:00 2001 From: Colin Basnett Date: Wed, 17 Jul 2024 01:38:32 -0700 Subject: [PATCH] Added support for collection exporters --- io_scene_psk_psa/psa/import_/operators.py | 2 +- io_scene_psk_psa/psk/builder.py | 138 +++++++++++++--------- io_scene_psk_psa/psk/export/operators.py | 127 +++++++++++++++++--- io_scene_psk_psa/psk/export/properties.py | 12 +- io_scene_psk_psa/psk/import_/operators.py | 3 +- io_scene_psk_psa/psk/writer.py | 4 + 6 files changed, 208 insertions(+), 78 deletions(-) diff --git a/io_scene_psk_psa/psa/import_/operators.py b/io_scene_psk_psa/psa/import_/operators.py index 0dd98c9..dad82fa 100644 --- a/io_scene_psk_psa/psa/import_/operators.py +++ b/io_scene_psk_psa/psa/import_/operators.py @@ -108,7 +108,6 @@ def load_psa_file(context, filepath: str): pg.psa_error = str(e) - def on_psa_file_path_updated(cls, context): load_psa_file(context, cls.filepath) @@ -261,6 +260,7 @@ class PSA_FH_import(FileHandler): bl_idname = 'PSA_FH_import' bl_label = 'File handler for Unreal PSA import' bl_import_operator = 'psa_import.import' + bl_export_operator = 'psa_export.export' bl_file_extensions = '.psa' @classmethod diff --git a/io_scene_psk_psa/psk/builder.py b/io_scene_psk_psa/psk/builder.py index f5953f7..b3297e0 100644 --- a/io_scene_psk_psa/psk/builder.py +++ b/io_scene_psk_psa/psk/builder.py @@ -3,7 +3,7 @@ from typing import Optional import bmesh import bpy import numpy as np -from bpy.types import Armature, Material +from bpy.types import Armature, Material, Collection, Context from .data import * from .properties import triangle_type_and_bit_flags_to_poly_flags @@ -20,30 +20,28 @@ class PskBuildOptions(object): def __init__(self): self.bone_filter_mode = 'ALL' self.bone_collection_indices: List[int] = [] - self.use_raw_mesh_data = True + self.object_eval_state = 'EVALUATED' self.materials: List[Material] = [] self.should_enforce_bone_name_restrictions = False -def get_psk_input_objects(context) -> PskInputObjects: - input_objects = PskInputObjects() - for selected_object in context.view_layer.objects.selected: - if selected_object.type == 'MESH': - input_objects.mesh_objects.append(selected_object) +def get_mesh_objects_for_collection(collection: Collection): + for obj in collection.all_objects: + if obj.type == 'MESH': + yield obj - if len(input_objects.mesh_objects) == 0: - raise RuntimeError('At least one mesh must be selected') - for mesh_object in input_objects.mesh_objects: - if len(mesh_object.data.materials) == 0: - raise RuntimeError(f'Mesh "{mesh_object.name}" must have at least one material') +def get_mesh_objects_for_context(context: Context): + for obj in context.view_layer.objects.selected: + if obj.type == 'MESH': + yield obj - # Ensure that there are either no armature modifiers (static mesh) - # or that there is exactly one armature modifier object shared between - # all selected meshes + +def get_armature_for_mesh_objects(mesh_objects: List[Object]) -> Optional[Object]: + # Ensure that there are either no armature modifiers (static mesh) or that there is exactly one armature modifier + # object shared between all meshes. armature_modifier_objects = set() - - for mesh_object in input_objects.mesh_objects: + for mesh_object in mesh_objects: modifiers = [x for x in mesh_object.modifiers if x.type == 'ARMATURE'] if len(modifiers) == 0: continue @@ -53,21 +51,46 @@ def get_psk_input_objects(context) -> PskInputObjects: if len(armature_modifier_objects) > 1: armature_modifier_names = [x.name for x in armature_modifier_objects] - raise RuntimeError(f'All selected meshes must have the same armature modifier, encountered {len(armature_modifier_names)} ({", ".join(armature_modifier_names)})') + raise RuntimeError( + f'All meshes must have the same armature modifier, encountered {len(armature_modifier_names)} ({", ".join(armature_modifier_names)})') elif len(armature_modifier_objects) == 1: - input_objects.armature_object = list(armature_modifier_objects)[0] + return list(armature_modifier_objects)[0] + else: + return None + + +def _get_psk_input_objects(mesh_objects: List[Object]) -> PskInputObjects: + if len(mesh_objects) == 0: + raise RuntimeError('At least one mesh must be selected') + + for mesh_object in mesh_objects: + if len(mesh_object.data.materials) == 0: + raise RuntimeError(f'Mesh "{mesh_object.name}" must have at least one material') + + input_objects = PskInputObjects() + input_objects.mesh_objects = mesh_objects + input_objects.armature_object = get_armature_for_mesh_objects(mesh_objects) return input_objects +def get_psk_input_objects_for_context(context: Context) -> PskInputObjects: + mesh_objects = list(get_mesh_objects_for_context(context)) + return _get_psk_input_objects(mesh_objects) + + +def get_psk_input_objects_for_collection(collection: Collection) -> PskInputObjects: + mesh_objects = list(get_mesh_objects_for_collection(collection)) + return _get_psk_input_objects(mesh_objects) + + class PskBuildResult(object): def __init__(self): self.psk = None self.warnings: List[str] = [] -def build_psk(context, options: PskBuildOptions) -> PskBuildResult: - input_objects = get_psk_input_objects(context) +def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions) -> PskBuildResult: armature_object: bpy.types.Object = input_objects.armature_object result = PskBuildResult() @@ -160,47 +183,48 @@ def build_psk(context, options: PskBuildOptions) -> PskBuildResult: material_indices = [material_names.index(material_slot.material.name) for material_slot in input_mesh_object.material_slots] # MESH DATA - if options.use_raw_mesh_data: - mesh_object = input_mesh_object - mesh_data = input_mesh_object.data - else: - # Create a copy of the mesh object after non-armature modifiers are applied. + match options.object_eval_state: + case 'ORIGINAL': + mesh_object = input_mesh_object + mesh_data = input_mesh_object.data + case 'EVALUATED': + # Create a copy of the mesh object after non-armature modifiers are applied. - # Temporarily force the armature into the rest position. - # We will undo this later. - old_pose_position = None - if armature_object is not None: - old_pose_position = armature_object.data.pose_position - armature_object.data.pose_position = 'REST' + # Temporarily force the armature into the rest position. + # We will undo this later. + old_pose_position = None + if armature_object is not None: + old_pose_position = armature_object.data.pose_position + armature_object.data.pose_position = 'REST' - depsgraph = context.evaluated_depsgraph_get() - bm = bmesh.new() - bm.from_object(input_mesh_object, depsgraph) - mesh_data = bpy.data.meshes.new('') - bm.to_mesh(mesh_data) - del bm - mesh_object = bpy.data.objects.new('', mesh_data) - mesh_object.matrix_world = input_mesh_object.matrix_world + depsgraph = context.evaluated_depsgraph_get() + bm = bmesh.new() + bm.from_object(input_mesh_object, depsgraph) + mesh_data = bpy.data.meshes.new('') + bm.to_mesh(mesh_data) + del bm + mesh_object = bpy.data.objects.new('', mesh_data) + mesh_object.matrix_world = input_mesh_object.matrix_world - scale = (input_mesh_object.scale.x, input_mesh_object.scale.y, input_mesh_object.scale.z) + scale = (input_mesh_object.scale.x, input_mesh_object.scale.y, input_mesh_object.scale.z) - # Negative scaling in Blender results in inverted normals after the scale is applied. However, if the scale - # is not applied, the normals will appear unaffected in the viewport. The evaluated mesh data used in the - # export will have the scale applied, but this behavior is not obvious to the user. - # - # In order to have the exporter be as WYSIWYG as possible, we need to check for negative scaling and invert - # the normals if necessary. If two axes have negative scaling and the third has positive scaling, the - # normals will be correct. We can detect this by checking if the number of negative scaling axes is odd. If - # it is, we need to invert the normals of the mesh by swapping the order of the vertices in each face. - should_flip_normals = sum(1 for x in scale if x < 0) % 2 == 1 + # Negative scaling in Blender results in inverted normals after the scale is applied. However, if the scale + # is not applied, the normals will appear unaffected in the viewport. The evaluated mesh data used in the + # export will have the scale applied, but this behavior is not obvious to the user. + # + # In order to have the exporter be as WYSIWYG as possible, we need to check for negative scaling and invert + # the normals if necessary. If two axes have negative scaling and the third has positive scaling, the + # normals will be correct. We can detect this by checking if the number of negative scaling axes is odd. If + # it is, we need to invert the normals of the mesh by swapping the order of the vertices in each face. + should_flip_normals = sum(1 for x in scale if x < 0) % 2 == 1 - # Copy the vertex groups - for vertex_group in input_mesh_object.vertex_groups: - mesh_object.vertex_groups.new(name=vertex_group.name) + # Copy the vertex groups + for vertex_group in input_mesh_object.vertex_groups: + mesh_object.vertex_groups.new(name=vertex_group.name) - # Restore the previous pose position on the armature. - if old_pose_position is not None: - armature_object.data.pose_position = old_pose_position + # Restore the previous pose position on the armature. + if old_pose_position is not None: + armature_object.data.pose_position = old_pose_position vertex_offset = len(psk.points) @@ -305,7 +329,7 @@ def build_psk(context, options: PskBuildOptions) -> PskBuildResult: w.weight = weight psk.weights.append(w) - if not options.use_raw_mesh_data: + if options.object_eval_state == 'EVALUATED': bpy.data.objects.remove(mesh_object) bpy.data.meshes.remove(mesh_data) del mesh_data diff --git a/io_scene_psk_psa/psk/export/operators.py b/io_scene_psk_psa/psk/export/operators.py index 04d5d39..56ab4f2 100644 --- a/io_scene_psk_psa/psk/export/operators.py +++ b/io_scene_psk_psa/psk/export/operators.py @@ -1,35 +1,40 @@ -from bpy.props import StringProperty -from bpy.types import Operator +from typing import List + +import bpy +from bpy.props import StringProperty, BoolProperty, EnumProperty +from bpy.types import Operator, Context, Object from bpy_extras.io_utils import ExportHelper -from ..builder import build_psk, PskBuildOptions, get_psk_input_objects +from .properties import object_eval_state_items +from ..builder import build_psk, PskBuildOptions, get_psk_input_objects_for_context, \ + get_psk_input_objects_for_collection from ..writer import write_psk from ...shared.helpers import populate_bone_collection_list def is_bone_filter_mode_item_available(context, identifier): - input_objects = get_psk_input_objects(context) + input_objects = get_psk_input_objects_for_context(context) armature_object = input_objects.armature_object if identifier == 'BONE_COLLECTIONS': if armature_object is None or armature_object.data is None or len(armature_object.data.collections) == 0: return False - # else if... you can set up other conditions if you add more options return True -def populate_material_list(mesh_objects, material_list): - material_list.clear() - +def get_materials_for_mesh_objects(mesh_objects: List[Object]): materials = [] for mesh_object in mesh_objects: for i, material_slot in enumerate(mesh_object.material_slots): material = material_slot.material - # TODO: put this in the poll arg? if material is None: raise RuntimeError('Material slot cannot be empty (index ' + str(i) + ')') if material not in materials: materials.append(material) + return materials +def populate_material_list(mesh_objects, material_list): + materials = get_materials_for_mesh_objects(mesh_objects) + material_list.clear() for index, material in enumerate(materials): m = material_list.add() m.material = material @@ -72,6 +77,84 @@ class PSK_OT_material_list_move_down(Operator): return {'FINISHED'} +class PSK_OT_export_collection(Operator, ExportHelper): + bl_idname = 'export.psk_collection' + bl_label = 'Export' + bl_options = {'INTERNAL'} + filename_ext = '.psk' + filter_glob: StringProperty(default='*.psk', options={'HIDDEN'}) + filepath: StringProperty( + name='File Path', + description='File path used for exporting the PSK file', + maxlen=1024, + default='', + subtype='FILE_PATH') + collection: StringProperty(options={'HIDDEN'}) + + object_eval_state: EnumProperty( + items=object_eval_state_items, + name='Object Evaluation State', + default='EVALUATED' + ) + should_enforce_bone_name_restrictions: BoolProperty( + default=False, + name='Enforce Bone Name Restrictions', + description='Enforce that bone names must only contain letters, numbers, spaces, hyphens and underscores.\n\n' + 'Depending on the engine, improper bone names might not be referenced correctly by scripts' + ) + + def execute(self, context): + collection = bpy.data.collections.get(self.collection) + + try: + input_objects = get_psk_input_objects_for_collection(collection) + except RuntimeError as e: + self.report({'ERROR_INVALID_CONTEXT'}, str(e)) + return {'CANCELLED'} + + options = PskBuildOptions() + options.bone_filter_mode = 'ALL' + options.object_eval_state = self.object_eval_state + options.materials = get_materials_for_mesh_objects(input_objects.mesh_objects) + options.should_enforce_bone_name_restrictions = self.should_enforce_bone_name_restrictions + + try: + result = build_psk(context, input_objects, options) + for warning in result.warnings: + self.report({'WARNING'}, warning) + write_psk(result.psk, self.filepath) + if len(result.warnings) > 0: + self.report({'WARNING'}, f'PSK export successful with {len(result.warnings)} warnings') + else: + self.report({'INFO'}, f'PSK export successful') + except RuntimeError as e: + self.report({'ERROR_INVALID_CONTEXT'}, str(e)) + return {'CANCELLED'} + + return {'FINISHED'} + + def draw(self, context: Context): + layout = self.layout + + # MESH + mesh_header, mesh_panel = layout.panel('Mesh', default_closed=False) + mesh_header.label(text='Mesh', icon='MESH_DATA') + if mesh_panel: + flow = mesh_panel.grid_flow(row_major=True) + flow.use_property_split = True + flow.use_property_decorate = False + flow.prop(self, 'object_eval_state', text='Data') + + # BONES + bones_header, bones_panel = layout.panel('Bones', default_closed=False) + bones_header.label(text='Bones', icon='BONE_DATA') + if bones_panel: + flow = bones_panel.grid_flow(row_major=True) + flow.use_property_split = True + flow.use_property_decorate = False + flow.prop(self, 'should_enforce_bone_name_restrictions') + + class PSK_OT_export(Operator, ExportHelper): bl_idname = 'export.psk' bl_label = 'Export' @@ -88,7 +171,7 @@ class PSK_OT_export(Operator, ExportHelper): def invoke(self, context, event): try: - input_objects = get_psk_input_objects(context) + input_objects = get_psk_input_objects_for_context(context) except RuntimeError as e: self.report({'ERROR_INVALID_CONTEXT'}, str(e)) return {'CANCELLED'} @@ -110,7 +193,7 @@ class PSK_OT_export(Operator, ExportHelper): @classmethod def poll(cls, context): try: - get_psk_input_objects(context) + get_psk_input_objects_for_context(context) except RuntimeError as e: cls.poll_message_set(str(e)) return False @@ -118,16 +201,20 @@ class PSK_OT_export(Operator, ExportHelper): def draw(self, context): layout = self.layout + pg = getattr(context.scene, 'psk_export') # MESH - mesh_header, mesh_panel = layout.panel('01_mesh', default_closed=False) + mesh_header, mesh_panel = layout.panel('Mesh', default_closed=False) mesh_header.label(text='Mesh', icon='MESH_DATA') if mesh_panel: - mesh_panel.prop(pg, 'use_raw_mesh_data') + flow = mesh_panel.grid_flow(row_major=True) + flow.use_property_split = True + flow.use_property_decorate = False + flow.prop(pg, 'object_eval_state', text='Data') # BONES - bones_header, bones_panel = layout.panel('02_bones', default_closed=False) + bones_header, bones_panel = layout.panel('Bones', default_closed=False) bones_header.label(text='Bones', icon='BONE_DATA') if bones_panel: bone_filter_mode_items = pg.bl_rna.properties['bone_filter_mode'].enum_items_static @@ -146,7 +233,7 @@ class PSK_OT_export(Operator, ExportHelper): bones_panel.prop(pg, 'should_enforce_bone_name_restrictions') # MATERIALS - materials_header, materials_panel = layout.panel('03_materials', default_closed=False) + materials_header, materials_panel = layout.panel('Materials', default_closed=False) materials_header.label(text='Materials', icon='MATERIAL') if materials_panel: row = materials_panel.row() @@ -157,16 +244,19 @@ class PSK_OT_export(Operator, ExportHelper): col.operator(PSK_OT_material_list_move_down.bl_idname, text='', icon='TRIA_DOWN') def execute(self, context): - pg = context.scene.psk_export + pg = getattr(context.scene, 'psk_export') + + input_objects = get_psk_input_objects_for_context(context) + options = PskBuildOptions() 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.use_raw_mesh_data = pg.use_raw_mesh_data + options.object_eval_state = pg.object_eval_state options.materials = [m.material for m in pg.material_list] options.should_enforce_bone_name_restrictions = pg.should_enforce_bone_name_restrictions try: - result = build_psk(context, options) + result = build_psk(context, input_objects, options) for warning in result.warnings: self.report({'WARNING'}, warning) write_psk(result.psk, self.filepath) @@ -185,4 +275,5 @@ classes = ( PSK_OT_material_list_move_up, PSK_OT_material_list_move_down, PSK_OT_export, + PSK_OT_export_collection, ) diff --git a/io_scene_psk_psa/psk/export/properties.py b/io_scene_psk_psa/psk/export/properties.py index 111cbf0..6dd31da 100644 --- a/io_scene_psk_psa/psk/export/properties.py +++ b/io_scene_psk_psa/psk/export/properties.py @@ -5,6 +5,12 @@ from ...shared.types import PSX_PG_bone_collection_list_item empty_set = set() + +object_eval_state_items = ( + ('EVALUATED', 'Evaluated', 'Use data from fully evaluated object'), + ('ORIGINAL', 'Original', 'Use data from original object with no modifiers applied'), +) + class PSK_PG_material_list_item(PropertyGroup): material: PointerProperty(type=Material) index: IntProperty() @@ -23,7 +29,11 @@ class PSK_PG_export(PropertyGroup): ) bone_collection_list: CollectionProperty(type=PSX_PG_bone_collection_list_item) bone_collection_list_index: IntProperty(default=0) - use_raw_mesh_data: BoolProperty(default=False, name='Raw Mesh Data', description='No modifiers will be evaluated as part of the exported mesh') + object_eval_state: EnumProperty( + items=object_eval_state_items, + name='Object Evaluation State', + default='EVALUATED' + ) material_list: CollectionProperty(type=PSK_PG_material_list_item) material_list_index: IntProperty(default=0) should_enforce_bone_name_restrictions: BoolProperty( diff --git a/io_scene_psk_psa/psk/import_/operators.py b/io_scene_psk_psa/psk/import_/operators.py index b8ce7c9..1d45f37 100644 --- a/io_scene_psk_psa/psk/import_/operators.py +++ b/io_scene_psk_psa/psk/import_/operators.py @@ -13,8 +13,9 @@ empty_set = set() class PSK_FH_import(FileHandler): bl_idname = 'PSK_FH_import' - bl_label = 'File handler for Unreal PSK/PSKX import' + bl_label = 'Unreal PSK' bl_import_operator = 'import_scene.psk' + bl_export_operator = 'export.psk_collection' bl_file_extensions = '.psk;.pskx' @classmethod diff --git a/io_scene_psk_psa/psk/writer.py b/io_scene_psk_psa/psk/writer.py index ff94b90..315e0d1 100644 --- a/io_scene_psk_psa/psk/writer.py +++ b/io_scene_psk_psa/psk/writer.py @@ -1,3 +1,4 @@ +import os from ctypes import Structure, sizeof from typing import Type @@ -34,6 +35,9 @@ def write_psk(psk: Psk, path: str): elif len(psk.bones) == 0: raise RuntimeError(f'At least one bone must be marked for export') + # Make the directory for the file if it doesn't exist. + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, 'wb') as fp: _write_section(fp, b'ACTRHEAD') _write_section(fp, b'PNTS0000', Vector3, psk.points)