qlsdk2 0.5.1__py3-none-any.whl → 0.6.0__py3-none-any.whl

qlsdk/rsc/device/base.py

@@ -1,4 +1,3 @@
-
 from multiprocessing import Queue
 from threading import Thread
 from time import sleep, time_ns
@@ -6,14 +5,30 @@ from typing import Any, Dict, Literal
 
 from loguru import logger
 import numpy as np
-from qlsdk.core.entity import RscPacket
+from qlsdk.core.entity import RscPacket, ImpedancePacket
 from qlsdk.core.utils import to_bytes
-from qlsdk.persist import RscEDFHandler
 from qlsdk.rsc.interface import IDevice, IParser
-from qlsdk.rsc.parser import TcpMessageParser
-from qlsdk.rsc.command import StartImpedanceCommand, StopImpedanceCommand, StartStimulationCommand, StopStimulationCommand, SetAcquisitionParamCommand, StartAcquisitionCommand, StopAcquisitionCommand
+from qlsdk.rsc.command import SetImpedanceParamCommand, StartImpedanceCommand, StopImpedanceCommand, StartStimulationCommand, StopStimulationCommand, SetAcquisitionParamCommand, StartAcquisitionCommand, StopAcquisitionCommand
+from qlsdk.rsc.paradigm import StimulationParadigm
+from qlsdk.rsc.parser.base import TcpMessageParser
 
+    
+def intersection_positions(A, B):
+    setB = set(B)
+    seen = set()
+    return [idx for idx, elem in enumerate(A) 
+            if elem in setB and elem not in seen and not seen.add(elem)]    
+    
 class QLBaseDevice(IDevice):
+    
+    __TRIGGER_MAPPING = {
+        0x3E8: "Start of stimulation",
+        0x3E9: "End of stimulation",
+        0x3EA: "Ascending end of stimulation",
+        0x3EB: "Descending start of stimulation ",
+        0x3EC: "刺激参数有误",
+        0x3ED: "End of stimulation (by force)",
+    }    
     def __init__(self, socket):
         self.socket = socket
         
@@ -63,52 +78,11 @@ class QLBaseDevice(IDevice):
             "channels": [],
         }
         
-        self._stim_param = {
-            "stim_type": 0,                 # 刺激类型：0-所有通道参数相同, 1: 通道参数不同
-            "channels": [],
-            "param": [{
-                    "channel_id": 0,        #通道号 从0开始                          -- 必填
-                    "waveform": 3,          #波形类型：0-直流，1-交流 2-方波 3-脉冲   -- 必填
-                    "current": 1,           #电流强度(mA)                            -- 必填
-                    "duration": 30,         #平稳阶段持续时间(s)                      -- 必填
-                    "ramp_up": 5,           #上升时间(s) 默认0
-                    "ramp_down": 5,         #下降时间(s) 默认0
-                    "frequency": 500,       #频率(Hz) -- 非直流必填
-                    "phase_position": 0,    #相位 -- 默认0
-                    "duration_delay": "0",  #延迟启动时间(s) -- 默认0
-                    "pulse_width": 0,       #脉冲宽度(us) -- 仅脉冲类型电流有效， 默认100us
-                },
-                {
-                    "channel_id": 1,        #通道号 从0开始                          -- 必填
-                    "waveform": 3,          #波形类型：0-直流，1-交流 2-方波 3-脉冲   -- 必填
-                    "current": 1,           #电流强度(mA)                            -- 必填
-                    "duration": 30,         #平稳阶段持续时间(s)                      -- 必填
-                    "ramp_up": 5,           #上升时间(s) 默认0
-                    "ramp_down": 5,         #下降时间(s) 默认0
-                    "frequency": 500,       #频率(Hz) -- 非直流必填
-                    "phase_position": 0,    #相位 -- 默认0
-                    "duration_delay": "0",  #延迟启动时间(s) -- 默认0
-                    "pulse_width": 0,       #脉冲宽度(us) -- 仅脉冲类型电流有效， 默认100us
-                }
-            ]
-        }
+        self._stim_param = None
         
