Multiple arms support / API router

pyproject.toml

@@ -25,7 +25,7 @@ Source = "https://github.com/SpesRobotics/teleop"
 Tracker = "https://github.com/SpesRobotics/teleop/issues"
 
 [project.optional-dependencies]
-utils = ["pin"]
+utils = ["pin", "ruckig"]
 
 [tool.setuptools.packages.find]
 include = ["teleop*"]

teleop/__init__.py

@@ -2,9 +2,9 @@
 import math
 import socket
 import logging
-from typing import Callable, List
+from typing import Callable, List, NamedTuple
 import uvicorn
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import APIRouter, FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 import transforms3d as t3d
@@ -15,6 +15,11 @@
 THIS_DIR = os.path.dirname(os.path.realpath(__file__))
 
 
+class FrontendMount(NamedTuple):
+    path: str
+    directory: str
+
+
 def get_local_ip():
     try:
         # Connect to an external address (doesn't actually send data)
@@ -147,7 +152,9 @@ class Teleop:
     Args:
         host (str, optional): The host IP address. Defaults to "0.0.0.0".
         port (int, optional): The port number. Defaults to 4443.
-        frontend_dir (str, optional): The directory containing the frontend assets. Defaults to THIS_DIR.
+        frontend_dir (str | list[tuple[str, str]], optional): The directory
+            containing the frontend assets, or a list of route prefix and
+            frontend directory pairs. Defaults to THIS_DIR.
     """
 
     def __init__(
@@ -157,6 +164,8 @@ def __init__(
         natural_phone_orientation_euler=None,
         natural_phone_position=None,
         frontend_dir=None,
+        offset_user_orientation=0,
+        current_pose_provider=None,
     ):
         self.__logger = logging.getLogger("teleop")
         self.__logger.setLevel(logging.INFO)
@@ -170,6 +179,8 @@ def __init__(
         self.__previous_received_pose = None
         self.__callbacks = []
         self.__pose = np.eye(4)
+        self.__previous_move = False
+        self.__current_pose_provider = current_pose_provider
 
         if natural_phone_orientation_euler is None:
             natural_phone_orientation_euler = [0, math.radians(-45), 0]
@@ -180,10 +191,9 @@ def __init__(
             t3d.euler.euler2mat(*natural_phone_orientation_euler),
             [1, 1, 1],
         )
+        self.__offset_user_orientation = t3d.affines.compose([0, 0, 0], t3d.euler.euler2mat(0, 0, offset_user_orientation), [1, 1, 1])
 
-        if frontend_dir is None:
-            frontend_dir = THIS_DIR
-        self.__frontend_dir = frontend_dir
+        self.__frontend_mounts = self.__normalize_frontend_mounts(frontend_dir)
 
         self.__app = FastAPI()
         self.__manager = ConnectionManager()
@@ -192,6 +202,52 @@ def __init__(
         logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
         self.__setup_routes()
 
+    def include_router(self, router: APIRouter, **kwargs) -> None:
+        self.__app.include_router(router, **kwargs)
+
+    def __normalize_frontend_mounts(self, frontend_dir):
+        if frontend_dir is None:
+            frontend_dir = THIS_DIR
+
+        if isinstance(frontend_dir, (str, os.PathLike)):
+            return [FrontendMount("/", os.fspath(frontend_dir))]
+
+        mounts = []
+        seen_paths = set()
+        for mount in frontend_dir:
+            try:
+                path, directory = mount
+            except (TypeError, ValueError):
+                raise ValueError(
+                    "frontend_dir entries must be (path, directory) pairs"
+                ) from None
+
+            path = self.__normalize_frontend_path(path)
+            if path in seen_paths:
+                raise ValueError(f"Duplicate frontend mount path: {path}")
+
+            seen_paths.add(path)
+            mounts.append(FrontendMount(path, os.fspath(directory)))
+
+        if not mounts:
+            raise ValueError("frontend_dir mount list cannot be empty")
+
+        return mounts
+
+    @staticmethod
+    def __normalize_frontend_path(path):
+        path = str(path).strip()
+        if not path:
+            raise ValueError("Frontend mount path cannot be empty")
+
+        if not path.startswith("/"):
+            path = f"/{path}"
+
+        if path != "/":
+            path = path.rstrip("/")
+
+        return path
+
     def set_pose(self, pose: np.ndarray) -> None:
         """
         Set the current pose of the end-effector.
@@ -222,6 +278,9 @@ def __update(self, message):
         position = message["position"]
         orientation = message["orientation"]
         scale = message.get("scale", 1.0)
+        #Detect rising edge of move signal
+        move_started = move and not self.__previous_move
+        self.__previous_move = move
 
         position = np.array([position["x"], position["y"], position["z"]])
         quat = np.array(
@@ -233,15 +292,21 @@ def __update(self, message):
             self.__absolute_pose_init = None
             self.__notify_subscribers(self.__pose, message)
             return
-
+        if move_started and self.__current_pose_provider is not None: # Rising edge
+            try:
+                current_pose = self.__current_pose_provider()
+                self.__pose = np.array(current_pose, dtype=float, copy=True)
+            except Exception:
+                self.__logger.warning("Failed to get current pose, using last known pose")  
+
         received_pose_rub = t3d.affines.compose(
             position, t3d.quaternions.quat2mat(quat), [1, 1, 1]
         )
         received_pose = TF_RUB2FLU @ received_pose_rub
         received_pose[:3, :3] = received_pose[:3, :3] @ np.linalg.inv(
             TF_RUB2FLU[:3, :3]
         )
-        received_pose = received_pose @ self.__natural_phone_pose
+        received_pose = self.__offset_user_orientation @ received_pose @ self.__natural_phone_pose
 
         # Pose jump protection
         if self.__previous_received_pose is not None:
@@ -285,14 +350,8 @@ def __update(self, message):
         self.__notify_subscribers(self.__pose, message)
 
     def __setup_routes(self):
-        # Mount static files directory
-        assets_dir = os.path.join(self.__frontend_dir, "assets")
-        self.__app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
-
-        @self.__app.get("/")
-        async def index():
-            self.__logger.debug("Serving the index.html file")
-            return FileResponse(os.path.join(self.__frontend_dir, "index.html"))
+        for mount_index, mount in enumerate(self.__frontend_mounts):
+            self.__mount_frontend(mount, mount_index)
 
         @self.__app.websocket("/ws")
         async def websocket_endpoint(websocket: WebSocket):
@@ -314,6 +373,43 @@ async def websocket_endpoint(websocket: WebSocket):
                 self.__manager.disconnect(websocket)
                 self.__logger.info("Client disconnected")
 
+    def __mount_frontend(self, mount: FrontendMount, mount_index: int):
+        assets_dir = os.path.join(mount.directory, "assets")
+        assets_path = "/assets" if mount.path == "/" else f"{mount.path}/assets"
+        self.__app.mount(
+            assets_path,
+            StaticFiles(directory=assets_dir),
+            name=f"frontend_assets_{mount_index}",
+        )
+
+        async def index():
+            self.__logger.debug(
+                f"Serving {mount.path} frontend index.html from {mount.directory}"
+            )
+            return FileResponse(os.path.join(mount.directory, "index.html"))
+
+        if mount.path == "/":
+            self.__app.add_api_route("/", index, methods=["GET"])
+            return
+
+        self.__app.add_api_route(mount.path, index, methods=["GET"])
+        self.__app.add_api_route(f"{mount.path}/", index, methods=["GET"])
+
+        @self.__app.get(f"{mount.path}/{{path:path}}")
+        async def prefixed_spa_index(path: str):
+            requested_path = os.path.abspath(os.path.join(mount.directory, path))
+            frontend_root = os.path.abspath(mount.directory)
+            if (
+                os.path.commonpath([frontend_root, requested_path]) == frontend_root
+                and os.path.isfile(requested_path)
+            ):
+                return FileResponse(requested_path)
+
+            self.__logger.debug(
+                f"Serving {mount.path} frontend index.html from {mount.directory}"
+            )
+            return FileResponse(os.path.join(mount.directory, "index.html"))
+
     def run(self) -> None:
         """
         Runs the teleop server. This method is blocking.