Initial web ui and control for robot

2026-04-30 21:35:18 +00:00
parent d64e1c24c8
commit 91f5f4d3ab
23 changed files with 2688 additions and 4 deletions
@@ -0,0 +1,105 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import FileResponse, JSONResponse
+
+from ros_node import get_node, start_ros, stop_ros
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+STATIC_DIR = Path(__file__).parent / 'static'
+
+
+class ConnectionManager:
+    def __init__(self) -> None:
+        self.connections: list[WebSocket] = []
+        self._loop: asyncio.AbstractEventLoop | None = None
+
+    def set_loop(self, loop: asyncio.AbstractEventLoop) -> None:
+        self._loop = loop
+
+    async def connect(self, ws: WebSocket) -> None:
+        await ws.accept()
+        self.connections.append(ws)
+
+    def disconnect(self, ws: WebSocket) -> None:
+        if ws in self.connections:
+            self.connections.remove(ws)
+
+    def broadcast_from_thread(self, data: dict) -> None:
+        if self._loop and self._loop.is_running():
+            asyncio.run_coroutine_threadsafe(self._broadcast(data), self._loop)
+
+    async def _broadcast(self, data: dict) -> None:
+        dead = []
+        for ws in self.connections:
+            try:
+                await ws.send_json(data)
+            except Exception:
+                dead.append(ws)
+        for ws in dead:
+            self.disconnect(ws)
+
+
+manager = ConnectionManager()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    manager.set_loop(asyncio.get_event_loop())
+    node = start_ros()
+    node.on_joint_states_callbacks.append(manager.broadcast_from_thread)
+    node.on_ultrasonic_callbacks.append(manager.broadcast_from_thread)
+    yield
+    stop_ros()
+
+
+app = FastAPI(lifespan=lifespan)
+
+
+@app.websocket('/ws')
+async def websocket_endpoint(ws: WebSocket):
+    await manager.connect(ws)
+    node = get_node()
+    try:
+        while True:
+            data = await ws.receive_json()
+            msg_type = data.get('type')
+            if msg_type == 'cmd_vel':
+                node.publish_cmd_vel(data.get('linear_x', 0.0), data.get('angular_z', 0.0))
+            elif msg_type == 'joint_command':
+                node.publish_joint_command(data.get('pan', 0.0), data.get('tilt', 0.0))
+    except WebSocketDisconnect:
+        pass
+    except Exception as e:
+        logger.error('WebSocket error: %s', e)
+    finally:
+        manager.disconnect(ws)
+        try:
+            node.publish_cmd_vel(0.0, 0.0)
+        except Exception:
+            pass
+
+
+@app.get('/{full_path:path}')
+async def serve_spa(full_path: str):
+    index = STATIC_DIR / 'index.html'
+    if not index.exists():
+        return JSONResponse(
+            {'detail': 'Frontend not built. In dev mode open http://localhost:5173'},
+            status_code=404,
+        )
+    candidate = (STATIC_DIR / full_path).resolve()
+    static_root = STATIC_DIR.resolve()
+    if str(candidate).startswith(str(static_root)) and full_path and candidate.is_file():
+        return FileResponse(candidate)
+    return FileResponse(index)
+
+
+if __name__ == '__main__':
+    import uvicorn
+    uvicorn.run(app, host='0.0.0.0', port=8080)
@@ -0,0 +1,3 @@
+fastapi>=0.115.0
+uvicorn[standard]>=0.32.0
+numpy
@@ -0,0 +1,76 @@
+import threading
+
+import rclpy
+from geometry_msgs.msg import Twist
+from rclpy.node import Node
+from sensor_msgs.msg import JointState, Range
+
+
+class RobotBridgeNode(Node):
+    def __init__(self):
+        super().__init__('webui_bridge')
+
+        self._cmd_vel_pub = self.create_publisher(Twist, '/cmd_vel', 10)
+        self._joint_cmd_pub = self.create_publisher(JointState, '/joint_command', 10)
+
+        self.create_subscription(JointState, '/joint_states', self._on_joint_states, 10)
+        self.create_subscription(Range, '/ultrasonic/range', self._on_ultrasonic, 10)
+
+        self.on_joint_states_callbacks: list = []
+        self.on_ultrasonic_callbacks: list = []
+
+    def publish_cmd_vel(self, linear_x: float, angular_z: float) -> None:
+        # print(f'Publishing cmd_vel: linear_x={linear_x}, angular_z={angular_z}')
+        msg = Twist()
+        msg.linear.x = float(linear_x)
+        msg.angular.z = float(angular_z)
+        try:
+            self._cmd_vel_pub.publish(msg)
+        except Exception as e:
+            print('Failed to publish cmd_vel:', e)
+        # self._cmd_vel_pub.publish(msg)
+
+    def publish_joint_command(self, pan: float, tilt: float) -> None:
+        msg = JointState()
+        msg.header.stamp = self.get_clock().now().to_msg()
+        msg.name = ['pan', 'tilt']
+        msg.position = [float(pan), float(tilt)]
+        self._joint_cmd_pub.publish(msg)
+
+    def _on_joint_states(self, msg: JointState) -> None:
+        data = {'type': 'joint_states', 'positions': dict(zip(msg.name, msg.position))}
+        for cb in self.on_joint_states_callbacks:
+            cb(data)
+
+    def _on_ultrasonic(self, msg: Range) -> None:
+        data = {'type': 'ultrasonic', 'range': msg.range}
+        for cb in self.on_ultrasonic_callbacks:
+            cb(data)
+
+
+_node: RobotBridgeNode | None = None
+_spin_thread: threading.Thread | None = None
+
+
+def start_ros() -> RobotBridgeNode:
+    global _node, _spin_thread
+    rclpy.init()
+    _node = RobotBridgeNode()
+    _spin_thread = threading.Thread(target=rclpy.spin, args=(_node,), daemon=True)
+    _spin_thread.start()
+    return _node
+
+
+def get_node() -> RobotBridgeNode:
+    assert _node is not None, "ROS node not started"
+    return _node
+
+
+def stop_ros() -> None:
+    global _node
+    try:
+        if _node:
+            _node.destroy_node()
+        rclpy.shutdown()
+    except Exception:
+        pass
@@ -0,0 +1,39 @@
+import asyncio
+import os
+
+import av
+import cv2
+from aiortc import VideoStreamTrack
+
+CAMERA_DEVICE = int(os.environ.get('CAMERA_DEVICE', '0'))
+
+
+class CameraVideoTrack(VideoStreamTrack):
+    def __init__(self):
+        super().__init__()
+        self._cap = cv2.VideoCapture(CAMERA_DEVICE)
+        self._cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
+        self._cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
+        self._cap.set(cv2.CAP_PROP_FPS, 30)
+        self._cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+
+    async def recv(self) -> av.VideoFrame:
+        pts, time_base = await self.next_timestamp()
+
+        loop = asyncio.get_event_loop()
+        ret, frame = await loop.run_in_executor(None, self._cap.read)
+
+        if not ret:
+            import numpy as np
+            frame = np.zeros((480, 640, 3), dtype=np.uint8)
+        else:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+        video_frame.pts = pts
+        video_frame.time_base = time_base
+        return video_frame
+
+    def __del__(self) -> None:
+        if self._cap:
+            self._cap.release()