update db5 script with optional majority voting

scripts/db5.py

@@ -121,6 +121,49 @@ def augment_train_data(
     return new_data, new_labels
 
 
+def resample_to_rate(
+    emg: np.ndarray,
+    label: np.ndarray,
+    rerep: np.ndarray,
+    original_fs: float,
+    target_fs: float,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Resamples EMG and aligns sample-wise labels/repetition indices.
+
+    EMG is resampled with a polyphase anti-imaging filter. Labels and repetition
+    indices are mapped with nearest-neighbor sampling so they remain discrete.
+
+    Args:
+        emg (np.ndarray): EMG array of shape (T, D).
+        label (np.ndarray): Sample-wise labels of shape (T,).
+        rerep (np.ndarray): Sample-wise repetition indices of shape (T,).
+        original_fs (float): Original sampling frequency in Hz.
+        target_fs (float): Target sampling frequency in Hz.
+
+    Returns:
+        Tuple[np.ndarray, np.ndarray, np.ndarray]: Resampled EMG, labels, and
+        repetition indices.
+    """
+    if target_fs <= 0:
+        raise ValueError("target_fs must be positive")
+    if np.isclose(original_fs, target_fs):
+        return emg, label, rerep
+
+    from fractions import Fraction
+
+    ratio = Fraction(target_fs / original_fs).limit_denominator(1000)
+    up, down = ratio.numerator, ratio.denominator
+    emg_resampled = signal.resample_poly(emg, up, down, axis=0)
+
+    new_len = emg_resampled.shape[0]
+    source_idx = np.rint(
+        np.arange(new_len, dtype=np.float64) * original_fs / target_fs
+    ).astype(np.int64)
+    source_idx = np.clip(source_idx, 0, len(label) - 1)
+
+    return emg_resampled, label[source_idx], rerep[source_idx]
+
+
 def notch_filter(
     data: np.ndarray, notch_freq: float = 50.0, Q: float = 30.0, fs: float = 200.0
 ) -> np.ndarray:
@@ -178,6 +221,7 @@ def process_emg_features(
     rerep: np.ndarray,
     window_size: int = 1024,
     stride: int = 512,
+    majority: bool = False,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Segments raw EMG signals into overlapping windows.
 
@@ -196,18 +240,24 @@ def process_emg_features(
     """
     segs, lbls, reps = [], [], []
     N = len(label)
-    for start in range(0, N, stride):
+    for start in range(0, N - window_size + 1, stride):
         end = start + window_size
-        if end > N:
-            cut = emg[start:N]
-            pad = np.zeros((end - N, emg.shape[1]))
-            win = np.vstack([cut, pad])
+        win = emg[start:end]
+        window_labels = label[start:end]
+        window_reps = rerep[start:end]
+
+        if majority:
+            # Majority voting over gestures
+            labels_u, labels_c = np.unique(window_labels, return_counts=True)
+            assigned_label = int(labels_u[np.argmax(labels_c)])
         else:
-            win = emg[start:end]
+            # Pick index 0 for labe
+            assigned_label = window_labels[0]
+        assigned_rep = window_reps[0]
 
         segs.append(win)
-        lbls.append(label[start])
-        reps.append(rerep[start])
+        lbls.append(assigned_label)
+        reps.append(assigned_rep)
     return np.array(segs), np.array(lbls), np.array(reps)
 
 
@@ -242,6 +292,18 @@ def main():
         default=3,
         help="Number of augmented versions to create for each training sample.",
     )
+    args.add_argument(
+        "--resample-2khz",
+        action="store_true",
+        help=(
+            "If set, resample EMG from 200 Hz to 2000 Hz before segmentation. "
+            "Labels and repetition indices are aligned with nearest-neighbor mapping."
+        ),
+    )
+    args.add_argument(
+        "--majority",
+        action="store_true",
+    )
     args = args.parse_args()
 
     data_dir = args.data_dir
@@ -261,11 +323,17 @@ def main():
             os.system(f"rm {data_dir}/s{i}.zip")
             print(f"Downloaded and unzipped subject {i}\n{data_dir}/s{i}.zip")
 
-    fs = 200.0  # original sampling rate
+    original_fs = 200.0
+    target_fs = 2000.0 if args.resample_2khz else original_fs
     window_size, stride = args.seq_len, args.stride
 
-    window_seconds = sequence_to_seconds(window_size, fs)
+    window_seconds = sequence_to_seconds(window_size, target_fs)
+    print(
+        f"Sampling rate: {target_fs:.0f} Hz"
+        + (f" (resampled from {original_fs:.0f} Hz)" if args.resample_2khz else "")
+    )
     print(f"Window size: {window_size} samples ({window_seconds:.2f} seconds)")
+    print(f"{args.majority=}")
 
     train_reps = [1, 3, 4, 6]
     val_reps = [2]
@@ -313,16 +381,24 @@ def main():
                 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)
-            emg_filt = notch_filter(emg_filt, 50, 30, fs=fs)
+            # Filter at the original acquisition rate. Upsampling afterward does not
+            # create new frequency content, but provides a 2 kHz sample grid when requested.
+            emg_filt = bandpass_filter_emg(emg, 20, 90, fs=original_fs)
+            emg_filt = notch_filter(emg_filt, 50, 30, fs=original_fs)
+
+            if args.resample_2khz:
+                emg_filt, label, rerep = resample_to_rate(
+                    emg_filt, label, rerep, original_fs, target_fs
+                )
 
             # z-score
-            emg_z = (emg_filt - emg_filt.mean(axis=0)) / emg_filt.std(axis=0, ddof=1)
+            channel_std = emg_filt.std(axis=0, ddof=1)
+            channel_std[channel_std == 0] = 1.0
+            emg_z = (emg_filt - emg_filt.mean(axis=0)) / channel_std
 
             # segment
             segs, lbls, reps = process_emg_features(
-                emg_z, label, rerep, window_size, stride
+                emg_z, label, rerep, window_size, stride, args.majority
             )
 
             # split by repetition index