Fixed the translation bug. TODO: Make sure the new animation curves are properly connected to their nodes.

fbx_handler.py

@@ -45,6 +45,15 @@ def make_ghost_markers(missing: int) -> np.array:
         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
     ])
 
@@ -60,15 +69,31 @@ def append_suffix(file_path: Path, suffix: str = '_INF'):
     return file_path.with_name(new_file_name)
 
 
+def append_zero(arr: np.ndarray) -> np.ndarray:
+    zeros = np.zeros((arr.shape[0], arr.shape[1], 1), dtype=float)
+    return np.concatenate((arr, zeros), axis=-1)
+
+
+def append_one(arr: np.ndarray) -> np.ndarray:
+    ones = np.ones((arr.shape[0], arr.shape[1], 1), dtype=float)
+    return np.concatenate((arr, ones), axis=-1)
+
+
 def merge_tdc(actor_classes: np.array,
               marker_classes: np.array,
               translation_vectors: np.array,
+              rotation_vectors: np.array,
+              scale_vectors: np.array,
               ordered: bool = True) -> np.array:
     # Actor and marker classes enter as shape (x, 1000), so use np.expand_dims to create an extra dimension at the end.
     # Return the concatenated array of shape (x, 1000, 5), which matches the original timeline dense cloud before
     # splitting it into sub arrays.
-    tdc = np.concatenate((np.expand_dims(actor_classes, -1), np.expand_dims(marker_classes, -1),
-                          translation_vectors), axis=2)
+
+    tdc = np.concatenate((np.expand_dims(actor_classes, -1),
+                          np.expand_dims(marker_classes, -1),
+                          append_zero(translation_vectors),
+                          append_zero(rotation_vectors),
+                          append_one(scale_vectors)), axis=2)
     if ordered:
         tdc = sort_cloud(tdc)
 
@@ -115,27 +140,6 @@ def create_keyframe(anim_curve: fbx.FbxAnimCurve, frame: int, value: float):
     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()):
@@ -167,7 +171,16 @@ def timeline_cloud_to_dict(data: np.array, start_frame: int = 0) -> dict:
         if actor_class == 0 or marker_class == 0:
             continue
 
-        translation_vector = data[frame, node, 2:]
+        # Just to be sure, forcing the last numbers of each array to be the correct values.
+        # Also check self.get_world_transform() for this.
+        translations = data[frame, node, 2:5] + np.array([0.0])
+        rotations = data[frame, node, 6:9] + np.array([0.0])
+        scales = data[frame, node, 10:13] + np.array([1.0])
+
+        world_matrix = fbx.FbxAMatrix()
+        world_matrix.SetT(fbx.FbxVector4(*translations))
+        world_matrix.SetR(fbx.FbxVector4(*rotations))
+        world_matrix.SetS(fbx.FbxVector4(*scales))
 
         # Create the actor dictionary if it doesn't exist.
         if actor_class not in result:
@@ -178,35 +191,31 @@ def timeline_cloud_to_dict(data: np.array, start_frame: int = 0) -> dict:
             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
+        result[actor_class][marker_class][frame + start_frame] = world_matrix
 
     return result
 
 
-def world_to_local_translation(node, parent_node, f):
+def world_to_local_transform(node, world_transform, frame):
     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)
+    t.SetFrame(frame)
+    if node.GetParent():
+        # Get the parent's world transform
+        parent_world_transform = node.GetParent().EvaluateGlobalTransform(t)
 
-    # 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 transform
+        parent_world_transform_inv = parent_world_transform.Inverse()
 
-    # Compute the inverse of the parent's world transformation matrix
-    parent_world_matrix_inv = np.linalg.inv(parent_world_matrix)
+        # Multiply the inverse of the parent's world transform by the world transform to get the local transform
+        lcl = parent_world_transform_inv * world_transform
+    else:
+        # If the node doesn't have a parent, the local transform is the same as the world transform
+        lcl = world_transform
 
