From 3a103d1e9e63f6a12bec6829d7f7441429588bfb Mon Sep 17 00:00:00 2001
From: kaatrasa <valtteri.kaatrasalo@gmail.com>
Date: Thu, 25 Apr 2024 16:04:31 +0300
Subject: [PATCH 1/2] Add option to use Python visualizer with sai-cli record
 oak

---
 python/cli/record/oak.py | 51 +++++++++++++++++++++++++++++-----------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/python/cli/record/oak.py b/python/cli/record/oak.py
index b8361bf..24cfad6 100644
--- a/python/cli/record/oak.py
+++ b/python/cli/record/oak.py
@@ -40,6 +40,7 @@ def define_args(p):
     p.add_argument("--gray", help="Record (rectified) gray video data", action="store_true")
     p.add_argument("--no_convert", help="Skip converting h265 video file", action="store_true")
     p.add_argument('--no_preview', help='Do not show a live preview', action="store_true")
+    p.add_argument('--preview3d', help='Use more advanced visualizer instead of matplotlib', action="store_true")
     p.add_argument('--no_slam', help='Record with SLAM module disabled', action="store_true")
     p.add_argument('--recording_only', help='Do not run VIO, may be faster', action="store_true")
     p.add_argument('--april_tag_path', help='Record with April Tags (path to tags.json)')
@@ -169,8 +170,27 @@ def record(args):
 
     config.internalParameters = internalParameters
 
+    if args.no_preview:
+        plotter = None
+        visualizer = None
+    elif args.preview3d:
+        from spectacularAI.cli.visualization.visualizer import Visualizer, VisualizerArgs
+        visArgs = VisualizerArgs()
+        visArgs.targetFps = 30
+        visualizer = Visualizer(visArgs)
+        plotter = None
+        config.parameterSets.append('point-cloud')
+    else:
+        from spectacularAI.cli.visualization.vio_visu import make_plotter
+        import matplotlib.pyplot as plt
+        plotter, anim = make_plotter()
+        visualizer = None
+
+    def on_mapping_output(mapperOutput):
+        visualizer.onMappingOutput(mapperOutput)
+
     # Enable recoding by setting recordingFolder option
-    vio_pipeline = spectacularAI.depthai.Pipeline(pipeline, config)
+    vio_pipeline = spectacularAI.depthai.Pipeline(pipeline, config, None if visualizer is None else on_mapping_output)
 
     # Optionally also record other video streams not used by the Spectacular AI SDK, these
     # can be used for example to render AR content or for debugging.
@@ -214,7 +234,7 @@ def create_rgb_camera_control(colorCameraNode):
         if vio_pipeline.monoRight: create_rgb_camera_control(vio_pipeline.monoRight)
 
     should_quit = threading.Event()
-    def main_loop(plotter=None):
+    def main_loop(plotter=None, visualizer=None):
         frame_number = 1
 
         deviceInfo = None
@@ -224,6 +244,9 @@ def createDevice():
                 return depthai.Device(pipeline, deviceInfo=deviceInfo, maxUsbSpeed=depthai.UsbSpeed.SUPER_PLUS)
             return depthai.Device(pipeline)
 
+        def onMappingOutput(mapperOutput):
+            visualizer.onMappingOutput(mapperOutput)
+
         with createDevice() as device, vio_pipeline.startSession(device) as vio_session:
 
             if args.ir_dot_brightness > 0:
@@ -259,7 +282,7 @@ def open_gray_video(name):
 
             print("Recording to '{0}'".format(config.recordingFolder))
             print("")
-            if plotter is not None:
+            if plotter is not None or visualizer is not None:
                 print("Close the visualization window to stop recording")
 
             while not should_quit.is_set():
@@ -280,6 +303,10 @@ def open_gray_video(name):
                 if vio_session.hasOutput():
                     out = vio_session.getOutput()
                     progress = True
+
+                    if visualizer is not None:
+                        visualizer.onVioOutput(out.getCameraPose(0), status=out.status)
+
                     if plotter is not None:
                         if not plotter(json.loads(out.asJson())): break
 
@@ -309,19 +336,15 @@ def open_gray_video(name):
                 print("Use ffmpeg to convert video into a viewable format:")
                 print("    " + ffmpegCommand)
 
-    if args.no_preview:
-        plotter = None
-    else:
-        from spectacularAI.cli.visualization.vio_visu import make_plotter
-        import matplotlib.pyplot as plt
-        plotter, anim = make_plotter()
-
-    reader_thread = threading.Thread(target = lambda: main_loop(plotter))
+    reader_thread = threading.Thread(target = lambda: main_loop(plotter, visualizer))
     reader_thread.start()
-    if plotter is None:
-        input("---- Press ENTER to stop recording ----")
-    else:
+    if visualizer is not None:
+        visualizer.run()
+    elif plotter is not None:
         plt.show()
+    else:
+        input("---- Press ENTER to stop recording ----")
+
     should_quit.set()
 
     reader_thread.join()

From 4c492ab467b215a279304c5f6eabfe9b153d4875 Mon Sep 17 00:00:00 2001
From: kaatrasa <valtteri.kaatrasalo@gmail.com>
Date: Thu, 25 Apr 2024 16:17:29 +0300
Subject: [PATCH 2/2] Clean code

---
 python/cli/record/oak.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/python/cli/record/oak.py b/python/cli/record/oak.py
index 24cfad6..f3a81d0 100644
--- a/python/cli/record/oak.py
+++ b/python/cli/record/oak.py
@@ -244,9 +244,6 @@ def createDevice():
                 return depthai.Device(pipeline, deviceInfo=deviceInfo, maxUsbSpeed=depthai.UsbSpeed.SUPER_PLUS)
             return depthai.Device(pipeline)
 
-        def onMappingOutput(mapperOutput):
-            visualizer.onMappingOutput(mapperOutput)
-
         with createDevice() as device, vio_pipeline.startSession(device) as vio_session:
 
             if args.ir_dot_brightness > 0: