446 lines
16 KiB
Python
446 lines
16 KiB
Python
import numpy as np
|
||
import zmq
|
||
import threading
|
||
import json
|
||
import queue
|
||
import time
|
||
from Device.SunnyLinker import SunnyLinker64, RingBuffer
|
||
from collections import deque
|
||
|
||
|
||
class zmqServer(threading.Thread):
|
||
def __init__(self, host='0.0.0.0', cmd_port=8099, data_port=8100):
|
||
threading.Thread.__init__(self)
|
||
self.host = host
|
||
self.cmd_port = cmd_port
|
||
self.data_port = data_port
|
||
self.running = False
|
||
self.get_Impedance = False
|
||
self.open_Impedance = None
|
||
self.StartDecode = False
|
||
self.StartTrain = False
|
||
self.state_mode = None
|
||
self.currentLabel = -1
|
||
self.IsExitApp = False
|
||
self.getReport = False
|
||
self.daemon = True
|
||
|
||
# ZMQ Context
|
||
self.context = zmq.Context()
|
||
|
||
# 指令通道 (8099) - ROUTER
|
||
self.cmd_socket = self.context.socket(zmq.ROUTER)
|
||
self.cmd_socket.setsockopt(zmq.RCVHWM, 1000)
|
||
self.cmd_socket.setsockopt(zmq.SNDHWM, 1000)
|
||
self.cmd_socket.bind(f"tcp://{self.host}:{cmd_port}")
|
||
|
||
# 数据通道 (8100)) - ROUTER
|
||
self.data_socket = self.context.socket(zmq.ROUTER)
|
||
self.data_socket.setsockopt(zmq.RCVHWM, 1000)
|
||
self.data_socket.setsockopt(zmq.RCVTIMEO, 50)
|
||
self.data_socket.bind(f"tcp://{self.host}:{data_port}")
|
||
|
||
self.targetFreqs = []
|
||
self.changeTarget = False
|
||
self.sunnyLinker = SunnyLinker64(None, None, None, None, None)
|
||
self.labels = [0x01, 0x02, 0x03]
|
||
|
||
self.decoder_switch = False
|
||
self.decoder_class = None
|
||
self.cmd_clients = set()
|
||
self.data_clients = set()
|
||
self.send_queue = queue.Queue()
|
||
|
||
# ========== 数据缓冲区 (RingBuffer) ==========
|
||
# 与 SunnyLinker 保持一致，使用 RingBuffer
|
||
# 66 = 64 EEG通道 + 1 事件通道(第65) + 1 标签序号通道(第66)
|
||
# 缓存约 10 秒数据 (250Hz * 10s = 2500 点)
|
||
self.n_chan = 66
|
||
self.t_buffer = 10.0 # 缓冲区时长(秒)
|
||
self.__ringBuffer = RingBuffer(self.n_chan, int(self.t_buffer * 250))
|
||
|
||
# 事件检测相关
|
||
self._event_lock = threading.Lock()
|
||
self._epoch_finished = False
|
||
self._event_inner_idx = -1
|
||
self.pack_contain_event = False
|
||
self.predict_event = 99
|
||
self.events = [1, 2, self.predict_event]
|
||
self.count_events = {}
|
||
self.latency = 50
|
||
self.train_latency = 50
|
||
|
||
# 当前事件标签序号 (从第66通道获取)
|
||
self.current_label_index = 0
|
||
|
||
# 初始化标志
|
||
self._interval_inited = False
|
||
self._currentLabel = -1
|
||
|
||
# 注册的客户端（兼容旧接口）
|
||
self.clients = set()
|
||
|
||
# ========== 事件属性：线程安全访问 ==========
|
||
@property
|
||
def epoch_finished(self):
|
||
with self._event_lock:
|
||
return self._epoch_finished
|
||
|
||
@epoch_finished.setter
|
||
def epoch_finished(self, value):
|
||
with self._event_lock:
|
||
self._epoch_finished = value
|
||
|
||
@property
|
||
def event_inner_idx(self):
|
||
with self._event_lock:
|
||
return self._event_inner_idx
|
||
|
||
@event_inner_idx.setter
|
||
def event_inner_idx(self, value):
|
||
with self._event_lock:
|
||
self._event_inner_idx = value
|
||
|
||
@property
|
||
def interval_inited(self):
|
||
return self._interval_inited
|
||
|
||
@interval_inited.setter
|
||
def interval_inited(self, value):
|
||
self._interval_inited = value
|
||
|
||
@property
|
||
def currentLabel(self):
|
||
return self._currentLabel
|
||
|
||
@currentLabel.setter
|
||
def currentLabel(self, value):
|
||
self._currentLabel = value
|
||
|
||
def broadcast_message(self, method, params):
|
||
"""Put message into queue to be sent to all connected clients"""
|
||
self.send_queue.put((method, params))
|
||
|
||
# ========== 数据缓冲区操作接口 ==========
|
||
def GetDataLenCount(self):
|
||
"""返回缓冲区当前数据点数"""
|
||
return self.__ringBuffer.nUpdate
|
||
|
||
def getData(self, count):
|
||
"""获取最新count个数据点，不消费（只读）"""
|
||
with self.__ringBuffer.RingBufferLock:
|
||
count = min(count, self.__ringBuffer.nUpdate)
|
||
if count == 0:
|
||
return np.zeros((self.n_chan, 0))
|
||
|
||
# 计算读取范围（从尾部取最新数据）
|
||
read_end = (self.__ringBuffer.currentPtr - 1) % self.__ringBuffer.n_points
|
||
read_start = (read_end - count + 1) % self.__ringBuffer.n_points
|
||
|
||
if self.__ringBuffer.currentPtr == 0:
|
||
read_start = self.__ringBuffer.n_points - count
|
||
read_end = self.__ringBuffer.n_points - 1
|
||
|
||
if read_start <= read_end:
|
||
data = self.__ringBuffer.buffer[:, read_start:read_end + 1]
|
||
else:
|
||
part1 = self.__ringBuffer.buffer[:, read_start:]
|
||
part2 = self.__ringBuffer.buffer[:, :read_end + 1]
|
||
data = np.concatenate((part1, part2), axis=1)
|
||
|
||
return data
|
||
|
||
def consumeData(self, count):
|
||
"""消费(丢弃)指定数量的数据点，从头部移除"""
|
||
with self.__ringBuffer.RingBufferLock:
|
||
count = min(count, self.__ringBuffer.nUpdate)
|
||
self.__ringBuffer.readPtr = (self.__ringBuffer.readPtr + count) % self.__ringBuffer.n_points
|
||
self.__ringBuffer.nUpdate -= count
|
||
|
||
def ResetAll(self):
|
||
"""重置缓冲区"""
|
||
with self.__ringBuffer.RingBufferLock:
|
||
self.__ringBuffer.resetAllPara()
|
||
with self._event_lock:
|
||
self._epoch_finished = False
|
||
self._event_inner_idx = -1
|
||
self.pack_contain_event = False
|
||
self.count_events.clear()
|
||
self.current_label_index = 0
|
||
|
||
def reset_data_buffer(self):
|
||
self.ResetAll()
|
||
|
||
def reset_state(self):
|
||
self.ResetAll()
|
||
|
||
def interval_init(self, decoder_class):
|
||
"""初始化事件检测参数"""
|
||
import ast
|
||
from PubLibrary.InifileHelper import IniRead
|
||
|
||
if decoder_class == 'ssmvep':
|
||
interval_epoch = ast.literal_eval(IniRead('system', 'SSMVEP_IntervalEpoch'))
|
||
self.interval_epoch = [int(i * 250) for i in interval_epoch]
|
||
self.train_epoch = [int(self.interval_epoch[0]),
|
||
int(self.interval_epoch[1] + 0.1 * 250)]
|
||
self.latency = (self.interval_epoch[1] + 0.1 * 250) // 5
|
||
self.train_latency = (self.train_epoch[1] + 0.1 * 250) // 5
|
||
|
||
elif decoder_class == 'mi':
|
||
interval_epoch = ast.literal_eval(IniRead('system', 'MI_IntervalEpoch'))
|
||
self.interval_epoch = [int(i * 250) for i in interval_epoch]
|
||
self.train_epoch = self.interval_epoch.copy()
|
||
self.latency = self.interval_epoch[1] // 5
|
||
self.train_latency = self.latency
|
||
|
||
self.count_events = {}
|
||
self._event_inner_idx = -1
|
||
self._epoch_finished = False
|
||
self.pack_contain_event = False
|
||
self.predict_event = 99
|
||
self.events = [1, 2, self.predict_event]
|
||
self._interval_inited = True
|
||
|
||
# ========== 事件检测 ==========
|
||
def detect_event(self, data_matrix):
|
||
"""
|
||
检测事件通道中的触发信号
|
||
|
||
@param data_matrix: shape (66, N) - N个采样点的数据
|
||
第65行(索引64) = 事件通道
|
||
第66行(索引65) = 标签通道
|
||
@return: 是否检测到事件
|
||
"""
|
||
if data_matrix.shape[1] == 0:
|
||
return False
|
||
|
||
self.pack_contain_event = False
|
||
event_channel = data_matrix[64, :] # 第65通道 = 标签值(event值)
|
||
label_channel = data_matrix[65, :] # 第66通道 = 标签序号(label index)
|
||
|
||
events = event_channel.tolist()
|
||
|
||
with self._event_lock:
|
||
self._event_inner_idx = -1
|
||
self.current_event_label = 0
|
||
|
||
for idx, event in enumerate(events):
|
||
if int(event) in self.events:
|
||
self._event_inner_idx = idx
|
||
self.current_label_index = int(label_channel[idx])
|
||
self.pack_contain_event = True
|
||
|
||
new_key = f"{event}_{time.time()}"
|
||
latency = self.latency if event == self.predict_event else self.train_latency
|
||
self.count_events[new_key] = latency + 1
|
||
|
||
# 延迟计数递减
|
||
drop_items = []
|
||
for key, value in self.count_events.items():
|
||
value = value - 1
|
||
if value == 0:
|
||
drop_items.append(key)
|
||
self.count_events[key] = value
|
||
for key in drop_items:
|
||
del self.count_events[key]
|
||
|
||
if drop_items:
|
||
self._epoch_finished = True
|
||
# 检测到事件时，清除RingBuffer中之前的数据，只保留当前包
|
||
if self.pack_contain_event:
|
||
self.__ringBuffer.resetAllPara()
|
||
return True
|
||
|
||
self._epoch_finished = False
|
||
return False
|
||
|
||
def run(self):
|
||
self.running = True
|
||
print(f"Server running - CMD: {self.cmd_port}, DATA: {self.data_port}")
|
||
|
||
cmd_poller = zmq.Poller()
|
||
cmd_poller.register(self.cmd_socket, zmq.POLLIN)
|
||
|
||
data_poller = zmq.Poller()
|
||
data_poller.register(self.data_socket, zmq.POLLIN)
|
||
|
||
try:
|
||
while self.running:
|
||
# --- 处理发送队列 (指令通道) ---
|
||
while not self.send_queue.empty():
|
||
method, params = self.send_queue.get()
|
||
if self.cmd_clients:
|
||
try:
|
||
msg = {'method': method, 'params': params}
|
||
msg_bytes = json.dumps(msg).encode('utf-8')
|
||
for client_id in list(self.cmd_clients):
|
||
try:
|
||
self.cmd_socket.send_multipart([client_id, b'', msg_bytes])
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
|
||
# --- 处理指令通道 ---
|
||
socks = dict(cmd_poller.poll(10))
|
||
if self.cmd_socket in socks:
|
||
self._handle_cmd_socket()
|
||
|
||
# --- 处理数据通道 ---
|
||
socks = dict(data_poller.poll(10))
|
||
if self.data_socket in socks:
|
||
self._handle_data_socket()
|
||
|
||
except Exception as e:
|
||
print(f"Server error: {e}")
|
||
finally:
|
||
self.running = False
|
||
self.cmd_socket.close()
|
||
self.data_socket.close()
|
||
self.context.term()
|
||
|
||
def _handle_cmd_socket(self):
|
||
"""处理指令通道消息"""
|
||
try:
|
||
frames = self.cmd_socket.recv_multipart()
|
||
if len(frames) < 3:
|
||
return
|
||
ident, _, message_bytes = frames[:3]
|
||
self.cmd_clients.add(ident)
|
||
self.clients.add(ident)
|
||
|
||
message = json.loads(message_bytes.decode('utf-8'))
|
||
method = message.get("method")
|
||
params = message.get("params")
|
||
|
||
print(f"[CMD] {method}: {params}")
|
||
|
||
if method == "sync":
|
||
self.state_mode = 'sync'
|
||
elif method == "targetFreqs":
|
||
if isinstance(params, list) and params != self.targetFreqs:
|
||
self.targetFreqs = params
|
||
self.changeTarget = True
|
||
elif method == "decoderClass":
|
||
if isinstance(params, str) and params != self.decoder_class:
|
||
self.decoder_class = params
|
||
self.decoder_switch = True
|
||
elif method == "getReport":
|
||
self.getReport = True
|
||
elif method == "train":
|
||
self.state_mode = 'train'
|
||
self.StartTrain = True
|
||
self.currentLabel = params
|
||
elif method == "predict":
|
||
self.state_mode = 'predict'
|
||
if params == 1:
|
||
self.StartDecode = True
|
||
elif params == 2:
|
||
self.IsExitApp = True
|
||
self.running = False
|
||
elif method == "rest":
|
||
self.state_mode = 'rest'
|
||
elif method == "impedance":
|
||
if params == 1:
|
||
self.open_Impedance = True
|
||
self.get_Impedance = True
|
||
elif params == 2:
|
||
self.open_Impedance = False
|
||
self.get_Impedance = False
|
||
|
||
except Exception as e:
|
||
print(f"CMD socket error: {e}")
|
||
|
||
def _handle_data_socket(self):
|
||
"""处理数据通道消息 (EEG数据)
|
||
|
||
上位机数据格式:
|
||
- 数据帧: [identity, '', meta_json, data_buffer]
|
||
data_buffer = [N, 66] float32 -> 转置为 [66, N]
|
||
"""
|
||
try:
|
||
frames = self.data_socket.recv_multipart()
|
||
if len(frames) < 4:
|
||
return
|
||
ident, _, message_bytes = frames[:3]
|
||
self.data_clients.add(ident)
|
||
|
||
meta = json.loads(message_bytes.decode('utf-8'))
|
||
|
||
# data: [N, 66] -> 转置 -> [66, N]
|
||
raw_data = np.frombuffer(frames[3], dtype=np.float32)
|
||
n_samples, n_channels = meta.get('shape', [5, 66])
|
||
data_matrix = raw_data.reshape(n_samples, n_channels).T.astype(np.float32)
|
||
|
||
# 写入 RingBuffer
|
||
with self.__ringBuffer.RingBufferLock:
|
||
self.__ringBuffer.appendBuffer(data_matrix)
|
||
|
||
# 事件检测
|
||
self.detect_event(data_matrix)
|
||
|
||
except Exception as e:
|
||
print(f"DATA socket error: {e}")
|
||
|
||
# ========== 各范式数据访问接口 ==========
|
||
def get_MIData(self):
|
||
"""获取MI导联数据 (21通道 + 事件)"""
|
||
data = self.getData(self.GetDataLenCount())
|
||
rows_to_extract = [8, 15, 12, 14, 18, 23, 16, 59, 50, 58, 17, 45, 29, 11, 10, 19, 20, 61, 51, 60, 21, 64, 65]
|
||
row_to_select = np.array(rows_to_extract)
|
||
if data.shape[1] > 0:
|
||
return data[row_to_select, :]
|
||
return np.zeros((len(rows_to_extract), 0))
|
||
|
||
def get_SSMVEPData(self):
|
||
"""获取SSMVEP导联数据 (8通道 + 事件)"""
|
||
data = self.getData(self.GetDataLenCount())
|
||
rows_to_extract = [13, 3, 2, 46, 9, 54, 47, 55, 64, 65]
|
||
row_to_select = np.array(rows_to_extract)
|
||
if data.shape[1] > 0:
|
||
return data[row_to_select, :]
|
||
return np.zeros((len(rows_to_extract), 0))
|
||
|
||
def getDataViaSSVEP(self, count):
|
||
"""获取SSVEP数据 (8通道 + 事件)"""
|
||
data = self.getData(count)
|
||
rows_to_extract = [13, 3, 2, 46, 9, 54, 47, 55, 64]
|
||
row_to_select = np.array(rows_to_extract)
|
||
if data.shape[1] > 0:
|
||
return data[row_to_select, :]
|
||
return np.zeros((len(rows_to_extract), 0))
|
||
|
||
def get_concentrateData(self, count):
|
||
"""获取专注力数据 (2通道)"""
|
||
data = self.getData(count)
|
||
rows_to_extract = [0, 1]
|
||
row_to_select = np.array(rows_to_extract)
|
||
if data.shape[1] > 0:
|
||
return data[row_to_select, :]
|
||
return np.zeros((len(rows_to_extract), 0))
|
||
|
||
def get_blinkData(self, count):
|
||
"""获取眨眼数据 (2通道)"""
|
||
data = self.getData(count)
|
||
rows_to_extract = [0, 1]
|
||
row_to_select = np.array(rows_to_extract)
|
||
if data.shape[1] > 0:
|
||
return data[row_to_select, :]
|
||
return np.zeros((len(rows_to_extract), 0))
|
||
|
||
def getImpedance(self, data, decoder_class):
|
||
"""计算阻抗（ZMQ模式下不可用）"""
|
||
return np.zeros(8)
|
||
|
||
def stop(self):
|
||
self.running = False
|
||
self.cmd_socket.close()
|
||
self.data_socket.close()
|
||
self.context.term()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
server = zmqServer()
|
||
server.start()
|