Removed old data extraction methods and documented all the new functions.

fbx_handler.py

@@ -33,37 +33,22 @@ def center_axis(a: Union[List[float], np.array]) -> np.array:
     return a
 
 
-def make_ghost_markers(missing: int) -> np.array:
+def append_zero(arr: np.ndarray) -> np.ndarray:
     """
-    Creates a np array containing enough rows to fill a point cloud up to 1000 points.
-    Ghost markers are always unlabeled markers, therefore their actor and marker class is 0.
-    :param missing: `int` amount of missing rows in the cloud that need to be filled with this function.
-    :return: multidimensional `np.array` with the shape: (missing, 5).
+    Appends a zero array to the end of an array.
+    :param arr: `np.array` to fill.
+    :return: `np.array` with a zero array appended to the end.
     """
-    return np.column_stack([
-        np.zeros((missing, 1), dtype=int),  # 0
-        np.zeros((missing, 1), dtype=int),  # 0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1),  # 0.0-1.0
-        np.random.rand(missing, 1)  # 0.0-1.0
-    ])
-
-
-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)
 
 
 def append_one(arr: np.ndarray) -> np.ndarray:
+    """
+    Appends a one array to the end of an array.
+    :param arr: `np.array` to fill.
+    :return: `np.array` with a one array appended to the end.
+    """
     ones = np.ones((arr.shape[0], arr.shape[1], 1), dtype=float)
     return np.concatenate((arr, ones), axis=-1)
 
@@ -74,10 +59,19 @@ 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 a new dimension at the end.
-    # Return the concatenated array of shape (x, 1000, 14), which matches the original timeline dense cloud before
+    """
+    Merges the given arrays into a single array of shape (n_frames, pc_size, 14).
+    :param actor_classes: `np.array` of actor class labels.
+    :param marker_classes: `np.array` of marker class labels.
+    :param translation_vectors: `np.array` of translation vectors (3 values each).
+    :param rotation_vectors: `np.array` of Euler rotation angles (3 values each).
+    :param scale_vectors: `np.array` of scale factors (3 values each).
+    :param ordered: Whether to sort the cloud by actor and marker classes.
+    :return: Merged `np.array` of shape (n_frames, pc_size, 14).
+    """
+    # Actor and marker classes enter as shape (x, 1024), so use np.expand_dims to create a new dimension at the end.
+    # Return the concatenated array of shape (x, 1024, 14), 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),
                           append_zero(translation_vectors),
@@ -111,7 +105,7 @@ 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.
+    :return: Sorted `np.array` point cloud.
     """
     # Extract the first two elements of the third dimension
     actor_classes = cloud[:, :, 0]
@@ -131,7 +125,14 @@ def sort_cloud(cloud: np.array) -> np.array:
     return sorted_tdc
 
 
-def create_keyframe(anim_curve: fbx.FbxAnimCurve, frame: int, value: float):
+def create_keyframe(anim_curve: fbx.FbxAnimCurve, frame: int, value: float) -> None:
+    """
+    Creates a keyframe at the given frame number on the given animation curve.
+    :param anim_curve: `fbx.FbxAnimCurve` node to add the keyframe to.
+    :param frame: `int` frame number at which to add the keyframe.
+    :param value: `float` value that the keyframe will have.
+    :return: True
+    """
     # Create an FbxTime object with the given frame number
     t = fbx.FbxTime()
     t.SetFrame(frame)
@@ -139,22 +140,43 @@ def create_keyframe(anim_curve: fbx.FbxAnimCurve, frame: int, value: float):
     # Create a new keyframe with the specified value
     key_index = anim_curve.KeyAdd(t)[0]
     anim_curve.KeySetValue(key_index, value)
-    return True
+    return
 
 
 def get_child_node_by_name(parent_node: fbx.FbxNode, name: str, ignore_namespace: bool = False) \
         -> Union[fbx.FbxNode, None]:
+    """
+    Gets the child node with the given name.
+    :param parent_node: `fbx.FbxNode` to get the child node from.
+    :param name: `str` name of the child node to get.
+    :param ignore_namespace: `bool` whether to ignore the namespace in the node name.
+    :return: `fbx.FbxNode` child node with the given name, if it exists, else None.
+    """
+    # Loop through all child nodes of the parent node.
     for c in range(parent_node.GetChildCount()):
+        # Get the child node.
         child = parent_node.GetChild(c)
+        # Check if the name of the child node matches the given name.
         if match_name(child, name, ignore_namespace):
+            # If it matches, return the child node.
             return child
+    # If no child node matches the given name, return None.
     return None
 
 
 def match_name(node: fbx.FbxNode, name: str, ignore_namespace: bool = True) -> bool:
+    """
+    Checks if the given node's name matches the given name.
+    :param node: `fbx.FbxNode` to check.
+    :param name: `str` name to match.
+    :param ignore_namespace: `bool` whether to ignore the namespace in the node name.
+    """
+    # get the name of the node
     node_name = node.GetName()
+    # if ignore_namespace is True, remove the namespace from the node name
     if ignore_namespace:
         node_name = node_name.split(':')[-1]
+    # return True if the node name matches the provided name, False otherwise
     return node_name == name
 
 
@@ -233,36 +255,6 @@ 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, 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 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(t)
-
-    # If axes is only the translation, we return a vector of (tx, ty, tz) only (useful for the training).
-    if axes == 't':
-        return np.array([matrix[i] for i in range(3)])
-
-    # Otherwise, we assemble the entire row depending on the axes.
-    world = []
-    if 't' in axes:
-        world += list(matrix.GetT())
-        world[3] = 0.0
-    if 'r' in axes:
-        world += list(matrix.GetR())
-        world[7] = 0.0
-    if 's' in axes:
-        world += list(matrix.GetS())
-        world[11] = 1.0
-
-    return np.array(world)
-
-
 def isolate_actor_from_tdc(tdc: np.array, actor: int) -> np.array:
     """
     Returns all markers of the given actor in the timeline dense cloud.