-        self.stim_paradigm = None
-                      
-        signal_info = {
-            "param" : None,
-            "start_time" : None,
-            "finished_time" : None,
-            "packet_total" : None,
-            "last_packet_time" : None,
-            "state" : 0
-        }
-        stim_info = {
-            
-        }
-        Impedance_info = {
-            
-        }
+        self._impedance_channels = []
+        
+        self.stim_paradigm: StimulationParadigm = None
         # 信号采集状态
         # 信号数据包总数（一个信号采集周期内）
         # 信号采集参数
@@ -129,6 +103,9 @@ class QLBaseDevice(IDevice):
         self.storage_enable = True 
         self._listening = False
         # self.ready()
+        self._signal_cache: Queue = None
+        self._recording = False
+        
         
     def parser(self) -> IParser:
         return self._parser
@@ -137,19 +114,55 @@ class QLBaseDevice(IDevice):
         self._record_duration = record_duration      
     
     # 数据包处理
-    def produce(self, data: RscPacket, type:Literal['signal', 'impedance']="signal"):
-        if data is None: return     
+    def produce(self, body: bytes, type:Literal['signal', 'impedance']="signal"):
+        if body is None: return      
         
-        # 处理信号数据
-        self._signal_wrapper(data)   
+        if type == "signal":
+            self._produce_signal(body)
+        elif type == "impedance":
+            self._produce_impedance(body)       
         
+    def _produce_impedance(self, body: bytes):
+        # 分发阻抗数据包给订阅者
+        if len(self._impedance_consumer) > 0:
+            packet = self._impedance_wrapper(body)
+            for topic, q in self._impedance_consumer.items():
+                try:
+                    # 队列满了就丢弃最早的数据
+                    if q.full(): 
+                        q.get()
+                    q.put(packet, timeout=1)
+                except Exception as e:
+                    logger.error(f"impedance data put to queue exception: {str(e)}")
+    
+    def _produce_signal(self, body: bytes):        
+        
+        # 处理信号数据
+        data = self._signal_wrapper(body)  
+        # logger.debug("pkg_id: {}, eeg len: {}".format(data.pkg_id, len(data.eeg)))
+        # 
+        trigger_positions = [index for index, value in enumerate(data.trigger) if value != 0]
+        if len(trigger_positions) > 0:
+            # logger.debug(f"Trigger触发点位置: {trigger_positions}, 触发点时间戳: {[data.time_stamp + int(pos * 1000 / data.sample_rate) for pos in trigger_positions]}")
+            for pos in trigger_positions:
+                self.trigger(self.trigger_info(data.trigger[pos]))
         # 存储
-        if self.storage_enable:
-            self._write_signal(data)
-                        
+        if self.storage_enable:            
+            # 确保记录线程启动
+            if self._recording is False:
+                self._start_recording() 
+                
+            # 写入文件的缓存队列
+            if self._signal_cache is None:
+                self._signal_cache = Queue(256 * 1024 * 1024)  # 256MB缓存
+            tmp = data.copy()
+            self._signal_cache.put(tmp)
+                                    
         if len(self.signal_consumers) > 0 :
+            logger.trace(f"dg eeg: {data.eeg}")
             # 信号数字值转物理值
             data.eeg = self.eeg2phy(np.array(data.eeg))
+            logger.trace(f"ph eeg: {data.eeg}")
             
             # 发送数据包到订阅者
             for q in list(self.signal_consumers.values()):
@@ -157,26 +170,59 @@ class QLBaseDevice(IDevice):
                 if q.full(): 
                     q.get()
                     
-                q.put(data)
+                q.put(data)    
                 
-    # 信号数据转换(默认不处理)
-    def _signal_wrapper(self, data: RscPacket):
-        pass
+    def _impedance_wrapper(self, body: bytes):        
+        packet = ImpedancePacket().transfer(body)    
+        if self._impedance_channels is not None and len(self._impedance_channels) > 0:
+            # 只保留设置的阻抗通道
+            channel_pos = intersection_positions(packet.channels, self._impedance_channels)
+            packet.impedance = [packet.impedance[i] for i in channel_pos]
+            packet.channels = [packet.channels[i] for i in channel_pos]
+            
+        return packet
+                
+    # 信号数据转换
+    def _signal_wrapper(self, body: bytes):
+        return RscPacket().transfer(body)
                 
-    def _write_signal(self, data: RscPacket):
+    def _write_signal(self):
         # 文件写入到edf
         if self._edf_handler is None:
             logger.debug("Initializing EDF handler for data storage")
             self.init_edf_handler()
         
