The app now allows you to select FBX files to process.

app.py

@@ -1,8 +1,66 @@
-import fbx
+# Import standard libs.
+import tempfile
+import os
+from pathlib import Path
+
+# Import primary libs.
 import streamlit as st
 
-st.title('Optical Motion Capture AI')
+# Import custom libs.
+from labeler import read_fbx
+
+
+def process_file(file: Path) -> bytes:
+    fbx_content = read_fbx.MarkerData(file)
+    return fbx_content.export(t='string')
+
+
+# Initialize session state variables if they don't exist
+if "uploaded_files" not in st.session_state:
+    st.session_state.uploaded_files = {}
+if "processed_files" not in st.session_state:
+    st.session_state.processed_files = {}
+
+st.title('Optical MoCap AI Processing')
+
+st.write('Select FBX files to upload and process. This will extract all marker animation data and turn it into a csv.')
+
+new_uploaded_files = st.file_uploader('Select FBX files', accept_multiple_files=True, type='fbx', label_visibility='collapsed')
+
+for uploaded_file in new_uploaded_files:
+    if uploaded_file.name not in st.session_state.processed_files.keys():
+        st.session_state.uploaded_files[uploaded_file.name] = uploaded_file
+
+if st.session_state.uploaded_files and st.button("Process Files"):
+    progress_bar = st.progress(0)
+    # Create a temporary directory to store the newly uploaded files
+    with tempfile.TemporaryDirectory() as temp_dir:
+        incr = 1. / len(st.session_state.uploaded_files)
+
+        for idx, (name, uploaded_file) in enumerate(st.session_state.uploaded_files.items()):
+            # Save the uploaded file to the temporary directory
+            temp_path = Path(os.path.join(temp_dir, name))
+            with open(temp_path, "wb") as f:
+                f.write(uploaded_file.getbuffer())
+                print(f'[LOAD FBX] Finished uploading {temp_path}.')
+                # Process the file and append the resulting DataFrame to dataframes
+                st.session_state.processed_files[name] = process_file(temp_path)
+                progress_bar.progress((idx+1) * incr, f'Processing {name}')
+
+        st.session_state.uploaded_files = {}
+        st.experimental_rerun()
+
+for name in list(st.session_state.processed_files.keys()):
+    new_file_name = name.replace('.fbx', '.csv')
+    if st.download_button(
+            label=f"Download {new_file_name}",
+            data=st.session_state.processed_files[name],
+            file_name=new_file_name,
+            mime='text/csv'):
+        del st.session_state.processed_files[name]
+        st.experimental_rerun()
 
-mng = fbx.FbxManager.Create()
-st.write('Successfully loaded the fbx module:')
-st.write(mng)
+if st.button('Delete cache', type='primary'):
+    st.session_state.uploaded_files = {}
+    st.session_state.processed_files = {}
+    st.experimental_rerun()

globals.py

@@ -0,0 +1,15 @@
+def get_time_modes():
+    return ('eDefaultMode', 'eFrames120', 'eFrames100', 'eFrames60', 'eFrames50', 'eFrames48',
+            'eFrames30', 'eFrames30Drop', 'eNTSCDropFrame', 'eNTSCFullFrame', 'ePAL',
+            'eFrames24', 'eFrames1000', 'eFilmFullFrame', 'eCustom', 'eFrames96', 'eFrames72',
+            'eFrames59dot94', 'eFrames119dot88')
+
+
+def get_marker_names():
+    return ('ARIEL', 'LFHD', 'LBHD', 'RFHD', 'RBHD', 'C7', 'T10', 'CLAV', 'STRN', 'LFSH', 'LBSH',
+            'LUPA', 'LELB', 'LIEL', 'LFRM', 'LIWR', 'LOWR', 'LIHAND', 'LOHAND', 'LTHM3', 'LTHM6',
+            'LIDX3', 'LIDX6', 'LMID0', 'LMID6', 'LRNG3', 'LRNG6', 'LPNK3', 'LPNK6', 'RFSH', 'RBSH',
+            'RUPA', 'RELB', 'RIEL', 'RFRM', 'RIWR', 'ROWR', 'RIHAND', 'ROHAND', 'RTHM3', 'RTHM6',
+            'RIDX3', 'RIDX6', 'RMID0', 'RMID6', 'RRNG3', 'RRNG6', 'RPNK3', 'RPNK6', 'LFWT', 'MFWT',
+            'RFWT', 'LBWT', 'MBWT', 'RBWT', 'LTHI', 'LKNE', 'LKNI', 'LSHN', 'LANK', 'LHEL', 'LMT5',
+            'LMT1', 'LTOE', 'RTHI', 'RKNE', 'RKNI', 'RSHN', 'RANK', 'RHEL', 'RMT5', 'RMT1', 'RTOE')

