Update hand_module/train_hand.py — 2026-07-07 18:04

code/train_hand.py

@@ -0,0 +1,308 @@
+# 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")