Moved things around for packaging on Blender extensions

psa/reader.py

@@ -0,0 +1,117 @@
+import ctypes
+
+import numpy as np
+
+from .data import *
+
+
+def _try_fix_cue4parse_issue_103(sequences) -> bool:
+    # Detect if the file was exported from CUE4Parse prior to the fix for issue #103.
+    # https://github.com/FabianFG/CUE4Parse/issues/103
+    # The issue was that the frame_start_index was not being set correctly, and was always being set to the same value
+    # as the frame_count.
+    # This fix will eventually be deprecated as it is only necessary for files exported prior to the fix.
+    if len(sequences) > 0 and sequences[0].frame_start_index == sequences[0].frame_count:
+        # Manually set the frame_start_index for each sequence. This assumes that the sequences are in order with
+        # no shared frames between sequences (all exporters that I know of do this, so it's a safe assumption).
+        frame_start_index = 0
+        for i, sequence in enumerate(sequences):
+            sequence.frame_start_index = frame_start_index
+            frame_start_index += sequence.frame_count
+        return True
+    return False
+
+
+class PsaReader(object):
+    """
+    This class reads the sequences and bone information immediately upon instantiation and holds onto a file handle.
+    The keyframe data is not read into memory upon instantiation due to its potentially very large size.
+    To read the key data for a particular sequence, call :read_sequence_keys.
+    """
+
+    def __init__(self, path):
+        self.keys_data_offset: int = 0
+        self.fp = open(path, 'rb')
+        self.psa: Psa = self._read(self.fp)
+
+    @property
+    def bones(self):
+        return self.psa.bones
+
+    @property
+    def sequences(self):
+        return self.psa.sequences
+
+    def read_sequence_data_matrix(self, sequence_name: str) -> np.ndarray:
+        """
+        Reads and returns the data matrix for the given sequence.
+        @param sequence_name: The name of the sequence.
+        @return: An FxBx7 matrix where F is the number of frames, B is the number of bones.
+        """
+        sequence = self.psa.sequences[sequence_name]
+        keys = self.read_sequence_keys(sequence_name)
+        bone_count = len(self.bones)
+        matrix_size = sequence.frame_count, bone_count, 7
+        matrix = np.zeros(matrix_size)
+        keys_iter = iter(keys)
+        for frame_index in range(sequence.frame_count):
+            for bone_index in range(bone_count):
+                matrix[frame_index, bone_index, :] = list(next(keys_iter).data)
+        return matrix
+
+    def read_sequence_keys(self, sequence_name: str) -> List[Psa.Key]:
+        """
+        Reads and returns the key data for a sequence.
+
+        @param sequence_name: The name of the sequence.
+        @return: A list of Psa.Keys.
+        """
+        # Set the file reader to the beginning of the keys data
+        sequence = self.psa.sequences[sequence_name]
+        data_size = sizeof(Psa.Key)
+        bone_count = len(self.psa.bones)
+        buffer_length = data_size * bone_count * sequence.frame_count
+        sequence_keys_offset = self.keys_data_offset + (sequence.frame_start_index * bone_count * data_size)
+        self.fp.seek(sequence_keys_offset, 0)
+        buffer = self.fp.read(buffer_length)
+        offset = 0
+        keys = []
+        for _ in range(sequence.frame_count * bone_count):
+            key = Psa.Key.from_buffer_copy(buffer, offset)
+            keys.append(key)
+            offset += data_size
+        return keys
+
+    @staticmethod
+    def _read_types(fp, data_class, section: Section, data):
+        buffer_length = section.data_size * section.data_count
+        buffer = fp.read(buffer_length)
+        offset = 0
+        for _ in range(section.data_count):
+            data.append(data_class.from_buffer_copy(buffer, offset))
+            offset += section.data_size
+
+    def _read(self, fp) -> Psa:
+        psa = Psa()
+        while fp.read(1):
+            fp.seek(-1, 1)
+            section = Section.from_buffer_copy(fp.read(ctypes.sizeof(Section)))
+            if section.name == b'ANIMHEAD':
+                pass
+            elif section.name == b'BONENAMES':
+                PsaReader._read_types(fp, Psa.Bone, section, psa.bones)
+            elif section.name == b'ANIMINFO':
+                sequences = []
+                PsaReader._read_types(fp, Psa.Sequence, section, sequences)
+                # Try to fix CUE4Parse bug, if necessary.
+                _try_fix_cue4parse_issue_103(sequences)
+                for sequence in sequences:
+                    psa.sequences[sequence.name.decode()] = sequence
+            elif section.name == b'ANIMKEYS':
+                # Skip keys on this pass. We will keep this file open and read from it as needed.
+                self.keys_data_offset = fp.tell()
+                fp.seek(section.data_size * section.data_count, 1)
+            else:
+                fp.seek(section.data_size * section.data_count, 1)
+                print(f'Unrecognized section in PSA: "{section.name}"')
+        return psa