| import os |
| from typing import Tuple, Optional |
|
|
| import h5py |
| import numpy as np |
| import scipy.io |
| import scipy.signal as signal |
| from joblib import Parallel, delayed |
| from scipy.signal import iirnotch |
| from tqdm import tqdm |
|
|
| def sequence_to_seconds(seq_len: int, fs: float) -> float: |
| """Converts a sequence length in samples to time in seconds. |
| |
| Args: |
| seq_len (int): The number of samples in the sequence. |
| fs (float): The sampling frequency in Hz. |
| |
| Returns: |
| float: The duration of the sequence in seconds. |
| """ |
| return seq_len / fs |
|
|
|
|
| _MATRIX_DOF2DOA_TRANSPOSED = np.array( |
| |
| |
| |
| |
| [ |
| [+0.6390, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.3830, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.0000, +1.0000, +0.0000, +0.0000, +0.0000], |
| [-0.6390, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.4000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.6000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.4000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.6000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.1667], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.3333], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.1667], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.3333], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.0000], |
| [-0.1900, +0.0000, +0.0000, +0.0000, +0.0000], |
| [+0.0000, +0.0000, +0.0000, +0.0000, +0.0000], |
| ], |
| dtype=np.float32, |
| ) |
|
|
| MATRIX_DOF2DOA = _MATRIX_DOF2DOA_TRANSPOSED.T |
|
|
|
|
| def notch_filter(data: np.ndarray, notch_freq: float = 50.0, Q: float = 30.0, fs: float = 1111.0) -> np.ndarray: |
| """Applies a notch filter to every channel of the input data independently. |
| |
| Args: |
| data (np.ndarray): The input signal array of shape (T, D). |
| notch_freq (float, optional): The frequency to be removed. Defaults to 50.0. |
| Q (float, optional): The quality factor. Defaults to 30.0. |
| fs (float, optional): The sampling frequency in Hz. Defaults to 1111.0. |
| |
| Returns: |
| np.ndarray: The filtered signal array. |
| """ |
| 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: np.ndarray, |
| lowcut: float = 20.0, |
| highcut: float = 90.0, |
| fs: float = 2000.0, |
| order: int = 4 |
| ) -> np.ndarray: |
| """Applies a Butterworth bandpass filter to the EMG signal. |
| |
| Args: |
| emg (np.ndarray): The input signal array of shape (T, D). |
| lowcut (float, optional): Lower bound of the passband in Hz. Defaults to 20.0. |
| highcut (float, optional): Upper bound of the passband in Hz. Defaults to 90.0. |
| fs (float, optional): The sampling frequency in Hz. Defaults to 2000.0. |
| order (int, optional): The order of the filter. Defaults to 4. |
| |
| Returns: |
| np.ndarray: The filtered signal array. |
| """ |
| nyq = 0.5 * fs |
| b, a = signal.butter(order, [lowcut / nyq, highcut / nyq], btype="bandpass") |
| out = np.zeros_like(emg) |
| for ch in range(emg.shape[1]): |
| out[:, ch] = signal.filtfilt(b, a, emg[:, ch]) |
| return out |
|
|
|
|
| def sliding_window_segment( |
| emg: np.ndarray, |
| label: np.ndarray, |
| window_size: int, |
| stride: int |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| """Segments EMG and label data using a sliding window. |
| |
| Args: |
| emg (np.ndarray): The raw EMG data of shape (T, n_ch). |
| label (np.ndarray): The corresponding labels/targets. |
| window_size (int): Number of samples per window. |
| stride (int): Number of samples to shift between windows. |
| |
| Returns: |
| Tuple[np.ndarray, np.ndarray]: A tuple containing: |
| - segmented EMG tokens (N, window_size, n_ch). |
| - segmented label tokens (N, window_size, target_dim). |
| """ |
| segments, labels = [], [] |
| n_samples = len(label) |
|
|
| for start in range(0, n_samples - window_size + 1, stride): |
| end = start + window_size |
| emg_segment = emg[start:end] |
| label_segment = label[start:end] |
| segments.append(emg_segment) |
| labels.append(label_segment) |
|
|
| return np.array(segments), np.array(labels) |
|
|
|
|
| def process_mat_file( |
| mat_path: str, |
| window_size: int, |
| stride: int, |
| fs: float |
| ) -> Optional[Tuple[str, np.ndarray, np.ndarray]]: |
| """Processes a single NinaPro DB8 .mat file. |
| |
| Loads the file, removes NaNs, normalizes EMG (Z-score), maps finger degrees |
| of freedom (DoF) to degrees of activation (DoA), and segments the data. |
| |
| Args: |
| mat_path (str): Absolute path to the .mat file. |
| window_size (int): Temporal window size in samples. |
| stride (int): Stride between windows in samples. |
| fs (float): Sampling frequency in Hz. |
| |
| Returns: |
| Optional[Tuple[str, np.ndarray, np.ndarray]]: A tuple of (split_name, segments, labels) |
| if the file is valid, else None. |
| """ |
| mat = scipy.io.loadmat(mat_path) |
| emg = mat["emg"] |
| label = mat["glove"] |
|
|
| |
| valid = ~np.isnan(label).any(axis=1) |
| emg = emg[valid] |
| label = label[valid] |
|
|
| |
| mu = emg.mean(axis=0) |
| sd = emg.std(axis=0, ddof=1) |
| sd[sd == 0] = 1.0 |
| emg = (emg - mu) / sd |
|
|
| |
| y_doa = (MATRIX_DOF2DOA @ label.T).T |
|
|
| |
| segs, labs = sliding_window_segment(emg, y_doa, window_size, stride) |
|
|
| |
| fname = os.path.basename(mat_path) |
| if "_A1" in fname: |
| split = "train" |
| elif "_A2" in fname: |
| split = "val" |
| elif "_A3" in fname: |
| split = "test" |
| else: |
| return None |
|
|
| return split, segs, labs |
|
|
|
|
| def main(): |
| import argparse |
|
|
| args = argparse.ArgumentParser(description="Process EMG data from DB8.") |
| args.add_argument("--download_data", action="store_true") |
| args.add_argument("--data_dir", type=str, required=True) |
| args.add_argument("--save_dir", type=str, required=True) |
| args.add_argument( |
| "--seq_len", type=int, help="Size of the window in samples for segmentation." |
| ) |
| args.add_argument( |
| "--stride", |
| type=int, |
| help="Step size between windows in samples for segmentation.", |
| ) |
| args.add_argument( |
| "--n_jobs", type=int, default=-1, help="Number of parallel jobs to run." |
| ) |
| args = args.parse_args() |
| data_dir = args.data_dir |
| os.makedirs(args.save_dir, exist_ok=True) |
|
|
| |
| if args.download_data: |
| |
| len_data = range(1, 13) |
| base_url = "https://ninapro.hevs.ch/files/DB8/" |
| |
| for i in len_data: |
| url_a = f"{base_url}S{i}_E1_A1.mat" |
| url_b = f"{base_url}S{i}_E1_A2.mat" |
| url_c = f"{base_url}S{i}_E1_A3.mat" |
| os.system(f"wget -P {data_dir} {url_a}") |
| os.system(f"wget -P {data_dir} {url_b}") |
| os.system(f"wget -P {data_dir} {url_c}") |
| print( |
| f"Downloaded subject {i}\n{data_dir}/S{i}_E1_A1.mat and {data_dir}/S{i}_E1_A2.mat and {data_dir}/S{i}_E1_A3.mat" |
| ) |
|
|
| fs = 2000.0 |
| window_size, stride = args.seq_len, args.stride |
|
|
| window_seconds = sequence_to_seconds(window_size, fs) |
| print(f"Window size: {window_size} samples ({window_seconds:.2f} seconds)") |
|
|
|
|
| |
| mat_paths = [ |
| os.path.join(args.data_dir, f) |
| for f in sorted(os.listdir(args.data_dir)) |
| if f.endswith(".mat") |
| ] |
|
|
| |
| results = Parallel(n_jobs=min(os.cpu_count(), args.n_jobs), verbose=5)( |
| delayed(process_mat_file)(mp, window_size, stride, fs) |
| for mp in mat_paths |
| ) |
|
|
| |
| splits = {k: {"data": [], "label": []} for k in ("train", "val", "test")} |
| for out in tqdm(results, desc="Processing files", unit="file"): |
| if out is None: |
| continue |
| split, segs, labs = out |
| splits[split]["data"].append(segs) |
| splits[split]["label"].append(labs) |
|
|
| |
| for split, d in tqdm(splits.items(), desc="Saving splits", unit="split"): |
| if not d["data"]: |
| continue |
|
|
| X = np.concatenate(d["data"], axis=0) |
| y = np.concatenate(d["label"], axis=0) |
|
|
| |
| X = X.transpose(0, 2, 1) |
|
|
| print(f"Split: {split}, X shape: {X.shape}, y shape: {y.shape}") |
| |
| with h5py.File(os.path.join(args.save_dir, f"{split}.h5"), "w") as hf: |
| hf.create_dataset("data", data=X) |
| hf.create_dataset("label", data=y) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