@@ -287,23 +279,56 @@ def split_tdc_into_actors(tdc: np.array) -> List[np.array]:
 
 
 def get_keyed_frames_from_curve(curve: fbx.FbxAnimCurve, length: int = -1) -> List[fbx.FbxAnimCurveKey]:
+    """
+    Returns a list of all the frames on the given curve.
+    :param curve: `fbx.FbxAnimCurve` to get frames from.
+    :param length: Desired amount of frame numbers to return. If this is more than there are keyframes on the curve,
+    it pads 0s to the end. Default -1, which will not add any 0s.
+    :return: List of all the frames on the given curve.
+    """
+    # Get all the frames on the curve.
     frames = [curve.KeyGet(i).GetTime().GetFrameCount() for i in range(curve.KeyGetCount())]
+    # Calculate the difference between desired length and actual length of frames.
     dif = length - len(frames)
+    # If desired length is greater than actual length and length is not -1, add 0s to the end.
     if dif > 0 and length != -1:
         frames += [0.] * dif
+    # Return the list of all frames on the curve.
     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]]:
+def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
+                         r: List[int], c: fbx.FbxAnimCurve, incl_keyed: int = 1) -> List[List[float]]:
+    """
+    For the given marker node, gets the world transform for each frame in r, and stores the translation, rotation
+    and scaling values as a list of lists. Stores the actor and marker classes at the start of this list of lists.
+    Optionally, if incl_keyed is 1, also stores the keyed frames as the last list.
+    Note: This function has to be passed the animation curve, because we need the animation layer to get the
+    animation curve. The animation layer is stored inside the FBXContainer class, to which we don't have access to here.
+    :param actor_idx: `int` actor class.
+    :param marker_idx: `int` marker class.
+    :param m: `fbx.FbxNode` to evaluate the world transform of at each frame.
+    :param r: `List[int]` list of frame numbers to evaluate the world transform at.
+    :param c: `fbx.FbxAnimCurve` node to read the keyframes from.
+    :param incl_keyed: `bool` whether to include if there was a key on a given frame or not. 0 if not.
+    :return:
+    """
+    # Create a list of zeros with the same length as r.
     zeros = [0.0 for _ in range(len(r))]
+    # Create a list of ones with the same length as r.
     ones = [1.0 for _ in range(len(r))]
 
+    # Create empty lists for each transformation parameter.
     tx, ty, tz, rx, ry, rz, sx, sy, sz = [], [], [], [], [], [], [], [], []
+    # Create a list of actor classes with the same length as r.
     actors = [actor_idx for _ in range(len(r))]
+    # Create a list of marker classes with the same length as r.
     markers = [marker_idx for _ in range(len(r))]
+    # Create a new FbxTime object without a frame set yet.
     t = fbx.FbxTime()
 
+    # For each frame in the given frame range (which does not need to start at 0),
+    # evaluate the world transform at each frame and store the relevant items into their respective lists.
     for f in r:
         t.SetFrame(f)
         wt = m.EvaluateGlobalTransform(t)
