File size: 4,718 Bytes
ca8e271 |
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 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import os
from pathlib import Path
import h5py
import numpy as np
import pandas as pd
import scipy.io
import scipy.signal as signal
from joblib import Parallel, delayed
from scipy.signal import iirnotch
from tqdm import tqdm
# ==== Filter functions (operate at original fs=2000) ====
def notch_filter(data, notch_freq=50.0, Q=30.0, fs=2000.0):
b, a = iirnotch(notch_freq, Q, fs)
out = np.zeros_like(data)
for ch in range(data.shape[1]):
out[:, ch] = signal.filtfilt(b, a, data[:, ch])
return out
def bandpass_filter_emg(emg, lowcut=20.0, highcut=90.0, fs=2000.0, order=4):
nyq = 0.5 * fs
low = lowcut / nyq
high = highcut / nyq
b, a = signal.butter(order, [low, high], btype="bandpass")
out = np.zeros_like(emg)
for c in range(emg.shape[1]):
out[:, c] = signal.filtfilt(b, a, emg[:, c])
return out
# ==== Window segmentation ====
def process_emg_features(emg, window_size=1000, stride=500):
segs, lbls = [], []
N = len(emg)
for start in range(0, N, stride):
end = start + window_size
if end > N: # skip the last segment if it is not complete
continue
win = emg[start:end]
segs.append(win)
return np.array(segs)
def process_one_recording(file_path, fs=2000.0, window_size=1000, stride=500):
"""
Process a single recording file to extract EMG features and labels
as to be used in the main pipeline with parallel processing.
"""
with h5py.File(file_path, "r") as f:
grp = f["emg2pose"]
data = grp["timeseries"]
emg = data["emg"][:].astype(np.float32)
# ==== Preprocessing EMG data ====
emg_filt = bandpass_filter_emg(emg, 20, 450, fs=fs)
emg_filt = notch_filter(emg_filt, 50, 30, fs=fs)
# z-score
mu = emg_filt.mean(axis=0)
sd = emg_filt.std(axis=0, ddof=1)
sd[sd == 0] = 1.0
emg_z = (emg_filt - mu) / sd
# segment
segs = process_emg_features(emg_z, window_size, stride)
return segs
# ==== Main pipeline ====
def main():
import argparse
args = argparse.ArgumentParser(description="Process EMG data from DB5.")
args.add_argument("--data_dir", type=str)
args.add_argument("--save_dir", type=str)
args.add_argument(
"--window_size", type=int, help="Size of the sliding window for segmentation."
)
args.add_argument(
"--stride", type=int, help="Stride for the sliding window segmentation."
)
args.add_argument(
"--subsample", type=float, default=1.0, help="Whether to subsample the data"
)
args.add_argument(
"--n_jobs",
type=int,
default=-1,
help="Number of parallel jobs to run. -1 means using all available cores.",
)
args.add_argument(
"--seed", type=int, default=42, help="Random seed for reproducibility."
)
args = args.parse_args()
data_dir = args.data_dir
save_dir = args.save_dir
os.makedirs(save_dir, exist_ok=True)
fs = 2000.0 # original sampling rate
window_size, stride = args.window_size, args.stride
df = pd.read_csv(os.path.join(data_dir, "metadata.csv"))
df = df.groupby("split").apply(
lambda x: (
x.sample(frac=args.subsample, random_state=args.seed)
if args.subsample < 1.0
else x
)
)
df.reset_index(drop=True, inplace=True)
splits = {}
for split, df_ in df.groupby("split"):
sessions = list(df_.filename)
splits[split] = [
Path(data_dir).expanduser().joinpath(f"{session}.hdf5")
for session in sessions
]
all_data = {"train": [], "val": [], "test": []}
for split, files in splits.items():
# Here we use joblib to parallelize the file processing, each file is processed independently as the task is embarrassingly parallel. We scale the processing across all available CPU cores since the number of files is around 25k (with training being 17k).
results = Parallel(n_jobs=args.n_jobs)(
delayed(process_one_recording)(file_path, fs, window_size, stride)
for file_path in tqdm(files, desc=f"Processing {split} files")
)
# Collect results
for segs in tqdm(results, desc=f"Collecting {split} data"):
all_data[split].append(segs)
# stack, augment train, transpose, save, and print stats
X = np.concatenate(all_data[split], axis=0) # [N, window_size, ch]
# transpose to [N, ch, window_size]
X = X.transpose(0, 2, 1)
# save
with h5py.File(os.path.join(save_dir, f"{split}.h5"), "w") as hf:
hf.create_dataset("data", data=X)
if __name__ == "__main__":
main()
|