diff --git a/.gitignore b/.gitignore index 8782aa4..61f48db 100644 --- a/.gitignore +++ b/.gitignore @@ -57,6 +57,11 @@ Thumbs.db credentials.json id_rsa - -mavlink-router/ -log + +mavlink-router/ +log + +*.swp + +# Local deep-dive docs (personal study notes) +docs_deep_dive/ diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..0a4bbea --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,72 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Setup + +```bash +git submodule update --init --recursive # clones src/modules/mavctl only +pip install -r requirements.txt +``` + +`dep/labeller` is NOT a registered submodule — that directory is empty. Do not import from `dep.labeller`; use `src.modules.imaging.detector` instead. + +## Commands + +```bash +./scripts/lint.sh # flake8, max-line-length 140 +./scripts/fmt.sh # yapf formatter +./scripts/type.sh # mypy on src/modules/imaging/ and test/ +./scripts/test.sh # pytest (uses PYTHONPATH=".:dep/labeller") + +# Run tests directly (preferred locally — avoids test_battery.py crash): +python -m pytest test/ --ignore=test/test_battery.py -v + +# Run a single test: +python -m pytest test/test_analysis.py::test_analysis_subscriber -v + +# Run a sample script: +samples/run.sh aruco_stream +``` + +CI runs lint and typecheck on every push; tests only run on push/PR to `main`. + +## Architecture + +Shepard runs on a Raspberry Pi or ODroid connected to a PixHawk via MAVLink. It serves an aiohttp web server that the Emu ground station connects to. + +### Data flow + +``` +Camera (RPiCamera / OakdCamera / DebugCamera) + └─► SharedFrameCamera # single capture thread, shared frame + ├─► VideoEmuStreamer # encodes JPEG, pushes to Emu /video endpoint + └─► ImageAnalysisDelegate # runs detector in background thread + └─► BaseDetector.predict() → BoundingBox + └─► get_object_location() → (x, y) direction vector + └─► subscribers(image, bounding_box, pos) +``` + +### Key abstractions + +- **`CameraProvider`** (`camera.py`): base class for all cameras. `capture()` returns `PIL.Image`. `SharedFrameCamera` wraps any provider and serves the latest frame to multiple consumers thread-safely. +- **`BaseDetector`** (`detector.py`): single method `predict(image) → Optional[BoundingBox]`. Implementations: `ArucoDetector` (cv2 DICT_4X4_50), `IrDetector` (brightness threshold), `BucketDetector` (YOLO via ultralytics). +- **`ImageAnalysisDelegate`** (`analysis.py`): runs `camera.capture() → detector.predict()` in a loop in a background thread. Subscribers receive `(image, bounding_box, pos)` — `pos` is currently a raw direction vector in meters (not lon/lat; `XY_To_LonLat` call is commented out in `inference_georeference.py`). +- **`Emu`** (`emu/emu.py`): aiohttp server. Routes: `/ws` (WebSocket for telemetry/commands), `/images/` (static), `/video` (latest JPEG frame). `send_video_frame(jpeg_bytes)` updates the frame; `VideoEmuStreamer` calls this from a background thread. +- **`VideoEmuStreamer`** (`video_emu_stream.py`): pulls frames from `SharedFrameCamera.capture()`, encodes to JPEG, calls `emu.send_video_frame()`. Frame rate set by `fps` param. +- **`Navigator`** (`autopilot/navigator.py`): wraps DroneKit vehicle for flight control via MAVLink. +- **`LocationProvider`** / **`MAVLinkDelegate`** (`location.py`, `mavlink.py`): provide GPS, heading, altitude, orientation from MAVLink messages. + +### Subscriber pattern + +`ImageAnalysisDelegate.subscribe(callback)` — callback signature: +```python +def on_detection(image: Image, bounding_box: Optional[BoundingBox], pos: Optional[tuple[float, float]]): + ... +``` + +### Testing without hardware + +Use `DebugCamera("res/test-image.jpeg")` and `DebugLocationProvider()`. Run `samples/aruco_stream.py` then open `http://localhost:8080/video` to verify streaming. + +`test_battery.py` is not a real test — it connects to a real drone at import time. Always `--ignore=test/test_battery.py` when running locally. diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..2266419 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +testpaths = test diff --git a/samples/aruco_stream.py b/samples/aruco_stream.py new file mode 100644 index 0000000..1113b55 --- /dev/null +++ b/samples/aruco_stream.py @@ -0,0 +1,52 @@ +import time + +from src.modules.emu import Emu +from src.modules.imaging.detector import ArucoDetector +from src.modules.imaging.camera import DebugCamera +from src.modules.imaging.location import DebugLocationProvider +from src.modules.imaging.analysis import ImageAnalysisDelegate +from src.modules.imaging.camera import SharedFrameCamera +from src.modules.imaging.video_emu_stream import VideoEmuStreamer + + +def on_detection(image, bounding_box, position): + if bounding_box is not None: + print(f"ArUco detected at position {position}") + + +def main(): + emu = Emu("tmp") + emu.start_comms() + time.sleep(1) + + # Base camera — swap DebugCamera for RPiCamera/OakdCamera on real hardware + base_camera = DebugCamera("res/test-image.jpeg") + + # SharedFrameCamera captures at 15fps; both video stream and analysis read from it + shared_cam = SharedFrameCamera(base_camera, fps=15) + shared_cam.start() + + # Video stream → EMU /video endpoint (browser polls: ) + video_streamer = VideoEmuStreamer(emu, shared_cam, fps=15, quality=70) + video_streamer.start() + + detector = ArucoDetector() + location_provider = DebugLocationProvider() + + analysis = ImageAnalysisDelegate(detector, shared_cam, location_provider) + analysis.subscribe(on_detection) + analysis.start() + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + pass + finally: + analysis.stop() + video_streamer.stop() + shared_cam.stop() + + +if __name__ == "__main__": + main() diff --git a/samples/slam_stream.py b/samples/slam_stream.py new file mode 100644 index 0000000..93f55ba --- /dev/null +++ b/samples/slam_stream.py @@ -0,0 +1,30 @@ +""" +Test SLAM streaming without any hardware. +Run this, open Emu, go to SLAM tab — you'll see a fake growing point cloud. +""" +import time + +from src.modules.emu import Emu +from src.modules.slam.debug_slam_provider import DebugSLAMProvider +from src.modules.slam.slam_emu_streamer import SLAMEmuStreamer + + +def main(): + emu = Emu("tmp") + emu.start_comms() + time.sleep(1) + + streamer = SLAMEmuStreamer(emu, DebugSLAMProvider()) + emu.register_slam_streamer(streamer) + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + pass + finally: + streamer.stop() + + +if __name__ == "__main__": + main() diff --git a/samples/slam_stream_real.py b/samples/slam_stream_real.py new file mode 100644 index 0000000..6489c68 --- /dev/null +++ b/samples/slam_stream_real.py @@ -0,0 +1,57 @@ +""" +Real SLAM stream using OAK-D depth camera + XM125 altimeter + MAVLink pose. +Run on the Pi/ODroid with hardware connected. +""" +import time + +from src.modules.autopilot.altimeter_xm125 import XM125 +from src.modules.emu import Emu +from src.modules.imaging.camera import OakdCamera +from src.modules.imaging.location import MAVLinkLocationProvider +from src.modules.imaging.mavlink import MAVLinkDelegate +from src.modules.slam.oakd_altimeter_slam_provider import OakdAltimeterSLAMProvider +from src.modules.slam.slam_emu_streamer import SLAMEmuStreamer + +MAVLINK_CONNECTION = "udp:127.0.0.1:14550" + + +def main(): + # MAVLink for drone pose + mavlink = MAVLinkDelegate(MAVLINK_CONNECTION) + location_provider = MAVLinkLocationProvider(mavlink) + + # XM125 radar altimeter (I2C bus 1, default address 0x52) + altimeter = XM125() + if not altimeter.begin(): + print("XM125 init failed — check wiring") + return + + # OAK-D depth camera + oakd = OakdCamera(fps=10) + oakd.start() + + # SLAM provider fuses the two + slam_provider = OakdAltimeterSLAMProvider(oakd, altimeter, location_provider) + + # Emu server + streamer + emu = Emu("tmp") + emu.start_comms() + time.sleep(1) + + streamer = SLAMEmuStreamer(emu, slam_provider) + emu.register_slam_streamer(streamer) + + print("Ready — open Emu and go to the SLAM tab") + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + pass + finally: + streamer.stop() + oakd.stop() + + +if __name__ == "__main__": + main() diff --git a/src/modules/emu/emu.py b/src/modules/emu/emu.py index b808e72..1445639 100644 --- a/src/modules/emu/emu.py +++ b/src/modules/emu/emu.py @@ -2,9 +2,8 @@ import queue import asyncio -from typing import Callable, List +from typing import Callable, List, Optional -from aiohttp import web from aiohttp import web import aiohttp @@ -28,6 +27,9 @@ def __init__(self, img_dir: str): self._on_connect = lambda: None self._is_connected = False + self._latest_video_frame: Optional[bytes] = None + self._video_lock = threading.Lock() + def start_comms(self): self._comms_thread = threading.Thread(target=self._start_comms_loop, daemon=True) self._comms_thread.start() @@ -38,16 +40,14 @@ def send_image(self, path: str): the path sent should be accessable from within self.img_dir so it can be accessed through /images/{filename} """ - print(path) img_url = "/images/" + path - print(img_url) content = { "type": "img", "value": img_url } self._send_queue.put(json.dumps(content)) - def send_log(self, message: str, severity: str="normal"): + def send_log(self, message: str, severity: str = "normal"): """ sends a log message to Emu message: string of flog @@ -66,6 +66,14 @@ def send_msg(self, message: str): """ self._send_queue.put(message) + def send_video_frame(self, jpeg_bytes: bytes): + """ + Update the latest video frame served at /video. + Call this from a background thread with raw JPEG bytes. + """ + with self._video_lock: + self._latest_video_frame = jpeg_bytes + def set_on_connect(self, func: Callable): self._on_connect = func @@ -75,19 +83,34 @@ def _start_comms_loop(self): """ print("start_comms loop") self.app = web.Application() - self.app.add_routes([web.static('/images', self.img_dir), - web.get('/ws', self.handle_websocket)]) + self.app.add_routes([ + web.static('/images', self.img_dir), + web.get('/ws', self.handle_websocket), + web.get('/video', self.handle_video_stream), + ]) web.run_app(self.app, handle_signals=False) def subscribe(self, subscriber: Callable): self._subscribers.append(subscriber) + def register_slam_streamer(self, streamer) -> None: + def _handle(raw_msg: str) -> None: + try: + msg = json.loads(raw_msg) + except json.JSONDecodeError: + return + if msg.get("type") == "slam": + if msg.get("command") == "start": + streamer.start() + elif msg.get("command") == "stop": + streamer.stop() + self.subscribe(_handle) + async def producer_handler(self, ws): """ handles sending messages to the client """ - event_loop = asyncio.get_running_loop() while not ws.closed: message = await asyncio.to_thread(self._send_queue.get) @@ -101,7 +124,20 @@ async def consumer_handler(self, ws): elif msg.type == aiohttp.WSMsgType.ERROR: print("WebSocket error:", ws.exception()) - + + async def handle_video_stream(self, request): + """ + Returns latest video frame as JPEG. Frame rate depends on how often + send_video_frame() is called (controlled by VideoEmuStreamer). + """ + with self._video_lock: + frame = self._latest_video_frame + + if frame is None: + raise web.HTTPNoContent() + + return web.Response(body=frame, content_type='image/jpeg') + async def handle_websocket(self, request): ws = web.WebSocketResponse() await ws.prepare(request) @@ -122,5 +158,5 @@ async def handle_websocket(self, request): print('websocket connection closed') self._is_connected = False - + return ws diff --git a/src/modules/imaging/.aruco_stream.py.swp b/src/modules/imaging/.aruco_stream.py.swp new file mode 100644 index 0000000..c618dce Binary files /dev/null and b/src/modules/imaging/.aruco_stream.py.swp differ diff --git a/src/modules/imaging/analysis.py b/src/modules/imaging/analysis.py index 94ad15a..bf4797a 100644 --- a/src/modules/imaging/analysis.py +++ b/src/modules/imaging/analysis.py @@ -1,4 +1,4 @@ -from typing import Callable, Optional, List, Callable, Any +from typing import Callable, Optional, List, Any, Tuple import threading # from multiprocessing import Process @@ -46,8 +46,8 @@ class ImageAnalysisDelegate: def __init__(self, detector: BaseDetector, camera: CameraProvider, - location_provider: LocationProvider = None, - navigation_provider: Navigator = None, + location_provider: Optional[LocationProvider] = None, + navigation_provider: Optional[Navigator] = None, debugger: Optional[ImageAnalysisDebugger] = None): self.detector = detector self.camera = camera @@ -59,9 +59,9 @@ def __init__(self, self.location_provider = location_provider self.navigation_provider = navigation_provider - self.subscribers: List[Callable[[Image.Image, float, float], Any]] = [] + self.subscribers: List[Callable[[Image.Image, Optional[BoundingBox], Optional[Tuple[float, float]]], Any]] = [] self.camera_attributes = CameraAttributes() - self.thread = None + self.thread: Optional[threading.Thread] = None self.loop = True def get_inference(self, bounding_box: BoundingBox) -> Inference: @@ -88,7 +88,8 @@ def start(self): def stop(self): self.loop = False - self.thread.join() + if self.thread: + self.thread.join() def _analyze_image(self): """ @@ -110,9 +111,9 @@ def _analyze_image(self): if inference: x, y = get_object_location(self.camera_attributes, inference) - subscriber(im, (x, y)) + subscriber(im, bounding_box, (x, y)) else: - subscriber(im, None) + subscriber(im, None, None) def _analysis_loop(self): """ diff --git a/src/modules/imaging/camera.py b/src/modules/imaging/camera.py index 3880002..f5dc61e 100644 --- a/src/modules/imaging/camera.py +++ b/src/modules/imaging/camera.py @@ -1,10 +1,12 @@ from typing import Tuple import pathlib +import threading +import time + from PIL import Image import numpy as np import cv2 -import depthai as dai from dataclasses import dataclass @@ -42,6 +44,7 @@ def caputure_as_ndarry(self) -> np.ndarray: """ return np.array(self.capture()) + @dataclass class DepthCapture: rgb: np.ndarray @@ -80,10 +83,13 @@ class OakdCamera(CameraProvider): """ def __init__(self, fps: int = 30): + import depthai as dai + self._dai = dai self._init_pipeline(fps) def _init_pipeline(self, fps: int): """Initialize the Depth AI pipeline (will be run on the OAK-D)""" + dai = self._dai pipeline = dai.Pipeline() camRgb = pipeline.create(dai.node.ColorCamera) @@ -145,15 +151,14 @@ def capture_with_depth(self) -> DepthCapture: return capture def capture(self) -> Image.Image: - capture = self.capture_with_depth + capture = self.capture_with_depth() img = Image.fromarray(capture.rgb, "RGB") return img - def start(self): """Start the depth-perception process on the OAK-D""" print("Starting OAK-D Connection") - self.device = dai.Device(self.pipeline) + self.device = self._dai.Device(self.pipeline) self.queue = self.device.getOutputQueue("out", maxSize=1, blocking=False) def stop(self): @@ -161,6 +166,7 @@ def stop(self): self.device.close() self.queue = None + class DebugCamera(CameraProvider): """ Debug camera source which always returns the same image loaded from @@ -210,7 +216,7 @@ def capture(self) -> Image.Image: self.index = (self.index + 1) % len(self.imgs) return Image.open(filename).resize(self.size) - + class GazeboCamera(CameraProvider): """ @@ -219,7 +225,7 @@ class GazeboCamera(CameraProvider): def __init__(self): self.port = 5600 - + gst_pipeline = ( "udpsrc address=127.0.0.1 port=5600 ! " "application/x-rtp, encoding-name=H264 ! " @@ -228,7 +234,7 @@ def __init__(self): "videoconvert ! " "appsink" ) - self.size = (640, 480) + self.size = (640, 480) self.cap = cv2.VideoCapture(gst_pipeline, cv2.CAP_GSTREAMER) if not self.cap.isOpened(): @@ -319,3 +325,44 @@ def capture(self) -> Image.Image: capture_result = self.camera.capture_array() image = Image.fromarray(capture_result) return image + + +class SharedFrameCamera(CameraProvider): + """ + Wraps a CameraProvider and shares the latest captured frame across + multiple consumers (e.g. video streamer + analysis pipeline) without + both threads calling camera.capture() simultaneously. + """ + + def __init__(self, camera: CameraProvider, fps: int = 15): + self._camera = camera + self._fps = fps + self._latest: Image.Image | None = None + self._lock = threading.Lock() + self._running = False + self._thread: threading.Thread | None = None + + def start(self): + self._running = True + self._thread = threading.Thread(target=self._capture_loop, daemon=True) + self._thread.start() + + def stop(self): + self._running = False + if self._thread: + self._thread.join() + + def capture(self) -> Image.Image: + while True: + with self._lock: + if self._latest is not None: + return self._latest + time.sleep(0.01) + + def _capture_loop(self): + interval = 1 / self._fps + while self._running: + frame = self._camera.capture() + with self._lock: + self._latest = frame + time.sleep(interval) diff --git a/src/modules/imaging/detector.py b/src/modules/imaging/detector.py index 65dd8fc..0a39b7b 100644 --- a/src/modules/imaging/detector.py +++ b/src/modules/imaging/detector.py @@ -1,10 +1,11 @@ -from functools import lru_cache +from functools import lru_cache, cached_property from typing import Optional +from dataclasses import dataclass +import math from PIL import Image import numpy as np import cv2 -from cv2 import aruco @dataclass @@ -56,35 +57,35 @@ def max(v1: 'Vec2', v2: 'Vec2') -> 'Vec2': class BoundingBox: -def __init__(self, position: Vec2, size: Vec2): - self.position = position - self.size = size + def __init__(self, position: Vec2, size: Vec2): + self.position = position + self.size = size -@lru_cache(maxsize=2) -def intersection(self, other: 'BoundingBox') -> float: - top_left = Vec2.max(self.position, other.position) - bottom_right = Vec2.min(self.position + self.size, - other.position + other.size) + @lru_cache(maxsize=2) + def intersection(self, other: 'BoundingBox') -> float: + top_left = Vec2.max(self.position, other.position) + bottom_right = Vec2.min(self.position + self.size, + other.position + other.size) - size = bottom_right - top_left + size = bottom_right - top_left - intersection = size.x * size.y - return max(intersection, 0) + intersection = size.x * size.y + return max(intersection, 0) -def union(self, other: 'BoundingBox') -> float: - intersection = self.intersection(other) - if intersection == 0: - return 0 + def union(self, other: 'BoundingBox') -> float: + intersection = self.intersection(other) + if intersection == 0: + return 0 - union = self.size.x * self.size.y + other.size.x * other.size.y - intersection - return union + union = self.size.x * self.size.y + other.size.x * other.size.y - intersection + return union -def intersection_over_union(self, pred: 'BoundingBox') -> Optional[float]: - intersection = self.intersection(pred) - if intersection == 0: - return 0 - iou = intersection / self.union(pred) - return iou + def intersection_over_union(self, pred: 'BoundingBox') -> Optional[float]: + intersection = self.intersection(pred) + if intersection == 0: + return 0 + iou = intersection / self.union(pred) + return iou class BaseDetector: @@ -98,7 +99,7 @@ def predict(self, image: Image.Image) -> Optional[BoundingBox]: img = np.array(image) gray_img = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY) - max_val = np.max(gray_img) # returns maximum value of brightness + max_val = int(np.max(gray_img)) # returns maximum value of brightness if max_val < 200: return None # lower threshold for intensity _, thresh = cv2.threshold(gray_img, max_val - 10, 255, cv2.THRESH_BINARY) @@ -115,20 +116,20 @@ def predict(self, image: Image.Image) -> Optional[BoundingBox]: return BoundingBox(Vec2(x, y), Vec2(w, h)) -class ArucoDetector(): +class ArucoDetector(BaseDetector): def predict(self, image: Image.Image) -> Optional[BoundingBox]: - img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50) - params = cv2.aruco.DetectorParemeters() + params = cv2.aruco.DetectorParameters() + + detector = cv2.aruco.ArucoDetector(aruco_dict, params) + corners, ids, rejected = detector.detectMarkers(img) - corners, ids, rejected = cv2.aruco.detectMarkers(img, aruco_dict, parameters=params) - if ids: for c in zip(corners, ids): - pts = c[0] x_min = pts[:, 0].min() @@ -140,6 +141,8 @@ def predict(self, image: Image.Image) -> Optional[BoundingBox]: y = (y_min + y_max) / 2 w = (x_max - x_min) h = (y_max - y_min) - + return BoundingBox(Vec2(x, y), Vec2(w, h)) + return None + diff --git a/src/modules/imaging/video_emu_stream.py b/src/modules/imaging/video_emu_stream.py new file mode 100644 index 0000000..d7a8c65 --- /dev/null +++ b/src/modules/imaging/video_emu_stream.py @@ -0,0 +1,40 @@ +import io +import threading +import time + +from src.modules.emu import Emu +from src.modules.imaging.camera import SharedFrameCamera + + +class VideoEmuStreamer: + """ + Continuously grabs frames from a SharedFrameCamera and pushes them + to EMU's MJPEG /video endpoint at the given fps. + """ + + def __init__(self, emu: Emu, shared_cam: SharedFrameCamera, fps: int = 15, quality: int = 70): + self.emu = emu + self.shared_cam = shared_cam + self.fps = fps + self.quality = quality + self._running = False + self._thread: threading.Thread | None = None + + def start(self): + self._running = True + self._thread = threading.Thread(target=self._stream_loop, daemon=True) + self._thread.start() + + def stop(self): + self._running = False + if self._thread: + self._thread.join() + + def _stream_loop(self): + interval = 1 / self.fps + while self._running: + frame = self.shared_cam.capture() + buf = io.BytesIO() + frame.save(buf, format="JPEG", quality=self.quality) + self.emu.send_video_frame(buf.getvalue()) + time.sleep(interval) diff --git a/src/modules/slam/__init__.py b/src/modules/slam/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/modules/slam/debug_slam_provider.py b/src/modules/slam/debug_slam_provider.py new file mode 100644 index 0000000..bc11027 --- /dev/null +++ b/src/modules/slam/debug_slam_provider.py @@ -0,0 +1,19 @@ +import time + +import numpy as np + +from src.modules.slam.slam_provider import SLAMFrame, SLAMPose, SLAMProvider + + +class DebugSLAMProvider(SLAMProvider): + """Returns a fake growing point cloud for testing the stream without hardware.""" + + def __init__(self): + self._t = 0 + + def get_frame(self) -> SLAMFrame: + time.sleep(0.1) + self._t += 1 + points = np.random.randn(50, 3) * 0.5 + np.array([self._t * 0.1, 0, 0]) + pose = SLAMPose(x=self._t * 0.1, y=0.0, z=1.5, roll=0.0, pitch=0.0, yaw=0.0) + return SLAMFrame(points=points, pose=pose) diff --git a/src/modules/slam/oakd_altimeter_slam_provider.py b/src/modules/slam/oakd_altimeter_slam_provider.py new file mode 100644 index 0000000..2717fb9 --- /dev/null +++ b/src/modules/slam/oakd_altimeter_slam_provider.py @@ -0,0 +1,90 @@ +import math + +import numpy as np + +from src.modules.autopilot.altimeter import Altimeter +from src.modules.imaging.camera import OakdCamera +from src.modules.imaging.location import LocationProvider +from src.modules.slam.slam_provider import SLAMFrame, SLAMPose, SLAMProvider + + +def _rotation_matrix(roll: float, pitch: float, yaw: float) -> np.ndarray: + """ + ZYX Euler rotation matrix (aerospace convention). + Rotates a vector from body/camera frame into world frame. + roll, pitch, yaw in radians. + """ + cr, sr = math.cos(roll), math.sin(roll) + cp, sp = math.cos(pitch), math.sin(pitch) + cy, sy = math.cos(yaw), math.sin(yaw) + + Rz = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]]) + Ry = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]]) + Rx = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]]) + + return Rz @ Ry @ Rx + + +class OakdAltimeterSLAMProvider(SLAMProvider): + """ + Fuses OAK-D depth camera + XM125 altimeter to produce a georeferenced + 3D point cloud and drone pose. + + Coordinate convention for output points: + X = right (camera X), Y = forward (camera Z), Z = up (altitude-anchored). + Each call to get_frame() blocks until a fresh depth frame is available. + """ + + def __init__( + self, + oakd: OakdCamera, + altimeter: Altimeter, + location_provider: LocationProvider, + ): + self._oakd = oakd + self._altimeter = altimeter + self._location = location_provider + + def get_frame(self) -> SLAMFrame: + # --- 1. Raw depth frame from OAK-D --- + depth = self._oakd.capture_with_depth() + + # point_cloud shape (N, 3), units mm, OAK-D camera frame (X right, Y down, Z forward) + pts_cam = depth.point_cloud.astype(np.float64) / 1000.0 # → metres + + # --- 2. Strip invalid points (zero or NaN) --- + valid = ~np.any(np.isnan(pts_cam), axis=1) & ~np.all(pts_cam == 0, axis=1) + pts_cam = pts_cam[valid] + + # --- 3. Altitude from XM125 (mm → m) --- + alt_mm = self._altimeter.get_distance_mm() + altitude = (alt_mm / 1000.0) if alt_mm is not None else 0.0 + + # --- 4. Drone orientation from MAVLink --- + orientation = self._location.orientation() + roll = math.radians(orientation.roll) + pitch = math.radians(orientation.pitch) + yaw = math.radians(orientation.yaw) + + # --- 5. Rotate points from camera frame to world frame --- + # OAK-D has Y pointing down; flip Y→-Y so Z is up before rotating + pts_cam[:, 1] *= -1 + + R = _rotation_matrix(roll, pitch, yaw) + pts_world = (R @ pts_cam.T).T + + # Shift Z so the drone sits at its altimeter-measured height + pts_world[:, 2] += altitude + + # --- 6. Build pose --- + location = self._location.location() + pose = SLAMPose( + x=float(location.lat), + y=float(location.lng), + z=altitude, + roll=orientation.roll, + pitch=orientation.pitch, + yaw=orientation.yaw, + ) + + return SLAMFrame(points=pts_world, pose=pose) diff --git a/src/modules/slam/recording_manager.py b/src/modules/slam/recording_manager.py new file mode 100644 index 0000000..02becb4 --- /dev/null +++ b/src/modules/slam/recording_manager.py @@ -0,0 +1,74 @@ +import json +import shutil +from pathlib import Path +from typing import Optional + +from src.modules.slam.slam_playback_provider import SLAMPlaybackProvider +from src.modules.slam.slam_provider import SLAMProvider +from src.modules.slam.slam_recorder import SLAMRecorder + + +class RecordingManager: + """ + Manages multiple SLAM recordings stored under a single directory. + + Each recording is a subdirectory with a metadata.json and a frames/ folder. + """ + + def __init__(self, recordings_dir: str = "recordings"): + self._dir = Path(recordings_dir) + self._dir.mkdir(parents=True, exist_ok=True) + + def list_recordings(self) -> list[dict]: + """Return metadata for all stored recordings, sorted by creation time.""" + recordings = [] + for path in sorted(self._dir.iterdir()): + meta_path = path / "metadata.json" + if path.is_dir() and meta_path.exists(): + with open(meta_path) as f: + meta = json.load(f) + meta["name"] = path.name + recordings.append(meta) + return recordings + + def new_recorder(self, name: str, provider: SLAMProvider) -> SLAMRecorder: + """Create a new recorder that wraps provider and saves to recordings/.""" + path = self._dir / name + if path.exists(): + raise FileExistsError(f"Recording '{name}' already exists") + return SLAMRecorder(provider, str(path)) + + def load_playback(self, name: str, loop: bool = False) -> SLAMPlaybackProvider: + """Load a recording by name and return a playback provider.""" + path = self._dir / name + if not path.exists(): + raise FileNotFoundError(f"Recording '{name}' not found in {self._dir}") + return SLAMPlaybackProvider(str(path), loop=loop) + + def delete_recording(self, name: str): + """Permanently delete a recording.""" + path = self._dir / name + if not path.exists(): + raise FileNotFoundError(f"Recording '{name}' not found") + shutil.rmtree(path) + print(f"Deleted recording '{name}'") + + def print_summary(self, recording: Optional[dict] = None): + """Print a summary table of all recordings, or one specific recording.""" + recordings = [recording] if recording else self.list_recordings() + if not recordings: + print("No recordings found.") + return + print(f"{'Name':<30} {'Frames':>8} {'Duration':>10} {'Created'}") + print("-" * 70) + for r in recordings: + import datetime + created = datetime.datetime.fromtimestamp( + r.get("created_at", 0) + ).strftime("%Y-%m-%d %H:%M") + print( + f"{r['name']:<30} " + f"{r.get('frame_count', '?'):>8} " + f"{r.get('duration_sec', 0):>9.1f}s " + f"{created}" + ) diff --git a/src/modules/slam/slam_emu_streamer.py b/src/modules/slam/slam_emu_streamer.py new file mode 100644 index 0000000..5dea155 --- /dev/null +++ b/src/modules/slam/slam_emu_streamer.py @@ -0,0 +1,49 @@ +import json +import threading + +from src.modules.emu import Emu +from src.modules.slam.slam_provider import SLAMProvider + + +class SLAMEmuStreamer: + """ + Continuously calls slam_provider.get_frame() and pushes each frame + to Emu over WebSocket as a JSON "slam" message. + Starts/stops on demand so bandwidth is only used while the SLAM tab is open. + """ + + def __init__(self, emu: Emu, slam_provider: SLAMProvider): + self.emu = emu + self.slam_provider = slam_provider + self._running = False + self._thread: threading.Thread | None = None + + def start(self): + if self._running: + return + self._running = True + self._thread = threading.Thread(target=self._stream_loop, daemon=True) + self._thread.start() + + def stop(self): + self._running = False + if self._thread: + self._thread.join() + self._thread = None + + def _stream_loop(self): + while self._running: + frame = self.slam_provider.get_frame() + msg = json.dumps({ + "type": "slam_data", + "point_cloud": frame.points.tolist(), + "pose": { + "x": frame.pose.x, + "y": frame.pose.y, + "z": frame.pose.z, + "roll": frame.pose.roll, + "pitch": frame.pose.pitch, + "yaw": frame.pose.yaw, + }, + }) + self.emu.send_msg(msg) diff --git a/src/modules/slam/slam_playback_provider.py b/src/modules/slam/slam_playback_provider.py new file mode 100644 index 0000000..c817f4d --- /dev/null +++ b/src/modules/slam/slam_playback_provider.py @@ -0,0 +1,75 @@ +import time +from pathlib import Path + +import numpy as np + +from src.modules.slam.slam_provider import SLAMFrame, SLAMPose, SLAMProvider + + +class SLAMPlaybackProvider(SLAMProvider): + """ + Replays a recorded SLAM session through the same SLAMProvider interface. + Maintains original timing between frames so replay feels real-time. + + Pass loop=True to repeat the recording indefinitely. + + Usage: + playback = SLAMPlaybackProvider("recordings/flight_01") + streamer = SLAMEmuStreamer(emu, playback) + streamer.start() + """ + + def __init__(self, recording_path: str, loop: bool = False): + self._path = Path(recording_path) + self._loop = loop + self._frame_idx = 0 + self._frame_files = sorted((self._path / "frames").glob("*.npz")) + self._prev_real_time: float | None = None + self._prev_frame_time: float | None = None + + if not self._frame_files: + raise FileNotFoundError(f"No frames found in {self._path / 'frames'}") + + @property + def frame_count(self) -> int: + return len(self._frame_files) + + def reset(self): + self._frame_idx = 0 + self._prev_real_time = None + self._prev_frame_time = None + + def get_frame(self) -> SLAMFrame: + if self._frame_idx >= len(self._frame_files): + if self._loop: + self.reset() + else: + raise StopIteration("Recording playback complete") + + data = np.load(self._frame_files[self._frame_idx]) + timestamp = float(data["timestamp"][0]) + points = data["points"] + pose_arr = data["pose"] + + pose = SLAMPose( + x=float(pose_arr[0]), + y=float(pose_arr[1]), + z=float(pose_arr[2]), + roll=float(pose_arr[3]), + pitch=float(pose_arr[4]), + yaw=float(pose_arr[5]), + ) + + # Reproduce original inter-frame timing + if self._prev_frame_time is not None and self._prev_real_time is not None: + frame_dt = timestamp - self._prev_frame_time + elapsed = time.time() - self._prev_real_time + sleep_time = frame_dt - elapsed + if sleep_time > 0: + time.sleep(sleep_time) + + self._prev_frame_time = timestamp + self._prev_real_time = time.time() + self._frame_idx += 1 + + return SLAMFrame(points=points, pose=pose, timestamp=timestamp) diff --git a/src/modules/slam/slam_provider.py b/src/modules/slam/slam_provider.py new file mode 100644 index 0000000..2bab468 --- /dev/null +++ b/src/modules/slam/slam_provider.py @@ -0,0 +1,28 @@ +import time +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + +import numpy as np + + +@dataclass +class SLAMPose: + x: float + y: float + z: float + roll: float + pitch: float + yaw: float + + +@dataclass +class SLAMFrame: + points: np.ndarray # shape (N, 3), metres, world frame + pose: SLAMPose + timestamp: float = field(default_factory=time.time) + + +class SLAMProvider(ABC): + @abstractmethod + def get_frame(self) -> SLAMFrame: + raise NotImplementedError() diff --git a/src/modules/slam/slam_recorder.py b/src/modules/slam/slam_recorder.py new file mode 100644 index 0000000..dc66cab --- /dev/null +++ b/src/modules/slam/slam_recorder.py @@ -0,0 +1,72 @@ +import json +import time +from pathlib import Path + +import numpy as np + +from src.modules.slam.slam_provider import SLAMFrame, SLAMProvider + + +class SLAMRecorder(SLAMProvider): + """ + Wraps any SLAMProvider and records every frame to disk while passing + it through unchanged. Acts as a transparent decorator — plug it in + between any provider and the streamer and recording happens automatically. + + Usage: + recorder = SLAMRecorder(real_provider, "recordings/flight_01") + recorder.start_recording() + streamer = SLAMEmuStreamer(emu, recorder) # recorder IS a SLAMProvider + streamer.start() + ... + recorder.stop_recording() + """ + + def __init__(self, provider: SLAMProvider, recording_path: str): + self._provider = provider + self._path = Path(recording_path) + self._frame_idx = 0 + self._start_time: float | None = None + self._recording = False + + def start_recording(self): + self._path.mkdir(parents=True, exist_ok=True) + (self._path / "frames").mkdir(exist_ok=True) + self._start_time = time.time() + self._frame_idx = 0 + self._recording = True + print(f"Recording started → {self._path}") + + def stop_recording(self): + self._recording = False + self._save_metadata() + print(f"Recording saved ({self._frame_idx} frames) → {self._path}") + + def get_frame(self) -> SLAMFrame: + frame = self._provider.get_frame() + if self._recording: + self._save_frame(frame) + return frame + + def _save_frame(self, frame: SLAMFrame): + frame_path = self._path / "frames" / f"{self._frame_idx:06d}.npz" + np.savez_compressed( + frame_path, + timestamp=np.array([frame.timestamp]), + points=frame.points, + pose=np.array([ + frame.pose.x, frame.pose.y, frame.pose.z, + frame.pose.roll, frame.pose.pitch, frame.pose.yaw, + ]), + ) + self._frame_idx += 1 + + def _save_metadata(self): + duration = (time.time() - self._start_time) if self._start_time else 0.0 + metadata = { + "frame_count": self._frame_idx, + "created_at": self._start_time, + "duration_sec": duration, + } + with open(self._path / "metadata.json", "w") as f: + json.dump(metadata, f, indent=2) diff --git a/test/test_analysis.py b/test/test_analysis.py index 080f1b8..32a2839 100644 --- a/test/test_analysis.py +++ b/test/test_analysis.py @@ -12,11 +12,10 @@ from src.modules.imaging.camera import DebugCamera from src.modules.imaging.location import DebugLocationProvider from src.modules.imaging.debug import ImageAnalysisDebugger -from dep.labeller.benchmarks.detector import LandingPadDetector, BoundingBox -from dep.labeller.loader.label import Vec2 +from src.modules.imaging.detector import BaseDetector, BoundingBox, Vec2 -class DebugLandingPadDetector(LandingPadDetector): +class DebugLandingPadDetector(BaseDetector): def __init__(self, vector: Optional[Vec2] = None, @@ -37,9 +36,12 @@ def test_analysis_subscriber(): global detected detected = None - def _callback(_image, lon, lat): + def _callback(_image, bounding_box, pos): global detected - detected = Vec2(lon, lat) + if pos is not None: + detected = Vec2(pos[0], pos[1]) + else: + detected = None analysis.subscribe(_callback) @@ -49,7 +51,7 @@ def _callback(_image, lon, lat): detector.bounding_box = BoundingBox(Vec2(20, 20), Vec2(50, 50)) analysis._analyze_image() assert (detected - - Vec2(-115.48873916832288, 5.483286467459389e-06)).norm < 0.01 + Vec2(0.4158184416499504, -0.574961758930409)).norm < 0.01 class MockImageAnlaysisDebugger(ImageAnalysisDebugger): @@ -87,7 +89,7 @@ def test_analysis_debugger(): location_provider = DebugLocationProvider() location_provider.set_altitude(1.0) analysis = ImageAnalysisDelegate(detector, camera, location_provider, - debug) + debugger=debug) def run_analysis(): detector.bounding_box = BoundingBox(Vec2(0, 0), Vec2(100, 100)) diff --git a/test/test_camera.py b/test/test_camera.py index 86a1f70..08aefda 100644 --- a/test/test_camera.py +++ b/test/test_camera.py @@ -26,7 +26,7 @@ def test_debug_camera(tmp_path): # Can save the image to a path im_path = tmp_path / "copy.jpeg" - cam.caputure_to(im_path) + cam.capture_to(im_path) im_md5 = md5sum(im_path) # Manually save a copy of 'test-image.jpeg'.