@@ -318,6 +343,7 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode, r: Lis
         sy.append(wts[1])
         sz.append(wts[2])
 
+    # If we don't need to include keyed frames, return the list of lists as is.
     if not incl_keyed:
         return [
             actors,
@@ -327,9 +353,14 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode, r: Lis
             sx, sy, sz, ones
         ]
 
+    # However, if we do need those keys, we first retrieve all the keyframed frame numbers from the curve.
+    # Note: We do this after returning the previous results, because the following lines are very slow
+    # and unnecessary for inference.
     keyed_frames = get_keyed_frames_from_curve(c)
+    # Then we check if any of the frame numbers are in the keyed frames, which means it had a keyframe and should be 1.
     keyed_bools = [1 if f in keyed_frames else 0 for f in r]
 
+    # Finally, return the complete lists of lists.
     return [
         actors,
         markers,
@@ -347,9 +378,7 @@ class FBXContainer:
                  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'):
+                 save_init: bool = True):
         """
         Class that stores references to important nodes in an FBX file.
         Offers utility functions to quickly load animation data.
@@ -396,14 +425,13 @@ class FBXContainer:
         self.pc_size = pc_size
 
         self.input_fbx = fbx_file
-        # self.output_fbx = append_suffix_to_fbx(fbx_file, '_INF')
         self.output_fbx = utils.append_suffix_to_file(fbx_file, '_INF')
         self.valid_frames = []
 
         # 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)
+            self.init()
 
     def __init_scene(self) -> None:
         """
@@ -498,7 +526,7 @@ class FBXContainer:
 
             self.markers.append(actor_markers)
 
-    def __init_unlabeled_markers(self, ignore_missing: bool = False) -> None:
+    def __init_unlabeled_markers(self, ignore_missing: bool = True) -> None:
         """
         Looks for the Unlabeled_Markers parent node under the root and stores references to all unlabeled marker nodes.
         """
@@ -514,47 +542,86 @@ class FBXContainer:
             raise ValueError('No unlabeled markers found.')
 
     def init_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> None:
-
+        """
+        Calls the init functions for the labeled and unlabeled world transforms.
+        :param r: Custom frame range to extract.
+        """
         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):
-
+                                      incl_keyed: int = 1) -> np.array:
+        """
+        For each actor, for each marker, stores a list for each element in the world transform for each frame
+        in r. This can later be used to recreate the world transform matrix.
+        :param r: Custom frame range to use.
+        :param incl_keyed: `bool` whether to check if the marker was keyed at the frame.
+        :return: `np.array` of shape (n_frames, n_markers, 14).
+        """
         r = self.convert_r(r)
         labeled_data = []
 
+        # Iterate through all actors.
         for actor_idx in range(self.actor_count):
+            # Initialize an empty list to store the results for this actor in.
             actor_data = []
+            # Iterate through all markers for this actor.
             for marker_idx, (n, m) in enumerate(self.markers[actor_idx].items()):
+                # Get this marker's local translation animation curve.
+                # This requires the animation layer, so we can't do it within the function itself.
                 curve = m.LclTranslation.GetCurve(self.anim_layer, 'X', True)
+                # Get a list of each world transform element for all frames.
                 marker_data = get_world_transforms(actor_idx + 1, marker_idx + 1, m, r, curve, incl_keyed)
+                # Add the result to actor_data.
                 actor_data.append(marker_data)
                 self._print(f'Actor {actor_idx} marker {marker_idx} done', 1)
+            # Add all this actor_data to the global labeled_data.
             labeled_data.append(actor_data)
 
+        # Convert the list to a np array. This will have all frames at the last dimension because of this order:
+        # Shape (n_actors, n_markers, 14/15, n_frames).
         wide_layout = np.array(labeled_data)
+        # Transpose the array so that the first dimension is the frames.
+        # Shape (n_frames, n_actors, n_markers, 14).
         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:
+        """
+        For all unlabeled markers, stores a list for each element in the world transform for each frame
+        in r. This can later be used to recreate the world transform matrix.
+        :param r: Custom frame range to use.
+        :return: `np.array` of shape (n_frames, n_unlabeled_markers, 14).
+        """
         r = self.convert_r(r)
-
         unlabeled_data = []
 
+        # Iterate through all unlabeled markers.
         for ulm in self.unlabeled_markers:
+            # Get this marker's local translation animation curve.
+            # This requires the animation layer, so we can't do it within the function itself.
             curve = ulm.LclTranslation.GetCurve(self.anim_layer, 'X', True)
+            # Get a list of each world transform element for all frames.
             marker_data = get_world_transforms(0, 0, ulm, r, curve, incl_keyed=0)
+            # Add the result to marker_data.
             unlabeled_data.append(marker_data)
             self._print(f'Unlabeled marker {ulm.GetName()} done', 1)
 
+        # Convert the list to a np array. This will have all frames at the last dimension because of this order:
+        # Shape (n_unlabeled_markers, 14/15, n_frames).
         wide_layout = np.array(unlabeled_data)
+        # Transpose the array so that the first dimension is the frames.
+        # Shape (n_frames, n_unlabeled_markers, 14).
         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:
+    def init(self, ignore_missing_labeled: bool = False, ignore_missing_unlabeled: bool = False) -> None:
+        """
+        Initializes the scene.
+        :param ignore_missing_labeled: `bool` whether to ignore errors for missing labeled markers.
+        :param ignore_missing_unlabeled: `bool` whether to ignore errors for missing unlabeled markers.
+        """
         self.__init_scene()
         self.__init_anim()
         self.__init_actors(ignore_missing=ignore_missing_labeled)
@@ -574,120 +641,6 @@ class FBXContainer:
         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):
-        """
-        Checks for each frame in the frame range, and for each marker, if there is a keyframe present
-        at that frame on LocalTranslation X.
-        If the keyframe is missing, removes that frame from the list of valid frames for that actor.
-        This eventually leaves a list of frames where each number is guaranteed to have a keyframe on all markers.
-        The list is appended to valid_frames, which can be indexed per actor.
-        Finally, stores a list of frames that is valid for all actors in common_frames.
-        :param actor: `int` index of the actor to find keyframes for.
-        """
-        # Make sure the actor index is in range.
-        self._check_actor(actor)
-
-        frames = self.get_frame_range()
-        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:
-                if frame not in keys:
-                    # If the frame is not present, that means there is no keyframe with that frame number,
-                    # so we don't want to use that frame because it is invalid, so we remove it from the list.
-                    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.
-        other_lists = [r for r in self.valid_frames if r]
-        # Make one list that contains all shared frame numbers.
-        self.common_frames = [num for num in self.get_frame_range()
-                              if all(num in other_list for other_list in other_lists)]
-
-    def _check_valid_frames(self, actor: int = 0):
-        """
-        Safety check to see if the given actor has any valid frames stored.
-        If not, calls _set_valid_frames_for_actor() for that actor.
-        :param actor: `int` actor index.
-        """
-        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_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.
-        :param actor: `int` actor index.
-        :param frame: `int` frame to evaluate the markers at.
-        :return: 1D list of `float` that contains the tx, ty and tz for each marker, in that order.
-        """
-        # Set new frame to evaluate at.
-        time = fbx.FbxTime()
-        time.SetFrame(frame)
-        # Prepare arrays for each axis.
-        x, y, z = [], [], []
-
-        # For each marker, store the x, y and z global position.
-        for n, m in self.markers[actor].items():
-            t = m.EvaluateGlobalTransform(time).GetT()
-            x += [t[0] * self.scale]
-            y += [t[1] * self.scale]
-            z += [t[2] * self.scale]
-
-        # Move the point cloud to the center of the x and y axes. This will put the actor in the middle.
-        x = center_axis(x)
-        z = center_axis(z)
-
-        # Move the actor to the middle of the volume floor by adding volume_dim/2 to x and z.
-        x += self.vol_x / 2.
-        z += self.vol_z / 2.
-
-        # Squeeze the actor into the 1x1 plane for the neural network by dividing the axes.
-        x /= self.vol_x
-        z /= self.vol_z
-        y = np.array(y) / self.vol_y
-
-        # 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.
-        pose = []
-        for i in range(len(x)):
-            pose += [x[i]]
-            pose += [y[i]]
-            pose += [z[i]]
-        return pose
-
     def get_frame_range(self) -> List[int]:
         """
         Replacement and improvement for:
