Save to c3d and convert trc to c3d (#92)

* make c3d

Pose2Sim/Utilities/trc_to_c3d.py

@@ -0,0 +1,73 @@
+"""
+Extracts marker data from a TRC file and creates a corresponding C3D file.
+
+Usage:
+    python trc_to_c3d.py --trc_path <path_to_trc_file> --f <frame_rate>
+
+--trc_path: Path to the TRC file.
+--f: Frame rate of the 3D data.
+c3d file will be saved in the same directory as the TRC file with the same name.
+
+"""
+
+import os
+import argparse
+import numpy as np
+import pandas as pd
+import c3d
+
+def extract_marker_data(trc_file):
+    with open(trc_file, 'r') as file:
+        lines = file.readlines()
+        marker_names_line = lines[3]
+        marker_names = marker_names_line.strip().split('\t')[2::3]
+
+    trc_data = pd.read_csv(trc_file, sep='\t', skiprows=5)
+    marker_coords = trc_data.iloc[:, 2:].to_numpy().reshape(-1, len(marker_names), 3)
+    # marker_coords = np.nan_to_num(marker_coords, nan=0.0)
+    marker_coords *= 1000  # Convert from meters to millimeters
+
+    return marker_names, marker_coords
+
+def create_c3d_file(trc_file, marker_names, marker_coords, frame_rate):
+    writer = c3d.Writer(point_rate=frame_rate, analog_rate=0, point_scale=1.0, point_units='mm', gen_scale=-1.0)
+    writer.set_point_labels(marker_names)
+    
+    markers_group = writer.point_group
+
+    for frame in marker_coords:
+        residuals = np.full((frame.shape[0], 1), 0.0)
+        cameras = np.zeros((frame.shape[0], 1))
+        points = np.hstack((frame, residuals, cameras))
+        writer.add_frames([(points, np.array([]))])
+
+    writer.set_start_frame(1)
+    writer._set_last_frame(len(marker_coords))
+
+    c3d_file_path = trc_file.replace('.trc', '.c3d')
+    with open(c3d_file_path, 'wb') as handle:
+        writer.write(handle)
+    print(f"Successfully created c3d file.")
+
+def trc_to_c3d(trc_file, frame_rate):
+    marker_names, marker_coords = extract_marker_data(trc_file)
+    create_c3d_file(trc_file, marker_names, marker_coords, frame_rate)
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert TRC files to C3D files.')
+    parser.add_argument('--trc_path', type=str, required=True, help='Path to the TRC file')
+    parser.add_argument('--f', type=int, required=True, help='Frame rate')
+
+    args = parser.parse_args()
+
+    trc_file = args.trc_path
+    frame_rate = args.f
+
+    if not os.path.isfile(trc_file):
+        print(f"Error: {trc_file} does not exist.")
+        return
+
+    trc_to_c3d(trc_file, frame_rate)
+
+if __name__ == '__main__':
+    main()

Pose2Sim/common.py

@@ -16,6 +16,10 @@ import json
 import numpy as np
 import re
 import cv2
+import os
+import pandas as pd
+import c3d
+import glob
 
 import matplotlib as mpl
 mpl.use('qt5agg')
@@ -431,4 +435,86 @@ class plotWindow():
         self.tab_handles.append(new_tab)
 
     def show(self):
-        self.app.exec_() 
+        self.app.exec_() 
+
+# Save c3d
+def trc_to_c3d(project_dir, frame_rate, called_from):
+    '''
+    Converts.trc files to.c3d files
+    INPUT:
+    - project_dir: path to project folder
+    - frame_rate: frame rate of the video
+    - called_from: string that determines which.trc files to convert.
+        'triangulation' for.trc files from triangulation step
+        'filtering' for.trc files from filtering step
+        
+    '''
+    # Determine the 3D pose folder
+    pose3d_dir = os.path.join(project_dir, 'pose-3d')
+
+    # Determine the .trc file name to read
+    trc_files = []
+    if called_from == 'triangulation':
+        trc_pattern = "*.trc"
+        trc_files = [os.path.basename(f) for f in glob.glob(os.path.join(pose3d_dir, trc_pattern)) if 'filt' not in f]
+    elif called_from == 'filtering':
+        trc_pattern = "*_filt_*.trc"
+        trc_files = [os.path.basename(f) for f in glob.glob(os.path.join(pose3d_dir, trc_pattern))]
+    else:
+        print("Invalid called_from value.")
+
+    for trc_file in trc_files:
+        # Extract marker names from the 4th row of the TRC file
+        with open(os.path.join(pose3d_dir, trc_file), 'r') as file:
+            lines = file.readlines()
+            marker_names_line = lines[3]
+            marker_names = marker_names_line.strip().split('\t')[2::3]
+
+        # Read the data frame (skiprows=5)
+        trc_data = pd.read_csv(os.path.join(pose3d_dir, trc_file), sep='\t', skiprows=5)
+
+        # Extract marker coordinates
+        marker_coords = trc_data.iloc[:, 2:].to_numpy().reshape(-1, len(marker_names), 3)
+        # marker_coords = np.nan_to_num(marker_coords, nan=0.0)
+
+        scale_factor = 1000
+        marker_coords = marker_coords * scale_factor
+
+        # Create a C3D writer
+        writer = c3d.Writer(point_rate=frame_rate, analog_rate=0, point_scale=1.0, point_units='mm', gen_scale=-1.0)
+
+        # Add marker parameters
+        writer.set_point_labels(marker_names)
+
+        # # Add marker descriptions (optional)
+        # marker_descriptions = [''] * len(marker_names)
+        # writer.point_group.add_param('DESCRIPTIONS', desc='Marker descriptions', 
+        #                             bytes_per_element=-1, dimensions=[len(marker_names)], 
+        #                             bytes=np.array(marker_descriptions, dtype=object))
+        
+        # # Set the data start parameter
+        # data_start = writer.header.data_block
+        # writer.point_group.add_param('DATA_START', desc='Data start parameter',
+        #                             bytes_per_element=2, dimensions=[], bytes=struct.pack('<H', data_start))
+        
+        # Create a C3D group for markers
+        markers_group = writer.point_group
+
+        # Add frame data
+        for frame in marker_coords:
+            # Add residual and camera columns
+            residuals = np.full((frame.shape[0], 1), 0.0)  # Set residuals to 0.0
+            cameras = np.zeros((frame.shape[0], 1))  # Set cameras to 0
+            points = np.hstack((frame, residuals, cameras))
+            writer.add_frames([(points, np.array([]))])
+
+        # Set the trial start and end frames
+        writer.set_start_frame(1)
+        writer._set_last_frame(len(marker_coords))
+
+        # Write the C3D file
+        c3d_file_path = trc_file.replace('.trc', '.c3d')
+        with open(os.path.join(pose3d_dir, c3d_file_path), 'wb') as handle:
+            writer.write(handle)
+        
+        print(f"-->c3d file saved: {c3d_file_path}")

Pose2Sim/filtering.py

@@ -37,6 +37,7 @@ from filterpy.kalman import KalmanFilter, rts_smoother
 from filterpy.common import Q_discrete_white_noise
 
 from Pose2Sim.common import plotWindow
+from Pose2Sim.common import trc_to_c3d
 
 ## AUTHORSHIP INFORMATION
 __author__ = "David Pagnon"
@@ -415,7 +416,7 @@ def recap_filter3d(config, trc_path):
     gaussian_filter_sigma_kernel = int(config.get('filtering').get('gaussian').get('sigma_kernel'))
     loess_filter_nb_values = config.get('filtering').get('LOESS').get('nb_values_used')
     median_filter_kernel_size = config.get('filtering').get('median').get('kernel_size')
-
+    
     # Recap
     filter_mapping_recap = {
         'kalman': f'--> Filter type: Kalman {kalman_filter_smooth_str}. Measurements trusted {kalman_filter_trustratio} times as much as previous data, assuming a constant acceleration process.', 
@@ -455,6 +456,8 @@ def filter_all(config):
     display_figures = config.get('filtering').get('display_figures')
     filter_type = config.get('filtering').get('type')
     seq_name = os.path.basename(os.path.realpath(project_dir))
+    make_c3d = config.get('filtering').get('make_c3d')
+    frame_rate = config.get('project').get('frame_rate')
     
     # Frames range
     pose_listdirs_names = next(os.walk(pose_dir))[1]
@@ -496,3 +499,8 @@ def filter_all(config):
 
         # Recap
         recap_filter3d(config, t_out)
+
+        # Save c3d
+        if make_c3d == True:
+            trc_to_c3d(project_dir, frame_rate, called_from='filtering')
+

Pose2Sim/triangulation.py

@@ -51,6 +51,7 @@ import logging
 
 from Pose2Sim.common import retrieve_calib_params, computeP, weighted_triangulation, \
     reprojection, euclidean_distance, sort_stringlist_by_last_number, zup2yup
+from Pose2Sim.common import trc_to_c3d
 from Pose2Sim.skeletons import *
 
 
@@ -702,6 +703,8 @@ def triangulate_all(config):
     interp_gap_smaller_than = config.get('triangulation').get('interp_if_gap_smaller_than')
     show_interp_indices = config.get('triangulation').get('show_interp_indices')
     undistort_points = config.get('triangulation').get('undistort_points')
+    make_c3d = config.get('triangulation').get('make_c3d')
+    frame_rate = config.get('project').get('frame_rate')
 
     calib_dir = [os.path.join(session_dir, c) for c in os.listdir(session_dir) if 'calib' in c.lower()][0]
     try:
@@ -922,3 +925,7 @@ def triangulate_all(config):
     
     # Recap message
     recap_triangulate(config, error_tot, nb_cams_excluded_tot, keypoints_names, cam_excluded_count, interp_frames, non_interp_frames, trc_paths)
+
+    # Save c3d
+    if make_c3d == True:
+            trc_to_c3d(project_dir, frame_rate, called_from='triangulation')