refactor: add data augmentation options for NinaPro DB5 processing and only 41 gestures setup

scripts/README.md

@@ -27,11 +27,13 @@ This guide provides commands to process raw EMG data into HDF5 format using slid
 
 | Dataset | Metric | Seq Len | Stride | Command |
 | :--- | :--- | :--- | :--- | :--- |
-| **NinaPro DB5** | Gesture | 200 (1s) | 50 | `python scripts/db5.py --data_dir $DATA_PATH/ninapro/DB5/ --save_dir $DATA_PATH/ninapro/DB5/h5_1sec/ --seq_len 200 --stride 50` |
-| **NinaPro DB5** | Gesture | 1000 (5s) | 250 | `python scripts/db5.py --data_dir $DATA_PATH/ninapro/DB5/ --save_dir $DATA_PATH/ninapro/DB5/h5_5sec/ --seq_len 1000 --stride 250` |
+| **NinaPro DB5** | Gesture | 200 (1s) | 50 | `python scripts/db5.py --data_dir $DATA_PATH/ninapro/DB5/ --save_dir $DATA_PATH/ninapro/DB5/h5_1sec/ --seq_len 200 --stride 50 --data-augment` |
+| **NinaPro DB5** | Gesture | 1000 (5s) | 250 | `python scripts/db5.py --data_dir $DATA_PATH/ninapro/DB5/ --save_dir $DATA_PATH/ninapro/DB5/h5_5sec/ --seq_len 1000 --stride 250 --data-augment` |
 | **EMG-EPN612** | Gesture | 200 (1s) | N/A | `python scripts/epn.py --data_dir $DATA_PATH/EPN612/ --source_training $DATA_PATH/EPN612/trainingJSON/ --source_testing $DATA_PATH/EPN612/testingJSON/ --dest_dir $DATA_PATH/EPN612/h5_1sec/ --seq_len 200` |
 | **EMG-EPN612** | Gesture | 1000 (5s) | N/A | `python scripts/epn.py --data_dir $DATA_PATH/EPN612/ --source_training $DATA_PATH/EPN612/trainingJSON/ --source_testing $DATA_PATH/EPN612/testingJSON/ --dest_dir $DATA_PATH/EPN612/h5_5sec/ --seq_len 1000` |
 | **UCI EMG** | Gesture | 200 (1s) | 50 | `python scripts/uci.py --data_dir $DATA_PATH/UCI_EMG/EMG_data_for_gestures-master/ --save_dir $DATA_PATH/UCI_EMG/EMG_data_for_gestures-master/h5_1sec/ --seq_len 200 --stride 50` |
 | **UCI EMG** | Gesture | 1000 (5s) | 250 | `python scripts/uci.py --data_dir $DATA_PATH/UCI_EMG/EMG_data_for_gestures-master/ --save_dir $DATA_PATH/UCI_EMG/EMG_data_for_gestures-master/h5_5sec/ --seq_len 1000 --stride 250` |
 | **NinaPro DB8** | Regression | 200 (0.1s) | 200 | `python scripts/db8.py --data_dir $DATA_PATH/ninapro/DB8/ --save_dir $DATA_PATH/ninapro/DB8/h5_100/ --seq_len 200 --stride 200` |
 | **NinaPro DB8** | Regression | 1000 (0.5s) | 1000 | `python scripts/db8.py --data_dir $DATA_PATH/ninapro/DB8/ --save_dir $DATA_PATH/ninapro/DB8/h5_500/ --seq_len 1000 --stride 1000` |
+
+>Note: For DB5, we used the `--data-augment` flag to augment the training data by a factor of 3 (see `--augment-factor` in `scripts/db5.py`).

scripts/db5.py

@@ -226,6 +226,22 @@ def main():
         type=int,
         help="Step size between windows in samples for segmentation.",
     )
+    args.add_argument(
+        "--only-41",
+        action="store_true",
+        help="If set, only the first 41 gestures (common between DB5, DB7) will be processed.",
+    )
+    args.add_argument(
+        "--data-augment",
+        action="store_true",
+        help="If set, data augmentation will be applied to the training set.",
+    )
+    args.add_argument(
+        "--augment-factor",
+        type=int,
+        default=3,
+        help="Number of augmented versions to create for each training sample.",
+    )
     args = args.parse_args()
 
     data_dir = args.data_dir
@@ -275,10 +291,27 @@ def main():
             rerep = dd["rerepetition"].ravel().astype(int)
 
             # label shift by exercise
-            if "E2" in mat:
-                label = np.where(label != 0, label + 12, 0)
-            elif "E3" in mat:
-                label = np.where(label != 0, label + 29, 0)
+            # E1: 12 gestures [1-12]
+            # E2: 17 gestures [1-17]
+            # E3: 23 gestures [1-23]
+
+            if args.only_41:
+                # Use only E2 and E3 and remap them to labels 1..17 and 18..40
+                # while keeping rest as label 0.
+                if "E1" in mat:
+                    continue
+                if "E2" in mat:
+                    label = np.where(label != 0, label, 0)
+                elif "E3" in mat:
+                    label = np.where(label != 0, label + 17, 0)
+                else:
+                    continue
+            else:
+                # Flatten the label space to avoid overlap across exercises
+                if "E2" in mat:
+                    label = np.where(label != 0, label + 12, 0)
+                elif "E3" in mat:
+                    label = np.where(label != 0, label + 29, 0)
 
             # filtering at original 200 Hz
             emg_filt = bandpass_filter_emg(emg, 20, 90, fs=fs)
@@ -310,8 +343,8 @@ def main():
         X = np.stack(all_data[split], axis=0)  # [N, window_size, ch]
         y = np.array(all_lbls[split], dtype=int)
 
-        if split == "train":
-            X, y = augment_train_data(X, y, factor=3)
+        if split == "train" and args.data_augment:
+            X, y = augment_train_data(X, y, factor=args.augment_factor)
 
         # transpose to [N, ch, window_size]
         X = X.transpose(0, 2, 1)