Fixed the math to align the poses, checked it in Houdini.

Added a method to remove clipping poses during train data extraction.

fbx_handler.py

@@ -27,7 +27,7 @@ def center_axis(a: Union[List[float], np.array]) -> np.array:
     # Find the centroid by subtracting the lowest value from the highest value.
     _min = np.min(a)
     _max = np.max(a)
-    _c = _max - _min
+    _c = _max - (_max - _min) * 0.5
     # Center the array by subtracting the centroid.
     a -= _c
     return a
@@ -342,7 +342,7 @@ def get_world_transforms(actor_idx: int, marker_idx: int, m: fbx.FbxNode, r: Lis
 
 class FBXContainer:
     def __init__(self, fbx_file: Path,
-                 volume_dims: Tuple[float] = (10., 4., 10.),
+                 volume_dims: Tuple[float] = (10., 10., 10.),
                  max_actors: int = 8,
                  pc_size: int = 1024,
                  scale: float = 0.01,
@@ -385,6 +385,9 @@ class FBXContainer:
         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
 
@@ -853,6 +856,16 @@ class FBXContainer:
         # 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)
+        mask_z2 = (arr[:, :, 4] > -self.hvol_z / self.scale).all(axis=1)
+        mask = mask_x1 & mask_x2 & mask_z1 & mask_z2
+        # print(mask.shape, mask)
+        return arr[mask]
+
     def extract_training_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None) -> np.array:
         if self.labeled_world_transforms is None:
             self.init_labeled_world_transforms(r=r, incl_keyed=1)
@@ -864,19 +877,21 @@ class FBXContainer:
         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).
+        # Returns shape of (n_frames * n_actors, 73, 15).
         mask = (flattened[..., -1] == 1)
         # We only need a filter for the first dimension, so use .all to check if all markers
         # have a keyframe. This results in shape (n_frames * n_actors,).
         mask = mask.all(axis=1)
         # Now isolate the right frames with the mask and remove the last element of the last dimension,
         # because it won't be useful anymore.
-        valid_poses = flattened[mask][..., :-1]
+        # Also remove any frames that cross the limits of the volume.
+        valid_poses = self.remove_clipping_poses(flattened[mask][..., :-1])
 
-        # Now we need to center the tx and tz axes.
+        # Now we need to center the tx and tz axes of each individual pose.
         for valid_pose in valid_poses:
             for axis in [2, 4]:
                 valid_pose[:, axis] = center_axis(valid_pose[:, axis])
+        # Finally, scale the data to the correct size by normalizing.
         return self.transform_translations(valid_poses)
 
     def extract_inf_translations(self, r: Union[int, Tuple[int, int], Tuple[int, int, int]] = None,
@@ -903,21 +918,23 @@ class FBXContainer:
         :param w: `np.array` that can either be a timeline dense cloud or translation vectors.
         :return: Modified `np.array`.
         """
-        if w.ndim != 3:
-            raise ValueError(f'Array does not have 3 dimensions: {w.ndim}/3.')
 
         # If the last dimension has 3 elements, it is a translation vector of shape (tx, ty, tz).
         # If it has 14 elements, it is a full marker row of shape (actor, marker, tx, ty, tz, tw, rx, ry, rz, tw, etc).
         start = 0 if w.shape[-1] == 3 else 2
 
-        # First multiply by self.scale, which turns meters to centimeters.
+        # 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] = np.clip(w[:, :, start + 0], -(self.vol_x * 0.5),
-                                     self.vol_x * 0.5) * self.scale / self.vol_x
-        w[:, :, start + 1] = np.clip(w[:, :, start + 1], -(self.vol_y * 0.5),
-                                     self.vol_y * 0.5) * self.scale / self.vol_y
-        w[:, :, start + 2] = np.clip(w[:, :, start + 2], -(self.vol_z * 0.5),
-                                     self.vol_z * 0.5) * self.scale / self.vol_z
+        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
+
+        # 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
+        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
 
         return w
 

preprocess_files.py

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

utils.py

@@ -48,13 +48,15 @@ def array4d_to_h5(array_4ds: Tuple, output_file: Path, group: str = None, datase
                 h5f.create_dataset(name=datasets[i], data=array_4ds[i], compression='gzip', compression_opts=9)
 
 
-def h5_to_array4d(input_file: Path, mode: str = 'train', pc_size: int = 1024) -> np.array:
+def h5_to_array4d(input_file: Path) -> np.array:
     with h5py.File(input_file, 'r') as h5f:
-        if mode == 'train':
-            return np.vstack([np.array(h5f[key]) for key in h5f.keys()])
+        return np.vstack([np.array(h5f[key]) for key in h5f.keys()])
 
+
+def combined_test_h5_to_array4d(input_file: Path, pc_size: int = 1024) -> np.array:
+    with h5py.File(input_file, 'r') as h5f:
         data = []
-        for grp_name in h5f.keys():
+        for grp_name in list(h5f.keys()):
             grp = h5f[grp_name]
             labeled = np.array(grp['labeled'])
             unlabeled = np.array(grp['unlabeled'])