-        if self._edf_handler:  
+        while self._recording:            
+            data = self._signal_cache.get()
             self._edf_handler.write(data)
+            if data is None:
+                break
+            
+        self._recording = False
+    def _start_recording(self):
+        if self.storage_enable is False:
+            logger.trace("Storage is disabled, will not start recording")
+            return  
+        
+        if self._signal_cache is None:
+            self._signal_cache = Queue(256 * 1024 * 1024)  # 256MB缓存
+            
+        try: 
+            self._recording = True
+            t = Thread(target=self._write_signal, daemon=True)
+            t.start()
+            logger.info(f"开启记录")
+        except Exception as e:
+            logger.error(f"开启记录失败: {str(e)}")
+            return
         
     def start_listening(self):        
         
         try:
+            # 启动消息解析器
             self.start_message_parser()
             
+            # 启动消息监听器
             self.start_message_listening()
         except Exception as e:
             logger.error(f"设备{self.device_no}准备失败: {str(e)}")
@@ -188,26 +234,50 @@ class QLBaseDevice(IDevice):
         logger.trace(f"设备{self.device_no}停止socket监听")
         self._listening = False        
         self._parser.stop()
-        
+     
+    def read_msg(self, size: int) -> bytes:
+        try:
+            self.socket.settimeout(2.0)
+            return self.socket.recv(size)
+        except Exception as e:
+            logger.error(f"read_msg exception: {str(e)}")
+            raise ValueError("read_msg exception") from e
+      
+    @classmethod
+    def from_parent(cls, parent:IDevice) -> IDevice:
+        rlt = cls(parent.socket)
+        rlt.device_id = parent.device_id
+        rlt._device_no = parent.device_no
+        return rlt
+         
     @property
     def device_type(self) -> int:
         return self._device_type
     
+    def set_impedance_channels(self, channels):
+        self._impedance_channels = channels
+        
+    def get_impedance_channels(self): 
+        return self._impedance_channels    
+    
     def start_message_parser(self) -> None:
         self._parser = TcpMessageParser(self)
         self._parser.start()
-        logger.debug("TCP消息解析器已启动")
+        logger.debug("RSC消息解析器已启动")
     
     def start_message_listening(self) -> None:
         def _accept():
             while self._listening:
-                # 缓冲区4M
-                data = self.socket.recv(4096*1024)
-                if not data:
-                    logger.warning(f"设备[{self.device_name}]连接结束")
+                try:
+                    # 缓冲区4M
+                    data = self.socket.recv(4096*1024)
+                    if not data:
+                        logger.warning(f"设备[{self.device_name}]连接结束")
+                        break
+                    self._parser.append(data)
+                except Exception as e:
+                    logger.debug(f"设备[{self.device_name}]接收数据异常: {str(e)}")
                     break
-                
-                self._parser.append(data)
     
         # 启动数据接收线程
         self._listening = True
@@ -252,11 +322,6 @@ class QLBaseDevice(IDevice):
     def init_edf_handler(self):
         logger.warning("init_edf_handler not implemented in base class, should be overridden in subclass")
         pass
-    #     self._edf_handler = RscEDFHandler(self.sample_rate,  self.sample_range * 1000 , - self.sample_range * 1000, self.resolution)  
-    #     self._edf_handler.set_device_type(self.device_type)
-    #     self._edf_handler.set_device_no(self.device_no)
-    #     self._edf_handler.set_storage_path(self._storage_path)
-    #     self._edf_handler.set_file_prefix(self._file_prefix)
         
     # eeg数字值转物理值
     def eeg2phy(self, digital:int):
@@ -266,7 +331,7 @@ class QLBaseDevice(IDevice):
     @property
     def edf_handler(self):
         if not self.storage_enable:
-            logger.warning("EDF storage is disabled, no edf handler available")
+            logger.info("已关闭文件记录，不会生成edf/bdf文件")
             return None
         
         if self._edf_handler is None:
@@ -300,8 +365,12 @@ class QLBaseDevice(IDevice):
         return self._impedance_consumer
     
     # 设置记录文件路径
-    def set_storage_path(self, path):
-        self._storage_path = path
+    def set_storage_path(self, dir: str):
+        import os
+        
+        abs_path = os.path.abspath(dir)
+        os.makedirs(abs_path, exist_ok=True)
+        self._storage_path = abs_path
        
     # 设置记录文件名称前缀 
     def set_file_prefix(self, prefix):
