refactor: improve Ninapro DB5 with progress indication for data augmentation

scripts/db5.py

@@ -6,6 +6,8 @@ import numpy as np
 import scipy.io
 import scipy.signal as signal
 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.
@@ -20,7 +22,9 @@ def sequence_to_seconds(seq_len: int, fs: float) -> float:
     return seq_len / fs
 
 
-def random_amplitude_scale(sig: np.ndarray, scale_range: Tuple[float, float] = (0.9, 1.1)) -> np.ndarray:
+def random_amplitude_scale(
+    sig: np.ndarray, scale_range: Tuple[float, float] = (0.9, 1.1)
+) -> np.ndarray:
     """Applies random amplitude scaling to the input signal.
 
     Args:
@@ -85,7 +89,9 @@ def augment_one_sample(seg: np.ndarray) -> np.ndarray:
     return out
 
 
-def augment_train_data(data: np.ndarray, labels: np.ndarray, factor: int = 3) -> Tuple[np.ndarray, np.ndarray]:
+def augment_train_data(
+    data: np.ndarray, labels: np.ndarray, factor: int = 3
+) -> Tuple[np.ndarray, np.ndarray]:
     """Augments the training dataset by creating multiple versions of each sample.
 
     Args:
@@ -104,7 +110,7 @@ def augment_train_data(data: np.ndarray, labels: np.ndarray, factor: int = 3) ->
     aug_segs = [data]
     aug_lbls = [labels]
     N = data.shape[0]
-    for i in range(N):
+    for i in tqdm(range(N), desc="Augmenting training data"):
         seg = data[i]  # [window_size, n_ch]
         lab = labels[i]
         for _ in range(factor):
@@ -115,7 +121,9 @@ def augment_train_data(data: np.ndarray, labels: np.ndarray, factor: int = 3) ->
     return new_data, new_labels
 
 
-def notch_filter(data: np.ndarray, notch_freq: float = 50.0, Q: float = 30.0, fs: float = 200.0) -> np.ndarray:
+def notch_filter(
+    data: np.ndarray, notch_freq: float = 50.0, Q: float = 30.0, fs: float = 200.0
+) -> np.ndarray:
     """Applies a notch filter to remove power line interference.
 
     Args:
@@ -140,7 +148,7 @@ def bandpass_filter_emg(
     lowcut: float = 20.0,
     highcut: float = 90.0,
     fs: float = 200.0,
-    order: int = 4
+    order: int = 4,
 ) -> np.ndarray:
     """Applies a Butterworth bandpass filter to the EMG signal.
 
@@ -169,7 +177,7 @@ def process_emg_features(
     label: np.ndarray,
     rerep: np.ndarray,
     window_size: int = 1024,
-    stride: int = 512
+    stride: int = 512,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Segments raw EMG signals into overlapping windows.