osc_recorder_pythonosc.py

"""
Muse OSC Recorder — records sensor streams to CSV files.

Requirements: pip install python-osc
Usage: python osc_recorder_pythonosc.py [port] [prefix] [output_dir] [--debug]

Examples:
  python osc_recorder_pythonosc.py 7888
  python osc_recorder_pythonosc.py 7888 /muse
  python osc_recorder_pythonosc.py 7888 /muse ./data --debug
"""
import sys
import csv
import time
import os
import threading
from datetime import datetime
from pythonosc import dispatcher, osc_server

debug = '--debug' in sys.argv
args = [a for a in sys.argv[1:] if a != '--debug']
port = int(args[0]) if len(args) > 0 else 7777
prefix = args[1] if len(args) > 1 else ''
output_dir = args[2] if len(args) > 2 else '.'

STREAMS = ['eeg', 'ppg', 'optics', 'acc', 'gyro', 'drlref', 'thermistor']

session_dir = os.path.join(output_dir, datetime.now().strftime('%Y%m%d_%H%M%S'))
os.makedirs(session_dir, exist_ok=True)

files = {}
writers = {}
lock = threading.Lock()
start_time = time.time()
sample_count = 0
started = False


def handler(address, *args):
    global sample_count, started
    name = address.strip('/').split('/')[-1].lower()
    if name not in STREAMS:
        return

    with lock:
        if name not in writers:
            f = open(os.path.join(session_dir, f'{name}.csv'), 'w', newline='')
            w = csv.writer(f)
            w.writerow(['timestamp'] + [f'ch{i}' for i in range(len(args))])
            files[name] = f
            writers[name] = w
            if debug:
                print(f"  New stream: {address} ({len(args)} ch)")

        writers[name].writerow([round(time.time(), 6)] +
                               [round(v, 6) if isinstance(v, float) else v for v in args])
        sample_count += 1

        if not started:
            started = True
            print("Recording...")

        if sample_count % 500 == 0:
            for f in files.values():
                f.flush()
            if debug:
                elapsed = round(time.time() - start_time, 1)
                print(f"[{elapsed}s] {sample_count} samples")


disp = dispatcher.Dispatcher()
disp.map("/*", handler)

print(f"Listening on port {port} — waiting for data...")
if prefix:
    print(f"Prefix: {prefix}")

server = osc_server.ThreadingOSCUDPServer(("0.0.0.0", port), disp)
try:
    server.serve_forever()
except KeyboardInterrupt:
    pass
finally:
    server.shutdown()
    for f in files.values():
        f.close()
    print(f"\nSaved to: {session_dir} ({sample_count} samples)")