@@ -316,23 +385,27 @@ class QLBaseDevice(IDevice):
         self.stim_paradigm = param
     
     # 设置采集参数
-    def set_acq_param(self, channels, sample_rate = 500, sample_range = 188):
+    def set_acq_param(self, channels, sample_rate:Literal[188, 375, 563, 750, 1125, 2250, 4500] = 500, sample_range:Literal[250, 500, 1000, 2000, 4000, 8000, 16000, 32000] = 188):
         self._acq_param["channels"] = channels 
         self._acq_param["sample_rate"] = sample_rate
         self._acq_param["sample_range"] = sample_range
         self._acq_channels = channels
         self._sample_rate = sample_rate
         self._sample_range = sample_range
-    
-    # 通用配置-TODO
-    def set_config(self, key:str, val: str):
-        pass      
         
     def start_impedance(self):
         logger.info(f"[设备-{self.device_no}]启动阻抗测量")  
-        msg = StartImpedanceCommand.build(self).pack()
-        logger.trace(f"start_impedance message is {msg.hex()}")
-        self.socket.sendall(msg)
+        # 设置数据采集参数
+        # set_param_msg = SetImpedanceParamCommand.build(self).pack()
+        device_id = bytes.fromhex(self.device_id)[::-1].hex() if self.device_id else '00000000'
+        set_param_msg = bytes.fromhex(f'5aa50239{device_id}3f0000001104ffffffffffffffff000000000000000000000000000000000000000000000000e8030000fa00000010000164000000745c5aa50239390045243a00000012040000000000000000000000000000000000000000000000000000000000000000000001000000000000004c2a')
+        logger.debug(f"set_param_msg message is {set_param_msg.hex()}")
+        self.socket.sendall(set_param_msg)
+        sleep(0.5)
+    
+        impedance_start_msg = StartImpedanceCommand.build(self).pack()
+        logger.debug(f"start_impedance message is {impedance_start_msg.hex()}")
+        self.socket.sendall(impedance_start_msg)
     
     def stop_impedance(self):
         logger.info(f"[设备{self.device_no}]停止阻抗测量")  
@@ -348,8 +421,11 @@ class QLBaseDevice(IDevice):
         msg = StartStimulationCommand.build(self).pack()
         logger.trace(f"start_stimulation message is {msg.hex()}")
         self.socket.sendall(msg)
-        t = Thread(target=self._stop_stimulation_trigger, args=(self.stim_paradigm.duration,), daemon=True)
-        t.start()
+        # t = Thread(target=self._stop_stimulation_trigger, args=(self.stim_paradigm.duration,), daemon=True)
+        # t.start()
+        
+    def get_stim_param(self) -> bytes:
+        return self.stim_paradigm.to_bytes()
         
     def _stop_stimulation_trigger(self, duration):
         delay = duration
@@ -357,23 +433,22 @@ class QLBaseDevice(IDevice):
             sleep(1)
             delay -= 1
         logger.debug(f"_stop_stimulation_trigger duration: {duration}")
-        if self._edf_handler:
-            self._edf_handler.trigger("stimulation should be stopped")
+        if self.edf_handler:
+            self.edf_handler.trigger("stimulation should be stopped")
         else:
             logger.warning("stop stim trigger fail. no edf writer alive")
         
     def stop_stimulation(self):
         logger.info(f"[设备-{self.device_no}]停止电刺激")  
-        msg = StopStimulationCommand.pack()
+        msg = StopStimulationCommand.build(self).pack()
         logger.trace(f"stop_stimulation message is {msg.hex()}")
         self.socket.sendall(msg)
         
     # 启动采集
-    def start_acquisition(self, recording = True):
+    def start_acquisition(self):
         logger.info(f"[设备-{self.device_no}]启动信号采集")  
-        self._recording = recording
-        # 初始化EDF处理器
-        self.init_edf_handler()
+        # 记录准备
+        self._start_recording()
         # 设置数据采集参数
         param_bytes = SetAcquisitionParamCommand.build(self).pack()
         # 启动数据采集
@@ -388,9 +463,10 @@ class QLBaseDevice(IDevice):
         msg = StopAcquisitionCommand.build(self).pack()
         logger.trace(f"stop_acquisition message is {msg.hex()}")
         self.socket.sendall(msg)
