refactor: processing in separate thread

src/__main__.py

@@ -0,0 +1,35 @@
+import logging
+import threading
+
+from . import ingest
+from . import config
+from . import visualise
+from .preprocess import Preprocessor
+
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(name)-40s %(levelname)-8s %(message)s",
+)
+
+receivers = [ingest.FeitReceiver(ip) for ip in config.RECEIVE_IP_ADDRESS_LIST]
+transmitter = ingest.FeitTransmitter()
+preprocess = Preprocessor()
+
+processor = ingest.CSIProcessor()
+
+# Start webapp in background thread
+visualise.aoa = processor.aoa
+webapp = threading.Thread(target=visualise.start)
+webapp.start()
+
+receiver_threads = [
+    threading.Thread(target=server.listen, args=(processor.add_data,))
+    for server in receivers
+]
+
+for thread in receiver_threads:
+    thread.start()
+
+processing_thread = threading.Thread(target=processor.process_forever)
+processing_thread.start()

src/ingest.py

@@ -0,0 +1,127 @@
+import logging
+import time
+import socket
+from typing import Callable
+import threading
+import struct
+from datetime import datetime
+import numpy as np
+
+from .csi import CSI
+from .aoa import AoA
+from .preprocess import Preprocessor
+from . import config
+from . import visualise
+
+
+class FeitTransmitter:
+    def __init__(self):
+        self.inject_server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.inject_server.connect((config.INJECT_IP_ADDRESS, config.FEITCSI_PORT))
+        inject_start_string = (
+            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
+            f"--channel-width {config.CHANNEL_WIDTH} "
+            f"--format {config.FRAME_FORMAT} "
+            f"--mode inject -s 1 --verbose "
+            f"--inject-delay {1_000_000 // config.SAMPLE_RATE}"
+        )
+        self.inject_server.send(b"stop\n")
+        self.inject_server.send(inject_start_string.encode())
+
+
+class FeitReceiver:
+    def __init__(self, host: str):
+        self.host = host
+        self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.server.connect((self.host, config.FEITCSI_PORT))
+        self.start_string = (
+            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
+            f"--channel-width {config.CHANNEL_WIDTH} "
+            f"--format {config.FRAME_FORMAT} "
+            f"--mode measure"
+        )
+        self.server.send(b"stop\n")
+        self.server.send(self.start_string.encode())
+        self.logger = logging.getLogger(
+            f"{__name__}.{self.__class__.__name__}-{self.host}"
+        )
+
+    def listen(self, callback: Callable[[str, CSI], None]):
+        prev_time = datetime.now()
+        while True:
+            # This is the max size of a UDP packet. The size of the actual CSI
+            # packet will depend on the frame format and channel width, which
+            # changes the number of subcarriers
+            data = self.server.recv(65535)
+            try:
+                csidata = CSI(data)
+                self.logger.debug(
+                    f"Received CSI data after {datetime.now() - prev_time}"
+                )
+                prev_time = datetime.now()
+                callback(self.host, csidata)
+            except struct.error:
+                self.logger.error("Failed to parse CSI data")
+
+
+class CSIProcessor:
+    def __init__(self):
+        self.pending_data: dict[str, tuple[datetime, CSI]] = {}
+        self.pending_data_lock = threading.Lock()
+        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
+        self.last_processed = datetime.now()
+
+        self.preprocess = Preprocessor()
+        self.aoa = AoA()
+
+    def add_data(self, host: str, data: CSI):
+        if (
+            host in self.pending_data
+            and self.last_processed < self.pending_data[host][0]
+        ):
+            self.logger.warning(
+                f"Skipping data from {host} at {self.pending_data[host][0]}"
+            )
+
+        with self.pending_data_lock:
+            self.pending_data[host] = (datetime.now(), data)
+
+            # Useful for figuring out the correct antenna order - RSSI values will decrease
+            # when the specific antenna is disconnected
+            rssis = [
+                (ip, csi.header.rssi1, csi.header.rssi2)
+                for ip, (_, csi) in sorted(self.pending_data.items())
+            ]
+            self.logger.debug("Antenna RSSI values: {}".format(rssis))
+
+    def is_ready(self):
+        for ip in config.RECEIVE_IP_ADDRESS_LIST:
+            if (
+                ip not in self.pending_data
+                or self.pending_data[ip][0] <= self.last_processed
+            ):
+                return False
+        return True
+
+    def process_data(self):
+        self.last_processed = datetime.now()
+        antenna_data = [
+            np.expand_dims(self.pending_data[ip][1].matrix[:, antenna], axis=2)
+            for ip, antenna in config.ANTENNA_ORDER
+        ]
+
+        # We have data from all servers
+        all_data = np.concat(antenna_data, axis=1)
+        self.logger.info(f"Got final CSI data with shape {all_data.shape}")
+
+        processed = self.preprocess.preprocess(all_data)
+        visualise.add_data(all_data, processed)
+        self.aoa.update(processed)
+
+    def process_forever(self):
+        while True:
+            if self.is_ready():
+                self.process_data()
+            else:
+                self.logger.debug("Not all data is ready")
+            time.sleep(0.001)

