57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import cv2
|
|
import numpy as np
|
|
import zmq
|
|
|
|
from utils.Msg import BMMsg, ImageArgMsg, KillMsg
|
|
from config import SOCKET1, VIDEO_WIDTH, VIDEO_HEIGHT
|
|
from nodes.Node import Node
|
|
from utils.RfFile import RfFrame, RfSequenceMeta
|
|
from utils.RfMat import RfMat
|
|
|
|
|
|
class Beamformer(Node):
|
|
topics = [ImageArgMsg]
|
|
|
|
def __init__(self):
|
|
super(Beamformer, self).__init__()
|
|
self.arg = ImageArgMsg('', 1400)
|
|
|
|
def process(self, data: bytes):
|
|
if data is None:
|
|
return
|
|
b = np.frombuffer(data, dtype=np.int16)
|
|
b = b.reshape((256, 1502)).astype(np.float32)
|
|
b = b[:, :self.arg.v1]
|
|
b -= b.min()
|
|
b /= b.max()
|
|
b *= 255
|
|
b = b.astype(np.uint8)
|
|
b = cv2.rotate(b, cv2.ROTATE_90_CLOCKWISE)
|
|
b = cv2.resize(b, (VIDEO_WIDTH, VIDEO_HEIGHT))
|
|
b = cv2.cvtColor(b, cv2.COLOR_GRAY2RGBA, b)
|
|
self.send(BMMsg(0, b.tobytes()))
|
|
|
|
def loop(self):
|
|
s2 = self.context.socket(zmq.PULL)
|
|
s2.setsockopt(zmq.CONFLATE, 1)
|
|
s2.connect(f"tcp://{SOCKET1}")
|
|
self.c.poller.register(s2, zmq.POLLIN)
|
|
buffer = None
|
|
while True:
|
|
socks = dict(self.c.poller.poll())
|
|
if s2 in socks and socks[s2] == zmq.POLLIN:
|
|
buffer = s2.recv()
|
|
|
|
RfMat(1, RfFrame.RfFrameMeta(
|
|
1, 1
|
|
), RfSequenceMeta())
|
|
|
|
if self.c.sub in socks and socks[self.c.sub] == zmq.POLLIN:
|
|
r = self.recv()
|
|
if isinstance(r, KillMsg):
|
|
if r.name == '':
|
|
return
|
|
if isinstance(r, ImageArgMsg):
|
|
self.arg = r
|
|
self.process(buffer)
|