@@ -697,7 +650,12 @@ 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):
+    def convert_r(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> List[int]:
+        """
+        Converts the value of r to a list of frame numbers, depending on what r is.
+        :param r: Custom frame range to use.
+        :return: List of `int` frame numbers (doesn't have to start at 0).
+        """
         # 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))
@@ -782,82 +740,12 @@ class FBXContainer:
         # 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,
-        number of missing frames, and the ratio of valid/total frames.
-        :param actor: `int` actor index.
-        :return: Tuple of `str` actor name, `int` total frames, `int` amount of valid frames, `float` valid frame ratio.
-        """
-        self._check_actor(actor)
-        self._check_valid_frames(actor)
-
-        len_valid = len(self.valid_frames[actor])
-        ratio = (len_valid / self.num_frames) * 100
-        print(f'Actor {self.actor_names[actor]}: Total: {self.num_frames}, valid: {len_valid}, missing: '
-              f'{self.num_frames - len_valid}, ratio: {ratio:.2f}% valid.')
-
-        return self.actor_names[actor], self.num_frames, len_valid, ratio
-
-    def get_valid_frames_for_actor(self, actor: int = 0) -> List[int]:
-        """
-        Collects the valid frames for the given actor.
-        :param actor: `int` actor index.
-        :return: List of `int` frame numbers that have a keyframe on tx for all markers.
+    def remove_clipping_poses(self, arr: np.array) -> np.array:
         """
-        self._check_valid_frames(actor)
-        return self.valid_frames[actor]
-
-    def extract_valid_translations_per_actor(self, actor: int = 0, t: str = 'np'):
+        Checks for each axis if it does not cross the volume limits. Returns an array without clipping poses.
+        :param arr: `np.array` to filter.
+        :return: Filtered `np.array` that only has non-clipping poses.
         """
-        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)
-
-        # Returns shape (n_valid_frames, 3, 73).
-        return np.vstack([self.get_transformed_pc(actor, frame) for frame in self.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]:
-        #     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, 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.
-        :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 remove_clipping_poses(self, arr: np.array) -> np.array:
-
         mask_x1 = (arr[:, :, 2] < self.hvol_x / self.scale).all(axis=1)
         mask_x2 = (arr[:, :, 2] > -self.hvol_x / self.scale).all(axis=1)
         mask_z1 = (arr[:, :, 4] < self.hvol_z / self.scale).all(axis=1)
@@ -867,6 +755,13 @@ class FBXContainer:
         return arr[mask]
 
     def extract_training_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
+        """
+        Manipulates the existing labeled world transform array into one that is suitable for training.
+        It does this through flattening the array to shape (n_frames, n_actors * 73, 15), then removing
+        all clipping frames and finally transforms the frames to the right location and scale.
+        :param r: Custom frame range to use if the labeled transforms are not stored yet.
+        :return: Transformed labeled world transforms.
+        """
         if self.labeled_world_transforms is None:
             self.init_labeled_world_transforms(r=r, incl_keyed=1)
 
@@ -896,12 +791,25 @@ class FBXContainer:
 
     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]]:
+        """
+        Manipulates the existing (un)labeled world transform arrays into arrays that are suitable for inference.
+        It does this through flattening the labeled world transforms to shape (n_frames, n_actors * 73, 14).
+        If merged is True, merges the unlabeled data with the labeled data.
+        :param r: Custom frame range to use if the transforms were not extracted yet.
+        :param merged: `bool` whether to merge both arrays into one or return separate arrays.
+        :return: If merged, returns one `np.array`, else flattened labeled `np.array` and unlabeled `np.array`.
+        """
+        # If either of the arrays is None, we can initialize them with r.
         if self.labeled_world_transforms is None:
+            # For inference, we don't need keyed frames, so incl_keyed is False.
             self.init_labeled_world_transforms(r=r, incl_keyed=0)
         if self.unlabeled_world_transforms is None:
+            # Note: Unlabeled data is already flattened.
             self.init_unlabeled_world_transforms(r=r)
 
+        # Returns (n_frames, n_actors, 73, 14).
         ls = self.labeled_world_transforms.shape
+        # Flatten the array, so we get a list of frames.
         # Returns shape (n_frames, 73 * n_actors, 14).
         flat_labeled = self.labeled_world_transforms.reshape(ls[0], -1, ls[-1])[..., :14]
 
