refactor: cleanup everything and ready to use as template
This commit is contained in:
173
main.py
173
main.py
@ -18,32 +18,29 @@ import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import streamlit as st
|
||||
import anyio
|
||||
from anyio.abc import TaskGroup
|
||||
from anyio import create_memory_object_stream
|
||||
from anyio.abc import TaskGroup, UDPSocket
|
||||
from anyio import create_memory_object_stream, create_udp_socket
|
||||
from anyio.streams.memory import MemoryObjectSendStream, MemoryObjectReceiveStream
|
||||
from aiomqtt import Client as MqttClient, Message as MqttMessage
|
||||
from threading import Thread
|
||||
from time import sleep
|
||||
from pydantic import BaseModel, computed_field
|
||||
from datetime import datetime
|
||||
import awkward as ak
|
||||
from awkward import Array as AwkwardArray, Record as AwkwardRecord
|
||||
import orjson
|
||||
from app.model import AlgoReport
|
||||
from collections import deque
|
||||
|
||||
|
||||
# https://handmadesoftware.medium.com/streamlit-asyncio-and-mongodb-f85f77aea825
|
||||
class AppState(TypedDict):
|
||||
worker_thread: Thread
|
||||
client: MqttClient
|
||||
message_queue: MemoryObjectReceiveStream[MqttMessage]
|
||||
message_queue: MemoryObjectReceiveStream[bytes]
|
||||
task_group: TaskGroup
|
||||
history: dict[str, AwkwardArray]
|
||||
history: deque[AlgoReport]
|
||||
|
||||
|
||||
MQTT_BROKER: Final[str] = "192.168.2.189"
|
||||
MQTT_BROKER_PORT: Final[int] = 1883
|
||||
UDP_LISTEN_PORT: Final[int] = 50_000
|
||||
MAX_LENGTH = 600
|
||||
TOPIC: Final[str] = "GwData"
|
||||
NDArray = np.ndarray
|
||||
|
||||
T = TypeVar("T")
|
||||
@ -57,146 +54,46 @@ def unwrap(value: Optional[T]) -> T:
|
||||
|
||||
@st.cache_resource
|
||||
def resource(params: Any = None):
|
||||
client: Optional[MqttClient] = None
|
||||
tx, rx = create_memory_object_stream[MqttMessage]()
|
||||
set_ev = anyio.Event()
|
||||
tx, rx = create_memory_object_stream[bytes]()
|
||||
tg: Optional[TaskGroup] = None
|
||||
|
||||
async def main():
|
||||
async def poll_task():
|
||||
nonlocal set_ev
|
||||
nonlocal tg
|
||||
nonlocal client
|
||||
tg = anyio.create_task_group()
|
||||
set_ev.set()
|
||||
async with tg:
|
||||
client = MqttClient(MQTT_BROKER, port=MQTT_BROKER_PORT)
|
||||
async with client:
|
||||
await client.subscribe(TOPIC)
|
||||
logger.info(
|
||||
"Subscribed {}:{} to topic {}", MQTT_BROKER, MQTT_BROKER_PORT, TOPIC
|
||||
)
|
||||
# https://aiomqtt.bo3hm.com/subscribing-to-a-topic.html
|
||||
async for message in client.messages:
|
||||
await tx.send(message)
|
||||
async with await create_udp_socket(
|
||||
local_port=UDP_LISTEN_PORT, reuse_port=True
|
||||
) as udp:
|
||||
async for packet, _ in udp:
|
||||
await tx.send(packet)
|
||||
|
||||
tr = Thread(target=anyio.run, args=(main,))
|
||||
tr = Thread(target=anyio.run, args=(poll_task,))
|
||||
tr.start()
|
||||
sleep(0.1)
|
||||
while not set_ev.is_set():
|
||||
sleep(0.01)
|
||||
logger.info("Poll task initialized")
|
||||
state: AppState = {
|
||||
"worker_thread": tr,
|
||||
"client": unwrap(client),
|
||||
"message_queue": rx,
|
||||
"task_group": unwrap(tg),
|
||||
"history": {},
|
||||
"history": deque(maxlen=MAX_LENGTH),
|
||||
}
|
||||
return state
|
||||
|
||||
|
||||
class GwMessage(TypedDict):
|
||||
v: int
|
||||
mid: int
|
||||
time: int
|
||||
ip: str
|
||||
mac: str
|
||||
devices: list[Any]
|
||||
rssi: int
|
||||
|
||||
|
||||
class DeviceMessage(BaseModel):
|
||||
mac: str
|
||||
"""
|
||||
Hex string, capital letters, e.g. "D6AF1CA9C491"
|
||||
"""
|
||||
service: str
|
||||
"""
|
||||
Hex string, capital letters, e.g. "180D"
|
||||
"""
|
||||
characteristic: str
|
||||
"""
|
||||
Hex string, capital letters, e.g. "2A37"
|
||||
"""
|
||||
value: str
|
||||
"""
|
||||
Hex string, capital letters, e.g. "0056"
|
||||
"""
|
||||
rssi: int
|
||||
|
||||
@property
|
||||
def value_bytes(self) -> bytes:
|
||||
return bytes.fromhex(self.value)
|
||||
|
||||
|
||||
def get_device_data(message: GwMessage) -> List[DeviceMessage]:
|
||||
"""
|
||||
devices
|
||||
|
||||
[[5,"D6AF1CA9C491","180D","2A37","0056",-58],[5,"A09E1AE4E710","180D","2A37","0055",-50]]
|
||||
|
||||
unknown, mac addr, service, characteristic, value (hex), rssi
|
||||
"""
|
||||
l: list[DeviceMessage] = []
|
||||
for d in message["devices"]:
|
||||
x, mac, service, characteristic, value, rssi = d
|
||||
l.append(
|
||||
DeviceMessage(
|
||||
mac=mac,
|
||||
service=service,
|
||||
characteristic=characteristic,
|
||||
value=value,
|
||||
rssi=rssi,
|
||||
)
|
||||
)
|
||||
return l
|
||||
|
||||
|
||||
def payload_to_hr(payload: bytes) -> int:
|
||||
"""
|
||||
ignore the first byte, parse the rest as a big-endian integer
|
||||
|
||||
Bit 0 (Heart Rate Format)
|
||||
0: Heart rate value is 8 bits
|
||||
1: Heart rate value is 16 bits
|
||||
Bit 3 (Energy Expended)
|
||||
Indicates whether energy expended data is present
|
||||
Bit 4 (RR Interval)
|
||||
Indicates whether RR interval data is present
|
||||
"""
|
||||
flags = payload[0]
|
||||
if flags & 0b00000001:
|
||||
return int.from_bytes(payload[1:3], "big")
|
||||
else:
|
||||
return payload[1]
|
||||
|
||||
|
||||
def main():
|
||||
state = resource()
|
||||
logger.info("Resource created")
|
||||
history = state["history"]
|
||||
|
||||
def push_new_message(message: GwMessage):
|
||||
dms = get_device_data(message)
|
||||
now = datetime.now()
|
||||
for dm in dms:
|
||||
rec = AwkwardRecord(
|
||||
{
|
||||
"time": now,
|
||||
"value": payload_to_hr(dm.value_bytes),
|
||||
"rssi": dm.rssi,
|
||||
}
|
||||
)
|
||||
if dm.mac not in history:
|
||||
history[dm.mac] = AwkwardArray([rec])
|
||||
else:
|
||||
history[dm.mac] = ak.concatenate([history[dm.mac], [rec]])
|
||||
if len(history[dm.mac]) > MAX_LENGTH:
|
||||
history[dm.mac] = AwkwardArray(history[dm.mac][-MAX_LENGTH:])
|
||||
|
||||
def on_export():
|
||||
now = datetime.now()
|
||||
filename = f"export-{now.strftime('%Y-%m-%d-%H-%M-%S')}.parquet"
|
||||
ak.to_parquet([history], filename)
|
||||
logger.info("Export to {}", filename)
|
||||
raise NotImplementedError
|
||||
|
||||
def on_clear():
|
||||
history.clear()
|
||||
logger.info("History cleared")
|
||||
raise NotImplementedError
|
||||
|
||||
st.button(
|
||||
"Export", help="Export the current data to a parquet file", on_click=on_export
|
||||
@ -208,25 +105,9 @@ def main():
|
||||
message = state["message_queue"].receive_nowait()
|
||||
except anyio.WouldBlock:
|
||||
continue
|
||||
m: str
|
||||
if isinstance(message.payload, str):
|
||||
m = message.payload
|
||||
elif isinstance(message.payload, bytes):
|
||||
m = message.payload.decode("utf-8")
|
||||
else:
|
||||
logger.warning("Unknown message type: {}", type(message.payload))
|
||||
continue
|
||||
d = cast(GwMessage, orjson.loads(m))
|
||||
push_new_message(d)
|
||||
|
||||
def to_scatter(key: str, dev_history: AwkwardArray):
|
||||
x = ak.to_numpy(dev_history["time"])
|
||||
y = ak.to_numpy(dev_history["value"])
|
||||
return go.Scatter(x=x, y=y, mode="lines+markers", name=key)
|
||||
|
||||
scatters = [to_scatter(k, el) for k, el in history.items()]
|
||||
fig = go.Figure(scatters)
|
||||
pannel.plotly_chart(fig)
|
||||
# TODO: plot
|
||||
# fig = go.Figure(scatters)
|
||||
# pannel.plotly_chart(fig)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user