import os, sys
|
from ctypes import sizeof, string_at
|
from socketserver import BaseRequestHandler, UDPServer
|
import warnings
|
from .components import *
|
from .constants import *
|
from ..protos import aerial_pb2 as pb
|
|
|
|
UDP_STATE_IDLE = 0
|
UDP_STATE_HAND_SHAKE = 1
|
UDP_STATE_DATA = 2
|
UDP_STATE_FOOTER = 3
|
|
class MicroWaveReceiver:
|
"""
|
micro wave data receiver
|
use AddBuf method add udp packet to receiver
|
when DataFrameReady == True, you can get data frame
|
"""
|
def __init__(self):
|
self.state = UDP_STATE_IDLE
|
self.PointCount = 0
|
self.TrackCount = 0
|
self.DataFrameReady = False
|
self.DataFrame = None
|
self.Header = None
|
self.Footer = None
|
|
def _is_handshake(self,data):
|
handshake = Handshake.from_buffer_copy(data)
|
if HAND_SHAKE_MAGIC == string_at(handshake.Magic, 8):
|
self.DataFrameReady = False
|
self.state = UDP_STATE_HAND_SHAKE
|
self.DataFrame = None
|
self.Header = None
|
self.Footer = None
|
return True
|
else:
|
return False
|
|
def _get_header(self, buf):
|
self.Header = HeaderBlock.from_buffer_copy(buf[:sizeof(HeaderBlock)])
|
if HEADER_MAGIC == string_at(self.Header.Magic, 8):
|
self.DataFrame = bytearray(buf)
|
N = self.Header.NumberOfDetection
|
Nt = self.Header.NumberOfTrack
|
self.PointCount = 0
|
self.TrackCount = 0
|
self.DataLength = N * 8 + Nt * sizeof(TrackerBlock)
|
self.RemainLength = self.DataLength + sizeof(FooterBlock)
|
self.data_buf = bytearray()
|
self.footer_buf = bytearray()
|
self.Points = []
|
self.ReceivedDataLength = 0
|
self.ReceivedFooterLength = 0
|
# the next packet should be data block
|
self.state = UDP_STATE_DATA
|
return True
|
else:
|
False
|
|
def _receive_data(self, data) -> None:
|
"""
|
receive data packet from udp
|
"""
|
self.data_buf.extend(bytes(data))
|
self.ReceivedDataLength += len(data)
|
if self.ReceivedDataLength == self.DataLength:
|
# if received data length equal expected data length
|
# then we can get footer block
|
self.state = UDP_STATE_FOOTER
|
return True
|
elif self.ReceivedDataLength > self.DataLength:
|
self.state = UDP_STATE_IDLE
|
warnings.warn(f"encounter an error when getting the detection data in udp packets ...")
|
return False
|
else:
|
return True
|
|
def _receive_footer(self, data) -> None:
|
# receive footer block
|
self.footer_buf.extend(bytes(data))
|
self.ReceivedFooterLength += len(data)
|
if self.ReceivedFooterLength == sizeof(FooterBlock):
|
self.Footer = FooterBlock.from_buffer_copy(self.footer_buf)
|
# after received footer block, the data frame is ready
|
self.DataFrameReady = True
|
# reset state to idle, ready for next data frame
|
self.state = UDP_STATE_IDLE
|
return
|
|
|
def AddBuf(self, buf: bytes) -> None:
|
# add a udp packet to receiver
|
buf_size = len(buf)
|
if buf_size == sizeof(Handshake):
|
self._is_handshake(buf)
|
|
if self.state == UDP_STATE_HAND_SHAKE:
|
# after received handshake packet, we can get header block
|
try:
|
# make sure to get a complete header
|
result = self._get_header(buf)
|
except Exception as e:
|
warnings.warn(f"encounter an unknow error {e} when getting the header data in udp packets ...")
|
result = False
|
|
if not result:
|
# waiting for the next handshake udp packet
|
self.state = UDP_STATE_IDLE
|
else:
|
# process the detection data in the first data packets
|
header_size = sizeof(HeaderBlock)
|
data_size = min(buf_size - header_size, self.DataLength)
|
|
if self._receive_data(buf[header_size: header_size + data_size]):
|
|
# 判断第一个数据包中是否包含Footer Block
|
if (header_size + data_size) < buf_size:
|
self._receive_footer(buf[header_size + data_size:])
|
|
elif self.state == UDP_STATE_DATA:
|
# receive data block
|
self.DataFrame.extend(bytes(buf))
|
remain_data_size = self.DataLength - self.ReceivedDataLength
|
data_size = min(buf_size, remain_data_size)
|
if self._receive_data(buf[: data_size]):
|
if data_size < buf_size:
|
self._receive_footer(buf[data_size: ])
|
|
elif self.state == UDP_STATE_FOOTER:
|
# receive footer block
|
self.DataFrame.extend(bytes(buf))
|
self._receive_footer(buf)
|
|
|
|
class MicroWaveRequestHandler(BaseRequestHandler):
|
"""
|
micro wave udp request handler
|
"""
|
def handle(self):
|
# Get message and client socket
|
buf, _ = self.request
|
receiver = self.server.Receiver
|
# add buffer to receiver
|
receiver.AddBuf(buf)
|
if receiver.DataFrameReady:
|
# dataframe ready, add to udp_queue
|
msg = pb.RequestCommand()
|
msg.sys_command = pb.SysCommand.MICRO_WAVE_COMMAND
|
msg.sub_command = pb.SubCommand.MICRO_WAVE_SUB_COMMAND
|
data1 = pb.MicroWaveMessage()
|
data1.length = len(receiver.DataFrame)
|
data1.data = bytes(receiver.DataFrame)
|
data1_str = data1.SerializeToString()
|
msg.data1.length = len(data1_str)
|
msg.data1.data = data1_str
|
self.server.msg_queue.put(msg)
|
|
|
class MicroWaveUDPServer(UDPServer):
|
"""
|
micro wave udp server
|
"""
|
def __init__(self, server_address, msg_queue):
|
super().__init__(server_address, MicroWaveRequestHandler)
|
# create receiver instance
|
self.Receiver = MicroWaveReceiver()
|
self.msg_queue = msg_queue
|
