File size: 5,125 Bytes
4caa42c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import numpy as np
import torch
from scipy.signal import butter, filtfilt
import secrets
def numpy_to_storage(labels, data, storage_file, datatype=None):
assert data.shape[1] == len(labels), "# labels doesn't match columns"
assert labels[0] == "time"
f = open(storage_file, 'w')
# Old style
if datatype is None:
f = open(storage_file, 'w')
f.write('name %s\n' % storage_file)
f.write('datacolumns %d\n' % data.shape[1])
f.write('datarows %d\n' % data.shape[0])
f.write('range %f %f\n' % (np.min(data[:, 0]), np.max(data[:, 0])))
f.write('endheader \n')
# New style
else:
if datatype == 'IK':
f.write('Coordinates\n')
elif datatype == 'ID':
f.write('Inverse Dynamics Generalized Forces\n')
elif datatype == 'GRF':
f.write('%s\n' % storage_file)
elif datatype == 'muscle_forces':
f.write('ModelForces\n')
f.write('version=1\n')
f.write('nRows=%d\n' % data.shape[0])
f.write('nColumns=%d\n' % data.shape[1])
if datatype == 'IK':
f.write('inDegrees=yes\n\n')
f.write('Units are S.I. units (second, meters, Newtons, ...)\n')
f.write(
"If the header above contains a line with 'inDegrees', this indicates whether rotational values are in degrees (yes) or radians (no).\n\n")
elif datatype == 'ID':
f.write('inDegrees=no\n')
elif datatype == 'GRF':
f.write('inDegrees=yes\n')
elif datatype == 'muscle_forces':
f.write('inDegrees=yes\n\n')
f.write('This file contains the forces exerted on a model during a simulation.\n\n')
f.write("A force is a generalized force, meaning that it can be either a force (N) or a torque (Nm).\n\n")
f.write('Units are S.I. units (second, meters, Newtons, ...)\n')
f.write('Angles are in degrees.\n\n')
f.write('endheader \n')
for i in range(len(labels)):
f.write('%s\t' % labels[i])
f.write('\n')
for i in range(data.shape[0]):
for j in range(data.shape[1]):
f.write('%20.8f\t' % data[i, j])
f.write('\n')
f.close()
def lowpass_filter(data, cutoff_cycles=2, num_samples=24, order=4):
nyquist = num_samples / 2 # Max frequency is Nyquist (12 cycles/stride for 24 samples)
normal_cutoff = cutoff_cycles / nyquist # Convert cycles per stride to normalized frequency
b, a = butter(order, normal_cutoff, btype='low', analog=False)
return filtfilt(b, a, data, axis=0)
# Define a function to get a random stride of the dataloader, save it to a mot file and print the metadata as well as the reconstruction
def save_random_stride(model, dataloader, joints, savepath, savepath_recon, device,
from_all_batches=False, filter=False):
# ----- choose a random example using secrets -----
if from_all_batches:
batch_data, batch_metadata = [], []
for batch in dataloader:
batch_data.append(batch["features"])
batch_metadata.extend(batch["metadata"])
batch_data = torch.cat(batch_data, dim=0)
stride_idx = secrets.randbelow(batch_data.shape[0]) # cryptographically strong
stride_data = batch_data[stride_idx]
stride_metadata = batch_metadata[stride_idx]
else:
ds = dataloader.dataset
ds_idx = secrets.randbelow(len(ds))
item = ds[ds_idx] # expects {"features": (S,T,D), "metadata": ...}
feats = item["features"] # (num_strides, T, D)
s_idx = secrets.randbelow(feats.shape[0])
stride_data = feats[s_idx]
meta = item["metadata"]
try:
stride_metadata = meta.iloc[s_idx] # per-stride metadata (DataFrame)
except Exception:
stride_metadata = meta # per-subject metadata
print(f"Stride Data Shape: {stride_data.shape}")
# ----- forward pass -----
model.eval()
with torch.no_grad():
reconstructed_data, _, mu, _, _ = model(stride_data.unsqueeze(0).to(device))
reconstructed_data = reconstructed_data.view(-1, 24, 32).cpu().numpy().squeeze()
stride_data = stride_data.cpu().numpy().squeeze()
# ----- keep requested joints -----
stride_data = stride_data[:, :len(joints)]
reconstructed_data = reconstructed_data[:, :len(joints)]
# ----- sampling freq -----
stride_time = stride_data[:, 0]
fs = 1.0 / np.mean(np.diff(stride_time))
print(f"Sampling Frequency: {fs} Hz")
# ----- optional filtering -----
if filter:
reconstructed_data = lowpass_filter(stride_data, cutoff_cycles=2, num_samples=24)
reconstructed_data[:, 0] = stride_time
joints = [j.replace('_ips', '_r').replace('_contra', '_l') for j in joints]
numpy_to_storage(joints, stride_data, savepath, datatype='IK')
numpy_to_storage(joints, reconstructed_data, savepath_recon, datatype='IK')
print(f"Stride Metadata: {stride_metadata}")
return mu # torch tensor on device
|