code/guided_test_hand.py

# guided_test_hand.py
# Real-time accuracy test for prosthetic hand model

import asyncio
import myo
from myo import ClassifierMode, EMGMode, IMUMode
import torch
import torch.nn as nn
import numpy as np
from scipy import signal
from collections import deque, Counter
import time
import json

FS          = 200
WIN_SAMPLES = 150

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',
}

GESTURE_INSTRUCTIONS = {
    0: 'Relax your hand completely — do not move anything',
    1: 'Close ALL fingers into a tight fist',
    2: 'Curl fingers into a C-shape — like holding a cup',
    3: 'Extend INDEX finger only — others closed',
    4: 'Extend MIDDLE finger only — others closed',
    5: 'Extend RING finger only — others closed',
    6: 'Extend PINKY finger only — others closed',
    7: 'Extend THUMB only — others closed',
    8: 'Rotate wrist so palm faces DOWN toward table',
    9: 'Rotate wrist so palm faces UP toward you',
}

TEST_SEQUENCE = [
    0, 1, 0, 2, 0, 3, 0, 4, 0, 5,
    0, 6, 0, 7, 0, 8, 0, 9, 0, 1,
    0, 3, 0, 5, 0, 7, 0, 2, 0, 4,
]

HOLD_SECONDS      = 5
COUNTDOWN_SECONDS = 3


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, 10)
        )
    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)


DEVICE = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
model  = EMG_CNN_LSTM().to(DEVICE)
model.load_state_dict(torch.load('hand_module/models/best_model_hand.pt',
                                  map_location=DEVICE))
model.eval()

NORM_MEAN = np.load('hand_module/models/hand_norm_mean.npy')
NORM_STD  = np.load('hand_module/models/hand_norm_std.npy')
print(f"✅ Model + normalization loaded")

nyq    = FS / 2
b,  a  = signal.butter(4, [20/nyq, 90/nyq], btype='band')
bn, an = signal.iirnotch(50, Q=30, fs=FS)


class State:
    emg_buffer      = deque(maxlen=WIN_SAMPLES)
    current_truth   = None
    predictions_log = []
    is_recording    = False

STATE = State()


def predict():
    if len(STATE.emg_buffer) < WIN_SAMPLES:
        return None, 0.0

    window = np.array(STATE.emg_buffer, dtype=np.float32)
    window = signal.filtfilt(b,  a,  window, axis=0)
    window = signal.filtfilt(bn, an, window, axis=0)
    window = (window - NORM_MEAN) / NORM_STD
    window = window.T.copy()

    x = torch.tensor(window, dtype=torch.float32).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        probs      = torch.softmax(model(x), dim=1)[0]
        confidence = probs.max().item()
        pred_label = probs.argmax().item()

    return pred_label, confidence


class TestClassifier(myo.MyoClient):

    async def on_emg_data(self, emg: myo.EMGData):
        for sample in [emg.sample1, emg.sample2]:
            STATE.emg_buffer.append(list(sample))

        if not STATE.is_recording:
            return

        pred_label, confidence = predict()
        if pred_label is None:
            return

        STATE.predictions_log.append({
            'truth': STATE.current_truth,
            'pred':  pred_label,
            'conf':  confidence,
        })

    async def on_imu_data(self, _): pass
    async def on_classifier_event(self, _): pass
    async def on_aggregated_data(self, _): pass
    async def on_emg_data_aggregated(self, _): pass
    async def on_fv_data(self, _): pass
    async def on_motion_event(self, _): pass


async def countdown(seconds, message):
    for i in range(seconds, 0, -1):
        print(f"\r  ⏳  {message}  —  {i}s   ", end='', flush=True)
        await asyncio.sleep(1)
    print(f"\r  ✅  GO!                                         ")


async def run_test():
    print("\n" + "═" * 64)
    print("  GUIDED TEST — Prosthetic Hand Model")
    print("═" * 64)
    print(f"\n  {len(TEST_SEQUENCE)} gestures  |  {HOLD_SECONDS}s each")
    print(f"  Keep your arm still — only hand/wrist moves\n")
    print("  Starting in 5 seconds...")
    await asyncio.sleep(5)

    all_results = []

    for idx, gesture_id in enumerate(TEST_SEQUENCE, 1):
        name        = GESTURE_NAMES[gesture_id]
        instruction = GESTURE_INSTRUCTIONS[gesture_id]

        print("\n" + "─" * 64)
        print(f"  [{idx}/{len(TEST_SEQUENCE)}]  {name.upper()}")
        print(f"  👉  {instruction}")

        await countdown(COUNTDOWN_SECONDS, f"Prepare for {name}")
        print(f"\n  🟢  Hold steady!\n")

        STATE.current_truth   = gesture_id
        STATE.predictions_log = []
        STATE.is_recording    = True

        start      = time.time()
        last_shown = None
        while time.time() - start < HOLD_SECONDS:
            await asyncio.sleep(0.1)
            if STATE.predictions_log:
                latest = STATE.predictions_log[-1]
                pred   = latest['pred']
                if pred != last_shown:
                    correct = "✅" if pred == gesture_id else "❌"
                    print(f"      {correct}  {GESTURE_NAMES[pred]:<20} "
                          f"(conf: {latest['conf']:.0%})")
                    last_shown = pred

        STATE.is_recording = False

        preds = [p['pred'] for p in STATE.predictions_log]
        if preds:
            correct_count = sum(1 for p in preds if p == gesture_id)
            acc           = correct_count / len(preds) * 100
            top3          = Counter(preds).most_common(3)
            print(f"\n      📊  Accuracy: {acc:.0f}%  ({correct_count}/{len(preds)})")
            print(f"      📊  Top predictions: "
                  f"{[(GESTURE_NAMES[k], v) for k,v in top3]}")

        all_results.append({'gesture': name, 'id': gesture_id, 'predictions': preds})

    # ── Final Summary ──
    print("\n" + "═" * 64)
    print("  FINAL SUMMARY")
    print("═" * 64)

    gesture_stats = {}
    for r in all_results:
        g  = r['gesture']
        gid = r['id']
        if g not in gesture_stats:
            gesture_stats[g] = {'correct': 0, 'total': 0, 'confusions': []}
        for p in r['predictions']:
            gesture_stats[g]['total'] += 1
            if p == gid:
                gesture_stats[g]['correct'] += 1
            else:
                gesture_stats[g]['confusions'].append(GESTURE_NAMES[p])

    print(f"\n  {'Gesture':<22} {'Accuracy':<12} {'Most Confused With'}")
    print("  " + "─" * 55)
    for g, stats in gesture_stats.items():
        acc = stats['correct'] / stats['total'] * 100 if stats['total'] else 0
        confusion = Counter(stats['confusions']).most_common(1)
        conf_str  = f"{confusion[0][0]} ({confusion[0][1]}x)" if confusion else "—"
        print(f"  {g:<22} {acc:>5.0f}%       {conf_str}")

    with open('hand_module/test_results_hand.json', 'w') as f:
        json.dump(all_results, f, indent=2)
    print(f"\n  💾  Saved: hand_module/test_results_hand.json")
    print("═" * 64)


async def main():
    print("🔍 Scanning for Myo Armband...")
    client = await TestClassifier.with_device()
    print(f"✅ Connected: {client.device.name}")

    await client.setup(
        classifier_mode=ClassifierMode.DISABLED,
        emg_mode=EMGMode.SEND_EMG,
        imu_mode=IMUMode.SEND_DATA,
    )
    await client.start()
    await run_test()
    await client.stop()
    await client.disconnect()


if __name__ == "__main__":
    asyncio.run(main())