218 lines
7.0 KiB
Python
218 lines
7.0 KiB
Python
import logging
|
|
import multiprocessing as mp
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable, NamedTuple
|
|
|
|
import numpy.typing as npt
|
|
|
|
from where_fi.collection import CSIMatrix, ingest
|
|
from where_fi.collection.csi_frame import CSI
|
|
from where_fi.collection.protocols import CSIProducer, MergedCSI
|
|
from where_fi.config import config
|
|
from where_fi.processing.preprocess import Preprocessor
|
|
from where_fi.visualise import server as visualise
|
|
|
|
|
|
class Receiver(NamedTuple):
|
|
ip: ingest.Host
|
|
receiver: ingest.FeitReceiver
|
|
thread: threading.Thread
|
|
|
|
|
|
class Scheduler:
|
|
"""
|
|
A simple scheduler that runs a function at a given rate.
|
|
"""
|
|
|
|
def __init__(self, rate: float, func: Callable[[], None]) -> None:
|
|
self.rate = rate
|
|
self.func = func
|
|
self.active = True
|
|
|
|
def run(self) -> None:
|
|
"""
|
|
Run the scheduler in a loop, calling the function at the given rate.
|
|
"""
|
|
while self.active:
|
|
self.func()
|
|
time.sleep(1 / self.rate)
|
|
|
|
def start(self) -> None:
|
|
"""
|
|
Start the scheduler in a separate thread.
|
|
"""
|
|
self.thread = threading.Thread(target=self.run)
|
|
self.thread.start()
|
|
|
|
def stop(self) -> None:
|
|
"""
|
|
Stop the scheduler.
|
|
"""
|
|
self.active = False
|
|
|
|
|
|
class CSIApplication:
|
|
"""
|
|
A high-level application that manages CSI data ingestion and processing.
|
|
|
|
This allows a simple interface to be used for the main CLI logic, without having to
|
|
deal with the threads and queues directly.
|
|
|
|
It can be used using decortors to register callbacks for different stages of the
|
|
processing pipeline:
|
|
|
|
- `on_pre_merge`: Called with the raw CSI data (including headers) from each
|
|
receiver, before it is merged into a single matrix.
|
|
- `on_sample`: Called once per sample with the merged CSI data.
|
|
- `on_process`: Called at the processing sample rate, without any data (this is
|
|
mostly used as a scheduler).
|
|
"""
|
|
|
|
def __init__(self, producer: CSIProducer, visualise_raw: bool = False) -> None:
|
|
self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
|
|
|
|
self.producer = producer
|
|
|
|
self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
|
|
config.processing_sample_rate
|
|
)
|
|
self.webapp = visualise.Webapp()
|
|
self.webapp_thread = threading.Thread(
|
|
target=self.webapp.start, args=(self.webapp_queue,)
|
|
)
|
|
self.visualise_raw = visualise_raw
|
|
|
|
self.pre_merge_callback = None
|
|
self.raw_csi_callback = None
|
|
self.preprocessed_csi_callback = None
|
|
self.processing_callback = None
|
|
|
|
self.preprocessor = Preprocessor()
|
|
|
|
def visualise_data(
|
|
self, data: npt.NDArray[Any], dtype: visualise.figures.Figure
|
|
) -> None:
|
|
"""
|
|
Update a visualisation with the given data. The available visualisations are as
|
|
per visualise.server.all_figures.
|
|
"""
|
|
if not self.webapp_queue.full():
|
|
self.webapp_queue.put(visualise.VisualiserData(data, dtype))
|
|
|
|
def on_raw(self, func: Callable[[CSIMatrix], None]) -> Callable[[CSIMatrix], None]:
|
|
"""
|
|
Decorator to register a callback for the sample preprocessing.
|
|
"""
|
|
|
|
def decorator(data: CSIMatrix) -> None:
|
|
func(data)
|
|
|
|
self.raw_csi_callback = decorator
|
|
return decorator
|
|
|
|
def on_sample(
|
|
self, func: Callable[[CSIMatrix], None]
|
|
) -> Callable[[CSIMatrix], None]:
|
|
"""
|
|
Decorator to register a callback for the sample preprocessing.
|
|
"""
|
|
|
|
def decorator(data: CSIMatrix) -> None:
|
|
self.visualise_data(data, visualise.figures.Figure.RAW_CSI)
|
|
func(data)
|
|
|
|
self.preprocessed_csi_callback = decorator
|
|
return decorator
|
|
|
|
def on_pre_merge(
|
|
self, func: Callable[[dict[ingest.Host, CSI]], None]
|
|
) -> Callable[[dict[ingest.Host, CSI]], None]:
|
|
"""
|
|
Decorator to register a callback for the samples before merging.
|
|
"""
|
|
self.pre_merge_callback = func
|
|
return func
|
|
|
|
def on_process(
|
|
self, func: Callable[[CSIMatrix], None]
|
|
) -> Callable[[CSIMatrix], None]:
|
|
"""
|
|
Decorator to register a callback for the sample processing.
|
|
"""
|
|
self.processing_callback = func
|
|
return func
|
|
|
|
def process_sample(self, sample: MergedCSI) -> None:
|
|
"""Data pipeline for processing a sample.
|
|
This is the entry point for a sample being received from the producer. Depending
|
|
on the callbacks which are registered, this will call the appropriate functions
|
|
to pre-process and visualise the sample.
|
|
"""
|
|
if self.visualise_raw:
|
|
self.visualise_data(sample.matrix, visualise.figures.Figure.RAW_CSI)
|
|
if self.raw_csi_callback is not None:
|
|
self.raw_csi_callback(sample.matrix)
|
|
if self.pre_merge_callback is not None:
|
|
self.pre_merge_callback(sample.frames)
|
|
processed = self.preprocessor.preprocess(sample.matrix)
|
|
if self.visualise_raw:
|
|
self.visualise_data(processed, visualise.figures.Figure.PROCESSED_CSI)
|
|
if self.preprocessed_csi_callback is not None:
|
|
self.preprocessed_csi_callback(processed)
|
|
|
|
def listen(self) -> None:
|
|
"""
|
|
Listen for incoming data from the selected ingestor, and call the appropriate
|
|
callback.
|
|
"""
|
|
data_gen = self.producer()
|
|
while True:
|
|
try:
|
|
sample = next(data_gen)
|
|
self.process_sample(sample)
|
|
except StopIteration:
|
|
break
|
|
except Exception as e:
|
|
self.logger.error(f"Error processing sample: {e}", exc_info=True)
|
|
break
|
|
|
|
def start(self) -> None:
|
|
# for receiver in self.receivers:
|
|
# receiver.thread.start()
|
|
self.buffer_thread = threading.Thread(
|
|
target=self.listen,
|
|
)
|
|
self.buffer_thread.start()
|
|
self.webapp_thread.start()
|
|
|
|
if self.processing_callback is not None:
|
|
|
|
def get_proc_sample() -> None:
|
|
sample = self.preprocessor.last_sample
|
|
if sample is not None and self.processing_callback is not None:
|
|
self.processing_callback(sample)
|
|
else:
|
|
self.logger.warning("No samples to process")
|
|
|
|
self.scheduler = Scheduler(config.processing_sample_rate, get_proc_sample)
|
|
self.scheduler.start()
|
|
try:
|
|
self.webapp_thread.join()
|
|
except KeyboardInterrupt:
|
|
self.stop()
|
|
|
|
def stop(self) -> None:
|
|
self.logger.info("Stopping application")
|
|
|
|
self.producer.stop()
|
|
|
|
if hasattr(self, "scheduler"):
|
|
self.scheduler.stop()
|
|
self.scheduler.thread.join()
|
|
self.buffer_thread.join()
|
|
|
|
self.webapp.active = False
|
|
self.webapp_thread.join()
|
|
self.logger.info("Application stopped")
|