From 15557bc6716551c051293ed97cc68999c39dd19a Mon Sep 17 00:00:00 2001 From: Christos Falas Date: Tue, 28 Jan 2025 20:27:56 +0000 Subject: [PATCH] Make heatmap on GPU Still has some weird artifacts, some axis ordering might be incorrect --- where_fi/cli/__init__.py | 11 ++- where_fi/config/models.py | 7 ++ where_fi/processing/__init__.py | 0 where_fi/processing/aoa.py | 118 +++++++++++++++++++++----------- where_fi/utils/__init__.py | 0 where_fi/visualise/__init__.py | 42 ++++-------- 6 files changed, 108 insertions(+), 70 deletions(-) create mode 100644 where_fi/processing/__init__.py create mode 100644 where_fi/utils/__init__.py diff --git a/where_fi/cli/__init__.py b/where_fi/cli/__init__.py index 5caa2b7..75641f0 100644 --- a/where_fi/cli/__init__.py +++ b/where_fi/cli/__init__.py @@ -38,9 +38,12 @@ def antennas() -> None: def heatmap() -> None: preprocessor = Preprocessor() aoa = AoA() - manager = mp.Manager() - webapp_queue: "mp.Queue[AoA]" = manager.Queue(config.sample_rate) + # Start webapp in background process + manager = mp.Manager() + webapp_queue: "mp.Queue[npt.NDArray[np.float32]]" = manager.Queue( + config.sample_rate + ) webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) webapp.start() @@ -50,8 +53,10 @@ def heatmap() -> None: processed_tensor = torch.tensor(processed, device=device) # visualise.add_data(all_data, processed) aoa.update(processed_tensor) + if not webapp_queue.full(): - webapp_queue.put(aoa) + heatmap = aoa.heatmap() + webapp_queue.put(heatmap) globals.csi_producer(csi_callback=callback) logger.info("Finished processing CSI data") diff --git a/where_fi/config/models.py b/where_fi/config/models.py index b052a8e..eaa937c 100644 --- a/where_fi/config/models.py +++ b/where_fi/config/models.py @@ -23,6 +23,13 @@ class MUSIC(BaseModel): eigval_threshold: int window_size: int + class Heatmap(BaseModel): + theta_resolution: int + tof_resolution: int + tof_max: float + + heatmap: Heatmap + class Antennas(BaseModel): spacing: float diff --git a/where_fi/processing/__init__.py b/where_fi/processing/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/where_fi/processing/aoa.py b/where_fi/processing/aoa.py index 0c8d2c8..872d121 100644 --- a/where_fi/processing/aoa.py +++ b/where_fi/processing/aoa.py @@ -68,53 +68,93 @@ class AoA: if self.historical_autocorr.shape[0] > WINDOW_SIZE: self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:] - # Is the moving average also Hermitian? - R = torch.mean(self.historical_autocorr, dim=0) + def steering_vector(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor: + assert theta.shape == tof.shape + assert len(theta.shape) == 1 + N = theta.shape[0] - # The smallest eigenvectors span the noise subspace, - # and the largest span the signal subspace. - eigvals, eigvecs = torch.linalg.eigh(R) - self.E_n = ( - eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold].cpu().numpy() - ) - - def steering_vector( - self, theta: float, tof: float - ) -> npt.NDArray[np.complexfloating]: - omega_t: npt.NDArray[np.complex64] = np.exp(-2j * np.pi * config.delta_f * tof) - phi_theta: npt.NDArray[np.complex64] = np.exp( + omega_t: torch.Tensor = torch.exp(-2j * np.pi * config.delta_f * tof) + phi_theta: torch.Tensor = torch.exp( 2j * np.pi * config.central_freq_hz * config.antennas.spacing - * (1 - np.cos(theta)) + * (1 - torch.cos(theta)) / 299_792_458 ) + assert omega_t.shape == phi_theta.shape == (N,) + print(omega_t, phi_theta) - omega_t = np.expand_dims(omega_t, axis=-1) - phi_theta = np.expand_dims(phi_theta, axis=-1) + omega_t = torch.unsqueeze(omega_t, dim=-1) + phi_theta = torch.unsqueeze(phi_theta, dim=-1) - antenna_v = omega_t ** torch.arange(self.N_subcarriers // 2) - phis = phi_theta ** torch.arange(self.N_rx // 2) - antenna_v = np.expand_dims(antenna_v, axis=-1) - print(antenna_v.shape, phis.shape) - steering = antenna_v[0] * phis - print(steering.shape) - return steering.T.reshape(-1) + assert omega_t.shape == phi_theta.shape == (N, 1) - def evaluate(self, theta: float, tof: float) -> float: - try: - steering = self.steering_vector(theta, tof) - steering_h = np.conj(steering).T - except Exception as e: - logger.exception(e) - return 0 + antenna_v = omega_t ** torch.arange( + self.N_subcarriers // 2, dtype=torch.float32 + ) + phis = phi_theta ** torch.arange(self.N_rx // 2, dtype=torch.float32) - assert isinstance(self.E_n, torch.Tensor) - E_n = self.E_n - E_n_H = np.conj(E_n).T - c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering)) - return np.abs(c.real).item() + assert antenna_v.shape == (N, self.N_subcarriers // 2) + assert phis.shape == (N, self.N_rx // 2) + + antenna_v = torch.unsqueeze(antenna_v, dim=1) + phis = torch.unsqueeze(phis, dim=-1) + + assert antenna_v.shape == (N, 1, self.N_subcarriers // 2) + assert phis.shape == (N, self.N_rx // 2, 1) + + steering = torch.bmm(phis, antenna_v) + + assert steering.shape == (N, self.N_rx // 2, self.N_subcarriers // 2) + return steering.reshape(N, -1) + + def evaluate(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor: + R = torch.mean(self.historical_autocorr, dim=0) + + # The smallest eigenvectors span the noise subspace, + # and the largest span the signal subspace. + eigvals, eigvecs = torch.linalg.eig(R) + assert isinstance(eigvals, torch.Tensor) + assert isinstance(eigvecs, torch.Tensor) + logger.info(f"Eigenvalues: {eigvals}") + E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold] + + logger.info(f"Signal subspace: {E_n.shape}") + steering = torch.unsqueeze(self.steering_vector(theta, tof), dim=-1) + steering_h = torch.conj(steering).permute(0, 2, 1) + + E_n = E_n.unsqueeze(0) + E_n_H = torch.conj(E_n).permute(0, 2, 1) + logger.info( + f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, {E_n_H.shape}, {steering.shape}" + ) + c: torch.Tensor = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering)) + return torch.abs(c.real) + + def heatmap(self) -> npt.NDArray[np.float32]: + thetas = np.linspace( + 0, np.pi, config.music.heatmap.theta_resolution, dtype=np.float32 + ) + tofs = np.linspace( + 0, + config.music.heatmap.tof_max, + config.music.heatmap.tof_resolution, + dtype=np.float32, + ) + thetas_mesh, tofs_mesh = np.meshgrid(thetas, tofs) + heatmap: npt.NDArray[np.float32] = ( + self.evaluate( + torch.tensor(thetas_mesh.reshape(-1)), + torch.tensor(tofs_mesh.reshape(-1)), + ) + .reshape( + config.music.heatmap.theta_resolution, + config.music.heatmap.tof_resolution, + ) + .numpy(force=True) + ) + return heatmap def test_smoothing() -> None: @@ -133,8 +173,8 @@ def test_smoothing() -> None: def test_steering_vector() -> None: aoa = AoA() aoa.N_subcarriers = 10 - aoa.N_rx = 2 - tau = 1 - theta = 0 + aoa.N_rx = 4 + tau = torch.Tensor([1, 0]) + theta = torch.Tensor([0, 1]) print(aoa.steering_vector(theta, tau)) assert False diff --git a/where_fi/utils/__init__.py b/where_fi/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/where_fi/visualise/__init__.py b/where_fi/visualise/__init__.py index 9407b6e..611e945 100644 --- a/where_fi/visualise/__init__.py +++ b/where_fi/visualise/__init__.py @@ -13,7 +13,7 @@ from flask import Flask, Response, render_template, request from flask_sock import Sock from simple_websocket import Server -from .. import config +from ..config import config from ..processing.aoa import AoA matplotlib.use("agg") @@ -92,18 +92,17 @@ def add_data( del subscriber_settings[subscriber] -def make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO: +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, 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 + 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 + X, Y = 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") + ax.pcolormesh(Y, X, heatmap, edgecolors="face") buf = io.BytesIO() fig.savefig(buf, format="jpeg") plt.close(fig) @@ -112,36 +111,23 @@ def make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO: return buf -def gather_aoa(max_tof: float) -> Generator[bytes, None, None]: +def gather_aoa() -> 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) + 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: - 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" - ) + return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame") -def start(conn: "mp.Queue[AoA]") -> None: +def start(conn: "mp.Queue[npt.NDArray[np.float32]]") -> None: global app, aoa_queue aoa_queue = conn