@@ -918,7 +826,6 @@ class FBXContainer:
         :param w: `np.array` that can either be a timeline dense cloud or translation vectors.
         :return: Modified `np.array`.
         """
-
         # 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, tw, rx, ry, rz, tw, etc).
         start = 0 if w.shape[-1] == 3 else 2
@@ -931,144 +838,15 @@ class FBXContainer:
 
         # Then move the x and z to the center of the volume. Y doesn't need to be done because pose needs to stand
         # on the floor.
-        # Finally, add 0.5 to the x and z to make the pose stand in the center of the normalized volume.
-        w[..., start + 0] = np.clip(w[..., start + 0], -0.5, 0.5) + 0.5
+        # We do not add 0.5 here to move the pose to the middle of the capture space,
+        # because in the Dataset we still need to randomly rotate it in world space.
+        # So we keep it centered here.
+        w[..., start + 0] = np.clip(w[..., start + 0], -0.5, 0.5)
         w[..., start + 1] = np.clip(w[..., start + 1], -0.5, 0.5)
-        w[..., start + 2] = np.clip(w[..., start + 2], -0.5, 0.5) + 0.5
+        w[..., start + 2] = np.clip(w[..., start + 2], -0.5, 0.5)
 
         return w
 
-    def is_kf_present(self, marker: fbx.FbxNode, time: fbx.FbxTime) -> bool:
-        """
-        Returns True if a keyframe is found on the given node's local translation x animation curve.
-        Else returns False.
-        :param marker: `fbx.FbxNode` marker node to evaluate.
-        :param time: `fbx.FbxTime` time to evaluate at.
-        :return: True if a keyframe was found, False otherwise.
-        """
-        curve = marker.LclTranslation.GetCurve(self.anim_layer, 'X')
-        return False if curve is None else curve.KeyFind(time) != -1
-
-    def get_sc(self, frame: int) -> np.array:
-        """
-        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.
-        :return: sparse point cloud as `np.array`.
-        """
-        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, *get_world_transform(m, time)]
-            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, *get_world_transform(m, time)]
-                for marker_class, (marker_name, m) in enumerate(
-                    self.markers[actor_idx].items(), start=1
-                )
-                # Only add the marker if it has a keyframe. Missing keyframes on these markers are potentially
-                # among the keyframes on the unlabeled markers. The job of the labeler AI is to predict which
-                # point (unlabeled or labeled) is which marker.
-                if self.is_kf_present(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_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
-        and returns the combined result.
-        :return: `np.array` that contains a sparse cloud for each frame in the frame range.
-        """
-        return np.array([self.get_sc(f) for f in self.get_frame_range()])
-
-    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,
-        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 r: tuple of `int` that indicates the frame range to get. Default is None,
-        resulting in the animation frame range.
-        :return: `np.array` that contains a dense point cloud for each frame,
-        with a shape of (self.num_frames, self.pc_size, 5).
-        """
-
-        r = self.convert_r(r)
-
-        # results = utils.parallel_process(r, self.get_dc)
-
-        return np.array([self.get_dc(f) for f in r])
-
-    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])
-
-        # 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.
-
-        # 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]:
-        """
-        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).
-        :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.
-        :return: Return tuple of `np.array` as (actor classes, marker classes, translation vectors).
-        """
-        if cloud is None:
-            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]}.")
-
-        # 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]:
         """
@@ -1102,17 +880,13 @@ class FBXContainer:
 
     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')
-
-        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':
+        """
+        Exports train data to an HDF5 file.
+        :param output_file: `Path` to the file.
+        :param r: Custom frame range to use.
+        :return: `np.array` of shape (n_poses, 73, 14) of train data.
+        """
+        if 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)
@@ -1120,10 +894,17 @@ class FBXContainer:
             return array_4d
 
         else:
-            raise ValueError('Invalid file extension. Must be .csv or .npy')
+            raise ValueError('Invalid file extension. Must be .h5')
 
     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]]:
+        """
+        Exports test data to an HDF5 file.
+        :param output_file: `Path` to the file.
+        :param r: Custom frame range to use.
+        :param merged: `bool` whether to merge the test data or output an unlabeled dataset and labeled dataset.
+        :return: `np.array` of the test data.
+        """
         # 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),
@@ -1281,6 +1062,12 @@ class FBXContainer:
         self.set_default_lcl_scaling(marker, lcl_s)
 
     def replace_keyframes_per_marker(self, marker: fbx.FbxNode, marker_keys: dict) -> None:
+        """
+        For the given marker, creates new keyframes on its local translation animation curves,
+        and sets the default local rotation values.
+        :param marker: `fbx.FbxNode` to set the default values on.
+        :param marker_keys: `dict` of keys that contain all info needed to create world transform matrices.
+        """
         # Initialize empty variables for the local rotation and scaling.
         # These will be filled at the first keyframe
         self.set_default_lcl_transforms(marker, marker_keys)