145 lines
4.2 KiB
Python
145 lines
4.2 KiB
Python
import asyncio
|
|
import dataclasses
|
|
import logging
|
|
from threading import Thread
|
|
|
|
from fastapi import FastAPI, Response
|
|
|
|
from flandre.nodes.Node import Node
|
|
from flandre.utils.Msg import ImageArgMsg, RobotRtsiMsg, SetSidMsg
|
|
from flandre.utils.RfSequence import RfSequence
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
import uvicorn
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
# class ConnectionManager:
|
|
# def __init__(self):
|
|
# # 存放激活的ws连接对象
|
|
# self.active_connections: List[WebSocket] = []
|
|
# self.router = APIRouter()
|
|
# self.router.add_api_route("/hello", self.hello, methods=["GET"])
|
|
#
|
|
#
|
|
# def hello(self):
|
|
# return {"Hello": self.name}
|
|
|
|
# async def connect(self, ws: WebSocket):
|
|
# # 等待连接
|
|
# await ws.accept()
|
|
# # 存储ws连接对象
|
|
# self.active_connections.append(ws)
|
|
#
|
|
# def disconnect(self, ws: WebSocket):
|
|
# # 关闭时 移除ws对象
|
|
# self.active_connections.remove(ws)
|
|
#
|
|
# @staticmethod
|
|
# async def send_personal_message(message: str, ws: WebSocket):
|
|
# # 发送个人消息
|
|
# await ws.send_text(message)
|
|
#
|
|
# async def broadcast(self, message: str):
|
|
# # 广播消息
|
|
# for connection in self.active_connections:
|
|
# await connection.send_text(message)
|
|
#
|
|
#
|
|
# manager = ConnectionManager()
|
|
#
|
|
#
|
|
# @app.websocket("/ws/{user}")
|
|
# async def websocket_endpoint(websocket: WebSocket, user: str):
|
|
# await manager.connect(websocket)
|
|
#
|
|
# await manager.broadcast(f"用户{user}进入聊天室")
|
|
#
|
|
# try:
|
|
# while True:
|
|
# data = await websocket.receive_text()
|
|
# await manager.send_personal_message(f"你说了: {data}", websocket)
|
|
# await manager.broadcast(f"用户:{user} 说: {data}")
|
|
#
|
|
# except WebSocketDisconnect:
|
|
# manager.disconnect(websocket)
|
|
# await manager.broadcast(f"用户-{user}-离开")
|
|
|
|
|
|
class Web(Node):
|
|
topics = [ImageArgMsg, RobotRtsiMsg, SetSidMsg]
|
|
|
|
def __init__(self, level=logging.INFO):
|
|
super(Web, self).__init__(level=level)
|
|
self.wss: list[WebSocket] = []
|
|
self.arg = ImageArgMsg(sender="frontend", t_end=123)
|
|
|
|
def custom_setup(self):
|
|
self.router = APIRouter()
|
|
self.router.add_api_route("/hello", self.hello, methods=["GET"])
|
|
self.router.add_websocket_route("/ws", self.websocket_endpoint)
|
|
self.router.add_api_route(
|
|
"/p",
|
|
self.get_image,
|
|
response_class=Response,
|
|
responses={200: {"content": {"image/png": {}}}},
|
|
)
|
|
|
|
def hello(self):
|
|
return {"Hello": "asd"}
|
|
|
|
def get_image(self, i: int):
|
|
rfs = RfSequence()
|
|
f = rfs.frames[i]
|
|
|
|
image_bytes: bytes = b""
|
|
return Response(content=image_bytes, media_type="image/png")
|
|
|
|
async def websocket_endpoint(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.wss.append(websocket)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
if data["type"] == "ImageArgMsg":
|
|
t = dataclasses.replace(self.arg)
|
|
t.sender = data["sender"]
|
|
t.t_start = data["t_start"]
|
|
t.g8 = data["g8"]
|
|
self.send(t)
|
|
except WebSocketDisconnect:
|
|
self.wss.remove(websocket)
|
|
|
|
def boardcast(self, data: dict):
|
|
try:
|
|
for ws in self.wss:
|
|
asyncio.run(ws.send_json(data))
|
|
except Exception as e:
|
|
logger.warning(e)
|
|
|
|
def wst(self):
|
|
app = FastAPI()
|
|
app.include_router(self.router)
|
|
uvicorn.run(app=app, host="0.0.0.0", port=8010)
|
|
|
|
def loop(self):
|
|
Thread(target=self.wst).start()
|
|
while True:
|
|
msg = self.recv()
|
|
if isinstance(msg, ImageArgMsg):
|
|
if msg.sender != "frontend":
|
|
self.arg = msg
|
|
self.boardcast(msg.dict)
|
|
elif isinstance(msg, RobotRtsiMsg):
|
|
self.boardcast(msg.dict)
|
|
elif isinstance(msg, SetSidMsg):
|
|
self.boardcast(msg.dict)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
Web()()
|