Moved the core functionality of FBXContainer to a new class called FBXContainerBase.

extract_inf_translations() will now replace the transforms of nodes without a keyframe with np.inf upon extraction.
This is so when we load the test data, we can easily replace this data with random data to serve as unlabeled markers.

fbx_handler.py

@@ -5,7 +5,6 @@ from typing import List, Union, Tuple
 import h5py
 
 # Import util libs.
-import contextlib
 import fbx
 import itertools
 
@@ -371,69 +370,32 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode,
     ]
 
 
-class FBXContainer:
-    def __init__(self, fbx_file: Path,
-                 volume_dims: Tuple[float] = (10., 10., 10.),
-                 max_actors: int = 8,
-                 pc_size: int = 1024,
-                 scale: float = 0.01,
-                 debug: int = -1,
-                 save_init: bool = True):
+def get_children_of_parent(parent: fbx.FbxNode) -> List[fbx.FbxNode]:
+    """
+    Returns a list of all the children of the given parent.
+    :param parent: `fbx.FbxNode` to get children from.
+    :return: List of `fbx.FbxNode` children.
+    """
+    return [parent.GetChild(i) for i in range(parent.GetChildCount())]
+
+
+class FBXContainerBase:
+    def __init__(self, fbx_file: Path, debug: int = -1) -> None:
         """
         Class that stores references to important nodes in an FBX file.
         Offers utility functions to quickly load animation data.
         :param fbx_file: `Path` to the file to load.
-        :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.input_fbx = fbx_file
         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.
-        self.marker_names = globals.get_marker_names()
 
         # Initiate empty lists to store references to nodes.
-        self.markers = []
-        self.actors = []
+        self.parents = []
+        self.children = []
         # 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]
-        self.vol_z = volume_dims[2]
-        self.hvol_x = volume_dims[0] / 2
-        self.hvol_y = volume_dims[1] / 2
-        self.hvol_z = volume_dims[2] / 2
-
-        self.scale = scale
-
-        self.max_actors = max_actors
-        # Maximum point cloud size = 73 * max_actors + unlabeled markers.
-        self.pc_size = pc_size
-
-        self.input_fbx = fbx_file
-        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()
+        self.parent_names = []
 
-    def __init_scene(self) -> None:
+    def _init_scene(self) -> None:
         """
         Stores scene, root, and time_mode properties.
         Destroys the importer to remove the reference to the loaded file.
@@ -456,7 +418,7 @@ class FBXContainer:
         # This will allow us to delete the uploaded file.
         importer.Destroy()
 
-    def __init_anim(self) -> None:
+    def _init_anim(self) -> None:
         """
         Stores the anim_stack, num_frames, start_frame, end_frame properties.
         """
@@ -474,12 +436,11 @@ class FBXContainer:
         self.start_frame = local_time_span.GetStart().GetFrameCount()
         self.end_frame = local_time_span.GetStop().GetFrameCount()
 
-    def __init_actors(self, ignore_missing: bool = False) -> None:
+    def _init_parents(self, names_to_look_for: List[str]) -> 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)
@@ -487,46 +448,234 @@ class FBXContainer:
         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:
             gen2_nodes = [gen1_node.GetChild(i) for i in
                           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 names_to_look_for):
-                self.actor_names.append(gen1_node.GetName())
-                self.actors.append(gen1_node)
+            # If the list of names_to_look_for are children of this parent, it must be a parent.
+            gen2_names = [node.GetName().split(':')[-1] for node in gen2_nodes]
+            if all(name in gen2_names for name in names_to_look_for):
+                self.parent_names.append(gen1_node.GetName())
+                self.parents.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: ' +
+        if len(self.parents) == 0:
+            raise ValueError('No parents found. A node is considered a parent ' +
+                             'if it has the following child nodes: ' +
                              ', '.join(names_to_look_for) + '.')
 
-        self.actor_count = len(self.actors)
-        self.valid_frames = [[] for _ in range(self.actor_count)]
+        self.parent_count = len(self.parents)
 
