diff --git a/python/oak/README.md b/python/oak/README.md index 27c0b3e..c27e60b 100644 --- a/python/oak/README.md +++ b/python/oak/README.md @@ -18,6 +18,7 @@ Also requires `PyOpenGL_accelerate` to be installed, see [`mixed_reality.py`](mixed_reality.py) for details. * **GNSS-VIO** example, reads external GNSS from standard input [`vio_gnss.py`](vio_gnss.py) (see also [these instructions](https://spectacularai.github.io/docs/pdf/GNSS-VIO_OAK-D_Python.pdf)) * **April Tag example**: Visualize detected April Tags [`python april_tag.py path/to/tags.json`](april_tag.py). See: https://spectacularai.github.io/docs/pdf/april_tag_instructions.pdf + * **Relocalize April Tag example**: Example of how to align a saved SLAM map with a recorded April Tag session [`python relocalize_apriltags.py path/to/recording/data`](relocalize_apriltags.py). For complete instructions on how to use this tool to align a digital twin in Unity with your SLAM map, see the https://github.com/SpectacularAI/unity-wrapper/tree/main/unity-examples/Assets/SpectacularAI/Examples/AprilTag/README.md. * **Remote visualization over SSH**. Can be achieved by combining the `vio_jsonl.py` and `vio_visu.py` scripts as follows: ssh user@example.org 'python -u /full/path/to/vio_jsonl.py' | python -u vio_visu.py --file=- diff --git a/python/oak/relocalize_apriltags.py b/python/oak/relocalize_apriltags.py new file mode 100644 index 0000000..c197948 --- /dev/null +++ b/python/oak/relocalize_apriltags.py @@ -0,0 +1,220 @@ +import argparse +import spectacularAI +import time +import threading +import os +import numpy as np +from enum import Enum +import json + +class State(Enum): + FindRelocalizeTime = 1 + FindAprilTagTime = 2 + FindRelocalizePose = 3 + +aprilTagFrameTime = None +aprilTagFramePose = None +relocalizedFrameTime = None +relocalizedFramePose = None +args = None +aprilTagReplay1 = None +relocalizeReplay1 = None +relocalizeReplay2 = None +state = None + +vio_config_yaml = """ +depthErrorScale: 0.1 +useSlam: True +keyframeCandidateInterval: 0 +alreadyRectified: True +parameterSets: [wrapper-base,oak-d,live] +""" + +def writeVioConfig(): + global vio_config_yaml + vio_config_path = os.path.join(args.dataFolder, "vio_config.yaml") + with open(vio_config_path, 'w') as vio_config_file: + vio_config_file.write(vio_config_yaml) + print(f"Updated VIO configuration written to: {vio_config_path}") + +def renameMapBin(): + slam_map_path = os.path.join(args.dataFolder, "slam_map.bin") + map_path = os.path.join(args.dataFolder, "map.bin") + + if os.path.exists(slam_map_path): + if os.path.exists(map_path): + os.remove(map_path) + os.rename(slam_map_path, map_path) + print(f"Renamed {slam_map_path} to {map_path}") + +def onRelocalizeExtendedOutput(output, frameSet): + global args, relocalizedFrameTime, relocalizedFramePose, relocalizeReplay1, relocalizeReplay2, pairedFrameTime, state + if not output.status == spectacularAI.TrackingStatus.TRACKING: return # not tracking yet + + for frame in frameSet: + if not frame.image: continue + if state == State.FindRelocalizeTime: + if not relocalizedFrameTime: + relocalizedFrameTime = output.pose.time + print("Tracking started with Relocalization at: {}".format(relocalizedFrameTime)) + # print("Close relocalize replay") + relocalizeReplay1.close() + break + elif state == State.FindRelocalizePose and output.pose.time >= aprilTagFrameTime: + if relocalizedFramePose is None: + relocalizedFramePose = frame.cameraPose.pose.asMatrix() + print("Camera to relocalized world at: {}\n{}".format(output.pose.time, relocalizedFramePose)) + # print("Close relocalize replay") + relocalizeReplay2.close() + break + +def onAprilTagOutput(output): + global args, aprilTagFrameTime, aprilTagFramePose, aprilTagReplay1, state + + if aprilTagFrameTime is not None: + # print("Close April Tag replay") + aprilTagReplay1.close() + return + + if not state == State.FindAprilTagTime: return # not looking for April Tag time + # if not output.status == spectacularAI.TrackingStatus.TRACKING: return # not tracking yet + if not output.map.keyFrames: return # empty map + + # Find latest keyframe + keyFrameId = max(output.map.keyFrames.keys()) + keyFrame = output.map.keyFrames[keyFrameId] + frameSet = keyFrame.frameSet + frame = frameSet.primaryFrame + frameTime = frame.cameraPose.pose.time + if not frameTime >= relocalizedFrameTime: return # not after relocalized time + + # find a frame with a marker + if frame is None or frame.image is None: return + if frame.visualMarkers is None or frame.visualMarkers == []: return + marker = frame.visualMarkers[0] + if not marker.hasPose: + raise Exception("Found April Tag but it's not in the tags.json file") + + aprilTagFrameTime = frameTime + aprilTagFramePose = marker.pose.asMatrix() + print("Camera to April Tag at: {}\n{}".format(frameTime, aprilTagFramePose)) + +def replayAprilTags(): + global args, aprilTagReplay1, state + print("[Replay April Tags] {}".format(state)) + + configInternal = { + "aprilTagPath": args.dataFolder + "/tags.json", + "extendParameterSets" : ["april-tags"] + } + + aprilTagReplay1 = spectacularAI.Replay(args.dataFolder, onAprilTagOutput, configuration=configInternal) + aprilTagReplay1.runReplay() + + +def replayRelocalize(): + global args, relocalizeReplay1, relocalizeReplay2, state + # print("[Replay Relocalize] {}".format(args)) + + configInternal = { + "mapLoadPath": args.dataFolder + "/map.bin", + "extendParameterSets" : ["relocalization"] + } + + replay = spectacularAI.Replay(args.dataFolder, configuration=configInternal) + replay.setExtendedOutputCallback(onRelocalizeExtendedOutput) + + if state == State.FindRelocalizeTime: + relocalizeReplay1 = replay + else: + relocalizeReplay2 = replay + + replay.runReplay() + +def compute_slam_to_dt(camera_to_slam: np.ndarray, + camera_to_tag: np.ndarray, + tag_to_dt: np.ndarray) -> np.ndarray: + """ + Computes the transformation from the SLAM (relocalization) frame to the Digital Twin frame. + + Parameters: + camera_to_slam (np.ndarray): 4x4 transformation matrix (T^(S)_(C)) representing + the camera pose in the SLAM coordinate frame. + camera_to_tag (np.ndarray): 4x4 transformation matrix (T^(C)_(A)) representing + the AprilTag pose in the camera coordinate frame. + tag_to_dt (np.ndarray): 4x4 transformation matrix (T^(D)_(A)) representing + the AprilTag pose in the Digital Twin coordinate frame. + + Returns: + np.ndarray: 4x4 transformation matrix (T^(D)_(S)) that maps a point from the + SLAM coordinate frame to the Digital Twin coordinate frame. + """ + # Derive the transformation from SLAM to AprilTag coordinates: + # T^(S)_(A) = T^(S)_(C) * inverse(T^(A)_(C)) + T_slam_apr = camera_to_slam @ np.linalg.inv(camera_to_tag) + + # The transformation from SLAM to Digital Twin coordinates is: + # T^(D)_(S) = T^(D)_(A) * inverse(T^(S)_(A)) + # Since inverse(T^(S)_(A)) = april_tag_pose @ inverse(relocalized_pose) + T_digitalTwin_slam = tag_to_dt @ camera_to_tag @ np.linalg.inv(camera_to_slam) + return T_digitalTwin_slam + +if __name__ == '__main__': + p = argparse.ArgumentParser(__doc__) + p.add_argument("dataFolder", nargs="?", help="Folder containing the recorded session and tags.json", default="data") + p.add_argument("--preview", help="Show latest primary image as a preview", action="store_true") + args = p.parse_args() + + writeVioConfig() + renameMapBin() + + print("Pose of first April Tag from tags.json:") + with open(args.dataFolder + "/tags.json", 'r') as tags_file: + tags_json = json.load(tags_file) + aprilTagToWorld = np.array(tags_json[0]['tagToWorld']) + print(aprilTagToWorld) + + print("Finding first relocalized frame time...") + state = State.FindRelocalizeTime + relocalizeThread = threading.Thread(target=replayRelocalize) + relocalizeThread.start() + while relocalizedFrameTime is None and relocalizeThread.is_alive(): + time.sleep(0.1) + if relocalizedFrameTime is None: + raise Exception("No valid frame found in Relocalize replay") + + state = State.FindAprilTagTime + print("Finding April Tag frame time (after Relocalized time)...") + aprilTagThread = threading.Thread(target=replayAprilTags) + aprilTagThread.start() + while aprilTagFrameTime is None and aprilTagThread.is_alive(): + time.sleep(0.1) + if aprilTagFrameTime is None: + raise Exception("No valid frame found in April Tags replay") + + print("Finding paired time in Relocalize replay...") + state = State.FindRelocalizePose + relocalizeThread = threading.Thread(target=replayRelocalize) + relocalizeThread.start() + while relocalizedFramePose is None and relocalizeThread.is_alive(): + time.sleep(0.1) + if relocalizedFramePose is None: + raise Exception("Paired frame not found in Relocalize replay") + + print("SLAM to Digital Twin transformation:") + slam_to_dt = compute_slam_to_dt(relocalizedFramePose, + aprilTagFramePose, + aprilTagToWorld) + print(slam_to_dt) + + slam_config_json = { + "slamToUnity": slam_to_dt.tolist() + } + slam_config_path = args.dataFolder + "/slam_config.json" + slam_config_json["slamToUnity"] = np.array(slam_config_json["slamToUnity"]).flatten().tolist() + with open(slam_config_path, 'w') as slam_config_file: + json.dump(slam_config_json, slam_config_file, indent=4) + print("Wrote transformation to: {}".format(slam_config_path)) + os._exit(0) + + \ No newline at end of file