51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import zmq
|
||
|
|
import cupy as cp
|
||
|
|
|
||
|
|
from Msg import BMMsg, ImageArgMsg, KillMsg
|
||
|
|
from config import socket1, video_width, video_height
|
||
|
|
from nodes.Node import Node
|
||
|
|
|
||
|
|
|
||
|
|
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()
|
||
|
|
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)
|