-    def __init_markers(self, ignore_missing: bool = False) -> None:
+    def _init_children(self, child_names: List[str]) -> None:
         """
-        Goes through all actor nodes and stores references to its marker nodes.
+        Goes through all parent nodes and stores references to its child nodes.
         """
-        for actor_node in self.actors:
-            actor_markers = {}
-            for marker_name in self.marker_names:
-                for actor_idx in range(actor_node.GetChildCount()):
-                    child = actor_node.GetChild(actor_idx)
+        for parent_node in self.parents:
+            family = {}
+            for child_name in child_names:
+                for parent_idx in range(parent_node.GetChildCount()):
+                    child = parent_node.GetChild(parent_idx)
                     # Child name might have namespaces in it like this: Vera:ARIEL
                     # 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
+                    if match_name(child, child_name, ignore_namespace=True):
+                        family[child_name] = child
+
+            if len(family) != len(child_names):
+                raise ValueError(f'{parent_node.GetName()} does not have all children.')
 
-            if len(actor_markers) != len(self.marker_names) and not ignore_missing:
-                raise ValueError(f'{actor_node.GetName()} does not have all markers.')
+            self.children.append(family)
 
-            self.markers.append(actor_markers)
+    def _print(self, txt: str, lvl: int = 0) -> None:
+        if lvl <= self.debug:
+            print(txt)
 
-    def __init_unlabeled_markers(self, ignore_missing: bool = True) -> None:
+    def get_frame_range(self) -> List[int]:
+        """
+        Replacement and improvement for:
+        `list(range(self.num_frames))`
+        If the animation does not start at frame 0, this will return a list that has the correct frames.
+        :return: List of `int` frame numbers that are between the start and end frame of the animation.
+        """
+        return list(range(self.start_frame, self.end_frame))
+
+    def _check_parent(self, parent: int = 0):
+        """
+        Safety check to see if the actor `int` is a valid number (to avoid out of range errors).
+        :param parent: `int` actor index, which should be between 0-max_actors.
+        """
+        if not 0 <= parent <= self.parent_count:
+            raise ValueError(f'Actor index must be between 0 and {self.parent_count - 1} ({parent}).')
+
+    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))
+
+        # 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 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_child_by_name(self, parent: int, name: str):
+        """
+        Returns the reference to the parent's direct child.
+        :param parent: `int` parent index.
+        :param name: `str` child name.
+        :return: `fbx.FbxNode` reference.
+        """
+        self._check_parent(parent)
+        return self.children[parent][name]
+
+    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 get_hierarchy(self, start: fbx.FbxNode, hierarchy: List = None) -> List[fbx.FbxNode]:
+        """
+        Returns the hierarchy under the start node.
+        :param hierarchy: Hierarchical `List` of `fbx.FbxNode` references.
+        :param start: `fbx.FbxNode` that will be used as starting point for finding the leaf nodes.
+        :return:
+        """
+        if hierarchy is None:
+            hierarchy = []
+        hierarchy.append(start)
+        children = get_children_of_parent(start)
+        for child in children:
+            self.get_hierarchy(child, hierarchy)
+
+        return hierarchy
+
+    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
+
+
+class FBXContainer(FBXContainerBase):
+    def __init__(self, fbx_file: Path,
+                 volume_dims: Tuple[float] = (10., 10., 10.),
+                 max_actors: int = 8,
+                 pc_size: int = 1024,
+                 scale: float = 0.01,
+                 debug: int = -1):
+        """
+        Class that stores references to important nodes in an FBX file.
+        Offers utility functions to quickly load animation data.
+        :param fbx_file: `Path` to the file to load.
+        :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.
+        """
+        super().__init__(fbx_file=fbx_file, debug=debug)
+
+        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.
+        self.child_names = globals.get_marker_names()
+
+        # Initiate empty lists to store references to nodes.
+        self.children = []
+
+        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]
+        self.vol_z = volume_dims[2]
+        self.hvol_x = volume_dims[0] / 2
+        self.hvol_y = volume_dims[1] / 2
+        self.hvol_z = volume_dims[2] / 2
+
+        self.scale = scale
+
+        self.max_actors = max_actors
+        # Maximum point cloud size = 73 * max_actors + unlabeled markers.
+        self.pc_size = pc_size
+
+        self.output_fbx = utils.append_suffix_to_file(fbx_file, '_INF')
+        self.valid_frames = []
+
+        self.init()
+
+    def __init_unlabeled_markers(self) -> None:
         """
         Looks for the Unlabeled_Markers parent node under the root and stores references to all unlabeled marker nodes.
         """
