# train_hand.py # Prosthetic hand gesture recognition — personal dataset import numpy as np import pandas as pd from scipy import signal from sklearn.metrics import classification_report, confusion_matrix from sklearn.utils.class_weight import compute_class_weight import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import matplotlib.pyplot as plt import seaborn as sns import os FS = 200 WIN_SAMPLES = 150 STEP = 75 N_CHANNELS = 8 N_CLASSES = 10 BATCH_SIZE = 128 EPOCHS = 100 DEVICE = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') GESTURE_NAMES = { 0: 'rest', 1: 'fist', 2: 'grasp', 3: 'index', 4: 'middle', 5: 'ring', 6: 'pinky', 7: 'thumb', 8: 'wrist_rotate_out', 9: 'wrist_rotate_in', } print(f"Device: {DEVICE}") # ── Load ── def load_sessions(sessions_dir="hand_module/sessions"): all_dfs = [] dirs = sorted([ d for d in os.listdir(sessions_dir) if os.path.isdir(f"{sessions_dir}/{d}") and os.path.exists(f"{sessions_dir}/{d}/emg_data.csv") ]) for i, d in enumerate(dirs): df = pd.read_csv(f"{sessions_dir}/{d}/emg_data.csv") df['session_id'] = i all_dfs.append(df) print(f" Session {i+1}: {len(df):,} samples — {d}") return pd.concat(all_dfs, ignore_index=True) print("\nLoading sessions...") df = load_sessions() df['block_id'] = (df['label'] != df['label'].shift()).cumsum() print(f"Total: {len(df):,} samples") print(f"Total blocks: {df['block_id'].nunique()}\n") for lbl, name in GESTURE_NAMES.items(): count = (df['label'] == lbl).sum() print(f" {name:<20}: {count:,}") # ── Preprocessing — global fixed normalization ── GLOBAL_STATS = {} def preprocess_global(df): EMG_COLS = [f'emg_{i}' for i in range(8)] emg_out = np.zeros((len(df), 8), dtype=np.float32) nyq = FS / 2 bb, aa = signal.butter(4, [20/nyq, 90/nyq], btype='band') bn, an = signal.iirnotch(50, Q=30, fs=FS) all_filtered = [] for sid in df['session_id'].unique(): mask = (df['session_id'] == sid).values emg = df.loc[mask, EMG_COLS].values.astype(np.float32) emg = signal.filtfilt(bb, aa, emg, axis=0) emg = signal.filtfilt(bn, an, emg, axis=0) emg_out[mask] = emg all_filtered.append(emg) all_concat = np.concatenate(all_filtered, axis=0) GLOBAL_STATS['mean'] = all_concat.mean(axis=0) GLOBAL_STATS['std'] = np.where( all_concat.std(axis=0) < 1e-8, 1e-8, all_concat.std(axis=0) ) emg_out = (emg_out - GLOBAL_STATS['mean']) / GLOBAL_STATS['std'] return emg_out print("\nPreprocessing...") emg_norm = preprocess_global(df) labels = df['label'].values.astype(np.int64) blocks = df['block_id'].values np.save('hand_module/models/hand_norm_mean.npy', GLOBAL_STATS['mean']) np.save('hand_module/models/hand_norm_std.npy', GLOBAL_STATS['std']) print(f" Saved normalization stats") # ── Windowing per block ── def extract_windows_per_block(emg, labels, blocks, win=WIN_SAMPLES, step=STEP): X, y, block_ids = [], [], [] for bid in np.unique(blocks): mask = blocks == bid e_blk = emg[mask] l_blk = labels[mask] if len(np.unique(l_blk)) != 1: continue lbl = l_blk[0] n = len(l_blk) i = 0 while i + win <= n: X.append(e_blk[i:i+win]) y.append(lbl) block_ids.append(bid) i += step return (np.array(X, dtype=np.float32), np.array(y, dtype=np.int64), np.array(block_ids)) print("Extracting windows...") X, y, block_ids = extract_windows_per_block(emg_norm, labels, blocks) print(f"Windows: {len(X):,} Shape: {X.shape}\n") for lbl, name in GESTURE_NAMES.items(): print(f" {name:<20}: {(y==lbl).sum():,}") # ── Split — window level per class (single session) ── np.random.seed(42) train_idx, test_idx = [], [] for lbl in range(N_CLASSES): lbl_idx = np.where(y == lbl)[0] np.random.shuffle(lbl_idx) n_test = max(1, int(len(lbl_idx) * 0.2)) test_idx.extend(lbl_idx[:n_test].tolist()) train_idx.extend(lbl_idx[n_test:].tolist()) train_idx = np.array(train_idx) test_idx = np.array(test_idx) X_train, y_train = X[train_idx], y[train_idx] X_test, y_test = X[test_idx], y[test_idx] print(f"\nTrain: {len(X_train):,} | Test: {len(X_test):,}") # ── Dataset ── class EMGDataset(Dataset): def __init__(self, X, y): self.X = torch.tensor(X.transpose(0, 2, 1), dtype=torch.float32) self.y = torch.tensor(y, dtype=torch.long) def __len__(self): return len(self.y) def __getitem__(self, i): return self.X[i], self.y[i] train_loader = DataLoader(EMGDataset(X_train, y_train), batch_size=BATCH_SIZE, shuffle=True, drop_last=True) test_loader = DataLoader(EMGDataset(X_test, y_test), batch_size=BATCH_SIZE, shuffle=False) # ── Model ── class EMG_CNN_LSTM(nn.Module): def __init__(self, n_channels=8, n_classes=10): super().__init__() self.cnn = nn.Sequential( nn.Conv1d(n_channels, 64, kernel_size=3, padding=1), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2), nn.Dropout(0.3), nn.Conv1d(128, 256, kernel_size=3, padding=1), nn.BatchNorm1d(256), nn.ReLU(), nn.MaxPool1d(2), nn.Dropout(0.3), ) self.lstm = nn.LSTM( input_size=256, hidden_size=128, num_layers=2, batch_first=True, dropout=0.3, bidirectional=True ) self.fc = nn.Sequential( nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.4), nn.Linear(128, n_classes) ) def forward(self, x): x = self.cnn(x) x = x.permute(0, 2, 1) x, _ = self.lstm(x) x = x[:, -1, :] return self.fc(x) # ── Training ── model = EMG_CNN_LSTM(n_classes=N_CLASSES).to(DEVICE) optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS) cw = compute_class_weight('balanced', classes=np.unique(y_train), y=y_train) criterion = nn.CrossEntropyLoss( weight=torch.tensor(cw, dtype=torch.float32).to(DEVICE), label_smoothing=0.05 ) print("\n" + "=" * 56) print(" TRAINING — Prosthetic Hand (10 gestures)") print("=" * 56) best_acc, best_epoch = 0.0, 0 train_losses, test_accs = [], [] for epoch in range(1, EPOCHS + 1): model.train() epoch_loss = 0 for xb, yb in train_loader: xb, yb = xb.to(DEVICE), yb.to(DEVICE) optimizer.zero_grad() loss = criterion(model(xb), yb) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() epoch_loss += loss.item() scheduler.step() model.eval() correct = total = 0 with torch.no_grad(): for xb, yb in test_loader: xb, yb = xb.to(DEVICE), yb.to(DEVICE) preds = model(xb).argmax(1) correct += (preds == yb).sum().item() total += len(yb) acc = correct / total avg_loss = epoch_loss / len(train_loader) if acc > best_acc: best_acc, best_epoch = acc, epoch torch.save(model.state_dict(), 'hand_module/models/best_model_hand.pt') train_losses.append(avg_loss) test_accs.append(acc) if epoch % 10 == 0 or epoch == 1: print(f" Epoch {epoch:3d}/{EPOCHS} | " f"Loss: {avg_loss:.4f} | " f"Acc: {acc:.3f} | " f"Best: {best_acc:.3f} (ep {best_epoch})") print(f"\n Best accuracy: {best_acc:.3f} at epoch {best_epoch}") # ── Evaluation ── model.load_state_dict(torch.load('hand_module/models/best_model_hand.pt')) model.eval() all_preds, all_true = [], [] with torch.no_grad(): for xb, yb in test_loader: preds = model(xb.to(DEVICE)).argmax(1).cpu().numpy() all_preds.extend(preds) all_true.extend(yb.numpy()) names = [GESTURE_NAMES[i] for i in range(N_CLASSES)] print("\n" + "=" * 56) print(" CLASSIFICATION REPORT — Prosthetic Hand") print("=" * 56) print(classification_report(all_true, all_preds, target_names=names)) cm = confusion_matrix(all_true, all_preds) plt.figure(figsize=(10, 8)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=names, yticklabels=names) plt.title(f'Confusion Matrix — Hand Model — Best Acc: {best_acc:.3f}') plt.ylabel('True'); plt.xlabel('Predicted') plt.tight_layout() plt.savefig('hand_module/results/confusion_matrix_hand.png', dpi=150) print(" Saved: hand_module/results/confusion_matrix_hand.png") # ── Training Curves ── fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) ax1.plot(train_losses); ax1.set_title('Training Loss'); ax1.set_xlabel('Epoch') ax2.plot(test_accs); ax2.set_title('Test Accuracy'); ax2.set_xlabel('Epoch') ax2.axhline(y=best_acc, color='r', linestyle='--', label=f'Best: {best_acc:.3f}') ax2.legend() plt.tight_layout() plt.savefig('hand_module/results/training_curves_hand.png', dpi=150) print(" Saved: hand_module/results/training_curves_hand.png") # ── Training Curves ── fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) ax1.plot(train_losses) ax1.set_title('Training Loss') ax1.set_xlabel('Epoch') ax2.plot(test_accs) ax2.set_title('Test Accuracy') ax2.set_xlabel('Epoch') ax2.axhline(y=best_acc, color='r', linestyle='--', label=f'Best: {best_acc:.3f}') ax2.legend() plt.tight_layout() plt.savefig('hand_module/results/training_curves_hand.png', dpi=150) print(" Saved: hand_module/results/training_curves_hand.png")