Copy vio_record.py from SDK examples an wrap in to a subcommand

oseiskar · Bercon · pekkaran · oseiskar · commit 758b0492db3d · 2023-12-11T15:45:05.000+02:00
Also copy vio_visu.py

Co-authored-by: Jerry Ylilammi &lt;jerry.ylilammi@gmail.com&gt;
Co-authored-by: Pekka Rantalankila &lt;40019233+pekkaran@users.noreply.github.com&gt;
Co-authored-by: Valtteri Kaatrasalo &lt;46484036+kaatrasa@users.noreply.github.com&gt;
diff --git a/python/cli/record/oak.py b/python/cli/record/oak.py
@@ -0,0 +1,251 @@
+r"""
+Record data for later playback from OAK devices
+
+Requirements:
+
+    ffmpeg must be installed.
+
+    On Linux you can install it with package manager
+    of your choise. For example with
+        ap-get:   sudo apt-get install ffmpeg
+        yuM:      sudo yum install ffmpeg
+
+On Windows, you must download and install it from https://www.ffmpeg.org and
+then update your environment Path variable to contain the binary path. To do
+this, press Windows Key, type Path and press Enter. Open Environment Settings,
+edit the row named Path and add location of the ffmpeg bin folder to the list,
+for example: "C:\Program Files\ffmpeg\bin". To check that it works, open
+command prompt and type ffmpeg, you should see version information.
+
+To view the depth video file, you must use ffplay, because normal video players
+cannot play 16bit grayscale video.
+
+Plug in the OAK-D and run:
+
+    sai-cli record oak
+
+.
+"""
+
+# --- The following mechanism allows using this both as a stand-alone
+# script and as a subcommand in sai-cli.
+
+def define_args(p):
+    p.add_argument("--output", help="Recording output folder", default="data")
+    p.add_argument('--auto_subfolders', action='store_true',
+        help='Create timestamp-named subfolders for each recording')
+    p.add_argument("--use_rgb", help="Use RGB data for tracking (OAK-D S2)", action="store_true")
+    p.add_argument("--mono", help="Use a single camera (not stereo)", action="store_true")
+    p.add_argument("--no_rgb", help="Disable recording RGB video feed", action="store_true")
+    p.add_argument("--no_inputs", help="Disable recording JSONL and depth", action="store_true")
+    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('--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('--disable_cameras', help='Prevents SDK from using cameras, for example to only record RGB camera and IMU', action="store_true")
+    p.add_argument('--no_usb_speed_check', help='Disable USB speed check', action="store_true")
+    # This can reduce CPU load while recording with the --no_feature_tracker option
+    # and the 800p resolution. See "ffmpeg -codecs" (and see "encoders" under h264)
+    # for options that might be available. On Raspberry Pi or Jetson, try "h264_v4l2m2m",
+    # and on Linux machines with Nvidia GPUs, try "h264_nvenc".
+    p.add_argument('--ffmpeg_codec', help="FFMpeg codec for host", default=None)
+    p.add_argument('--map', help='Record SLAM map', action="store_true")
+    p.add_argument('--no_feature_tracker', help='Disable on-device feature tracking', action="store_true")
+    p.add_argument('--vio_auto_exposure', help='Enable SpectacularAI auto exposure which optimizes exposure parameters for VIO performance (BETA)', action="store_true")
+    p.add_argument('--ir_dot_brightness', help='OAK-D Pro (W) IR laser projector brightness (mA), 0 - 1200', type=float, default=0)
+    p.add_argument("--resolution", help="Gray input resolution (gray)",
+        default='400p',
+        choices=['400p', '800p'])
+
+    return p
+
+def define_subparser(subparsers):
+    import argparse
+    sub = subparsers.add_parser('oak',
+                description="Record data for later playback from OAK devices",
+                epilog=__doc__,
+                formatter_class=argparse.RawDescriptionHelpFormatter)
+    sub.set_defaults(func=record)
+    return define_args(sub)
+
+def record(args):
+    import depthai
+    import spectacularAI
+    import subprocess
+    import os
+    import json
+    import threading
+    import time
+
+    config = spectacularAI.depthai.Configuration()
+    pipeline = depthai.Pipeline()
+
+    config.useSlam = True
+    config.inputResolution = args.resolution
+    outputFolder = args.output
+    if args.auto_subfolders:
+        import datetime
+        autoFolderName = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
+        outputFolder = os.path.join(outputFolder, autoFolderName)
+
+    if not args.no_inputs:
+        config.recordingFolder = outputFolder
+    if args.map:
+        try: os.makedirs(outputFolder) # SLAM only
+        except: pass
+        config.mapSavePath = os.path.join(outputFolder, 'slam_map._')
+    if args.no_slam:
+        assert args.map == False
+        config.useSlam = False
+    if args.no_feature_tracker:
+        config.useFeatureTracker = False
+    if args.vio_auto_exposure:
+        config.useVioAutoExposure = True
+    if args.use_rgb:
+        config.useColor = True
+    if args.mono:
+        config.useStereo = False
+    if args.recording_only:
+        config.recordingOnly = True
+    if args.disable_cameras:
+        config.disableCameras = True
+    if args.ffmpeg_codec is not None:
+        config.internalParameters = { 'ffmpegVideoCodec': args.ffmpeg_codec + ' -b:v 8M' }
+        print(config.internalParameters)
+    if args.no_usb_speed_check:
+        config.ensureSufficientUsbSpeed = False
+
+    # Enable recoding by setting recordingFolder option
+    vio_pipeline = spectacularAI.depthai.Pipeline(pipeline, config)
+
+    # 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.
+    rgb_as_video = not args.no_rgb and not args.use_rgb
+    if rgb_as_video:
+        import numpy # Required by frame.getData(), otherwise it hangs indefinitely
+        camRgb = pipeline.create(depthai.node.ColorCamera)
+        videoEnc = pipeline.create(depthai.node.VideoEncoder)
+        xout = pipeline.create(depthai.node.XLinkOut)
+        xout.setStreamName("h265-rgb")
+        camRgb.setBoardSocket(depthai.CameraBoardSocket.CAM_A)
+        camRgb.setResolution(depthai.ColorCameraProperties.SensorResolution.THE_1080_P)
+        # no need to set input resolution anymore (update your depthai package if this does not work)
+        videoEnc.setDefaultProfilePreset(30, depthai.VideoEncoderProperties.Profile.H265_MAIN)
+        camRgb.video.link(videoEnc.input)
+        videoEnc.bitstream.link(xout.input)
+
+    if args.gray:
+        def create_gray_encoder(node, name):
+            videoEnc = pipeline.create(depthai.node.VideoEncoder)
+            xout = pipeline.create(depthai.node.XLinkOut)
+            xout.setStreamName("h264-" + name)
+            videoEnc.setDefaultProfilePreset(30, depthai.VideoEncoderProperties.Profile.H264_MAIN)
+            node.link(videoEnc.input)
+            videoEnc.bitstream.link(xout.input)
+
+        create_gray_encoder(vio_pipeline.stereo.rectifiedLeft, 'left')
+        create_gray_encoder(vio_pipeline.stereo.rectifiedRight, 'right')
+
+    should_quit = threading.Event()
+    def main_loop(plotter=None):
+        frame_number = 1
+
+        with depthai.Device(pipeline) as device, \
+            vio_pipeline.startSession(device) as vio_session:
+
+            if args.ir_dot_brightness > 0:
+                device.setIrLaserDotProjectorBrightness(args.ir_dot_brightness)
+
+            def open_gray_video(name):
+                grayVideoFile = open(outputFolder + '/rectified_' + name + '.h264', 'wb')
+                queue = device.getOutputQueue(name='h264-' + name, maxSize=10, blocking=False)
+                return (queue, grayVideoFile)
+
+            grayVideos = []
+            if args.gray:
+                grayVideos = [
+                    open_gray_video('left'),
+                    open_gray_video('right')
+                ]
+
+            if rgb_as_video:
+                videoFile = open(outputFolder + "/rgb_video.h265", "wb")
+                rgbQueue = device.getOutputQueue(name="h265-rgb", maxSize=30, blocking=False)
+
+            print("Recording!")
+            print("")
+            if plotter is not None:
+                print("Close the visualization window to stop recording")
+
+            while not should_quit.is_set():
+                progress = False
+                if rgb_as_video:
+                    if rgbQueue.has():
+                        frame = rgbQueue.get()
+                        vio_session.addTrigger(frame.getTimestampDevice().total_seconds(), frame_number)
+                        frame.getData().tofile(videoFile)
+                        frame_number += 1
+                        progress = True
+
+                for (grayQueue, grayVideoFile) in grayVideos:
+                    if grayQueue.has():
+                        grayQueue.get().getData().tofile(grayVideoFile)
+                        progress = True
+
+                if vio_session.hasOutput():
+                    out = vio_session.getOutput()
+                    progress = True
+                    if plotter is not None:
+                        if not plotter(json.loads(out.asJson())): break
+
+                if not progress:
+                    time.sleep(0.01)
+
+        videoFileNames = []
+
+        if rgb_as_video:
+            videoFileNames.append(videoFile.name)
+            videoFile.close()
+
+        for (_, grayVideoFile) in grayVideos:
+            videoFileNames.append(grayVideoFile.name)
+            grayVideoFile.close()
+
+        for fn in videoFileNames:
+            if not args.no_convert:
+                withoutExt = fn.rpartition('.')[0]
+                ffmpegCommand = "ffmpeg -framerate 30 -y -i  \"{}\" -avoid_negative_ts make_zero -c copy \"{}.mp4\"".format(fn, withoutExt)
+
+                result = subprocess.run(ffmpegCommand, shell=True)
+                if result.returncode == 0:
+                    os.remove(fn)
+            else:
+                print('')
+                print("Use ffmpeg to convert video into a viewable format:")
+                print("    " + ffmpegCommand)
+
+    if args.no_preview:
+        plotter = None
+    else:
+        from 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.start()
+    if plotter is None:
+        input("---- Press ENTER to stop recording ----")
+    else:
+        plt.show()
+    should_quit.set()
+
+    reader_thread.join()
+
+if __name__ == '__main__':
+    def parse_args(args):
+        import argparse
+        parser = argparse.ArgumentParser(description=__doc__.strip())
+        parser = define_args(parser)
+        return parser.parse_args()
+    record(parse_args())
diff --git a/python/cli/visualization/vio_visu.py b/python/cli/visualization/vio_visu.py
@@ -0,0 +1,132 @@
+"""
+Simple VIO result visualizer Python. Reads outputs from the
+Spectacular AI OAK-D plugin and plots them in real time.
+Plug in the OAK-D to an USB3 port using an USB3 cable before running.
+
+Can also visualize pre-recorded results using Replay API, or from a JSONL file or from a pipe.
+The device does not have to be attached in this case. (See vio_record.py)
+"""
+import time
+import json
+import threading
+import matplotlib.pyplot as plt
+
+def live_vio_reader():
+    import depthai
+    import spectacularAI
+    pipeline = depthai.Pipeline()
+    vio_pipeline = spectacularAI.depthai.Pipeline(pipeline)
+
+    with depthai.Device(pipeline) as device, \
+        vio_pipeline.startSession(device) as vio_session:
+
+        while True:
+            out = vio_session.waitForOutput()
+            yield(json.loads(out.asJson()))
+
+def replay_vio_reader(replay):
+    outputs = []
+    def onOutput(out):
+        outputs.append(out.asJson())
+
+    replay.setOutputCallback(onOutput)
+    replay.startReplay()
+
+    while True:
+        if outputs:
+            out = outputs.pop(0)
+            yield(json.loads(out))
+        time.sleep(0.01)
+
+def file_vio_reader(in_stream):
+    while True:
+        line = in_stream.readline()
+        if not line: break
+        try:
+            d = json.loads(line)
+            if 'position' not in d and 'pose' not in d: continue
+            yield(d)
+        except:
+            # Ignore all lines that aren't valid json
+            pass
+
+def make_plotter():
+    import numpy as np
+    from mpl_toolkits.mplot3d import Axes3D
+
+    fig = plt.figure()
+    ax = Axes3D(fig)
+    fig.add_axes(ax)
+
+    ax_bounds = (-0.5, 0.5) # meters
+    ax.set(xlim=ax_bounds, ylim=ax_bounds, zlim=ax_bounds)
+    ax.view_init(azim=-140) # initial plot orientation
+
+    vio_plot = ax.plot(
+        xs=[], ys=[], zs=[],
+        linestyle="-",
+        marker=""
+    )
+    ax.set_xlabel("x (m)")
+    ax.set_ylabel("y (m)")
+    ax.set_zlabel("z (m)")
+
+    title = ax.set_title("VIO trajectory")
+
+    data = { c: [] for c in 'xyz' }
+
+    control = { 'close': False }
+    fig.canvas.mpl_connect('close_event', lambda _: control.update({'close': True}))
+
+    def update_data(vio_out):
+        if control['close']: return False
+        # supports two slightly different JSONL formats
+        if 'pose' in vio_out: vio_out = vio_out['pose']
+        # SDK < 0.12 does not expose the TRACKING status
+        is_tracking = vio_out.get('status', 'TRACKING') == 'TRACKING'
+        for c in 'xyz':
+            val = vio_out['position'][c]
+            if not is_tracking: val = np.nan
+            data[c].append(val)
+        return True
+
+    def update_graph(frames):
+        x, y, z = [np.array(data[c]) for c in 'xyz']
+        vio_plot[0].set_data(x, y)
+        vio_plot[0].set_3d_properties(z)
+        return (vio_plot[0],)
+
+    from matplotlib.animation import FuncAnimation
+    anim = FuncAnimation(fig, update_graph, interval=15, blit=True)
+    return update_data, anim
+
+if __name__ == '__main__':
+    plotter, anim = make_plotter()
+    import argparse
+    parser = argparse.ArgumentParser(__doc__)
+    parser.add_argument("--dataFolder", help="Instead of running live mapping session, replay session from this folder")
+    parser.add_argument('--file', type=argparse.FileType('r'),
+        help='Read data from a JSONL file or pipe instead of displaying it live',
+        default=None)
+
+    args = parser.parse_args()
+
+    def reader_loop():
+        replay = None
+        if args.dataFolder:
+            import spectacularAI
+            replay = spectacularAI.Replay(args.dataFolder)
+            vio_source = replay_vio_reader(replay)
+        elif args.file:
+            vio_source = file_vio_reader(args.file)
+        else:
+            vio_source = live_vio_reader()
+
+        for vio_out in vio_source:
+            if not plotter(vio_out): break
+        if replay: replay.close()
+
+    reader_thread = threading.Thread(target = reader_loop)
+    reader_thread.start()
+    plt.show()
+    reader_thread.join()