@@ -538,9 +687,6 @@ 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:
         """
         Calls the init functions for the labeled and unlabeled world transforms.
@@ -562,11 +708,11 @@ class FBXContainer:
         labeled_data = []
 
         # Iterate through all actors.
-        for actor_idx in range(self.actor_count):
+        for actor_idx in range(self.parent_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()):
+            for marker_idx, (n, m) in enumerate(self.children[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)
@@ -586,11 +732,13 @@ class FBXContainer:
         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:
+    def init_unlabeled_world_transforms(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
+                                        incl_keyed: int = 0) -> 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.
+        :param incl_keyed: `bool` whether to check if the marker was keyed at the frame.
         :return: `np.array` of shape (n_frames, n_unlabeled_markers, 14).
         """
         r = self.convert_r(r)
@@ -602,7 +750,7 @@ class FBXContainer:
             # 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)
+            marker_data = get_world_transforms(0, 0, ulm, r, curve, incl_keyed=incl_keyed)
             # Add the result to marker_data.
             unlabeled_data.append(marker_data)
             self._print(f'Unlabeled marker {ulm.GetName()} done', 1)
@@ -616,71 +764,16 @@ class FBXContainer:
         # 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) -> None:
+    def init(self) -> 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)
-        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.
-        """
-        if not 0 <= actor <= self.actor_count:
-            raise ValueError(f'Actor index must be between 0 and {self.actor_count - 1} ({actor}).')
-
-    def get_frame_range(self) -> List[int]:
-        """
-        Replacement and improvement for:
-        `list(range(self.num_frames))`
-        If the animation does not start at frame 0, this will return a list that has the correct frames.
-        :return: List of `int` frame numbers that are between the start and end frame of the animation.
-        """
-        return list(range(self.start_frame, self.end_frame))
-
-    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))
-
-        # 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
+        self._init_scene()
+        self._init_anim()
+        self._init_parents(list(self.child_names[:4]))
+        self._init_children(self.child_names)
+        self.__init_unlabeled_markers()
+        self._print('Init done', 0)
 
     def columns_from_joints(self) -> List[str]:
         """
@@ -688,58 +781,11 @@ class FBXContainer:
         :return: List of column names, in the form of [node1_tx, node1_ty, node1_tz, node2_tx, node2_ty, node2_tz..].
         """
         columns = []
-        for name in self.marker_names:
+        for name in self.child_names:
             columns += [f'{name}x', f'{name}y', f'{name}z']
 
         return columns
 
-    def get_marker_by_name(self, actor: int, name: str):
-        """
-        Returns the reference to the actor's marker.
-        :param actor: `int` actor index.
-        :param name: `str` marker name.
-        :return: `fbx.FbxNode` reference.
-        """
-        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 remove_clipping_poses(self, arr: np.array) -> np.array:
         """
         Checks for each axis if it does not cross the volume limits. Returns an array without clipping poses.
@@ -767,9 +813,11 @@ class FBXContainer:
 
         # 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, 15).
@@ -777,6 +825,7 @@ class FBXContainer:
         # 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.
         # Also remove any frames that cross the limits of the volume.
@@ -802,16 +851,29 @@ class FBXContainer:
         # 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)
+            self.init_labeled_world_transforms(r=r, incl_keyed=1)
         if self.unlabeled_world_transforms is None:
             # Note: Unlabeled data is already flattened.
