BDK code commit

io_scene_psk_psa/__init__.py

@@ -43,9 +43,34 @@ else:
     from .psa import importer as psa_importer
 
 import bpy
-from bpy.props import PointerProperty
+from bpy.props import CollectionProperty, PointerProperty, StringProperty, IntProperty
+from bpy.types import AddonPreferences, PropertyGroup
 
-classes = (psx_types.classes +
+
+class MaterialPathPropertyGroup(PropertyGroup):
+    path: StringProperty(name='Path', subtype='DIR_PATH')
+
+
+class PskPsaAddonPreferences(AddonPreferences):
+    bl_idname = __name__
+
+    material_path_list: CollectionProperty(type=MaterialPathPropertyGroup)
+    material_path_index: IntProperty()
+
+    def draw_filter(self, context, layout):
+        pass
+
+    def draw(self, context: bpy.types.Context):
+        self.layout.label(text='Material Paths')
+        row = self.layout.row()
+        row.template_list('PSX_UL_MaterialPathList', '', self, 'material_path_list', self, 'material_path_index')
+        column = row.column()
+        column.operator(psx_types.PSX_OT_MaterialPathAdd.bl_idname, icon='ADD', text='')
+        column.operator(psx_types.PSX_OT_MaterialPathRemove.bl_idname, icon='REMOVE', text='')
+
+
+classes = ((MaterialPathPropertyGroup, PskPsaAddonPreferences) +
+           psx_types.classes +
            psk_importer.classes +
            psk_exporter.classes +
            psa_exporter.classes +
@@ -64,6 +89,10 @@ def psa_export_menu_func(self, context):
     self.layout.operator(psa_exporter.PsaExportOperator.bl_idname, text='Unreal PSA (.psa)')
 
 
+def psa_import_menu_func(self, context):
+    self.layout.operator(psa_importer.PsaImportOperator.bl_idname, text='Unreal PSA (.psa)')
+
+
 def register():
     for cls in classes:
         bpy.utils.register_class(cls)
@@ -72,14 +101,12 @@ def register():
     bpy.types.TOPBAR_MT_file_export.append(psa_export_menu_func)
     bpy.types.TOPBAR_MT_file_import.append(psa_import_menu_func)
     bpy.types.Scene.psa_import = PointerProperty(type=psa_importer.PsaImportPropertyGroup)
-    bpy.types.Scene.psk_import = PointerProperty(type=psk_importer.PskImportPropertyGroup)
     bpy.types.Scene.psa_export = PointerProperty(type=psa_exporter.PsaExportPropertyGroup)
     bpy.types.Scene.psk_export = PointerProperty(type=psk_exporter.PskExportPropertyGroup)
 
 
 def unregister():
     del bpy.types.Scene.psa_import
-    del bpy.types.Scene.psk_import
     del bpy.types.Scene.psa_export
     del bpy.types.Scene.psk_export
     bpy.types.TOPBAR_MT_file_export.remove(psk_export_menu_func)

io_scene_psk_psa/bdk.py

@@ -0,0 +1,39 @@
+import re
+from typing import Optional
+
+
+class UReference:
+    type_name: str
+    package_name: str
+    group_name: Optional[str]
+    object_name: str
+
+    def __init__(self, type_name: str, package_name: str, object_name: str, group_name: Optional[str] = None):
+        self.type_name = type_name
+        self.package_name = package_name
+        self.object_name = object_name
+        self.group_name = group_name
+
+    @staticmethod
+    def from_string(string: str) -> Optional['UReference']:
+        if string == 'None':
+            return None
+        pattern = r'(\w+)\'([\w\.\d\-\_]+)\''
+        match = re.match(pattern, string)
+        if match is None:
+            print(f'BAD REFERENCE STRING: {string}')
+            return None
+        type_name = match.group(1)
+        object_name = match.group(2)
+        pattern = r'([\w\d\-\_]+)'
+        values = re.findall(pattern, object_name)
+        package_name = values[0]
+        object_name = values[-1]
+        return UReference(type_name, package_name, object_name, group_name=None)
+
+    def __repr__(self):
+        s = f'{self.type_name}\'{self.package_name}'
+        if self.group_name:
+            s += f'.{self.group_name}'
+        s += f'.{self.object_name}'
+        return s

io_scene_psk_psa/psk/data.py

@@ -92,6 +92,10 @@ class Psk(object):
     def has_vertex_normals(self):
         return len(self.vertex_normals) > 0
 
+    @property
+    def has_material_references(self):
+        return len(self.material_references) > 0
+
     def __init__(self):
         self.points: List[Vector3] = []
         self.wedges: List[Psk.Wedge] = []
@@ -102,3 +106,4 @@ class Psk(object):
         self.extra_uvs: List[Vector2] = []
         self.vertex_colors: List[Color] = []
         self.vertex_normals: List[Vector3] = []
+        self.material_references: List[str] = []

io_scene_psk_psa/psk/importer.py

@@ -1,34 +1,38 @@
 import os
 import sys
 from math import inf
+from pathlib import Path
 from typing import Optional, List
 
 import bmesh
 import bpy
 import numpy as np
 from bpy.props import BoolProperty, EnumProperty, FloatProperty, StringProperty
-from bpy.types import Operator, PropertyGroup, VertexGroup
+from bpy.types import Operator, VertexGroup
 from bpy_extras.io_utils import ImportHelper
 from mathutils import Quaternion, Vector, Matrix
 
 from .data import Psk
 from .reader import read_psk
+from ..bdk import UReference
 from ..helpers import rgb_to_srgb
 
 
-class PskImportOptions(object):
+class PskImportOptions:
     def __init__(self):
         self.name = ''
         self.should_import_mesh = True
+        self.should_reuse_materials = True
         self.should_import_vertex_colors = True
         self.vertex_color_space = 'sRGB'
         self.should_import_vertex_normals = True
         self.should_import_extra_uvs = True
         self.should_import_skeleton = True
         self.bone_length = 1.0
+        self.should_import_materials = True
 
 
-class ImportBone(object):
+class ImportBone:
     """
     Intermediate bone type for the purpose of construction.
     """
@@ -51,6 +55,30 @@ class PskImportResult:
         self.warnings: List[str] = []
 
 
+def load_bdk_material(reference: UReference):
+    if reference is None:
+        return None
+    asset_libraries = bpy.context.preferences.filepaths.asset_libraries
+    asset_library_name = 'bdk-library'
+    try:
+        asset_library = next(filter(lambda x: x.name == asset_library_name, asset_libraries))
+    except StopIteration:
+        return None
+    asset_library_path = Path(asset_library.path)
+    # TODO: going to be very slow for automation!
+    blend_files = [fp for fp in asset_library_path.glob(f'**/{reference.package_name}.blend') if fp.is_file()]
+    if len(blend_files) == 0:
+        return None
+    blend_file = str(blend_files[0])
+    with bpy.data.libraries.load(blend_file, link=True, relative=False, assets_only=True) as (data_in, data_out):
+        if reference.object_name in data_in.materials:
+            data_out.materials = [reference.object_name]
+        else:
+            return None
+    material = bpy.data.materials[reference.object_name]
+    return material
+
+
 def import_psk(psk: Psk, context, options: PskImportOptions) -> PskImportResult:
     result = PskImportResult()
     armature_object = None
@@ -125,10 +153,19 @@ def import_psk(psk: Psk, context, options: PskImportOptions) -> PskImportResult:
         mesh_object = bpy.data.objects.new(options.name, mesh_data)
 
         # MATERIALS
-        for material in psk.materials:
+        if options.should_import_materials:
-            # TODO: re-use of materials should be an option
+            for material_index, psk_material in enumerate(psk.materials):
-            bpy_material = bpy.data.materials.new(material.name.decode('utf-8'))
+                material_name = psk_material.name.decode('utf-8')
-            mesh_data.materials.append(bpy_material)
+                if options.should_reuse_materials and material_name in bpy.data.materials:
+                    # Material already exists, just re-use it.
+                    material = bpy.data.materials[material_name]
+                elif psk.has_material_references:
+                    # Material does not yet exist, attempt to load it using BDK.
+                    reference = UReference.from_string(psk.material_references[material_index])
+                    material = load_bdk_material(reference)
+                else:
+                    material = None
+                mesh_data.materials.append(material)
 
         bm = bmesh.new()
 
@@ -249,7 +286,19 @@ def import_psk(psk: Psk, context, options: PskImportOptions) -> PskImportResult:
 empty_set = set()
 
 
-class PskImportPropertyGroup(PropertyGroup):
+class PskImportOperator(Operator, ImportHelper):
+    bl_idname = 'import_scene.psk'
+    bl_label = 'Import'
+    bl_options = {'INTERNAL', 'UNDO', 'PRESET'}
+    __doc__ = 'Load a PSK file'
+    filename_ext = '.psk'
+    filter_glob: StringProperty(default='*.psk;*.pskx', options={'HIDDEN'})
+    filepath: StringProperty(
+        name='File Path',
+        description='File path used for exporting the PSK file',
+        maxlen=1024,
+        default='')
+
     should_import_vertex_colors: BoolProperty(
         default=True,
         options=empty_set,
@@ -284,6 +333,17 @@ class PskImportPropertyGroup(PropertyGroup):
         options=empty_set,
         description='Import mesh'
     )
+    should_import_materials: BoolProperty(
+        default=True,
+        name='Import Materials',
+        options=empty_set,
+    )
+    should_reuse_materials: BoolProperty(
+        default=True,
+        name='Reuse Materials',
+        options=empty_set,
+        description='Existing materials with matching names will be reused when available'
+    )
     should_import_skeleton: BoolProperty(
         default=True,
         name='Import Skeleton',
@@ -300,34 +360,19 @@ class PskImportPropertyGroup(PropertyGroup):
         description='Length of the bones'
     )
 
-
-class PskImportOperator(Operator, ImportHelper):
-    bl_idname = 'import.psk'
-    bl_label = 'Import'
-    bl_options = {'INTERNAL', 'UNDO'}
-    __doc__ = 'Load a PSK file'
-    filename_ext = '.psk'
-    filter_glob: StringProperty(default='*.psk;*.pskx', options={'HIDDEN'})
-    filepath: StringProperty(
-        name='File Path',
-        description='File path used for exporting the PSK file',
-        maxlen=1024,
-        default='')
-
     def execute(self, context):
-        pg = getattr(context.scene, 'psk_import')
-
         psk = read_psk(self.filepath)
 
         options = PskImportOptions()
         options.name = os.path.splitext(os.path.basename(self.filepath))[0]
-        options.should_import_mesh = pg.should_import_mesh
+        options.should_import_mesh = self.should_import_mesh
-        options.should_import_extra_uvs = pg.should_import_extra_uvs
+        options.should_import_extra_uvs = self.should_import_extra_uvs
-        options.should_import_vertex_colors = pg.should_import_vertex_colors
+        options.should_import_vertex_colors = self.should_import_vertex_colors
-        options.should_import_vertex_normals = pg.should_import_vertex_normals
+        options.should_import_vertex_normals = self.should_import_vertex_normals
-        options.vertex_color_space = pg.vertex_color_space
+        options.vertex_color_space = self.vertex_color_space
-        options.should_import_skeleton = pg.should_import_skeleton
+        options.should_import_skeleton = self.should_import_skeleton
-        options.bone_length = pg.bone_length
+        options.bone_length = self.bone_length
+        options.should_import_materials = self.should_import_materials
 
         result = import_psk(psk, context, options)
 
@@ -341,27 +386,26 @@ class PskImportOperator(Operator, ImportHelper):
         return {'FINISHED'}
 
     def draw(self, context):
-        pg = getattr(context.scene, 'psk_import')
         layout = self.layout
-        layout.prop(pg, 'should_import_mesh')
+        layout.prop(self, 'should_import_materials')
+        layout.prop(self, 'should_import_mesh')
         row = layout.column()
         row.use_property_split = True
         row.use_property_decorate = False
-        if pg.should_import_mesh:
+        if self.should_import_mesh:
-            row.prop(pg, 'should_import_vertex_normals')
+            row.prop(self, 'should_import_vertex_normals')
-            row.prop(pg, 'should_import_extra_uvs')
+            row.prop(self, 'should_import_extra_uvs')
-            row.prop(pg, 'should_import_vertex_colors')
+            row.prop(self, 'should_import_vertex_colors')
-            if pg.should_import_vertex_colors:
+            if self.should_import_vertex_colors:
-                row.prop(pg, 'vertex_color_space')
+                row.prop(self, 'vertex_color_space')
-        layout.prop(pg, 'should_import_skeleton')
+        layout.prop(self, 'should_import_skeleton')
         row = layout.column()
         row.use_property_split = True
         row.use_property_decorate = False
-        if pg.should_import_skeleton:
+        if self.should_import_skeleton:
-            row.prop(pg, 'bone_length')
+            row.prop(self, 'bone_length')
 
 
 classes = (
     PskImportOperator,
-    PskImportPropertyGroup,
 )

io_scene_psk_psa/psk/reader.py

@@ -1,4 +1,8 @@
 import ctypes
+import os
+import re
+import warnings
+from pathlib import Path
 
 from .data import *
 
@@ -12,8 +16,22 @@ def _read_types(fp, data_class, section: Section, data):
         offset += section.data_size
 
 
+def _read_material_references(path: str) -> List[str]:
+    property_file_path = Path(path).with_suffix('.props.txt')
+    if not property_file_path.is_file():
+        # Property file does not exist.
+        return []
+    # Do a crude regex match to find the Material list entries.
+    contents = property_file_path.read_text()
+    pattern = r"Material\s*=\s*([^\s^,]+)"
+    return re.findall(pattern, contents)
+
+
 def read_psk(path: str) -> Psk:
+
     psk = Psk()
+
+    # Read the PSK file sections.
     with open(path, 'rb') as fp:
         while fp.read(1):
             fp.seek(-1, 1)
@@ -46,5 +64,14 @@ def read_psk(path: str) -> Psk:
             elif section.name == b'VTXNORMS':
                 _read_types(fp, Vector3, section, psk.vertex_normals)
             else:
-                raise RuntimeError(f'Unrecognized section "{section.name} at position {15:fp.tell()}"')
+                # Section is not handled, skip it.
+                fp.seek(section.data_size * section.data_count, os.SEEK_CUR)
+                warnings.warn(f'Unrecognized section "{section.name} at position {fp.tell():15}"')
+
+    '''
+    UEViewer exports a sidecar file (*.props.txt) with fully-qualified reference paths for each material
+    (e.g., Texture'Package.Group.Object').
+    '''
+    psk.material_references = _read_material_references(path)
+
     return psk

io_scene_psk_psa/types.py

@@ -1,5 +1,6 @@
+import bpy.props
 from bpy.props import StringProperty, IntProperty, BoolProperty
-from bpy.types import PropertyGroup, UIList, UILayout, Context, AnyType
+from bpy.types import PropertyGroup, UIList, UILayout, Context, AnyType, Operator
 
 
 class PSX_UL_BoneGroupList(UIList):
@@ -11,6 +12,56 @@ class PSX_UL_BoneGroupList(UIList):
         row.label(text=str(getattr(item, 'count')), icon='BONE_DATA')
 
 
+class PSX_OT_MaterialPathAdd(Operator):
+    bl_idname = 'psx.material_paths_add'
+    bl_label = 'Add Material Path'
+    bl_options = {'INTERNAL'}
+
+    directory: bpy.props.StringProperty(subtype='DIR_PATH', options={'HIDDEN'})
+    filter_folder: bpy.props.BoolProperty(default=True, options={'HIDDEN'})
+
+    def invoke(self, context: 'Context', event: 'Event'):
+        context.window_manager.fileselect_add(self)
+        return {'RUNNING_MODAL'}
+
+    def execute(self, context: 'Context'):
+        m = context.preferences.addons[__package__].preferences.material_path_list.add()
+        m.path = self.directory
+        return {'FINISHED'}
+
+
+class PSX_OT_MaterialPathRemove(Operator):
+    bl_idname = 'psx.material_paths_remove'
+    bl_label = 'Remove Material Path'
+    bl_options = {'INTERNAL'}
+
+    @classmethod
+    def poll(cls, context: 'Context'):
+        preferences = context.preferences.addons[__package__].preferences
+        return preferences.material_path_index >= 0
+
+    def execute(self, context: 'Context'):
+        preferences = context.preferences.addons[__package__].preferences
+        preferences.material_path_list.remove(preferences.material_path_index)
+        return {'FINISHED'}
+
+
+class PSX_UL_MaterialPathList(UIList):
+
+    def draw_item(self,
+                  context: 'Context',
+                  layout: 'UILayout',
+                  data: 'AnyType',
+                  item: 'AnyType',
+                  icon: int,
+                  active_data: 'AnyType',
+                  active_property: str,
+                  index: int = 0,
+                  flt_flag: int = 0):
+        row = layout.row()
+        row.label(text=getattr(item, 'path'))
+
+
 class BoneGroupListItem(PropertyGroup):
     name: StringProperty()
     index: IntProperty()
@@ -21,4 +72,7 @@ class BoneGroupListItem(PropertyGroup):
 classes = (
     BoneGroupListItem,
     PSX_UL_BoneGroupList,
+    PSX_UL_MaterialPathList,
+    PSX_OT_MaterialPathAdd,
+    PSX_OT_MaterialPathRemove
 )