code/quick_calibration.py

# This is the best model it worked perfect for me and safer after callibrration
# quick_calibration.py
# Fine-tune last layer only for new user — 2 min recording + 30s training

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

# ── Config ──
FS          = 200
WIN_SAMPLES = 150
STEP        = 75
N_CHANNELS  = 8
N_CLASSES   = 10
DEVICE      = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

MODEL_PATH    = "hand_module/models/best_model_hand.pt"
NORM_MEAN     = np.load("hand_module/models/hand_norm_mean.npy")
NORM_STD      = np.load("hand_module/models/hand_norm_std.npy")

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',
    1: 'Close ALL fingers into a tight fist',
    2: 'Curl fingers — like holding a cup',
    3: 'Extend INDEX finger only',
    4: 'Extend MIDDLE finger only',
    5: 'Extend RING finger only',
    6: 'Extend PINKY finger only',
    7: 'Extend THUMB only',
    8: 'Rotate wrist — palm faces DOWN',
    9: 'Rotate wrist — palm faces UP',
}

CALIBRATION_REPS  = 3
HOLD_SECONDS      = 5
COUNTDOWN_SECONDS = 3
FINETUNE_EPOCHS   = 30


# ── 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)


# ── Load Model ──
model = EMG_CNN_LSTM(N_CHANNELS, N_CLASSES).to(DEVICE)
model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.eval()
print(f"✅ Model loaded — Device: {DEVICE}")

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


# ── State ──
class State:
    emg_buffer   = deque(maxlen=WIN_SAMPLES)
    is_recording = False
    recorded_emg = []
    calibrated   = False
    pred_history = deque(maxlen=5)
    last_pred    = 0
    last_print   = 0

STATE = State()


def preprocess(window):
    window = signal.filtfilt(b,  a,  window, axis=0)
    window = signal.filtfilt(bn, an, window, axis=0)
    window = (window - NORM_MEAN) / NORM_STD
    return window


def predict(window):
    w = preprocess(window.copy())
    x = torch.tensor(w.T.copy(), dtype=torch.float32).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        probs = torch.softmax(model(x), dim=1)[0]
        conf  = probs.max().item()
        pred  = probs.argmax().item()
    return pred, conf


# ── Myo Client ──
class CalibrationClient(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 STATE.is_recording:
                STATE.recorded_emg.append(list(sample))

    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, msg):
    for i in range(seconds, 0, -1):
        print(f"\r  ⏳  {msg}  —  {i}s   ", end='', flush=True)
        await asyncio.sleep(1)
    print(f"\r  ✅  GO!                              ")


async def calibrate():
    print("\n" + "═"*60)
    print("  QUICK CALIBRATION")
    print("═"*60)
    print(f"\n  {len(GESTURE_NAMES)} gestures × {CALIBRATION_REPS} reps × {HOLD_SECONDS}s")
    print(f"  Total recording: ~{len(GESTURE_NAMES)*CALIBRATION_REPS*8//60} minutes")
    print(f"  Fine-tuning: ~30 seconds\n")
    print("  Starting in 5 seconds...")
    await asyncio.sleep(5)

    all_X, all_y = [], []

    for gesture_id in range(N_CLASSES):
        name        = GESTURE_NAMES[gesture_id]
        instruction = GESTURE_INSTRUCTIONS[gesture_id]

        print(f"\n{'─'*60}")
        print(f"  GESTURE: {name.upper()}")
        print(f"  {instruction}")

        for rep in range(1, CALIBRATION_REPS + 1):
            print(f"\n  Rep {rep}/{CALIBRATION_REPS}")
            await countdown(COUNTDOWN_SECONDS, f"Prepare for {name}")
            print(f"  🟢  HOLD STEADY!\n")

            STATE.recorded_emg = []
            STATE.is_recording = True

            start = time.time()
            while time.time() - start < HOLD_SECONDS:
                await asyncio.sleep(0.1)
                elapsed = time.time() - start
                bar     = '█' * int(elapsed/HOLD_SECONDS*20) + '░' * (20-int(elapsed/HOLD_SAMPLES*20)) if False else ''
                print(f"\r  Recording... {elapsed:.1f}s/{HOLD_SECONDS}s  "
                      f"({len(STATE.recorded_emg)} samples)",
                      end='', flush=True)

            STATE.is_recording = False
            print()

            emg = np.array(STATE.recorded_emg, dtype=np.float32)
            if len(emg) < WIN_SAMPLES:
                continue

            # Extract windows
            j = 0
            while j + WIN_SAMPLES <= len(emg):
                window = preprocess(emg[j:j+WIN_SAMPLES].copy())
                all_X.append(window.T.copy())
                all_y.append(gesture_id)
                j += STEP

            print(f"  ✅  {len(all_X)} total windows collected")
            await asyncio.sleep(1)

    # ── Fine-tune last layer only ──
    print(f"\n{'═'*60}")
    print(f"  FINE-TUNING on your data...")
    print(f"  Windows: {len(all_X)}")

    # Freeze all layers except last fc layer
    for param in model.parameters():
        param.requires_grad = False
    for param in model.fc[-1].parameters():
        param.requires_grad = True

    model.train()

    X_tensor = torch.tensor(np.array(all_X), dtype=torch.float32).to(DEVICE)
    y_tensor = torch.tensor(np.array(all_y), dtype=torch.long).to(DEVICE)

    optimizer = torch.optim.Adam(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=1e-3
    )
    criterion = nn.CrossEntropyLoss()

    dataset = torch.utils.data.TensorDataset(X_tensor, y_tensor)
    loader  = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

    for epoch in range(1, FINETUNE_EPOCHS + 1):
        epoch_loss = 0
        correct = total = 0
        for xb, yb in loader:
            optimizer.zero_grad()
            out  = model(xb)
            loss = criterion(out, yb)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
            correct    += (out.argmax(1) == yb).sum().item()
            total      += len(yb)

        if epoch % 10 == 0:
            acc = correct / total
            print(f"  Epoch {epoch:2d}/{FINETUNE_EPOCHS} | "
                  f"Loss: {epoch_loss/len(loader):.4f} | Acc: {acc:.3f}")

    model.eval()
    STATE.calibrated = True

    print(f"\n  ✅  Calibration complete!")
    print(f"  Model fine-tuned on YOUR data")
    print("═"*60)


async def realtime():
    print("\n" + "═"*60)
    print("  REAL-TIME INFERENCE")
    print("  Try any gesture!")
    print("  Press Ctrl+C to stop")
    print("═"*60 + "\n")

    count = 0
    while True:
        await asyncio.sleep(0.05)
        count += 1
        if count % 10 != 0:
            continue
        if len(STATE.emg_buffer) < WIN_SAMPLES:
            continue

        window     = np.array(STATE.emg_buffer, dtype=np.float32)
        pred, conf = predict(window)

        STATE.pred_history.append(pred)
        top_pred = Counter(STATE.pred_history).most_common(1)[0][0]

        now = time.time()
        if top_pred != STATE.last_pred or (now - STATE.last_print) > 1.5:
            name = GESTURE_NAMES[top_pred]
            print(f"\r  🖐  {name:<22} (conf: {conf:.0%})        ",
                  end='', flush=True)
            STATE.last_pred  = top_pred
            STATE.last_print = now


async def main():
    print("🔍 Scanning for Myo Armband...")
    client = await CalibrationClient.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()

    try:
        await calibrate()
        await realtime()
    except (KeyboardInterrupt, EOFError):
        pass
    finally:
        print("\n\n  Stopping...")
        await client.stop()
        await client.disconnect()


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