298146e8b0
The idea is to allow relatively flexible visualisations in the browser, while minimising rendering load on the server.
135 lines
3.7 KiB
Python
135 lines
3.7 KiB
Python
import io
|
|
import logging
|
|
import multiprocessing as mp
|
|
from typing import Generator
|
|
|
|
import matplotlib
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
from flask import Flask, Response, render_template
|
|
from flask_sock import Sock
|
|
from simple_websocket import Server
|
|
|
|
from ..config import config
|
|
from ..processing.aoa import AoA
|
|
|
|
matplotlib.use("agg")
|
|
|
|
app = Flask(__name__)
|
|
sock = Sock(app)
|
|
aoa_queue = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
data: npt.NDArray[np.complex128] = np.array([], dtype=complex)
|
|
aoa: AoA = AoA()
|
|
|
|
|
|
@app.route("/preprocessed")
|
|
def preprocessed() -> str:
|
|
return render_template("preprocessed.html")
|
|
|
|
|
|
Subscriber = Server
|
|
subscriber_settings: dict[Subscriber, tuple[int, int, int]] = {}
|
|
|
|
|
|
@sock.route("/data")
|
|
def get_data(sock: Subscriber) -> None:
|
|
while True:
|
|
msg = sock.receive()
|
|
if len(msg.split()) != 3:
|
|
break
|
|
subcarrier, rx, tx = map(int, msg.split())
|
|
subscriber_settings[sock] = (subcarrier, rx, tx)
|
|
|
|
|
|
def add_data(
|
|
raw_data: npt.NDArray[np.complex128], new_data: npt.NDArray[np.complex128]
|
|
) -> None:
|
|
if config.VISUALISE_RAW:
|
|
magn = np.abs(raw_data)
|
|
phase = np.angle(raw_data)
|
|
|
|
new_mag = np.abs(new_data)
|
|
new_phase = np.angle(new_data)
|
|
|
|
fig, axs = plt.subplots(2, 2)
|
|
axs[0, 0].plot(magn[:, 0, 0], c="b")
|
|
axs[0, 0].plot(magn[:, 1, 0], c="orange")
|
|
axs[1, 0].plot(new_mag[:, 0, 0], c="b")
|
|
axs[1, 0].plot(new_mag[:, 1, 0], c="orange")
|
|
axs[0, 1].plot(phase[:, 0, 0], c="b")
|
|
axs[0, 1].plot(phase[:, 1, 0], c="orange")
|
|
axs[1, 1].plot(new_phase[:, 0, 0], c="b")
|
|
axs[1, 1].plot(new_phase[:, 1, 0], c="orange")
|
|
fig.savefig("/tmp/plot.png")
|
|
plt.close(fig)
|
|
|
|
global data
|
|
if data.size == 0:
|
|
data = np.expand_dims(new_data, axis=0)
|
|
else:
|
|
data = np.concat([data, np.expand_dims(new_data, axis=0)], axis=0)
|
|
|
|
# Only keep latest 100 entries
|
|
if data.shape[0] > 100:
|
|
data = data[-100:]
|
|
to_remove: list[Subscriber] = []
|
|
for subscriber in subscriber_settings:
|
|
try:
|
|
subcarrier, rx, tx = subscriber_settings[subscriber]
|
|
subscriber.send(new_data.real[subcarrier, rx, tx])
|
|
except Exception as e:
|
|
to_remove.append(subscriber)
|
|
print(e)
|
|
|
|
for subscriber in to_remove:
|
|
del subscriber_settings[subscriber]
|
|
|
|
|
|
def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO:
|
|
logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
|
|
fig = plt.figure()
|
|
ax = fig.add_axes([0.1, 0.1, 0.9, 0.9]) # , polar=True)
|
|
r = np.linspace(
|
|
0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution
|
|
)
|
|
theta = np.linspace(0, np.pi, config.music.heatmap.theta_resolution) # Angle values
|
|
|
|
mesh = ax.pcolormesh(theta, r, heatmap, edgecolors="face", vmin=0, vmax=50)
|
|
|
|
fig.colorbar(mesh, ax=ax)
|
|
buf = io.BytesIO()
|
|
fig.savefig(buf, format="jpeg")
|
|
plt.close(fig)
|
|
|
|
buf.seek(0)
|
|
return buf
|
|
|
|
|
|
def gather_aoa() -> Generator[bytes, None, None]:
|
|
assert aoa_queue is not None
|
|
|
|
while True:
|
|
logger.debug("Receiving from aoa pipe")
|
|
heatmap = aoa_queue.get()
|
|
buf = plot_heatmap(heatmap)
|
|
yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
|
|
buf.close()
|
|
|
|
|
|
@app.route("/aoa_tof")
|
|
def aoa_tof() -> Response:
|
|
return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame")
|
|
|
|
|
|
def start(conn: "mp.Queue[npt.NDArray[np.float32]]") -> None:
|
|
global app, aoa_queue
|
|
|
|
aoa_queue = conn
|
|
app.run(debug=True, use_reloader=False, host="0.0.0.0")
|
|
logger.info("Visualisation server shut down")
|