Changed the way transforms are extracted and added functions to export them to HDF5 (.h5) files.

app.py

@@ -10,9 +10,9 @@ import streamlit as st
 import fbx_handler
 
 
-def process_file(file: Path) -> bytes:
+def process_file(file: Path) -> int:
     fbx_content = fbx_handler.FBXContainer(file)
-    return fbx_content.export(t='string')
+    return 1
 
 
 # Initialize session state variables if they don't exist

fbx_handler.md

@@ -8,26 +8,35 @@ input_file = Path('/path/to/file.fbx')
 container = FBXContainer(input_file)
 ```
 
+## Preprocess data:
+```python
+container.init_world_transforms(r=...)
+train_raw_data = container.extract_training_translations()
+test_raw_data = container.extract_inf_translations()
+```
+
 ## Training workflow:
 ```python
-# Get dataframe with all valid translation numbers.
-df = container.extract_all_valid_translations()
+# Load file.
+container = FBXContainer(input_file)
+# Get np.array with all valid translation numbers.
+actors_train, markers_train, t_test, _, _ = container.get_split_transforms(mode='train')
 # Convert to dataset...
 ...
 ```
 
 ## Testing workflow:
 ```python
-# Get timeline dense cloud.
-tdc = container.get_tdc()  # wrap in shuffle_tdc() to shuffle nodes.
-# Split array into subarrays.
-actors_test, markers_test, t_test, r_test, s_test = container.split_tdc(tdc)
+# Load file.
+container = FBXContainer(input_file)
+# Get splitted original data (no transforms applied).
+actors_test, markers_test, t_test, r_test_, s_test = container.get_split_transforms(mode='test')
 # Predict the new actors and classes...
 actors_pred, markers_pred = Labeler(container.transform_translations(t_test))
 # Merge the new labels with their original translations.
 merged = merge_tdc(actors_pred, markers_pred, t_test, r_test, s_test)
 # Convert the full cloud into a dict structured for easy keyframes.
-new_dict = tsc_to_dict(merged)
+new_dict = array_to_dict(merged)
 # Replace the old translation keyframes with the new values.
 container.replace_keyframes_for_all_actors(new_dict)
 # Export file.

fbx_handler.py

@@ -1,10 +1,8 @@
-# Import core libs.
-from pprint import pprint
-
 import pandas as pd
 import numpy as np
 from pathlib import Path
 from typing import List, Union, Tuple
+import h5py
 
 # Import util libs.
 import contextlib
@@ -13,16 +11,18 @@ import itertools
 
 # Import custom data.
 import globals
+import utils
 
 
-def center_axis(a: List[float]) -> np.array:
+def center_axis(a: Union[List[float], np.array]) -> np.array:
     """
     Centers a list of floats.
     :param a: List of floats to center.
     :return: The centered list as a `np.array`.
     """
     # Turn list into np array for optimized math.
-    a = np.array(a)
+    if not isinstance(a, np.ndarray):
+        a = np.array(a)
 
     # Find the centroid by subtracting the lowest value from the highest value.
     _min = np.min(a)
@@ -58,17 +58,6 @@ def make_ghost_markers(missing: int) -> np.array:
     ])
 
 
-def append_suffix(file_path: Path, suffix: str = '_INF'):
-    """
-    Adds a suffix to the given file path.
-    :param file_path: `Path` object to the original file.
-    :param suffix: `str` suffix to add to the end of the original file name.
-    :return: Updated `Path`.
-    """
-    new_file_name = file_path.stem + suffix + file_path.suffix
-    return file_path.with_name(new_file_name)
-
-
 def append_zero(arr: np.ndarray) -> np.ndarray:
     zeros = np.zeros((arr.shape[0], arr.shape[1], 1), dtype=float)
     return np.concatenate((arr, zeros), axis=-1)
@@ -85,8 +74,8 @@ def merge_tdc(actor_classes: np.array,
               rotation_vectors: np.array,
               scale_vectors: np.array,
               ordered: bool = True) -> np.array:
