code/collect_hand_data.py

# collect_hand_data.py
# Guided EMG data collection for prosthetic hand control
# 10 gestures x 10 rounds x 5 seconds

import asyncio
import myo
from myo import ClassifierMode, EMGMode, IMUMode
import csv
import os
import time
from datetime import datetime

FS = 200

GESTURES = {
    0: {
        'name': 'rest',
        'instruction': 'Rest your hand completely flat — do not move anything',
    },
    1: {
        'name': 'fist',
        'instruction': 'Close ALL fingers into a tight fist',
    },
    2: {
        'name': 'grasp',
        'instruction': 'Curl fingers into a C-shape — like holding a cup or bottle. NOT a full fist, leave space in the middle',
    },
    3: {
        'name': 'index',
        'instruction': 'Extend INDEX finger only — keep all others closed in a fist',
    },
    4: {
        'name': 'middle',
        'instruction': 'Extend MIDDLE finger only — keep all others closed in a fist',
    },
    5: {
        'name': 'ring',
        'instruction': 'Extend RING finger only — keep all others closed. Take your time, this is hard',
    },
    6: {
        'name': 'pinky',
        'instruction': 'Extend PINKY finger only — keep all others closed in a fist',
    },
    7: {
        'name': 'thumb',
        'instruction': 'Extend THUMB only — keep all other fingers closed in a fist',
    },
    8: {
        'name': 'wrist_rotate_out',
        'instruction': 'Rotate wrist so palm faces DOWN toward the table — arm stays still',
    },
    9: {
        'name': 'wrist_rotate_in',
        'instruction': 'Rotate wrist so palm faces UP toward you — arm stays still',
    },
}

ROUNDS        = 10
HOLD_SECONDS  = 5
REST_SECONDS  = 4
COUNTDOWN     = 4

class State:
    emg_buffer = []
    is_recording = False
    sample_count = 0

STATE = State()


class Collector(myo.MyoClient):
    async def on_emg_data(self, emg: myo.EMGData):
        for sample in [emg.sample1, emg.sample2]:
            if STATE.is_recording:
                STATE.emg_buffer.append(list(sample))
                STATE.sample_count += 1

    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

def save_session(session_dir, all_data):
    os.makedirs(session_dir, exist_ok=True)
    filepath = f"{session_dir}/emg_data.csv"

    with open(filepath, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow([
            'emg_0', 'emg_1', 'emg_2', 'emg_3',
            'emg_4', 'emg_5', 'emg_6', 'emg_7',
            'label', 'gesture', 'timestamp'
        ])
        for row in all_data:
            writer.writerow(row)

    print(f"\n  Saved {len(all_data):,} samples → {filepath}\n")
    return filepath

async def countdown_display(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_collection():
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    session_dir = f"hand_module/sessions/session_{timestamp}"

    print("\n" + "═" * 64)
    print("  PROSTHETIC HAND — EMG DATA COLLECTION")
    print("═" * 64)
    print(f"\n  Gestures : {len(GESTURES)}")
    print(f"  Rounds   : {ROUNDS} per gesture")
    print(f"  Hold     : {HOLD_SECONDS}s each")
    print(f"  Rest     : {REST_SECONDS}s between gestures")
    print(f"\n  Place your arm comfortably on the table.")
    print(f"  Only your hand/wrist moves — keep your arm still.\n")
    print("  Starting in 8 seconds — get ready!")
    await asyncio.sleep(8)

    all_data    = []
    total_steps = len(GESTURES) * ROUNDS
    step        = 0

    for gesture_id, gesture_info in GESTURES.items():
        name        = gesture_info['name']
        instruction = gesture_info['instruction']

        print(f"\n{'═'*64}")
        print(f"  GESTURE: {name.upper()}")
        print(f"  {instruction}")
        print(f"{'═'*64}")
        print(f"  Study this gesture — starting in {COUNTDOWN} seconds...")

        await countdown_display(COUNTDOWN, f"Prepare for {name.upper()}")

        for round_num in range(1, ROUNDS + 1):
            step += 1
            progress = f"[{step}/{total_steps}]"

            print(f"\n  {progress}  Round {round_num}/{ROUNDS}  —  {name.upper()}")
            print(f"  👉  {instruction}\n")

            # countdown قبل التسجيل
            await countdown_display(3, "Get ready")

            # ابدأ التسجيل
            print(f"  🟢  RECORDING — hold steady for {HOLD_SECONDS} seconds!\n")
            STATE.emg_buffer   = []
            STATE.sample_count = 0
            STATE.is_recording = True

            start = time.time()
            while time.time() - start < HOLD_SECONDS:
                elapsed  = time.time() - start
                progress_bar = int(elapsed / HOLD_SECONDS * 20)
                bar = '█' * progress_bar + '░' * (20 - progress_bar)
                print(f"\r  [{bar}] {elapsed:.1f}s / {HOLD_SECONDS}s  "
                      f"({STATE.sample_count} samples)",
                      end='', flush=True)
                await asyncio.sleep(0.1)

            STATE.is_recording = False
            print()

            # احفظ الـ samples مع الـ label
            ts = datetime.now().isoformat()
            for sample in STATE.emg_buffer:
                row = sample + [gesture_id, name, ts]
                all_data.append(row)

            samples = len(STATE.emg_buffer)
            print(f"  ✅  Captured {samples} samples "
                  f"({samples/FS:.1f}s)")

            # راحة بين الـ rounds (إلا آخر round في كل gesture)
            if round_num < ROUNDS:
                print(f"\n  😌  Rest...")
                await countdown_display(REST_SECONDS, "Relax your hand")

        # راحة أطول بين الـ gestures
        if gesture_id < len(GESTURES) - 1:
            print(f"\n  💤  Gesture complete! Rest for 6 seconds before next gesture.")
            await countdown_display(6, "Relax completely")

    # احفظ كل الداتا
    filepath = save_session(session_dir, all_data)

    # ملخص
    print("\n" + "═" * 64)
    print("  SESSION COMPLETE")
    print("═" * 64)

    print(f"\n  {'Gesture':<20} {'Samples':<12} {'Duration'}")
    print("  " + "─" * 45)
    for gesture_id, info in GESTURES.items():
        name    = info['name']
        samples = sum(1 for r in all_data if r[8] == gesture_id)
        dur     = samples / FS
        print(f"  {name:<20} {samples:<12,} {dur:.1f}s")

    total = len(all_data)
    print(f"\n  Total: {total:,} samples ({total/FS:.0f}s)")
    print(f"  Saved: {filepath}")
    print("═" * 64)


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

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

    try:
        await run_collection()
    except KeyboardInterrupt:
        print("\n\n  Interrupted — saving collected data so far...")
    finally:
        await client.stop()
        await client.disconnect()
        print("  Done.")


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