-    return np.dot(parent_world_matrix_inv, child_world_matrix)
+    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)]
 
 
 class FBXContainer:
-    # TODO: Model is currently built for training. Add testing mode.
     def __init__(self, fbx_file: Path,
                  volume_dims: Tuple[float] = (10., 4., 10.),
                  max_actors: int = 10,
@@ -443,7 +452,7 @@ class FBXContainer:
         z /= self.vol_z
         y = np.array(y) / self.vol_y
 
-        # TODO: Optionally: Add any extra modifications to the point cloud here.
+        # TODO: Optional: Add any extra modifications to the point cloud here.
 
         # Append all values to a new array, one axis at a time.
         # This way it will match the column names order.
@@ -592,7 +601,7 @@ class FBXContainer:
 
         return pd.DataFrame(all_poses, columns=columns)
 
-    def get_worldspace(self, m: fbx.FbxNode, time: fbx.FbxTime, apply_transform: bool = True) -> List[float]:
+    def get_world_transform(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.
@@ -601,17 +610,23 @@ class FBXContainer:
         :param apply_transform: `bool` Whether to transform the translation or not.
         :return: Vector in the form: [tx, ty, tz].
         """
-        t = m.EvaluateGlobalTransform(time).GetT()
+        world = m.EvaluateGlobalTransform(time)
+        world = list(world.GetT()) + list(world.GetR()) + list(world.GetS())
+        # Make sure that the last numbers of each row are the correct values.
+        world[3] = 0.0
+        world[7] = 0.0
+        world[11] = 1.0
+
         if not apply_transform:
-            return [t[i] for i in range(3)]
+            return world
 
         # 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
+        world[2] = np.clip(world[2], -(self.vol_x * 0.5), self.vol_x * 0.5) * self.scale / self.vol_x
+        world[3] = np.clip(world[3], -(self.vol_y * 0.5), self.vol_y * 0.5) * self.scale / self.vol_y
+        world[4] = np.clip(world[4], -(self.vol_z * 0.5), self.vol_z * 0.5) * self.scale / self.vol_z
 
-        return [x, y, z]
+        return world
 
     def is_kf_present(self, marker: fbx.FbxNode, time: fbx.FbxTime) -> bool:
         """
@@ -638,7 +653,7 @@ class FBXContainer:
         # 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)]
+            [0, 0, *self.get_world_transform(m, time, apply_transform)]
             for m in self.unlabeled_markers
             if self.is_kf_present(m, time)
         ]
@@ -649,7 +664,7 @@ class FBXContainer:
                 # 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_worldspace(m, time, apply_transform)]
+                [actor_idx + 1, marker_class, *self.get_world_transform(m, time, apply_transform)]
                 for marker_class, (marker_name, m) in enumerate(
                     self.markers[actor_idx].items(), start=1
                 )
@@ -673,17 +688,17 @@ class FBXContainer:
         """
         return np.array([self.get_sc(f, apply_transform) for f in self.get_frame_range()])
 
-    def get_tdc(self, shuffle: bool = False,
-                r: Union[int, Tuple[int, int]] = None,
+    def get_tdc(self, r: Union[int, Tuple[int, int]] = None,
+                shuffle: bool = False,
                 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 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 shuffle: If `True`, shuffles the dense point cloud of each frame.
         :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).
@@ -722,13 +737,14 @@ class FBXContainer:
         return np.array(clouds)
 
     def split_tdc(self, cloud: np.array = None, shuffle: bool = False, apply_transform: bool = True) \
-            -> Tuple[np.array, np.array, np.array]:
+            -> 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, 3).
+        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.
         :param shuffle: `bool` whether to shuffle the generated cloud if no cloud was given.
@@ -736,13 +752,15 @@ class FBXContainer:
         :return: Return tuple of `np.array` as (actor classes, marker classes, translation vectors).
         """
         if cloud is None:
-            cloud = self.get_tdc(shuffle, apply_transform=apply_transform)
+            cloud = self.get_tdc(shuffle=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."
+        if cloud.shape[1] != 1000:
+            raise ValueError(f"Dense cloud doesn't have enough points. {cloud.shape[1]}/1000.")
+        if cloud.shape[2] != 14:
+            raise ValueError(f"Dense cloud is missing columns: {cloud.shape[2]}/14.")
 
-        # Return np arrays as (actor classes, marker classes, translation vectors).
-        return cloud[:, :, 0], cloud[:, :, 1], cloud[:, :, -3:]
+        # 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):
         """
@@ -841,19 +859,66 @@ class FBXContainer:
         self.remove_unlabeled_markers()
         self.remove_system()
 
+    def replace_animation_curves(self, node, curve_types):
+        anim_curve_dict = {
+            't': [
+                ("X", node.LclTranslation),
+                ("Y", node.LclTranslation),
+                ("Z", node.LclTranslation)
+            ],
+            'r': [
+                ("X", node.LclRotation),
+                ("Y", node.LclRotation),
+                ("Z", node.LclRotation)
+            ],
+            's': [
+                ("X", node.LclScaling),
+                ("Y", node.LclScaling),
+                ("Z", node.LclScaling)
+            ]
+        }
+
+        anim_curves = []
+        for curve_type in curve_types:
+            if curve_type in anim_curve_dict:
+                for anim_curve_name, p in anim_curve_dict[curve_type]:
+                    # Disconnect and remove existing animation curve, if any
+                    existing_anim_curve = p.GetCurve(self.anim_layer, anim_curve_name, False)
+                    # TODO: Make sure a new anim curve is properly connected.
+                    if existing_anim_curve:
+                        existing_anim_curve.KeyClear()
+                        anim_curves.append(existing_anim_curve)
+                    #     p.DisconnectSrcObject(existing_anim_curve)
+                    #     existing_anim_curve.Destroy()
+                    # del existing_anim_curve
+                    #
+                    # # Create a new animation curve and connect it to the node and animation layer
+                    # new_anim_curve = fbx.FbxAnimCurve.Create(self.manager, anim_curve_name)
+                    # p.ConnectSrcObject(new_anim_curve)
+                    # new_anim_curve.ConnectDstObject(self.anim_layer)
+                    #
+                    # anim_curves.append(new_anim_curve)
+
+        return anim_curves
+
     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)
+
+        # Collect lcl transform curves.
+        curves = self.replace_animation_curves(marker, 't')
+
+        # TODO: Only set translation keys. Set rotation and scale as property values instead of curves.
+        for axis, curve in enumerate(curves):
+
             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)
+            # The dict has frames mapped to world matrices.
+            # The world_transform here is that full matrix, so we only need to convert this to local space.
+            for frame, world_transform in marker_keys.items():
+                # Convert world to local transform at the given frame.
+                lcl_t, lcl_r, lcl_s = world_to_local_transform(marker, world_transform, frame)
+                # Only for translations, set keyframes.
+                create_keyframe(curve, frame, lcl_t[axis])
 
-                    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:
@@ -866,48 +931,12 @@ class FBXContainer:
 
 
 # d = FBXContainer(Path('G:/Firestorm/mocap-ai/data/fbx/dowg/TAKE_01+1_ALL_001.fbx'))
-# 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()
+# # cloud = d.get_tdc(apply_transform=False)
+# actors_train, markers_train, t_train, r_train, s_train = d.split_tdc(apply_transform=True)
+# actors_test, markers_test, t_test, r_test, s_test = d.split_tdc(apply_transform=False)
+# # splits = d.split_tdc(apply_transform=False)
+# merged = merge_tdc(actors_train, markers_train, t_test, r_test, s_test)
+# pc_dict = timeline_cloud_to_dict(merged, d.start_frame)
+# d.replace_keyframes_for_all_actors(pc_dict)
+# # d.cleanup()
 # d.export_fbx(Path('G:/Firestorm/mocap-ai/data/fbx/export/TAKE_01+1_ALL_001.fbx'))