150 lines
4.2 KiB
Python
150 lines
4.2 KiB
Python
import io
|
|
import logging
|
|
import multiprocessing as mp
|
|
import time
|
|
from datetime import datetime
|
|
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, request
|
|
from flask_sock import Sock
|
|
from simple_websocket import Server
|
|
|
|
from .. 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 make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO:
|
|
logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
|
|
fig = plt.figure()
|
|
ax = fig.add_axes([0, 0, 1, 1], polar=True)
|
|
r = np.linspace(0, max_tof, 100) # Radius values
|
|
theta = np.linspace(0, np.pi, 50) # Angle values
|
|
R, Theta = np.meshgrid(r, theta) # Create a 2D grid of r and theta
|
|
|
|
# Compute the function values
|
|
Z = np.log(np.vectorize(aoa.evaluate)(Theta, R))
|
|
|
|
ax.pcolormesh(Theta, R, Z, edgecolors="face")
|
|
buf = io.BytesIO()
|
|
fig.savefig(buf, format="jpeg")
|
|
plt.close(fig)
|
|
|
|
buf.seek(0)
|
|
return buf
|
|
|
|
|
|
def gather_aoa(max_tof: float) -> Generator[bytes, None, None]:
|
|
assert aoa_queue is not None
|
|
|
|
prev_frame = datetime.now()
|
|
while True:
|
|
while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
|
|
time.sleep(0.01)
|
|
while not aoa_queue.empty():
|
|
logger.debug("Receiving from aoa pipe")
|
|
aoa = aoa_queue.get()
|
|
prev_frame = datetime.now()
|
|
logger.debug(f"Generating heatmap of time {aoa.timestamp}")
|
|
buf = make_heatmap(aoa, max_tof)
|
|
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:
|
|
max_tof_str = request.args.get("max_tof")
|
|
try:
|
|
max_tof = float(max_tof_str)
|
|
except Exception:
|
|
max_tof = 5e-8
|
|
return Response(
|
|
gather_aoa(max_tof), mimetype="multipart/x-mixed-replace; boundary=frame"
|
|
)
|
|
|
|
|
|
def start(conn: "mp.Queue[AoA]") -> 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")
|