src/main.py

@@ -1,134 +0,0 @@
-import logging
-import socket
-from typing import Callable
-import threading
-import struct
-from datetime import datetime
-import numpy as np
-
-from .csi import CSI
-from .aoa import AoA
-from .preprocess import Preprocessor
-from . import config
-from . import visualise
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s %(name)-8s %(levelname)-8s %(message)s",
-)
-
-logger = logging.getLogger(__name__)
-
-
-class FeitTransmitter:
-    def __init__(self):
-        self.inject_server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        self.inject_server.connect((config.INJECT_IP_ADDRESS, config.FEITCSI_PORT))
-        inject_start_string = (
-            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
-            f"--channel-width {config.CHANNEL_WIDTH} "
-            f"--format {config.FRAME_FORMAT} "
-            f"--mode inject -s 1 --verbose "
-            f"--inject-delay {1_000_000 // config.SAMPLE_RATE}"
-        )
-        self.inject_server.send(b"stop\n")
-        self.inject_server.send(inject_start_string.encode())
-
-
-class FeitServer:
-    def __init__(self, host: str):
-        self.host = host
-        self.server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        self.server.connect((self.host, config.FEITCSI_PORT))
-        self.start_string = (
-            f"feitcsi --frequency {config.CENTRAL_FREQUENCY_MHZ} "
-            f"--channel-width {config.CHANNEL_WIDTH} "
-            f"--format {config.FRAME_FORMAT} "
-            f"--mode measure"
-        )
-        self.server.send(b"stop\n")
-        self.server.send(self.start_string.encode())
-
-    def listen(self, callback: Callable[[str, CSI], None]):
-        prev_time = datetime.now()
-        while True:
-            # This is the max size of a UDP packet. The size of the actual CSI
-            # packet will depend on the frame format and channel width, which
-            # changes the number of subcarriers
-            data = self.server.recv(65535)
-            try:
-                csidata = CSI(data)
-                logger.debug(f"Received CSI data after {datetime.now() - prev_time}")
-                prev_time = datetime.now()
-                callback(self.host, csidata)
-            except struct.error:
-                logger.error("Failed to parse CSI data")
-
-
-pending_data: dict[str, tuple[datetime, CSI]] = {}
-pending_data_lock = threading.Lock()
-
-
-def process_data(host: str, data: CSI):
-    global pending_data
-    current_time = datetime.now()
-
-    logger.debug(f"Got CSI data from {host}")
-
-    with pending_data_lock:
-        if host in pending_data and (
-            current_time - pending_data[host][0]
-        ).total_seconds() < 1 / (2 * config.SAMPLE_RATE):
-            logger.warning(f"Received CSI data from {host} too quickly")
-            return
-        elif host in pending_data:
-            logger.warning(f"Skipping existing CSI data from {host}")
-        pending_data[host] = (datetime.now(), data)
-
-        for ip in config.RECEIVE_IP_ADDRESS_LIST:
-            if ip not in pending_data:
-                return
-
-        logger.debug(
-            "Antenna RSSI values:"
-            + str(
-                [
-                    (ip, csi.header.rssi1, csi.header.rssi2)
-                    for ip, (_, csi) in sorted(pending_data.items())
-                ]
-            ),
-        )
-
-        antenna_data = [
-            np.expand_dims(pending_data[ip][1].matrix[:, antenna], axis=2)
-            for ip, antenna in config.ANTENNA_ORDER
-        ]
-
-        # We have data from all servers
-        all_data = np.concat(antenna_data, axis=1)
-        logger.info(f"Got final CSI data with shape {all_data.shape}")
-        pending_data = {}
-
-    # logging.info(f"Got CSI frame with shape {data.matrix.shape}")
-    processed = preprocess.preprocess(all_data)
-    visualise.add_data(data.matrix, processed)
-    aoa.update(processed)
-
-
-if __name__ == "__main__":
-    servers = [FeitServer(ip) for ip in config.RECEIVE_IP_ADDRESS_LIST]
-    transmitter = FeitTransmitter()
-    preprocess = Preprocessor()
-    aoa = AoA()
-
-    # Start webapp in background thread
-    visualise.aoa = aoa
-    webapp = threading.Thread(target=visualise.start)
-    webapp.start()
-
-    server_threads = [
-        threading.Thread(target=server.listen, args=(process_data,))
-        for server in servers
-    ]
-    for thread in server_threads:
-        thread.start()