labeler/read_fbx.py

@@ -0,0 +1,230 @@
+import pandas as pd
+import numpy as np
+from pathlib import Path
+
+import contextlib
+import fbx
+from typing import List, Union
+
+# Import custom data.
+import globals
+
+
+class MarkerData:
+    # TODO: Model is currently built for training. Add testing mode.
+    def __init__(self, fbx_file: Path):
+        """
+        Class that stores references to important nodes in an FBX file.
+        Offers utility functions to quickly load animation data.
+        :param fbx_file: `str` Path to the file to load.
+        """
+        self.time_modes = globals.get_time_modes()
+        self.marker_names = globals.get_marker_names()
+
+        self.markers = []
+        self.actor_names = []
+        self.actors = []
+
+        self.volume_dim_x = 10.
+        self.volume_dim_y = 4.
+
+        self.fbx_file = fbx_file
+        self.valid_frames = []
+
+        self.__init_scene()
+        self.__init_anim()
+        self.__init_actors()
+        self.__init_markers()
+
+    def __init_scene(self):
+        # Create an FBX manager and importer
+        manager = fbx.FbxManager.Create()
+        importer = fbx.FbxImporter.Create(manager, '')
+
+        # Import the FBX file
+        importer.Initialize(str(self.fbx_file))
+        self.scene = fbx.FbxScene.Create(manager, '')
+        importer.Import(self.scene)
+        self.root = self.scene.GetRootNode()
+        self.time_mode = self.scene.GetGlobalSettings().GetTimeMode()
+
+        # Destroy importer to remove reference to imported file.
+        # This will allow us to delete the uploaded file.
+        importer.Destroy()
+
+    def __init_anim(self):
+        # Get the animation stack and layer.
+        anim_stack = self.scene.GetCurrentAnimationStack()
+        self.anim_layer = anim_stack.GetSrcObject(fbx.FbxCriteria.ObjectType(fbx.FbxAnimLayer.ClassId), 0)
+
+        # Find the total number of frames to expect from the local time span.
+        local_time_span = anim_stack.GetLocalTimeSpan()
+        self.num_frames = int(local_time_span.GetDuration().GetFrameCount(self.time_mode))
+
+    def __init_actors(self):
+
+        # 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() for node in gen2_nodes] for name in self.marker_names[:4]):
+                self.actor_names.append(gen1_node.GetName())
+                self.actors.append(gen1_node)
+
+        self.actor_count = len(self.actors)
+        self.valid_frames = [[] for _ in range(self.actor_count)]
+
+    def __init_markers(self):
+        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)
+                    child_name = child.GetName()
+                    if child_name == marker_name:
+                        actor_markers[child_name] = child
+
+            assert len(actor_markers) == len(self.marker_names), f'{actor_node.GetName()} does not have all markers.'
+
+            self.markers.append(actor_markers)
+
+    def _check_actor(self, actor: int = 0):
+        assert 0 <= actor <= self.actor_count, f'Actor number must be between 0 and {self.actor_count - 1}. ' \
+                                               f'It is {actor}.'
+
+    def _set_valid_frames_for_actor(self, actor: int = 0):
+        self._check_actor(actor)
+
+        frames = list(range(self.num_frames))
+        for marker_name in self.marker_names:
+            marker = self.markers[actor][marker_name]
+            t_curve = marker.LclTranslation.GetCurve(self.anim_layer, 'X')
+            keys = [t_curve.KeyGet(i).GetTime().GetFrameCount(self.time_mode) for i in range(t_curve.KeyGetCount())]
+            for frame in frames:
+                if frame not in keys:
+                    with contextlib.suppress(ValueError):
+                        frames.remove(frame)
+
+        self.valid_frames[actor] = frames
+
+    def _check_valid_frames(self, actor: int = 0):
+        if not len(self.valid_frames[actor]):
+            self._set_valid_frames_for_actor(actor)
+
+    def _modify_pose(self, actor, frame) -> List[float]:
+        # 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).GetRow(3)
+            x += [t[0] * 0.01]
+            y += [t[1] * 0.01]
+            z += [t[2] * 0.01]
+
+        # Move the point cloud to the center of the x and y axes. This will put the actor in the middle.
+        x = self.center_axis(x)
+        z = self.center_axis(z)
+
+        # Move the actor to the middle of the volume floor by adding volume_dim_x/2 to x and z.
+        x += self.volume_dim_x / 2.
+        z += self.volume_dim_x / 2.
+
+        # Squeeze the actor into the 1x1 plane for the neural network by dividing the axes.
+        x /= self.volume_dim_x
+        z /= self.volume_dim_x
+        y = np.array(y) / self.volume_dim_y
+
+        # TODO: Optionally: 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.
+        pose = []
+        for i in range(len(x)):
+            pose += [x[i]]
+            pose += [y[i]]
+            pose += [z[i]]
+        return pose
+
+    def get_marker_by_name(self, actor: int, name: str):
+        self._check_actor(actor)
+        return self.markers[actor][name]
+
+    def get_valid_frames_for_actor(self, actor: int = 0):
+        self._check_valid_frames(actor)
+        return self.valid_frames[actor]
+
+    def print_valid_frames_stats_for_actor(self, actor: int = 0):
+        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 columns_from_joints(self):
+        columns = []
+        for name in self.marker_names:
+            columns += [f'{name}x', f'{name}y', f'{name}z']
+
+        return columns
+
+    @staticmethod
+    def center_axis(a) -> np.array:
+        a = np.array(a)
+        _min = np.min(a)
+        _max = np.max(a)
+
+        _c = _max - _min
+        a -= _c
+        return a
+
+    def extract_translations_per_actor(self, actor: int = 0):
+        self._check_actor(actor)
+        self._check_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]:
+            # Get the centered point cloud as an array.
+            pose_at_frame = self._modify_pose(actor, frame)
+            poses.append(pose_at_frame)
+
+        return poses
+
+    def extract_all_translations(self) -> pd.DataFrame:
+
+        columns = self.columns_from_joints()
+
+        all_poses = []
+
+        for i in range(self.actor_count):
+            all_poses.extend(self.extract_translations_per_actor(i))
+
+        return pd.DataFrame(all_poses, columns=columns)
+
+    def export(self, t: str = 'csv', output_file: Path = None) -> Union[bytes, Path]:
+        # Get the dataframe with all animation data.
+        df = self.extract_all_translations()
+
+        if t == 'string':
+            return df.to_csv(index=False).encode('utf-8')
+
+        if output_file is None:
+            output_file = self.fbx_file.with_suffix('.csv')
+
+        if output_file.suffix != '.csv':
+            raise ValueError(f'{output_file} needs to be a .csv file.')
+
+        df.to_csv(output_file, index=False)
+        return output_file

requirements.txt

@@ -1 +1,3 @@
-streamlit
+streamlit~=1.21.0
+pandas~=1.3.5
+numpy~=1.21.5