cam_test: add -raw for streaming/saving raw frames (+ISP YUV)

Author: alex-luxonis

utilities/cam_test.py

@@ -36,6 +36,7 @@
 #os.environ["DEPTHAI_LEVEL"] = "debug"
 
 import cv2
+import numpy as np
 import argparse
 import collections
 import time
@@ -80,6 +81,8 @@ def socket_type_pair(arg):
                     help="Path to custom camera tuning database")
 parser.add_argument('-d', '--device', default="", type=str,
                     help="Optional MX ID of the device to connect to.")
+parser.add_argument('-raw', '--enable-raw', default=False, action="store_true",
+                    help='Enable the RAW camera streams')
 
 parser.add_argument('-ctimeout', '--connection-timeout', default=30000,
                     help="Connection timeout in ms. Default: %(default)s (sets DEPTHAI_CONNECTION_TIMEOUT environment variable)")
@@ -178,9 +181,12 @@ def get(self):
 
 cam = {}
 xout = {}
+xout_raw = {}
+streams = []
 for c in cam_list:
     xout[c] = pipeline.createXLinkOut()
     xout[c].setStreamName(c)
+    streams.append(c)
     if cam_type_color[c]:
         cam[c] = pipeline.createColorCamera()
         cam[c].setResolution(color_res_opts[args.color_resolution])
@@ -206,6 +212,15 @@ def get(self):
     cam[c].setFps(args.fps)
     cam[c].setIsp3aFps(args.isp3afps)
 
+    if args.enable_raw:
+        raw_name = 'raw_' + c
+        xout_raw[c] = pipeline.create(dai.node.XLinkOut)
+        xout_raw[c].setStreamName(raw_name)
+        if args.enable_raw:
+            streams.append(raw_name)
+        cam[c].raw.link(xout_raw[c].input)
+        # to be added cam[c].setRawOutputPacked(False)
+
 if args.camera_tuning:
     pipeline.setCameraTuningBlobPath(str(args.camera_tuning))
 
@@ -232,6 +247,8 @@ def exit_cleanly(signum, frame):
         print('auto ' if p.hasAutofocus else 'fixed', '- ', end='')
         print(*[type.name for type in p.supportedTypes])
         cam_name[p.socket.name] = p.sensorName
+        if args.enable_raw:
+            cam_name['raw_'+p.socket.name] = p.sensorName
 
     print('USB speed:', device.getUsbSpeed().name)
 
@@ -240,7 +257,7 @@ def exit_cleanly(signum, frame):
     q = {}
     fps_host = {}  # FPS computed based on the time we receive frames in app
     fps_capt = {}  # FPS computed based on capture timestamps from device
-    for c in cam_list:
+    for c in streams:
         q[c] = device.getOutputQueue(name=c, maxSize=4, blocking=False)
         # The OpenCV window resize may produce some artifacts
         if args.resizable_windows:
@@ -299,7 +316,7 @@ def exit_cleanly(signum, frame):
 
     capture_list = []
     while True:
-        for c in cam_list:
+        for c in streams:
             try:
                 pkt = q[c].tryGet()
             except Exception as e:
@@ -308,25 +325,44 @@ def exit_cleanly(signum, frame):
             if pkt is not None:
                 fps_host[c].update()
                 fps_capt[c].update(pkt.getTimestamp().total_seconds())
+                width, height = pkt.getWidth(), pkt.getHeight()
                 frame = pkt.getCvFrame()
-                if c in capture_list:
-                    width, height = pkt.getWidth(), pkt.getHeight()
-                    capture_file_name = ('capture_' + c + '_' + cam_name[cam_socket_opts[c].name]
-                                         + '_' + str(width) + 'x' + str(height)
-                                         + '_exp_' +
-                                         str(int(
-                                             pkt.getExposureTime().total_seconds()*1e6))
-                                         + '_iso_' + str(pkt.getSensitivity())
-                                         + '_lens_' +
-                                         str(pkt.getLensPosition())
-                                         + '_' + capture_time
-                                         + '_' + str(pkt.getSequenceNum())
-                                         + ".png"
-                                         )
-                    print("\nSaving:", capture_file_name)
-                    cv2.imwrite(capture_file_name, frame)
+                capture = c in capture_list
+                if capture:
+                    capture_file_info = ('capture_' + c + '_' + cam_name[c]
+                         + '_' + str(width) + 'x' + str(height)
+                         + '_exp_' + str(int(pkt.getExposureTime().total_seconds()*1e6))
+                         + '_iso_' + str(pkt.getSensitivity())
+                         + '_lens_' + str(pkt.getLensPosition())
+                         + '_' + capture_time
+                         + '_' + str(pkt.getSequenceNum())
+                        )
                     capture_list.remove(c)
-
+                    print()
+                if c.startswith('raw_'):
+                    if capture:
+                        filename = capture_file_info + '_10bit.bw'
+                        print('Saving:', filename)
+                        frame.tofile(filename)
+                    # Full range for display, use bits [15:6] of the 16-bit pixels
+                    frame = frame * (1 << 6)
+                    # Debayer color for preview/png
+                    if cam_type_color[c.split('_')[-1]]:
+                        # See this for the ordering, at the end of page:
+                        # https://docs.opencv.org/4.5.1/de/d25/imgproc_color_conversions.html
+                        # TODO add bayer order to ImgFrame getType()
+                        frame = cv2.cvtColor(frame, cv2.COLOR_BayerGB2BGR)
+                else:
+                    # Save YUV too, but only when RAW is also enabled (for tuning purposes)
+                    if capture and args.enable_raw:
+                        payload = pkt.getData()
+                        filename = capture_file_info + '_P420.yuv'
+                        print('Saving:', filename)
+                        payload.tofile(filename)
+                if capture:
+                    filename = capture_file_info + '.png'
+                    print('Saving:', filename)
+                    cv2.imwrite(filename, frame)
                 cv2.imshow(c, frame)
         print("\rFPS:",
               *["{:6.2f}|{:6.2f}".format(fps_host[c].get(),
@@ -337,7 +373,7 @@ def exit_cleanly(signum, frame):
         if key == ord('q'):
             break
         elif key == ord('c'):
-            capture_list = cam_list.copy()
+            capture_list = streams.copy()
             capture_time = time.strftime('%Y%m%d_%H%M%S')
         elif key == ord('t'):
             print("Autofocus trigger (and disable continuous)")