First attempt at recalculating the correct world spaces back to local spaces.

fbx_handler.py

@@ -1,4 +1,6 @@
 # Import core libs.
+from pprint import pprint
+
 import pandas as pd
 import numpy as np
 from pathlib import Path
@@ -7,6 +9,7 @@ from typing import List, Union, Tuple
 # Import util libs.
 import contextlib
 import fbx
+import itertools
 
 # Import custom data.
 import globals
@@ -46,6 +49,162 @@ 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 merge_tdc(actor_classes: np.array,
+              marker_classes: np.array,
+              translation_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
+    # splitting it into sub arrays.
+    tdc = np.concatenate((np.expand_dims(actor_classes, -1), np.expand_dims(marker_classes, -1),
+                          translation_vectors), axis=2)
+    if ordered:
+        tdc = sort_cloud(tdc)
+
+    return tdc
+
+
+def sort_cloud(cloud: np.array) -> np.array:
+    """
+    Convenience function to sort a timeline dense cloud by actor and marker classes.
+    Not required.
+    :param cloud: `np.array` point cloud to sort.
+    :return: sorted `np.array` point cloud.
+    """
+    # Extract the first two elements of the third dimension
+    actor_classes = cloud[:, :, 0]
+    marker_classes = cloud[:, :, 1]
+
+    # Create an empty array with the same shape as the input array
+    sorted_tdc = np.empty_like(cloud)
+
+    # Sort the input array row by row using the first two elements of the third dimension
+    for i in range(cloud.shape[0]):
+        # Get the sorting indices for the current row
+        sorted_indices = np.lexsort((marker_classes[i], actor_classes[i]))
+
+        # Sort the current row using the sorted indices
+        sorted_tdc[i] = cloud[i, sorted_indices]
+
+    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()
+    t.SetFrame(frame)
+
+    # Create a new keyframe with the specified value
+    key_index = anim_curve.KeyAdd(t)[0]
+    anim_curve.KeySetValue(key_index, value)
+    return True
+
+
+def replace_keyframes_on_curve(anim_curve: fbx.FbxAnimCurve, frames: List[int], values: np.array):
+    # Check if the anim_curve is of type FbxAnimCurve
+    if not isinstance(anim_curve, fbx.FbxAnimCurve):
+        print("Input is not an FbxAnimCurve instance.")
+        return False
+
+    # Check if the frames and values lists have the same length
+    if len(frames) != len(values):
+        print("Frames and values lists have different lengths.")
+        return False
+
+    # Remove all existing keyframes
+    anim_curve.KeyClear()
+
+    # Create new keyframes with the given frames and values
+    for frame, value in zip(frames, values):
+        create_keyframe(anim_curve, frame, value)
+
+    return True
+
+
+def get_child_node_by_name(parent_node: fbx.FbxNode, name: str, ignore_namespace: bool = False) \
+        -> Union[fbx.FbxNode, None]:
+    for c in range(parent_node.GetChildCount()):
+        child = parent_node.GetChild(c)
+        if match_name(child, name, ignore_namespace):
+            return child
+    return None
+
+
+def match_name(node: fbx.FbxNode, name: str, ignore_namespace: bool = True) -> bool:
+    node_name = node.GetName()
+    if ignore_namespace:
+        node_name = node_name.split(':')[-1]
+    return node_name == name
+
+
+def timeline_cloud_to_dict(data: np.array, start_frame: int = 0) -> dict:
+    # Initialize an empty dictionary.
+    result = {}
+
+    # Iterate over the first dimension (frames) and second dimension (markers).
+    for frame, node in itertools.product(range(data.shape[0]), range(data.shape[1])):
+
+        # Extract the actor class, node class, and translation vector.
+        actor_class = int(data[frame, node, 0])
+        marker_class = int(data[frame, node, 1])
+        # If actor or marker class is predicted to be 0 (unlabeled marker), then skip adding it to the dict,
+        # because we only want to keyframe labeled markers.
+        if actor_class == 0 or marker_class == 0:
+            continue
+
+        translation_vector = data[frame, node, 2:]
+
+        # Create the actor dictionary if it doesn't exist.
+        if actor_class not in result:
+            result[actor_class] = {}
+
+        # Create the node dictionary if it doesn't exist.
+        if marker_class not in result[actor_class]:
+            result[actor_class][marker_class] = {}
+
+        # Add the frame number and translation vector to the node dictionary.
+        result[actor_class][marker_class][frame + start_frame] = translation_vector
+
+    return result
+
+
+def world_to_local_translation(node, parent_node, f):
+    t = fbx.FbxTime()
+    t.SetFrame(f)
+    child_world_matrix = node.EvaluateGlobalTransform(t)
+    child_world_matrix = np.array([
+        [child_world_matrix.Get(i, j) for j in range(4)] for i in range(4)
+    ])
+    # Convert the world_translation vector to a homogeneous 4D vector by appending a 1
+    # world_translation_homogeneous = np.append(world_translation, 1)
+
+    # Get the parent's world transformation matrix as a numpy array
+    parent_world_matrix = parent_node.EvaluateGlobalTransform(t)
+    parent_world_matrix = np.array([
+        [parent_world_matrix.Get(i, j) for j in range(4)] for i in range(4)
+    ])
+
+    # Compute the inverse of the parent's world transformation matrix
+    parent_world_matrix_inv = np.linalg.inv(parent_world_matrix)
+
+    return np.dot(parent_world_matrix_inv, child_world_matrix)
+
+
 class FBXContainer:
     # TODO: Model is currently built for training. Add testing mode.
     def __init__(self, fbx_file: Path,
@@ -80,14 +239,15 @@ class FBXContainer:
         self.vol_x = volume_dims[0]
         self.vol_y = volume_dims[1]
         self.vol_z = volume_dims[2]
-        
+
         self.scale = scale
 
         self.max_actors = max_actors
         # Maximum point cloud size = 73 * max_actors + unlabeled markers.
         self.pc_size = pc_size
 
-        self.fbx_file = fbx_file
+        self.input_fbx = fbx_file
+        self.output_fbx = append_suffix(fbx_file, '_INF')
         self.valid_frames = []
 
         self.__init_scene()
@@ -102,15 +262,16 @@ class FBXContainer:
         Destroys the importer to remove the reference to the loaded file.
         """
         # Create an FBX manager and importer.
-        manager = fbx.FbxManager.Create()
-        importer = fbx.FbxImporter.Create(manager, '')
+        self.manager = fbx.FbxManager.Create()
+        importer = fbx.FbxImporter.Create(self.manager, '')
 
         # Import the FBX file.
-        importer.Initialize(str(self.fbx_file))
-        self.scene = fbx.FbxScene.Create(manager, '')
+        importer.Initialize(str(self.input_fbx))
+        self.scene = fbx.FbxScene.Create(self.manager, '')
         importer.Import(self.scene)
         self.root = self.scene.GetRootNode()
         self.time_mode = self.scene.GetGlobalSettings().GetTimeMode()
+        fbx.FbxTime.SetGlobalTimeMode(self.time_mode)
 
         # Destroy importer to remove reference to imported file.
         # This will allow us to delete the uploaded file.
@@ -126,9 +287,9 @@ class FBXContainer:
 
         # 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(self.time_mode))
-        self.start_frame = local_time_span.GetStart().GetFrameCount(self.time_mode)
-        self.end_frame = local_time_span.GetStop().GetFrameCount(self.time_mode)
+        self.num_frames = int(local_time_span.GetDuration().GetFrameCount())
+        self.start_frame = local_time_span.GetStart().GetFrameCount()
+        self.end_frame = local_time_span.GetStop().GetFrameCount()
 
     def __init_actors(self):
         """
@@ -161,9 +322,8 @@ class FBXContainer:
                 for actor_idx in range(actor_node.GetChildCount()):
                     child = actor_node.GetChild(actor_idx)
                     # Child name might have namespaces in it like this: Vera:ARIEL
-                    # We want to match only on the actual name, so strip everything before off.
-                    child_name = child.GetName().split(':')[-1]
-                    if child_name == marker_name:
+                    # We want to match only on the actual name, so ignore namespaces.
+                    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.'
@@ -177,7 +337,7 @@ class FBXContainer:
         # Find the Unlabeled_Markers parent node.
         for i in range(self.root.GetChildCount()):
             gen1_node = self.root.GetChild(i)
-            if gen1_node.GetName() == 'Unlabeled_Markers':
+            if match_name(gen1_node, 'Unlabeled_Markers'):
                 self.unlabeled_markers_parent = gen1_node
                 self.unlabeled_markers = [gen1_node.GetChild(um) for um in range(gen1_node.GetChildCount())]
                 return
@@ -214,7 +374,7 @@ class FBXContainer:
                 return
 
             # Get all keyframes on the animation curve and store their frame numbers.
-            keys = [t_curve.KeyGet(i).GetTime().GetFrameCount(self.time_mode) for i in range(t_curve.KeyGetCount())]
+            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:
                 if frame not in keys:
@@ -334,6 +494,43 @@ class FBXContainer:
         self._check_actor(actor)
         return self.markers[actor][name]
 
+    def get_parent_node_by_name(self, parent_name: str, ignore_namespace: bool = True) -> Union[fbx.FbxNode, None]:
+        """
+        Utility function to get a parent node reference by name.
+        :param parent_name: `str` name that will be looked for in the node name.
+        :param ignore_namespace: `bool` Whether to ignore namespaces in a node's name.
+        :return:
+        """
+        # Find all parent nodes (/System, /Unlabeled_Markers, /Actor1, etc).
+        parent_nodes = [self.root.GetChild(i) for i in range(self.root.GetChildCount())]
+
+        return next(
+            (
+                parent_node
+                for parent_node in parent_nodes
+                if match_name(parent_node, parent_name, ignore_namespace=ignore_namespace)
+            ),
+            None,
+        )
+
+    def get_node_by_path(self, path: str) -> fbx.FbxNode:
+        """
+        Utility function to retrieve a node reference from a path like /Actor1/Hips/UpperLeg_l.
+        :param path: `str` path with forward slashes to follow.
+        :return: `fbx.FbxNode` reference to that node.
+        """
+        # Split the path into node names.
+        node_names = [x for x in path.split('/') if x]
+        # Start the list of node references with the parent node that has its own function.
+        nodes = [self.get_parent_node_by_name(node_names[0], False)]
+        # Extend the list with each following child node of the previous parent.
+        nodes.extend(
+            get_child_node_by_name(nodes[idx], node_name)
+            for idx, node_name in enumerate(node_names[1:])
+        )
+        # Return the last node in the chain, which will be the node we were looking for.
+        return nodes[-1]
+
     def print_valid_frames_stats_for_actor(self, actor: int = 0):
         """
         Prints: actor name, total amount of frames in the animation, amount of valid frames for the given actor,
@@ -395,18 +592,24 @@ class FBXContainer:
 
         return pd.DataFrame(all_poses, columns=columns)
 
-    def get_transformed_worldspace(self, m: fbx.FbxNode, time: fbx.FbxTime) -> List[float]:
+    def get_worldspace(self, m: fbx.FbxNode, time: fbx.FbxTime, apply_transform: bool = True) -> List[float]:
         """
         Evaluates the world translation of the given marker 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 apply_transform: `bool` Whether to transform the translation or not.
         :return: Vector in the form: [tx, ty, tz].
         """
         t = m.EvaluateGlobalTransform(time).GetT()
-        x = t[0] * self.scale / self.vol_x
-        y = t[1] * self.scale / self.vol_y
-        z = t[2] * self.scale / self.vol_z
+        if not apply_transform:
+            return [t[i] for i in range(3)]
+
+        # First multiply by self.scale, which turns meters to centimeters.
+        # Then divide by volume dimensions, to normalize to the total area of the capture volume.
+        x = np.clip(t[0], -(self.vol_x * 0.5), self.vol_x * 0.5) * self.scale / self.vol_x
+        y = np.clip(t[1], -(self.vol_y * 0.5), self.vol_y * 0.5) * self.scale / self.vol_y
+        z = np.clip(t[2], -(self.vol_z * 0.5), self.vol_z * 0.5) * self.scale / self.vol_z
 
         return [x, y, z]
 
@@ -421,21 +624,32 @@ class FBXContainer:
         curve = marker.LclTranslation.GetCurve(self.anim_layer, 'X')
         return False if curve is None else curve.KeyFind(time) != -1
 
-    def get_sparse_cloud(self, time: fbx.FbxTime) -> np.array:
+    def get_sc(self, frame: int, apply_transform: bool = True) -> np.array:
         """
-        For each actor,
-        :param time:
-        :return:
+        For each actor at the given time, find all markers with keyframes and add their values to a point cloud.
+        :param frame: `fbx.FbxTime` time at which to evaluate the marker.
+        :param apply_transform: `bool` Whether to transform the translation or not.
+        :return: sparse point cloud as `np.array`.
         """
-        cloud = []
+        time = fbx.FbxTime()
+        time.SetFrame(frame)
+        # Start with a cloud of unlabeled markers, which will use actor and marker class 0.
+        # It is important to start with these before the labeled markers,
+        # because by adding the labeled markers after (which use classes 1-74),
+        # we eventually return an array that doesn't need to be sorted anymore.
+        cloud = [
+            [0, 0, *self.get_worldspace(m, time, apply_transform)]
+            for m in self.unlabeled_markers
+            if self.is_kf_present(m, time)
+        ]
+
         # Iterate through all actors to get their markers' world translations and add them to the cloud list.
         for actor_idx in range(self.actor_count):
-
             cloud.extend(
                 # This actor's point cloud is made up of all markers that have a keyframe at the given time.
                 # For each marker, we create this row: [actor class (index+1), marker class (index+1), tx, ty, tz].
                 # We use index+1 because the unlabeled markers will use index 0 for both classes.
-                [actor_idx + 1, marker_class, *self.get_transformed_worldspace(m, time)]
+                [actor_idx + 1, marker_class, *self.get_worldspace(m, time, apply_transform)]
                 for marker_class, (marker_name, m) in enumerate(
                     self.markers[actor_idx].items(), start=1
                 )
@@ -445,56 +659,58 @@ class FBXContainer:
                 if self.is_kf_present(m, time)
             )
 
-        # Unlabeled markers are their own 'actor', so we only need one loop here.
-        for m in self.unlabeled_markers:
-            if self.is_kf_present(m, time):
-                # Unlabeled markers use actor class 0 and marker class 0.
-                cloud.extend([[0, 0, *self.get_transformed_worldspace(m, time)]])
-
         # If the data is extremely noisy, it might have only a few labeled markers and a lot of unlabeled markers.
         # The returned point cloud is not allowed to be bigger than the maximum size (self.pc_size),
         # so return the cloud as a np array that cuts off any excessive markers.
         return np.array(cloud)[:self.pc_size]
 
-    def get_timeline_sparse_cloud(self) -> np.array:
+    def get_tsc(self, apply_transform: bool = True) -> np.array:
         """
         Convenience method that calls self.get_sparse_cloud() for all frames in the frame range
         and returns the combined result.
+        :param apply_transform: `bool` Whether to transform the translation or not.
         :return: `np.array` that contains a sparse cloud for each frame in the frame range.
         """
-        # We need time objects in a list to do list comprehension in the return line.
-        # The SetFrame() method is a void/in-place method, so we can't use list comprehension
-        # to create a list of fbx.FbxTime()s.
-        time = fbx.FbxTime()
-        times = []
-        for f in self.get_frame_range():
-            # Use in-place function to update the time. This returns None, so we can't append this directly.
-            time.SetFrame(f)
-            times.append(time)
+        return np.array([self.get_sc(f, apply_transform) for f in self.get_frame_range()])
 
-        return np.array([self.get_sparse_cloud(t) for t in times])
-
-    def get_timeline_dense_cloud(self, shuffle: bool = False) -> np.array:
+    def get_tdc(self, shuffle: bool = False,
+                r: Union[int, Tuple[int, int]] = None,
+                apply_transform: bool = True) -> 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,
         and adds them together to create a dense cloud whose shape is always (self.pc_size, 5).
         Optionally shuffles this dense cloud before adding it to the final list.
-        :param shuffle: If `True`, shuffles the dense point cloud before appending it to the overall list.
+        :param shuffle: If `True`, shuffles the dense point cloud of each frame.
+        :param r: tuple of `int` that indicates the frame range to get. Default is None,
+        resulting in the animation frame range.
+        :param apply_transform: `bool` Whether to transform the translation or not.
         :return: `np.array` that contains a dense point cloud for each frame,
         with a shape of (self.num_frames, self.pc_size, 5).
         """
-        time = fbx.FbxTime()
+
         clouds = []
-        for frame in self.get_frame_range():
-            time.SetFrame(frame)
-            cloud = self.get_sparse_cloud(time)
+
+        # 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()
+
+        for frame in r:
+            cloud = self.get_sc(frame, apply_transform)
             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:
                 ghost_cloud = make_ghost_markers(missing)
-                cloud = np.vstack([cloud, ghost_cloud])
+                cloud = np.vstack([ghost_cloud, cloud])
 
             # Shuffle the rows if needed. Because each row contains all dependent and independent variables,
             # shuffling won't mess up the labels.
@@ -505,7 +721,7 @@ class FBXContainer:
 
         return np.array(clouds)
 
-    def split_timeline_dense_cloud(self, cloud: np.array = None, shuffle: bool = False) \
+    def split_tdc(self, cloud: np.array = None, shuffle: bool = False, apply_transform: bool = True) \
             -> Tuple[np.array, np.array, np.array]:
         """
         Splits a timeline dense cloud with shape (self.num_frames, self.pc_size, 5) into 3 different
@@ -516,10 +732,11 @@ class FBXContainer:
         :param cloud: `np.array` of shape (self.num_frames, self.pc_size, 5) that contains a dense point cloud
         (self.pc_size, 5) per frame in the frame range.
         :param shuffle: `bool` whether to shuffle the generated cloud if no cloud was given.
+        :param apply_transform: `bool` Whether to transform the translation or not.
         :return: Return tuple of `np.array` as (actor classes, marker classes, translation vectors).
         """
         if cloud is None:
-            cloud = self.get_timeline_dense_cloud(shuffle)
+            cloud = self.get_tdc(shuffle, apply_transform=apply_transform)
 
         assert cloud.shape[1] == 1000, f"Dense cloud doesn't have enough points. {cloud.shape[1]}/1000."
         assert cloud.shape[2] == 5, f"Dense cloud is missing columns: {cloud.shape[2]}/5."
@@ -533,7 +750,7 @@ class FBXContainer:
         :param c: `float` actor class index.
         :return: `str` actor name.
         """
-        return 'UNLABELED' if c == 0. else self.actor_names[int(c) - 1]
+        return 'UNLABELED' if int(c) == 0 else self.actor_names[int(c) - 1]
 
     def convert_class_to_marker(self, c: float = 0):
         """
@@ -541,7 +758,7 @@ class FBXContainer:
         :param c: `float` marker class index.
         :return: `str` marker name.
         """
-        return 'UNLABELED' if c == 0. else self.marker_names[int(c) - 1]
+        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.
@@ -551,7 +768,7 @@ class FBXContainer:
             return df.to_csv(index=False).encode('utf-8')
 
         if output_file is None:
-            output_file = self.fbx_file.with_suffix('.csv')
+            output_file = self.input_fbx.with_suffix('.csv')
 
         if output_file.suffix != '.csv':
             raise ValueError(f'{output_file} needs to be a .csv file.')
@@ -559,6 +776,138 @@ class FBXContainer:
         df.to_csv(output_file, index=False)
         return output_file
 
+    def export_fbx(self, output_file: Path = None) -> bool:
+        """
+        Exports the entire scene to the output file. If output_file is None, it uses the automatic output_fbx property.
+        :param output_file: `Path` path to an FBX file to export the contents to.
+        :return: `True` if successful, otherwise `False`.
+        """
+        if output_file is None:
+            output_file = self.output_fbx
+
+        # Create an exporter using the manager
+        exporter = fbx.FbxExporter.Create(self.manager, "")
+
+        # 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}'.")
+            return False
+
+        # Export the scene
+        result = exporter.Export(self.scene)
+        if not result:
+            print(f"Failed to export the scene to file '{output_file}'.")
+            return False
+
+        # Clean up the manager and exporter
+        exporter.Destroy()
+
+        return True
+
+    def remove_node(self, node: fbx.FbxNode, recursive: bool = False) -> bool:
+        """
+        Removes a node by reference from the scene.
+        :param node: `fbx.FbxNode` to remove.
+        :param recursive: `bool` Apply deletion recursively.
+        :return: True if success.
+        """
+        if recursive:
+            children = [node.GetChild(c) for c in range(node.GetChildCount())]
+            for child in children:
+                self.remove_node(child, True)
+        # Disconnect the marker node from its parent
+        node.GetParent().RemoveChild(node)
+
+        # Remove the marker node from the scene
+        self.scene.RemoveNode(node)
+
+        return True
+
+    def remove_unlabeled_markers(self) -> None:
+        """
+        Uses self.remove_node() to delete all unlabeled markers from the scene.
+        """
+        for m in self.unlabeled_markers:
+            self.remove_node(m)
+
+        self.remove_node(self.unlabeled_markers_parent)
+
+    def remove_system(self) -> None:
+        system_node = self.get_parent_node_by_name('System')
+        self.remove_node(system_node, recursive=True)
+
+    def cleanup(self) -> None:
+        self.remove_unlabeled_markers()
+        self.remove_system()
+
+    def replace_keyframes_per_marker(self, marker: fbx.FbxNode, marker_keys: dict) -> None:
+        parent = marker.GetParent()
+        for axis in ['X', 'Y', 'Z']:
+            curve = marker.LclTranslation.GetCurve(self.anim_layer, axis)
+            curve.KeyModifyBegin()
+            if curve is not None:
+                curve.KeyClear()
+
+                for frame, world_translation in marker_keys.items():
+                    local_translation = world_to_local_translation(world_translation, parent)
+
+                    for axis_idx in range(3):
+                        create_keyframe(curve, frame, local_translation[axis_idx])
+            curve.KeyModifyEnd()
+
+    def replace_keyframes_per_actor(self, actor: int, actor_keys: dict) -> None:
+        for marker_class, (marker_name, marker) in enumerate(self.markers[actor].items(), start=1):
+            self.replace_keyframes_per_marker(marker, actor_keys[marker_class])
+
+    def replace_keyframes_for_all_actors(self, key_dict: dict) -> None:
+        for actor_idx in range(self.actor_count):
+            self.replace_keyframes_per_actor(actor_idx, key_dict[actor_idx + 1])
+
 
 # d = FBXContainer(Path('G:/Firestorm/mocap-ai/data/fbx/dowg/TAKE_01+1_ALL_001.fbx'))
-# TODO: Make functions to write new class predictions to the fbx file.
+# t = fbx.FbxTime()
+# t.SetFrame(0)
+#
+# one = list(d.unlabeled_markers[0].EvaluateGlobalTransform(t).GetT())
+# # one = np.array([
+# #         [one.Get(i, j) for j in range(4)] for i in range(4)
+# #     ])
+# two = world_to_local_translation(d.unlabeled_markers[0], d.unlabeled_markers_parent, 0)[3, :3]
+# print(one)
+# print(two)
+# for e in zip(one, two):
+#     print(e)
+# f = d.get_sc(0, False)
+# f = f[f[:, 1] == 1.]
+# print(f)
+# print(d.convert_class_to_actor(f[0]))
+# print(d.convert_class_to_marker(f[1]))
+# print(f[2:5])
+# lt = world_to_local_translation(f[2:5], d.unlabeled_markers_parent)
+# print(lt)
+
+# for u in sc:
+#     print(d.convert_class_to_actor(u[0]), d.convert_class_to_marker(u[1]))
+# train_cloud = d.get_tdc()
+# train_actors, train_markers, train_X = d.split_tdc(train_cloud)
+# test_cloud = d.get_tdc(r=1, apply_transform=False)
+# for row in test_cloud:
+#     for m in row[:5]:
+#         t = m[2:5]
+#         print(t)
+#         lt = world_to_local_translation()
+#     print()
+#     print(row[:5, 2:5])
+# print(test_cloud.shape)
+# print(test_cloud[0, :, :2])
+# print(d.convert_class_to_actor(test_cloud[0, 0, 0]))
+# print(d.convert_class_to_marker(test_cloud[0, 0, 1]))
+# test_actors, test_markers, test_X = d.split_tdc(test_cloud, apply_transform=False)
+# Predict...
+# merged_preds = merge_tdc(train_actors, train_markers, test_X, ordered=False)
+# di = timeline_cloud_to_dict(test_cloud)
+
+# d.replace_keyframes_for_all_actors(di)
+# d.cleanup()
+# d.export_fbx(Path('G:/Firestorm/mocap-ai/data/fbx/export/TAKE_01+1_ALL_001.fbx'))