"""Manages the movement and use of data files.""" import os import warnings from scipy.spatial.transform import Rotation as R import numpy as np from numpy.lib.recfunctions import append_fields class TRCFile(object): """A plain-text file format for storing motion capture marker trajectories. TRC stands for Track Row Column. The metadata for the file is stored in attributes of this object. See http://simtk-confluence.stanford.edu:8080/display/OpenSim/Marker+(.trc)+Files for more information. """ def __init__(self, fpath=None, **kwargs): #path=None, #data_rate=None, #camera_rate=None, #num_frames=None, #num_markers=None, #units=None, #orig_data_rate=None, #orig_data_start_frame=None, #orig_num_frames=None, #marker_names=None, #time=None, #): """ Parameters ---------- fpath : str Valid file path to a TRC (.trc) file. """ self.marker_names = [] if fpath != None: self.read_from_file(fpath) else: for k, v in kwargs.items(): setattr(self, k, v) def read_from_file(self, fpath): # Read the header lines / metadata. # --------------------------------- # Split by any whitespace. # TODO may cause issues with paths that have spaces in them. f = open(fpath) # These are lists of each entry on the first few lines. first_line = f.readline().split() # Skip the 2nd line. f.readline() third_line = f.readline().split() fourth_line = f.readline().split() f.close() # First line. if len(first_line) > 3: self.path = first_line[3] else: self.path = '' # Third line. self.data_rate = float(third_line[0]) self.camera_rate = float(third_line[1]) self.num_frames = int(third_line[2]) self.num_markers = int(third_line[3]) self.units = third_line[4] self.orig_data_rate = float(third_line[5]) self.orig_data_start_frame = int(third_line[6]) self.orig_num_frames = int(third_line[7]) # Marker names. # The first and second column names are 'Frame#' and 'Time'. self.marker_names = fourth_line[2:] len_marker_names = len(self.marker_names) if len_marker_names != self.num_markers: warnings.warn('Header entry NumMarkers, %i, does not ' 'match actual number of markers, %i. Changing ' 'NumMarkers to match actual number.' % ( self.num_markers, len_marker_names)) self.num_markers = len_marker_names # Load the actual data. # --------------------- col_names = ['frame_num', 'time'] # This naming convention comes from OpenSim's Inverse Kinematics tool, # when it writes model marker locations. for mark in self.marker_names: col_names += [mark + '_tx', mark + '_ty', mark + '_tz'] dtype = {'names': col_names, 'formats': ['int'] + ['float64'] * (3 * self.num_markers + 1)} usecols = [i for i in range(3 * self.num_markers + 1 + 1)] self.data = np.loadtxt(fpath, delimiter='\t', skiprows=5, dtype=dtype, usecols=usecols) self.time = self.data['time'] # Check the number of rows. n_rows = self.time.shape[0] if n_rows != self.num_frames: warnings.warn('%s: Header entry NumFrames, %i, does not ' 'match actual number of frames, %i, Changing ' 'NumFrames to match actual number.' % (fpath, self.num_frames, n_rows)) self.num_frames = n_rows def __getitem__(self, key): """See `marker()`. """ return self.marker(key) def units(self): return self.units def time(self): this_dat = np.empty((self.num_frames, 1)) this_dat[:, 0] = self.time return this_dat def marker(self, name): """The trajectory of marker `name`, given as a `self.num_frames` x 3 array. The order of the columns is x, y, z. """ this_dat = np.empty((self.num_frames, 3)) this_dat[:, 0] = self.data[name + '_tx'] this_dat[:, 1] = self.data[name + '_ty'] this_dat[:, 2] = self.data[name + '_tz'] return this_dat def add_marker(self, name, x, y, z): """Add a marker, with name `name` to the TRCFile. Parameters ---------- name : str Name of the marker; e.g., 'R.Hip'. x, y, z: array_like Coordinates of the marker trajectory. All 3 must have the same length. """ if (len(x) != self.num_frames or len(y) != self.num_frames or len(z) != self.num_frames): raise Exception('Length of data (%i, %i, %i) is not ' 'NumFrames (%i).', len(x), len(y), len(z), self.num_frames) self.marker_names += [name] self.num_markers += 1 if not hasattr(self, 'data'): self.data = np.array(x, dtype=[('%s_tx' % name, 'float64')]) self.data = append_fields(self.data, ['%s_t%s' % (name, s) for s in 'yz'], [y, z], usemask=False) else: self.data = append_fields(self.data, ['%s_t%s' % (name, s) for s in 'xyz'], [x, y, z], usemask=False) def marker_at(self, name, time): x = np.interp(time, self.time, self.data[name + '_tx']) y = np.interp(time, self.time, self.data[name + '_ty']) z = np.interp(time, self.time, self.data[name + '_tz']) return [x, y, z] def marker_exists(self, name): """ Returns ------- exists : bool Is the marker in the TRCFile? """ return name in self.marker_names def write(self, fpath): """Write this TRCFile object to a TRC file. Parameters ---------- fpath : str Valid file path to which this TRCFile is saved. """ f = open(fpath, 'w') # Line 1. f.write('PathFileType 4\t(X/Y/Z) %s\n' % os.path.split(fpath)[0]) # Line 2. f.write('DataRate\tCameraRate\tNumFrames\tNumMarkers\t' 'Units\tOrigDataRate\tOrigDataStartFrame\tOrigNumFrames\n') # Line 3. f.write('%.1f\t%.1f\t%i\t%i\t%s\t%.1f\t%i\t%i\n' % ( self.data_rate, self.camera_rate, self.num_frames, self.num_markers, self.units, self.orig_data_rate, self.orig_data_start_frame, self.orig_num_frames)) # Line 4. f.write('Frame#\tTime\t') for imark in range(self.num_markers): f.write('%s\t\t\t' % self.marker_names[imark]) f.write('\n') # Line 5. f.write('\t\t') for imark in np.arange(self.num_markers) + 1: f.write('X%i\tY%s\tZ%s\t' % (imark, imark, imark)) f.write('\n') # Line 6. f.write('\n') # Data. for iframe in range(self.num_frames): f.write('%i' % (iframe + 1)) f.write('\t%.7f' % self.time[iframe]) for mark in self.marker_names: idxs = [mark + '_tx', mark + '_ty', mark + '_tz'] f.write('\t%.7f\t%.7f\t%.7f' % tuple( self.data[coln][iframe] for coln in idxs)) f.write('\n') f.close() def add_noise(self, noise_width): """ add random noise to each component of the marker trajectory The noise mean will be zero, with the noise_width being the standard deviation. noise_width : int """ for imarker in range(self.num_markers): components = ['_tx', '_ty', '_tz'] for iComponent in range(3): # generate noise noise = np.random.normal(0, noise_width, self.num_frames) # add noise to each component of marker data. self.data[self.marker_names[imarker] + components[iComponent]] += noise def rotate(self, axis, value): """ rotate the data. axis : rotation axis value : angle in degree """ for imarker in range(self.num_markers): temp = np.zeros((self.num_frames, 3)) temp[:,0] = self.data[self.marker_names[imarker] + '_tx'] temp[:,1] = self.data[self.marker_names[imarker] + '_ty'] temp[:,2] = self.data[self.marker_names[imarker] + '_tz'] r = R.from_euler(axis, value, degrees=True) temp_rot = r.apply(temp) self.data[self.marker_names[imarker] + '_tx'] = temp_rot[:,0] self.data[self.marker_names[imarker] + '_ty'] = temp_rot[:,1] self.data[self.marker_names[imarker] + '_tz'] = temp_rot[:,2] def offset(self, axis, value): """ offset the data. axis : rotation axis value : offset in m """ for imarker in range(self.num_markers): if axis.lower() == 'x': self.data[self.marker_names[imarker] + '_tx'] += value elif axis.lower() == 'y': self.data[self.marker_names[imarker] + '_ty'] += value elif axis.lower() == 'z': self.data[self.marker_names[imarker] + '_tz'] += value else: raise ValueError("Axis not recognized")