-    # Actor and marker classes enter as shape (x, 1000), so use np.expand_dims to create an extra dimension at the end.
-    # Return the concatenated array of shape (x, 1000, 5), which matches the original timeline dense cloud before
+    # Actor and marker classes enter as shape (x, 1000), so use np.expand_dims to create a new dimension at the end.
+    # Return the concatenated array of shape (x, 1000, 14), which matches the original timeline dense cloud before
     # splitting it into sub arrays.
 
     tdc = np.concatenate((np.expand_dims(actor_classes, -1),
@@ -111,7 +100,7 @@ def shuffle_tdc(tdc: np.array) -> np.array:
     if tdc.ndim != 3:
         raise ValueError(f'Array does not have 3 dimensions: {tdc.ndim}/3.')
 
-    # Shuffle the node rows.
+    # Shuffle the frames.
     for i in range(tdc.shape[0]):
         np.random.shuffle(tdc[i])
     return tdc
@@ -142,10 +131,6 @@ def sort_cloud(cloud: np.array) -> np.array:
     return sorted_tdc
 
 
-def isolate_labeled_markers_from_tdc(tdc: np.array) -> np.array:
-    return np.stack([tdc[i, tdc[i, :, 0] > 0.] for i in range(tdc.shape[0])], axis=0)
-
-
 def create_keyframe(anim_curve: fbx.FbxAnimCurve, frame: int, value: float):
     # Create an FbxTime object with the given frame number
     t = fbx.FbxTime()
@@ -173,7 +158,7 @@ def match_name(node: fbx.FbxNode, name: str, ignore_namespace: bool = True) -> b
     return node_name == name
 
 
-def tsc_to_dict(tsc: np.array, start_frame: int = 0) -> dict:
+def array_to_dict(tsc: np.array, start_frame: int = 0) -> dict:
     """
     Converts an `np.array` timeline sparse cloud to a dictionary structured for keyframed animation.
     :param tsc: `np.array` timeline sparse cloud to process.
@@ -248,16 +233,16 @@ def world_to_local_transform(node: fbx.FbxNode, world_transform: fbx.FbxAMatrix,
     return [lcl.GetT()[t] for t in range(3)], [lcl.GetR()[r] for r in range(3)], [lcl.GetS()[s] for s in range(3)]
 
 
-def get_world_transform(m: fbx.FbxNode, time: fbx.FbxTime, axes: str = 'trs') -> np.array:
+def get_world_transform(m: fbx.FbxNode, t: fbx.FbxTime, axes: str = 'trs') -> np.array:
     """
     Evaluates the world translation of the given node at the given time,
     scales it down by scale and turns it into a vector list.
     :param m: `fbx.FbxNode` marker to evaluate the world translation of.
-    :param time: `fbx.FbxTime` time to evaluate at.
+    :param t: `fbx.FbxTime` time to evaluate at.
     :param axes: `str` that contains types of info to include. Options are a combination of t, r, and s.
     :return: Vector in the form: [tx, ty, etc..].
     """
-    matrix = m.EvaluateGlobalTransform(time)
+    matrix = m.EvaluateGlobalTransform(t)
 
     # If axes is only the translation, we return a vector of (tx, ty, tz) only (useful for the training).
     if axes == 't':
@@ -301,12 +286,70 @@ def split_tdc_into_actors(tdc: np.array) -> List[np.array]:
     return [isolate_actor_from_tdc(tdc, i) for i in range(1, actor_count + 1)]
 
 
+def get_keyed_frames_from_curve(curve: fbx.FbxAnimCurve, length: int = -1) -> List[fbx.FbxAnimCurveKey]:
+    frames = [curve.KeyGet(i).GetTime().GetFrameCount() for i in range(curve.KeyGetCount())]
+    dif = length - len(frames)
+    if dif > 0 and length != -1:
+        frames += [0.] * dif
+    return frames
+
+
+def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode, r: List[int], c, incl_keyed: int = 1) \
+        -> List[List[float]]:
+    zeros = [0.0 for _ in range(len(r))]
+    ones = [1.0 for _ in range(len(r))]
+
+    tx, ty, tz, rx, ry, rz, sx, sy, sz = [], [], [], [], [], [], [], [], []
+    actors = [actor_idx for _ in range(len(r))]
+    markers = [marker_idx for _ in range(len(r))]
+    t = fbx.FbxTime()
+
+    for f in r:
+        t.SetFrame(f)
+        wt = m.EvaluateGlobalTransform(t)
+        wtt, wtr, wts = wt.GetT(), wt.GetR(), wt.GetS()
+        tx.append(wtt[0])
+        ty.append(wtt[1])
+        tz.append(wtt[2])
+        rx.append(wtr[0])
+        ry.append(wtr[1])
+        rz.append(wtr[2])
+        sx.append(wts[0])
+        sy.append(wts[1])
+        sz.append(wts[2])
+
+    if not incl_keyed:
+        return [
+            actors,
+            markers,
+            tx, ty, tz, zeros,
+            rx, ry, rz, zeros,
+            sx, sy, sz, ones
+        ]
+
+    keyed_frames = get_keyed_frames_from_curve(c)
+    keyed_bools = [1 if f in keyed_frames else 0 for f in r]
+
+    return [
+        actors,
+        markers,
+        tx, ty, tz, zeros,
+        rx, ry, rz, zeros,
+        sx, sy, sz, ones,
+        keyed_bools
+    ]
+
+
 class FBXContainer:
     def __init__(self, fbx_file: Path,
                  volume_dims: Tuple[float] = (10., 4., 10.),
-                 max_actors: int = 10,
-                 pc_size: int = 1000,
-                 scale: float = 0.01):
+                 max_actors: int = 8,
+                 pc_size: int = 1024,
+                 scale: float = 0.01,
+                 debug: int = -1,
+                 save_init: bool = True,
+                 r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
+                 mode: str = 'train'):
         """
         Class that stores references to important nodes in an FBX file.
         Offers utility functions to quickly load animation data.
@@ -314,11 +357,16 @@ class FBXContainer:
         :param volume_dims: `tuple` of `float` that represent the dimensions of the capture volume in meters.
         :param max_actors: `int` maximum amount of actors to expect in a point cloud.
         :param pc_size: `int` amount of points in a point cloud.
+        :param debug: If higher than -1, will print out debugging statements.
+        :param save_init: If the file is guaranteed to have all data, set to True to automatically call self.init().
+        :param r: Optional frame range that will be passed to init_transforms.
+        :param mode: `str` to indicate whether to store world transforms for inference only. Default 'train'.
         """
         if pc_size < max_actors * 73:
             raise ValueError('Point cloud size must be large enough to contain the maximum amount of actors * 73'
                              f' markers: {pc_size}/{max_actors * 73}.')
 
+        self.debug = debug
         # Python ENUM of the C++ time modes.
         self.time_modes = globals.get_time_modes()
         # Ordered list of marker names. Note: rearrange this in globals.py.
@@ -330,6 +378,9 @@ class FBXContainer:
         # Store names of the actors (all parent nodes that have the first 4 markers as children).
         self.actor_names = []
 
+        self.labeled_world_transforms = None
+        self.unlabeled_world_transforms = None
+
         # Split the dimensions tuple into its axes for easier access.
         self.vol_x = volume_dims[0]
         self.vol_y = volume_dims[1]
@@ -342,16 +393,16 @@ class FBXContainer:
         self.pc_size = pc_size
 
         self.input_fbx = fbx_file
-        self.output_fbx = append_suffix(fbx_file, '_INF')
+        # self.output_fbx = append_suffix_to_fbx(fbx_file, '_INF')
+        self.output_fbx = utils.append_suffix_to_file(fbx_file, '_INF')
         self.valid_frames = []
 
-        self.__init_scene()
-        self.__init_anim()
-        self.__init_actors()
-        self.__init_markers()
-        self.__init_unlabeled_markers()
+        # If we know that the input file has valid data,
+        # we can automatically call the init function and ignore missing data.
+        if save_init:
+            self.init(r=r)
 
-    def __init_scene(self):
+    def __init_scene(self) -> None:
         """
         Stores scene, root, and time_mode properties.
         Destroys the importer to remove the reference to the loaded file.
@@ -365,6 +416,8 @@ class FBXContainer:
         self.scene = fbx.FbxScene.Create(self.manager, '')
         importer.Import(self.scene)
         self.root = self.scene.GetRootNode()
+        if self.root is None:
+            raise ValueError('No root node found.')
         self.time_mode = self.scene.GetGlobalSettings().GetTimeMode()
         fbx.FbxTime.SetGlobalTimeMode(self.time_mode)
 
@@ -372,27 +425,38 @@ class FBXContainer:
         # This will allow us to delete the uploaded file.
         importer.Destroy()
 
-    def __init_anim(self):
+    def __init_anim(self) -> None:
         """
         Stores the anim_stack, num_frames, start_frame, end_frame properties.
         """
         # Get the animation stack and layer.
         anim_stack = self.scene.GetCurrentAnimationStack()
         self.anim_layer = anim_stack.GetSrcObject(fbx.FbxCriteria.ObjectType(fbx.FbxAnimLayer.ClassId), 0)
+        if self.anim_layer is None:
+            raise ValueError('No animation layer found.')
 
         # Find the total number of frames to expect from the local time span.
         local_time_span = anim_stack.GetLocalTimeSpan()
         self.num_frames = int(local_time_span.GetDuration().GetFrameCount())
+        if self.num_frames == 0:
+            raise ValueError('Number of animated frames is 0.')
         self.start_frame = local_time_span.GetStart().GetFrameCount()
         self.end_frame = local_time_span.GetStop().GetFrameCount()
 
-    def __init_actors(self):
+    def __init_actors(self, ignore_missing: bool = False) -> None:
         """
         Goes through all root children (generation 1).
         If a child has 4 markers as children, it is considered an actor (Shogun subject) and appended to actors
         and actor_names list properties.
         Also initializes an empty valid_frames list for each found actor.
         """
+        ts = fbx.FbxTime()
+        ts.SetFrame(self.start_frame)
+
+        te = fbx.FbxTime()
+        te.SetFrame(self.end_frame)
+
+        names_to_look_for = list(self.marker_names[:4])
         # Find all parent nodes (/System, /Unlabeled_Markers, /Actor1, etc).
         gen1_nodes = [self.root.GetChild(i) for i in range(self.root.GetChildCount())]
         for gen1_node in gen1_nodes:
@@ -400,14 +464,19 @@ class FBXContainer:
                           range(gen1_node.GetChildCount())]  # Actor nodes (/Mimi/Hips, /Mimi/ARIEL, etc)
 
             # If the first 3 marker names are children of this parent, it must be an actor.
-            if all(name in [node.GetName().split(':')[-1] for node in gen2_nodes] for name in self.marker_names[:4]):
+            if all(name in [node.GetName().split(':')[-1] for node in gen2_nodes] for name in names_to_look_for):
                 self.actor_names.append(gen1_node.GetName())
                 self.actors.append(gen1_node)
 
+        if len(self.actors) == 0 and not ignore_missing:
+            raise ValueError('No actors/subjects found. A node is considered an actor ' +
+                             'if it has the following children nodes: ' +
+                             ', '.join(names_to_look_for) + '.')
+
         self.actor_count = len(self.actors)
         self.valid_frames = [[] for _ in range(self.actor_count)]
 
-    def __init_markers(self):
+    def __init_markers(self, ignore_missing: bool = False) -> None:
         """
         Goes through all actor nodes and stores references to its marker nodes.
         """
@@ -421,11 +490,12 @@ class FBXContainer:
                     if match_name(child, marker_name, ignore_namespace=True):
                         actor_markers[marker_name] = child
 
-            assert len(actor_markers) == len(self.marker_names), f'{actor_node.GetName()} does not have all markers.'
+            if len(actor_markers) != len(self.marker_names) and not ignore_missing:
+                raise ValueError(f'{actor_node.GetName()} does not have all markers.')
 
             self.markers.append(actor_markers)
 
-    def __init_unlabeled_markers(self):
+    def __init_unlabeled_markers(self, ignore_missing: bool = False) -> None:
         """
         Looks for the Unlabeled_Markers parent node under the root and stores references to all unlabeled marker nodes.
         """
@@ -437,13 +507,69 @@ class FBXContainer:
                 self.unlabeled_markers = [gen1_node.GetChild(um) for um in range(gen1_node.GetChildCount())]
                 return
 
+        if not ignore_missing:
+            raise ValueError('No unlabeled markers found.')
+
+    def init_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> None:
+
+        self.init_labeled_world_transforms(r=r, incl_keyed=1)
+        self.init_unlabeled_world_transforms(r=r)
+
+    def init_labeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
+                                      incl_keyed: int = 1):
+
+        r = self.convert_r(r)
+        labeled_data = []
+
+        for actor_idx in range(self.actor_count):
+            actor_data = []
+            for marker_idx, (n, m) in enumerate(self.markers[actor_idx].items()):
+                curve = m.LclTranslation.GetCurve(self.anim_layer, 'X', True)
+                marker_data = get_world_transforms(actor_idx + 1, marker_idx + 1, m, r, curve, incl_keyed)
+                actor_data.append(marker_data)
+                self._print(f'Actor {actor_idx} marker {marker_idx} done', 1)
+            labeled_data.append(actor_data)
+
+        wide_layout = np.array(labeled_data)
+        self.labeled_world_transforms = np.transpose(wide_layout, axes=(3, 0, 1, 2))
+        return self.labeled_world_transforms
+
+    def init_unlabeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
+        r = self.convert_r(r)
+
+        unlabeled_data = []
+
+        for ulm in self.unlabeled_markers:
+            curve = ulm.LclTranslation.GetCurve(self.anim_layer, 'X', True)
+            marker_data = get_world_transforms(0, 0, ulm, r, curve, incl_keyed=0)
+            unlabeled_data.append(marker_data)
+            self._print(f'Unlabeled marker {ulm.GetName()} done', 1)
+
+        wide_layout = np.array(unlabeled_data)
+        self.unlabeled_world_transforms = np.transpose(wide_layout, axes=(2, 0, 1))
+        # Returns shape (n_frames, n_unlabeled_markers, 14).
+        return self.unlabeled_world_transforms
+
+    def init(self, ignore_missing_labeled: bool = False, ignore_missing_unlabeled: bool = False,
+             r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> None:
+        self.__init_scene()
+        self.__init_anim()
+        self.__init_actors(ignore_missing=ignore_missing_labeled)
+        self.__init_markers(ignore_missing=ignore_missing_labeled)
+        self.__init_unlabeled_markers(ignore_missing=ignore_missing_unlabeled)
+        self._print('Init done', 0)
+
+    def _print(self, txt: str, lvl: int = 0) -> None:
+        if lvl <= self.debug:
+            print(txt)
+
     def _check_actor(self, actor: int = 0):
         """
         Safety check to see if the actor `int` is a valid number (to avoid out of range errors).
         :param actor: `int` actor index, which should be between 0-max_actors.
         """
-        assert 0 <= actor <= self.actor_count, f'Actor index must be between 0 and {self.actor_count - 1}. ' \
-                                               f'It is {actor}.'
+        if not 0 <= actor <= self.actor_count:
+            raise ValueError(f'Actor index must be between 0 and {self.actor_count - 1} ({actor}).')
 
     def _set_valid_frames_for_actor(self, actor: int = 0):
         """
@@ -459,16 +585,18 @@ class FBXContainer:
         self._check_actor(actor)
 
         frames = self.get_frame_range()
-        for _, marker in self.markers[actor].items():
+        for n, marker in self.markers[actor].items():
             # Get the animation curve for local translation x.
             t_curve = marker.LclTranslation.GetCurve(self.anim_layer, 'X')
             # If an actor was recorded but seems to have no animation curves, we set their valid frames to nothing.
             # Then we return, because there is no point in further checking non-existent keyframes.
             if t_curve is None:
                 self.valid_frames[actor] = []
+                self._print('Found no animation curve', 2)
                 return
 
             # Get all keyframes on the animation curve and store their frame numbers.
+            self._print(f'Checking keyframes for {n}', 2)
             keys = [t_curve.KeyGet(i).GetTime().GetFrameCount() for i in range(t_curve.KeyGetCount())]
             # Check for each frame in frames if it is present in the list of keyframed frames.
             for frame in frames:
@@ -478,6 +606,7 @@ class FBXContainer:
                     with contextlib.suppress(ValueError):
                         frames.remove(frame)
 
+        self._print(f'Found {len(frames)}/{self.num_frames} valid frames for {self.actor_names[actor]}', 1)
         self.valid_frames[actor] = frames
 
         # Store all frame lists that have at least 1 frame.
@@ -486,13 +615,6 @@ class FBXContainer:
         self.common_frames = [num for num in self.get_frame_range()
                               if all(num in other_list for other_list in other_lists)]
 
-    def set_valid_frames(self):
-        """
-        For each actor, calls _set_valid_frames_for_actor().
-        """
-        for i in range(self.actor_count):
-            self._set_valid_frames_for_actor(i)
-
     def _check_valid_frames(self, actor: int = 0):
         """
         Safety check to see if the given actor has any valid frames stored.
@@ -502,9 +624,10 @@ class FBXContainer:
         self._check_actor(actor)
 
         if not len(self.valid_frames[actor]):
+            self._print(f'Getting missing valid frames for {self.actor_names[actor]}', 1)
             self._set_valid_frames_for_actor(actor)
 
-    def get_transformed_pc(self, actor: int = 0, frame: int = 0) -> List[float]:
+    def get_transformed_axes(self, actor: int = 0, frame: int = 0) -> Tuple[np.array, np.array, np.array]:
         """
         Evaluates all marker nodes for the given actor and modifies the resulting point cloud,
         so it is centered and scaled properly for training.
@@ -538,7 +661,20 @@ class FBXContainer:
         z /= self.vol_z
         y = np.array(y) / self.vol_y
 
-        # TODO: Optional: Add any extra modifications to the point cloud here.
+        # EXTRA: Add any extra modifications to the point cloud here.
+
+        return x, y, z
+
+    def get_transformed_pc(self, actor: int = 0, frame: int = 0, t: str = 'np') -> Union[np.array, List[float]]:
+
+        x, y, z = self.get_transformed_axes(actor, frame)
+        # If we need to return a numpy array, simply vstack the axes to get a shape of (3, 73).
+        # This is in preparation for PyTorch's CNN layers that use input shape (batch_size, C, H, W).
+        if t == 'np':
+            # Exports shape of (3, 9, 9).
+            # return make_pc_ghost_markers(np.vstack((x, y, z)))
+            # Exports shape of (1, 3, 73).
+            return np.vstack((x, y, z))[None, ...]
 
         # Append all values to a new array, one axis at a time.
         # This way it will match the column names order.
@@ -549,16 +685,6 @@ class FBXContainer:
             pose += [z[i]]
         return pose
 
-    def extract_scaled_translation(self, m: fbx.FbxNode, time: fbx.FbxTime) -> List[float]:
-        """
-        Evaluates a node's world translation at the given time and scales the vector down by a factor of self.scale.
-        :param m: `fbx.FbxNode` node that needs to be evaluated.
-        :param time: `fbx.FbxTime` at which frame/time the node needs to be evaluated.
-        :return: Translation vector as a list of floats.
-        """
-        t = m.EvaluateGlobalTransform(time).GetT()
-        return [t[i] * self.scale for i in range(3)]
-
     def get_frame_range(self) -> List[int]:
         """
         Replacement and improvement for:
@@ -568,6 +694,33 @@ class FBXContainer:
         """
         return list(range(self.start_frame, self.end_frame))
 
+    def convert_r(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None):
+        # If r is one int, use 0 as start frame. If r is higher than the total frames, limit the range.
+        if isinstance(r, int):
+            r = list(range(self.num_frames)) if r > self.num_frames else list(range(r))
+
+        # A tuple of 2 indicates a frame range without step.
+        elif isinstance(r, tuple) and len(r) == 2:
+            # If the requested frame range is longer than the total frames, limit the range.
+            if r[1] - r[0] > self.num_frames:
+                r = list(range(r[0], r[0] + self.num_frames))
+            else:
+                r = list(range(r[0], r[1]))
+
+        # A tuple of 3 indicates a frame range with step.
+        elif isinstance(r, tuple) and len(r) == 3:
+            # If the requested frame range is longer than the total frames, limit the range.
+            if r[1] - r[0] > self.num_frames:
+                r = list(range(r[0], r[0] + self.num_frames, r[2]))
+            else:
+                r = list(range(r[0], r[1], r[2]))
+
+        # If r is None, return the default frame range.
+        else:
+            r = self.get_frame_range()
+
+        return r
+
     def columns_from_joints(self) -> List[str]:
         """
         Generates a list of column names based on the (order of the) marker names.
@@ -652,40 +805,97 @@ class FBXContainer:
         self._check_valid_frames(actor)
         return self.valid_frames[actor]
 
-    def extract_valid_translations_per_actor(self, actor: int = 0) -> List[List[float]]:
+    def extract_valid_translations_per_actor(self, actor: int = 0, t: str = 'np'):
         """
         Assembles the poses for the valid frames for the given actor as a 2D list where each row is a pose.
         :param actor: `int` actor index.
+        :param t: If 'np', returns a (3, -1) `np.array`. Otherwise returns a list of floats.
         :return: List of poses, where each pose is a list of `float` translations.
         """
         # Ensure the actor index is within range.
         self._check_actor(actor)
+        self._check_valid_frames(actor)
 
-        poses = []
-        # Go through all valid frames for this actor.
-        # Note that these frames can be different per actor.
-        for frame in self.valid_frames[actor]:
-            # Get the centered point cloud as a 1D list.
-            pose_at_frame = self.get_transformed_pc(actor, frame)
-            poses.append(pose_at_frame)
+        # Returns shape (n_valid_frames, 3, 73).
+        return np.vstack([self.get_transformed_pc(actor, frame) for frame in self.valid_frames[actor]])
 
-        return poses
+        # poses = []
+        # # Go through all valid frames for this actor.
+        # # Note that these frames can be different per actor.
+        # for frame in self.valid_frames[actor]:
+        #     self._print(f'  Extracting frame: {frame}', 1)
+        #     # Get the centered point cloud as a 1D list.
+        #     pose_at_frame = self.get_transformed_pc(actor, frame, t)
+        #     poses.append(pose_at_frame)
+        #
+        # return np.array(poses) if t == 'np' else poses
 
-    def extract_all_valid_translations(self) -> pd.DataFrame:
+    def extract_all_valid_translations(self, t: str = 'np') -> Union[np.array, pd.DataFrame]:
         """
         Convenience method that calls self.extract_valid_translations_per_actor() for all actors
         and returns a `DataFrame` containing all poses after each other.
-        :return: `DataFrame` where each row is a pose.
-        """
-        # Note that the column names are/must be in the same order as the markers.
-        columns = self.columns_from_joints()
-
-        all_poses = []
-        # For each actor, add their valid poses to all_poses.
-        for i in range(self.actor_count):
-            all_poses.extend(self.extract_valid_translations_per_actor(i))
-
-        return pd.DataFrame(all_poses, columns=columns)
+        :param t: If 'np', returns a `np.array`. Otherwise, returns a DataFrame.
+        :return: `np.array` or `DataFrame` where each row is a pose.
+        """
+        # Returns shape (n_total_valid_frames, 3, 73).
+        return np.vstack([self.extract_valid_translations_per_actor(i) for i in range(self.actor_count)])
+        # all_poses = []
+        # # For each actor, add their valid poses to all_poses.
+        # for i in range(self.actor_count):
+        #     self._print(f'Extracting actor {self.actor_names[i]}', 0)
+        #     all_poses.extend(self.extract_valid_translations_per_actor(i, t))
+        #
+        # self._print('Extracting finished')
+        # # Note that the column names are/must be in the same order as the markers.
+        # if t == 'np':
+        #     # Shape: (n_poses, 3, 73).
+        #     return np.array(all_poses)
+        # else:
+        #     return pd.DataFrame(all_poses, columns=self.columns_from_joints())
+
+    def extract_training_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
+        if self.labeled_world_transforms is None:
+            self.init_labeled_world_transforms(r=r, incl_keyed=1)
+
+        # Returns (n_frames, n_actors, 73, 15).
+        l_shape = self.labeled_world_transforms.shape
+        # Flatten the array, so we get a list of frames.
+        # Reshape to (n_frames * n_actors, 73, 15).
+        flattened = self.labeled_world_transforms.reshape(-1, l_shape[2], l_shape[3])
+        # Isolates the poses with all keyframes present by checking the last elements.
+        # Start with the mask.
+        # Returns shape of (n_frames * n_actors, 73).
+        mask = (flattened[..., -1] == 1)
+        # We only need a filter for the first dimension, so use .all to check if all markers
+        # have a keyframe. This results in shape (n_frames * n_actors,).
+        mask = mask.all(axis=1)
+        # Now isolate the right frames with the mask and remove the last element of the last dimension,
+        # because it won't be useful anymore.
+        valid_poses = flattened[mask][..., :-1]
+
+        # Now we need to center the tx and tz axes.
+        for valid_pose in valid_poses:
+            for axis in [2, 4]:
+                valid_pose[:, axis] = center_axis(valid_pose[:, axis])
+        return self.transform_translations(valid_poses)
+
+    def extract_inf_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
+                                 merged: bool = True) -> Union[np.array, Tuple[np.array, np.array]]:
+        if self.labeled_world_transforms is None:
+            self.init_labeled_world_transforms(r=r, incl_keyed=0)
+        if self.unlabeled_world_transforms is None:
+            self.init_unlabeled_world_transforms(r=r)
+
+        ls = self.labeled_world_transforms.shape
+        # Returns shape (n_frames, 73 * n_actors, 14).
+        flat_labeled = self.labeled_world_transforms.reshape(ls[0], -1, ls[-1])[..., :14]
+
+        if merged:
+            return utils.merge_labeled_and_unlabeled_data(labeled=flat_labeled,
+                                                          unlabeled=self.unlabeled_world_transforms,
+                                                          pc_size=self.pc_size)
+        else:
+            return flat_labeled, self.unlabeled_world_transforms
 
     def transform_translations(self, w: np.array) -> np.array:
         """
@@ -697,7 +907,7 @@ class FBXContainer:
             raise ValueError(f'Array does not have 3 dimensions: {w.ndim}/3.')
 
         # If the last dimension has 3 elements, it is a translation vector of shape (tx, ty, tz).
-        # If it has 14 elements, it is a full marker row of shape (actor, marker, tx, ty, tz, rx, ry, rz, etc).
+        # If it has 14 elements, it is a full marker row of shape (actor, marker, tx, ty, tz, tw, rx, ry, rz, tw, etc).
         start = 0 if w.shape[-1] == 3 else 2
 
         # First multiply by self.scale, which turns meters to centimeters.
@@ -761,6 +971,21 @@ class FBXContainer:
         # so return the cloud as a np array that cuts off any excessive markers.
         return np.array(cloud)[:self.pc_size]
 
+    def get_dc(self, frame: int = 0) -> np.array:
+        self._print(f'Getting sparse cloud for frame {frame}', 2)
+        cloud = self.get_sc(frame)
+        missing = self.pc_size - cloud.shape[0]
+
+        # Only bother creating ghost markers if there are any missing rows.
+        # If we need to add ghost markers, add them before the existing cloud,
+        # so that the cloud will remain a sorted array regarding the actor and marker classes.
+        if missing > 0:
+            self._print('Making ghost markers', 2)
+            ghost_cloud = make_ghost_markers(missing)
+            cloud = np.vstack([ghost_cloud, cloud])
+
+        return cloud
+
     def get_tsc(self) -> np.array:
         """
         Convenience method that calls self.get_sparse_cloud() for all frames in the frame range
@@ -769,7 +994,7 @@ class FBXContainer:
         """
         return np.array([self.get_sc(f) for f in self.get_frame_range()])
 
-    def get_tdc(self, r: Union[int, Tuple[int, int]] = None) -> np.array:
+    def get_tdc(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
         """
         For each frame in the frame range, collects the point cloud that is present in the file.
         Then it creates a ghost cloud of random markers that are treated as unlabeled markers,
@@ -781,32 +1006,27 @@ class FBXContainer:
         with a shape of (self.num_frames, self.pc_size, 5).
         """
 
-        clouds = []
+        r = self.convert_r(r)
 
-        # If r is one int, use 0 as start frame.
-        if isinstance(r, int):
-            r = list(range(r))
-        # If r is two ints, use that as specific frame range.
-        elif isinstance(r, tuple) and len(r) >= 2:
-            r = list(range(r[0], r[1]))
-        # If r is empty, use the animation frame range.
-        else:
-            r = self.get_frame_range()
+        # results = utils.parallel_process(r, self.get_dc)
 
-        for frame in r:
-            cloud = self.get_sc(frame)
-            missing = self.pc_size - cloud.shape[0]
+        return np.array([self.get_dc(f) for f in r])
 
-            # Only bother creating ghost markers if there are any missing rows.
-            # If we need to add ghost markers, add them before the existing cloud,
-            # so that the cloud will remain a sorted array regarding the actor and marker classes.
-            if missing > 0:
-                ghost_cloud = make_ghost_markers(missing)
-                cloud = np.vstack([ghost_cloud, cloud])
+    def modify_actor_pose(self, actor: np.array) -> np.array:
+        # Scale to cm.
+        actor[:, 2:5] *= self.scale
+        # Move the point cloud to the center of the x and y axes. This will put the actor in the middle.
+        for axis in range(2, 5):
+            actor[:, axis] = center_axis(actor[:, axis])
 
-            clouds.append(cloud)
+        # Move the actor to the middle of the volume floor by adding volume_dim/2 to x and z.
+        actor[:, 2] += self.vol_x / 2.
+        actor[:, 4] += self.vol_z / 2.
 
-        return np.array(clouds)
+        # Squeeze the actor into the 1x1 plane for the neural network by dividing the axes.
+        actor[:, 2] /= self.vol_x
+        actor[:, 3] /= self.vol_y
+        actor[:, 4] /= self.vol_z
 
     def split_tdc(self, cloud: np.array = None) \
             -> Tuple[np.array, np.array, np.array, np.array, np.array]:
@@ -822,13 +1042,28 @@ class FBXContainer:
         :return: Return tuple of `np.array` as (actor classes, marker classes, translation vectors).
         """
         if cloud is None:
-            cloud = self.get_tdc()
+            cloud = self.extract_inf_translations()
+
+        if cloud.shape[1] != self.pc_size:
+            raise ValueError(f"Dense cloud doesn't have enough points. {cloud.shape[1]}/{self.pc_size}.")
+        if cloud.shape[2] < 14:
+            raise ValueError(f"Dense cloud is missing columns: {cloud.shape[2]}.")
 
-        if cloud.shape[1] != 1000:
-            raise ValueError(f"Dense cloud doesn't have enough points. {cloud.shape[1]}/1000.")
-        if cloud.shape[2] != 14:
-            raise ValueError(f"Dense cloud is missing columns: {cloud.shape[2]}/14.")
+        # Return np arrays as (actor classes, marker classes, translation vectors, rotation vectors, scale vectors).
+        return cloud[:, :, 0], cloud[:, :, 1], cloud[:, :, 2:5], cloud[:, :, 6:9], cloud[:, :, 10:13]
 
+    def get_split_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
+                             mode: str = 'train') -> Tuple[np.array, np.array, np.array, np.array, np.array]:
+        """
+        Splits a timeline dense cloud with shape (self.num_frames, self.pc_size, 5) into 3 different
+        arrays:
+        1. A `np.array` with the actor classes as shape (self.num_frames, self.pc_size, 1).
+        2. A `np.array` with the marker classes as shape (self.num_frames, self.pc_size, 1).
+        3. A `np.array` with the translation floats as shape (self.num_frames, self.pc_size, 4).
+        4. A `np.array` with the rotation Euler angles as shape (self.num_frames, self.pc_size, 3).
+        :return: Return tuple of `np.array` as (actor classes, marker classes, translation vectors).
+        """
+        cloud = self.extract_training_translations(r) if mode == 'train' else self.extract_inf_translations(r)
         # Return np arrays as (actor classes, marker classes, translation vectors, rotation vectors, scale vectors).
         return cloud[:, :, 0], cloud[:, :, 1], cloud[:, :, 2:5], cloud[:, :, 6:9], cloud[:, :, 10:13]
 
@@ -848,21 +1083,50 @@ class FBXContainer:
         """
         return 'UNLABELED' if int(c) == 0 else self.marker_names[int(c) - 1]
 
-    def export(self, t: str = 'csv', output_file: Path = None) -> Union[bytes, Path]:
-        # Get the dataframe with all animation data.
-        df = self.extract_all_valid_translations()
-
-        if t == 'string':
+    def export_train_data(self, output_file: Path, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) \
+            -> Union[bytes, pd.DataFrame, np.array]:
+        if output_file is None:
+            df = pd.DataFrame(self.extract_training_translations(r))
             return df.to_csv(index=False).encode('utf-8')
 
-        if output_file is None:
-            output_file = self.input_fbx.with_suffix('.csv')
+        elif output_file.suffix == '.npy':
+            array_4d = self.extract_training_translations(r)
+            np.save(str(output_file), array_4d)
+            self._print(f'Exported train data to {output_file}', 0)
+            return array_4d
+
+        elif output_file.suffix == '.h5':
+            array_4d = self.extract_training_translations(r)
+            with h5py.File(output_file, 'w') as h5f:
+                h5f.create_dataset('array_data', data=array_4d, compression='gzip', compression_opts=9)
+                self._print(f'Exported train data to {output_file}', 0)
+            return array_4d
+
+        else:
+            raise ValueError('Invalid file extension. Must be .csv or .npy')
+
+    def export_test_data(self, output_file: Path, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
+                         merged: bool = True) -> Union[np.array, Tuple[np.array, np.array]]:
+        # Retrieve the clean world transforms.
+        # If merged is True, this will be one array of shape (n_frames, pc_size, 14).
+        # If merged is False, this will be two arrays, one of shape (n_frames, 73 * n_actors, 14),
+        # and one of shape (n_frames, n_unlabeled_markers, 14).
+        array_4d = self.extract_inf_translations(r, merged=merged)
 
-        if output_file.suffix != '.csv':
-            raise ValueError(f'{output_file} needs to be a .csv file.')
+        if output_file.suffix == '.h5':
+            with h5py.File(output_file, 'w') as h5f:
+                if merged:
+                    # If merged, this can be one dataset.
+                    h5f.create_dataset('merged_data', data=array_4d, compression='gzip', compression_opts=9)
 
-        df.to_csv(output_file, index=False)
-        return output_file
+                else:
+                    # If not merged, we split it up because array_4d is a tuple of 2.
+                    h5f.create_dataset('labeled', data=array_4d[0], compression='gzip', compression_opts=9)
+                    h5f.create_dataset('unlabeled', data=array_4d[1], compression='gzip', compression_opts=9)
+
+                self._print(f'Exported test data to {output_file}', 0)
+
+        return array_4d
 
     def export_fbx(self, output_file: Path = None) -> bool:
         """
@@ -879,18 +1143,20 @@ class FBXContainer:
         # Initialize the exporter with the output file path
         result = exporter.Initialize(str(output_file))
         if not result:
-            print(f"Failed to initialize the exporter for file '{output_file}'.")
+            self._print(f"Failed to initialize the exporter for file '{output_file}'.", 0)
             return False
 
         # Export the scene
         result = exporter.Export(self.scene)
         if not result:
-            print(f"Failed to export the scene to file '{output_file}'.")
+            self._print(f"Failed to export the scene to file '{output_file}'.", 0)
             return False
 
         # Clean up the manager and exporter
         exporter.Destroy()
 
+        self._print('Export finished', 0)
+
         return True
 
     def remove_node(self, node: fbx.FbxNode, recursive: bool = False) -> bool:
@@ -1024,6 +1290,7 @@ class FBXContainer:
         for marker_class, (marker_name, marker) in enumerate(self.markers[actor].items(), start=1):
             marker_keys = actor_keys.get(marker_class)
             if marker_keys:
+                self._print(f'Replacing keys for {marker_name}', 1)
                 self.replace_keyframes_per_marker(marker, marker_keys)
 
     def replace_keyframes_for_all_actors(self, key_dict: dict) -> None:
@@ -1032,23 +1299,7 @@ class FBXContainer:
         :param key_dict: `dict` with all actor keyframes.
         """
         for actor_idx in range(self.actor_count):
-            actor_dict = key_dict.get(actor_idx+1)
+            actor_dict = key_dict.get(actor_idx + 1)
             if actor_dict:
+                self._print(f'Replacing keys for actor {actor_idx}', 1)
                 self.replace_keyframes_per_actor(actor_idx, actor_dict)
-
-
-# d = FBXContainer(Path('G:/Firestorm/mocap-ai/data/fbx/dowg/TAKE_01+1_ALL_001.fbx'))
-# og_cloud = d.get_tdc()
-# # print(og_cloud[0, -10:, 2:5])
-# di = tsc_to_dict(og_cloud)
-# d.replace_keyframes_for_all_actors(di)
-# # new_cloud = d.get_tdc(r=100)
-# # print(new_cloud[0, -10:, 2:5])
-# # actors_train, markers_train, t_train, r_train, s_train = d.split_tdc(cloud)
-# # # t_train_transformed = d.transform_translations(t_train)
-# # # splits = d.split_tdc(apply_transform=False)
-# # merged = merge_tdc(actors_train, markers_train, t_train, r_train, s_train)
-# # pc_dict = tsc_to_dict(merged, d.start_frame)
-# # d.replace_keyframes_for_all_actors(pc_dict)
-# # # d.cleanup()
-# d.export_fbx(Path('G:/Firestorm/mocap-ai/data/fbx/export/TAKE_01+1_ALL_001.fbx'))

labeler/data_setup.py

@@ -0,0 +1,161 @@
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch.utils import data
+import math
+
+
+def apply_random_y_rotation(point_cloud_data: torch.Tensor) -> torch.Tensor:
+    # Convert the random angle from degrees to radians
+    angle = (torch.rand(1).item() * 2 - 1) * 180 * torch.tensor(math.pi / 180, device='cuda')
+
+    # Create the rotation matrix for the y-axis
+    rotation_matrix = torch.tensor([[torch.cos(angle), 0, torch.sin(angle)],
+                                    [0, 1, 0],
+                                    [-torch.sin(angle), 0, torch.cos(angle)]], device='cuda')
+
+    # Apply the rotation to the point cloud data
+    return torch.matmul(point_cloud_data, rotation_matrix.T)
+
+
+class PointCloudDataset(data.Dataset):
+    def __init__(self, file: Path,
+                 n_samples=100,
+                 max_actors: int = 8,
+                 translation_factor=0.1,
+                 max_overlap: Tuple[float] = (0.2, 0.2, 0.2)):
+        point_clouds_np = torch.tensor(np.load(str(file)), dtype=torch.float32, device='cuda')
+        self.sparse_point_clouds = point_clouds_np
+        self.n_samples = n_samples
+        self.max_actors = max_actors
+        self.translation_factor = translation_factor
+        self.max_overlap = max_overlap
+
+        # Generate a random permutation of indices.
+        self.indices = torch.randperm(len(self.sparse_point_clouds))
+
+        dataset = []
+        for _ in range(n_samples):
+            accumulated_cloud = []
+            # TODO: Get a random number up to the max of actors.
+            # TODO: Transform one row of the available rows, and check if it doesn't overlap.
+            # TODO: Accumulate all actors into one point cloud and append that to dataset.
+            # TODO: __getitem__() needs to get one of these point cloud rows.
+            for i in range(max_actors):
+
+                # Get a point cloud from the tensor using the shuffled index, shape (1, 1024).
+                point_cloud = self.sparse_point_clouds[self.indices[index]]
+
+                point_cloud_data = point_cloud[:, 2:5]  # returns shape: (1024, 3)
+
+                valid_transform = False
+                while not valid_transform:
+
+                    point_cloud = point_cloud_data.clone()
+                    # Randomly translate the point cloud along the x and z axes
+
+                    self.apply_random_translation(point_cloud)
+                    # Apply random rotation around the y-axis
+                    rotated_point_cloud_data = apply_random_y_rotation(point_cloud)
+
+                    if not does_overlap(accumulated_cloud, point_cloud, self.max_overlap):
+                        accumulated_cloud.append(point_cloud)
+                        valid_transform = True
+
+    def apply_random_translation(self, point_cloud: torch.Tensor) -> None:
+        x_translation = (torch.rand(1).item() * 2 - 1) * self.translation_factor
+        z_translation = (torch.rand(1).item() * 2 - 1) * self.translation_factor
+        point_cloud[:, [0, 2]] += torch.tensor([x_translation, z_translation], device='cuda')
+
+    def fill_point_cloud(self, point_cloud):
+        target_num_points = 73 * self.max_actors
+        current_num_points = point_cloud.shape[1]
+
+        if current_num_points < target_num_points:
+            num_points_to_add = target_num_points - current_num_points
+            random_indices = torch.randint(0, current_num_points, (num_points_to_add,))
+            additional_points = point_cloud[:, random_indices, :]
+
+            filled_point_cloud = torch.cat((point_cloud, additional_points), dim=1)
+        else:
+            filled_point_cloud = point_cloud
+
+        return filled_point_cloud
+
+    def __getitem__(self, index):
+
+        point_cloud = np.vstack(accumulated_cloud)
+        # Separate the labels from the point cloud data
+        actor_labels = point_cloud[:, :, 0]  # shape: (1024,)
+        marker_labels = point_cloud[:, :, 1]  # shape: (1024,)
+
+        return actor_labels, marker_labels, rotated_point_cloud_data
+
+    def __len__(self):
+        return len(self.sparse_point_clouds)
+
+
+def does_overlap(accumulated_point_cloud, new_point_cloud, overlap_thresholds=(0.2, 0.2, 0.2)):
+    def project_to_axis(point_cloud, axis):
+        projected_points = point_cloud.clone()
+        projected_points[:, axis] = 0
+        return projected_points
+
+    def get_bounding_box_2d(points):
+        min_values, _ = torch.min(points, dim=0)
+        max_values, _ = torch.max(points, dim=0)
+        return min_values, max_values
+
+    def check_surface_area_overlap(bb1_min, bb1_max, bb2_min, bb2_max, axis, overlap_threshold):
+        bb1_area = (bb1_max[axis] - bb1_min[axis]) * (bb1_max[1] - bb1_min[1])
+        bb2_area = (bb2_max[axis] - bb2_min[axis]) * (bb2_max[1] - bb2_min[1])
+
+        overlap_min = torch.max(bb1_min, bb2_min)
+        overlap_max = torch.min(bb1_max, bb2_max)
+
+        overlap_area = (overlap_max[axis] - overlap_min[axis]) * (overlap_max[1] - overlap_min[1])
+        overlap_area = torch.max(torch.tensor(0.0, device='cuda'), overlap_area)  # Clamp to 0 if negative
+
+        overlap_percentage = overlap_area / torch.min(bb1_area, bb2_area)
+
+        return overlap_percentage >= overlap_threshold
+
+    new_point_cloud_xz = project_to_axis(new_point_cloud, 1)  # Project to xz-plane (remove y-axis values)
+    new_point_cloud_min, new_point_cloud_max = get_bounding_box_2d(new_point_cloud_xz)
+
+    overlaps = []
+
+    for pc in accumulated_point_cloud:
+        for axis in range(len(overlap_thresholds)):
+            pc_xz = project_to_axis(pc, axis)  # Project to xz-plane (remove y-axis values)
+            pc_min, pc_max = get_bounding_box_2d(pc_xz)
+
+            if all(
+                    check_surface_area_overlap(
+                        new_point_cloud_min,
+                        new_point_cloud_max,
+                        pc_min,
+                        pc_max,
+                        axis,
+                        overlap_thresholds[axis],
+                    )
+                    for axis in range(len(overlap_thresholds))
+            ):
+                return True
+
+    return False
+
+
+class NoOverlapDataLoader(data.DataLoader):
+    def __init__(self, dataset: data.Dataset, max_overlap: Tuple[float] = (0.2, 0.2, 0.2), *args, **kwargs):
+        super().__init__(dataset, *args, **kwargs)
+        self.max_overlap = max_overlap
+
+    def __iter__(self):
+        accumulated_point_clouds = []
+        for actor_labels, marker_labels, point_cloud_data in super().__iter__():
+            if not does_overlap(accumulated_point_clouds, point_cloud_data, self.max_overlap):
+                accumulated_point_clouds.append(point_cloud_data)
+                yield actor_labels, marker_labels, point_cloud_data

preprocess_files.py

@@ -0,0 +1,71 @@
+from pathlib import Path
+import shutil
+
+# Import custom libs.
+import fbx_handler
+import utils
+
+
+def process_fbx_files(source_folder: Path, train_folder: Path, test_folder: Path, v: int = 1):
+    # Delete the existing folders and make them again, because the array4d_to_h5 function will append
+    # new data to any existing files.
+    shutil.rmtree(train_folder)
+    shutil.rmtree(test_folder)
+    train_folder.mkdir(parents=True, exist_ok=True)
+    test_folder.mkdir(parents=True, exist_ok=True)
+
+    files = list(source_folder.glob('*.fbx'))
+    # files = [Path('G:/Firestorm/mocap-ai/data/fbx/mes-1/HangoutSpot_1_003.fbx')]
+
+    # Create Paths to new files that will contain all data.
+    train_all = train_folder / 'ALL.h5'
+    test_all = test_folder / 'ALL.h5'
+
+    with utils.Timer('Extracting took'):
+        # Iterate through all .fbx files in the source folder
+        for idx, fbx_file in enumerate(files):
+            print(f'{idx + 1}/{len(files)}: {fbx_file}')
+            # Create a new class object with the file path.
+            my_obj = fbx_handler.FBXContainer(fbx_file, max_actors=4, pc_size=296, debug=0, save_init=True)
+            # Init world transforms for labeled and unlabeled data. This will store all relevant transform info.
+            with utils.Timer('Getting world transforms took'):
+                my_obj.init_world_transforms()
+            # Define the export file path with the same file name but in the export folder
+            export_train_path = train_folder / fbx_file.with_suffix('.h5').name
+            export_test_path = test_folder / fbx_file.with_suffix('.h5').name
+
+            # Get the train data as an array of shape (n_valid_frames, 73, 14).
+            # This will also export it to a h5 file just in case.
+            train_data = my_obj.export_train_data(export_train_path)
+            print(f'Train shape:          {train_data.shape}')
+            # Append the array to the existing ALL file.
+            utils.array4d_to_h5(array_4ds=(train_data,),
+                                output_file=train_all,
+                                datasets=(fbx_file.stem,))
+
+            # Do the same thing for the test data.
+            test_data = my_obj.export_test_data(export_test_path, merged=False)
+            print(f'Test labeled shape:   {test_data[0].shape}')
+            print(f'Test unlabeled shape: {test_data[1].shape}')
+            print(f'Minimum cloud size:   {test_data[0].shape[1] + test_data[1].shape[1]}')
+            utils.array4d_to_h5(array_4ds=test_data,
+                                output_file=test_all,
+                                group=fbx_file.stem,
+                                datasets=('labeled', 'unlabeled'))
+
+    print('--- FINAL ---')
+    # Just to be sure, print the shapes of the final results.
+    with utils.Timer('Loading training data took'):
+        print(f"Final train shape:    {utils.h5_to_array4d(train_all, mode='train').shape}")
+
+    with utils.Timer('Loading testing data took'):
+        print(f"Final test shape:     {utils.h5_to_array4d(test_all, mode='test').shape}")
+
+
+if __name__ == '__main__':
+    source = Path('G:/Firestorm/mocap-ai/data/fbx/mes-1/')
+    train = Path('G:/Firestorm/mocap-ai/data/h5/mes-1/train')
+    test = Path('G:/Firestorm/mocap-ai/data/h5/mes-1/test')
+
+    with utils.Timer('Full execution took'):
+        process_fbx_files(source, train, test)

requirements.txt

@@ -1,3 +1,5 @@
 streamlit~=1.21.0
 pandas~=1.3.5
-numpy~=1.21.5
+numpy~=1.21.5
+torch~=1.13.1
+h5py

utils.py

@@ -0,0 +1,104 @@
+import cProfile
+import pstats
+import time
+from pathlib import Path
+from typing import List, Tuple
+
+import h5py
+import numpy as np
+
+
+def append_suffix_to_file(file_path: Path, suffix: str = '_INF', ext: str = None):
+    """
+    Adds a suffix to the given file path.
+    :param file_path: `Path` object to the original file.
+    :param suffix: `str` suffix to add to the end of the original file name.
+    :param ext: `str` potential new file extension.
+    :return: Updated `Path`.
+    """
+    if ext:
+        file_path = file_path.with_suffix(ext)
+    new_file_name = file_path.stem + suffix + file_path.suffix
+    return file_path.with_name(new_file_name)
+
+
+def is_int_in_list(n: int, l: List[int]) -> int:
+    if l[0] > n:
+        return 0
+
+    for e in l:
+        if e == n:
+            return 1
+        elif e > n:
+            return 0
+
+    return 0
+
+
+def array4d_to_h5(array_4ds: Tuple, output_file: Path, group: str = None, datasets: Tuple = 'array_data'):
+    if len(array_4ds) != len(datasets):
+        raise ValueError(f'Amount of arrays {len(array_4ds)} must match amount of dataset names {len(datasets)}.')
+    with h5py.File(output_file, 'a') as h5f:
+        if group is not None:
+            grp = h5f.create_group(group)
+            for i in range(len(array_4ds)):
+                grp.create_dataset(name=datasets[i], data=array_4ds[i], compression='gzip', compression_opts=9)
+        else:
+            for i in range(len(array_4ds)):
+                h5f.create_dataset(name=datasets[i], data=array_4ds[i], compression='gzip', compression_opts=9)
+
+
+def h5_to_array4d(input_file: Path, mode: str = 'train', pc_size: int = 1024) -> np.array:
+    with h5py.File(input_file, 'r') as h5f:
+        if mode == 'train':
+            return np.vstack([np.array(h5f[key]) for key in h5f.keys()])
+
+        data = []
+        for grp_name in h5f.keys():
+            grp = h5f[grp_name]
+            labeled = np.array(grp['labeled'])
+            unlabeled = np.array(grp['unlabeled'])
+            data.append(merge_labeled_and_unlabeled_data(labeled, unlabeled, pc_size=pc_size))
+
+        return np.vstack(data)
+
+
+def merge_labeled_and_unlabeled_data(labeled: np.array, unlabeled: np.array, pc_size: int) -> np.array:
+    missing = pc_size - (labeled.shape[1] + unlabeled.shape[1])
+    if missing <= 0:
+        # Returns shape (n_frames, self.pc_size, 14).
+        return np.concatenate((unlabeled, labeled), axis=1)[:, -pc_size:]
+
+    missing_markers = np.random.rand(labeled.shape[0], missing, labeled.shape[-1])
+    missing_markers[:, :, 0] = 0.
+    missing_markers[:, :, 1] = 0.
+
+    # Returns shape (n_frames, self.pc_size, 14).
+    return np.concatenate((missing_markers,
+                           unlabeled,
+                           labeled), axis=1)
+
+
+class Timer:
+    def __init__(self, txt: str = 'Execution time: ', profiler: bool = False):
+        self.txt = txt
+        self.profiler = profiler
+
+    def __enter__(self):
+        self.start_time = time.time()
+        if self.profiler:
+            self.p = cProfile.Profile()
+            self.p.enable()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end_time = time.time()
+        dif = self.end_time - self.start_time
+        print(f"{self.txt}: {dif:.4f} seconds")
+
+        if self.profiler:
+            self.p.disable()
+            stats = pstats.Stats(self.p).sort_stats('time')
+            stats.print_stats()
+
+