-            self.init_unlabeled_world_transforms(r=r)
+            self.init_unlabeled_world_transforms(r=r, incl_keyed=1)
 
-        # Returns (n_frames, n_actors, 73, 14).
+        # Returns (n_frames, n_actors, 73, 15).
         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]
+        # Returns shape (n_frames, 73 * n_actors, 15).
+        flat_labeled = self.labeled_world_transforms.reshape(ls[0], -1, ls[-1])
+
+        # Find all labeled markers that have their keyed value set to 0 (which means they had no keyframe on tx),
+        # and set their transforms to np.inf.
+        mask = flat_labeled[..., -1] == 0
+        flat_labeled[mask, 2:] = np.inf
+
+        # Do the same for the unlabeled markers.
+        mask = self.unlabeled_world_transforms[..., -1] == 0
+        self.unlabeled_world_transforms[mask, 2:] = np.inf
+
+        # Remove the last element of the last dimension of both arrays, because it won't be useful anymore.
+        flat_labeled = flat_labeled[..., :-1]
+        self.unlabeled_world_transforms = self.unlabeled_world_transforms[..., :-1]
 
         if merged:
             return utils.merge_labeled_and_unlabeled_data(labeled=flat_labeled,
@@ -832,9 +894,9 @@ class FBXContainer:
 
         # First multiply by self.scale, which turns centimeters to meters.
         # Then divide by volume dimensions, to normalize to the total area of the capture volume.
-        w[..., start + 0] = w[..., start + 0] * self.scale / self.hvol_x
-        w[..., start + 1] = w[..., start + 1] * self.scale / self.hvol_y
-        w[..., start + 2] = w[..., start + 2] * self.scale / self.hvol_z
+        w[..., start + 0] = w[..., start + 0] * self.scale / self.vol_x
+        w[..., start + 1] = w[..., start + 1] * self.scale / self.vol_y
+        w[..., start + 2] = w[..., start + 2] * self.scale / self.vol_z
 
         # 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.
@@ -862,21 +924,21 @@ class FBXContainer:
         # 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 convert_class_to_actor(self, c: float = 0):
+    def convert_class_to_parent(self, c: float = 0):
         """
         Returns the actor name based on the class value.
         :param c: `float` actor class index.
         :return: `str` actor name.
         """
-        return 'UNLABELED' if int(c) == 0 else self.actor_names[int(c) - 1]
+        return 'UNLABELED' if int(c) == 0 else self.parent_names[int(c) - 1]
 
-    def convert_class_to_marker(self, c: float = 0):
+    def convert_class_to_child(self, c: float = 0):
         """
         Returns the marker name based on the class value.
         :param c: `float` marker class index.
         :return: `str` marker name.
         """
-        return 'UNLABELED' if int(c) == 0 else self.marker_names[int(c) - 1]
+        return 'UNLABELED' if int(c) == 0 else self.child_names[int(c) - 1]
 
     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]:
@@ -957,25 +1019,6 @@ class FBXContainer:
 
         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.
@@ -1091,7 +1134,7 @@ class FBXContainer:
         :param actor: `int` actor index to apply to. Index starts at 0.
         :param actor_keys: `dict` with all marker keys for this actor.
         """
-        for marker_class, (marker_name, marker) in enumerate(self.markers[actor].items(), start=1):
+        for marker_class, (marker_name, marker) in enumerate(self.children[actor].items(), start=1):
             marker_keys = actor_keys.get(marker_class)
             if marker_keys:
                 self._print(f'Replacing keys for {marker_name}', 1)
@@ -1102,7 +1145,7 @@ class FBXContainer:
         For all actors, uses self.replace_keyframes_per_actor() to set keyframes on each actor's marker nodes.
         :param key_dict: `dict` with all actor keyframes.
         """
-        for actor_idx in range(self.actor_count):
+        for actor_idx in range(self.parent_count):
             actor_dict = key_dict.get(actor_idx + 1)
             if actor_dict:
                 self._print(f'Replacing keys for actor {actor_idx}', 1)