-        if self._edf_handler:
+        # 结束标识
+        if self._signal_cache:
             # 发送结束标识
-            self.edf_handler.write(None)
+            self._signal_cache.put(None)
         
     '''
         订阅数据
@@ -408,7 +484,7 @@ class QLBaseDevice(IDevice):
             
         # 数据队列
         if q is None:
-            q = Queue(maxsize=1000)
+            q = Queue(maxsize=100 * 1024 * 1024)
         
         # 订阅生理电信号数据
         if type == "signal":
@@ -428,11 +504,12 @@ class QLBaseDevice(IDevice):
         return topic, q
     
     def trigger(self, desc):
-        if self._edf_handler:
+        if self.edf_handler:
             self.edf_handler.trigger(desc)
         else:
-            logger.warning("没有开启文件记录时，无法记录trigger信息")
+            logger.info("已关闭文件记录，不会记录trigger信息")
         
+    # 设置信号采集参数
     def gen_set_acquirement_param(self) -> bytes:
         
         body = to_bytes(self.acq_channels) 
@@ -443,13 +520,45 @@ class QLBaseDevice(IDevice):
         body += bytes.fromhex('00')     
         
         return body
+    # 设置阻抗测量参数
+    def gen_set_impedance_param(self) -> bytes:
+        
+        # 仅通道生效 32字节，其他不生效-272字节，实际73字节
+        body = to_bytes(self._impedance_channels) 
+        # 100 bytes
+        # body += bytes.fromhex('00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000')     
+        # # 100 bytes
+        # body += bytes.fromhex('00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000')     
+        # 73 bytes
+        body += bytes.fromhex('00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000')     
+        
+        return bytes.fromhex('ffffffffffffffff000000000000000000000000000000000000000000000000e8030000fa00000010000164000000745c5aa50239390045243a0000001204000000000000000000000000000000000000000000000000000000000000000000000100000000000000')
+    def disconnect(self):
+        logger.info(f"[断开设备-{self.device_no}]的连接...")  
+        self._listening = False
+        try:
+            sleep(0.1)
+            self.socket.shutdown(2)
+            self.socket.close()
+            logger.info(f"[设备-{self.device_no}]设备连接已断开")  
+        except Exception as e:
+            logger.error(f"断开设备连接异常: {str(e)}")
+         
+        # 关闭解析器
+        self._parser.stop()   
+        
+    def enable_storage(self, enable: bool = True):
+        self.storage_enable = enable
+        
+    def trigger_info(self, code: int) -> str:
+        return QLBaseDevice.__TRIGGER_MAPPING.get(code, hex(code))    
     
     def __str__(self):
         return f''' 
             Device: 
                 Name: {self.device_name}, 
                 Type: {hex(self.device_type) if self.device_type else None}, 
-                ID: {self.device_id if self.device_id else None}, 
+                ID: {self.device_no if self.device_no else None}, 
                 Software: {self.software_version}, 
                 Hardware: {self.hardware_version}, 
                 Connect Time: {self.connect_time},
@@ -459,8 +568,8 @@ class QLBaseDevice(IDevice):
                 Battery Total: {str(self.battery_total) + "%" if self.battery_total else None}
             '''
 
-    def __eq__(self, other):
-        return self.device_name == other.device_name and self.device_type == other.device_type and self.device_id == other.device_id
+    def __eq__(self, other:IDevice):
+        return self.device_type == other.device_type and self.device_no == other.device_no
 
     def __hash__(self):
-        return hash((self.device_name, self.device_type, self.device_id))    
+        return hash((self.device_type, self.device_no))    

qlsdk/rsc/device/c16_rs.py

@@ -115,11 +115,7 @@ class C16RS(QLBaseDevice):
         self._edf_handler.set_device_no(self.device_no)
         self._edf_handler.set_storage_path(self._storage_path)
         self._edf_handler.set_file_prefix(self._file_prefix if self._file_prefix else 'C16R')
-            
-    # 设置刺激参数
-    def set_stim_param(self, param):
-        self.stim_paradigm = param
-    
+                
     # 设置采集参数
     def set_acq_param(self, channels, sample_rate = 500, sample_range = 188):
         # 保存原始通道参数