Support multiple cam streams from a single spot

include/spot-observer.h

@@ -9,7 +9,9 @@
 
 #include "IUnityInterface.h"
 
+#ifdef __cplusplus
 extern "C" {
+#endif
 
 // Define a function pointer type for logging
 typedef void (*SOb_LogCallback)(const char* message);
@@ -39,11 +41,16 @@ bool UNITY_INTERFACE_API SOb_DisconnectFromSpot(int32_t robot_id);
 
 // Start reading spot camera feeds. Runs in a separate threads.
 UNITY_INTERFACE_EXPORT
-bool UNITY_INTERFACE_API SOb_ReadCameraFeeds(int32_t robot_id, uint32_t cam_bitmask);
+int32_t UNITY_INTERFACE_API SOb_CreateCameraStream(int32_t robot_id, uint32_t cam_bitmask);
+
+UNITY_INTERFACE_EXPORT
+bool UNITY_INTERFACE_API SOb_DestroyCameraStream(int32_t robot_id, int32_t cam_stream_id);
+
 
 UNITY_INTERFACE_EXPORT
 bool UNITY_INTERFACE_API SOb_GetNextImageSet(
     int32_t robot_id,
+    int32_t cam_stream_id,
     int32_t n_images_requested,
     uint8_t** images,
     float** depths
@@ -52,6 +59,7 @@ bool UNITY_INTERFACE_API SOb_GetNextImageSet(
 UNITY_INTERFACE_EXPORT
 bool UNITY_INTERFACE_API SOb_RegisterUnityReadbackBuffers(
     int32_t robot_id,
+    int32_t cam_stream_id,
     uint32_t cam_bit,         // Single bit only
     void* out_img_buf,        // ID3D12Resource* (aka tensor)
     void* out_depth_buf,      // ID3D12Resource* (aka tensor)
@@ -63,22 +71,27 @@ UNITY_INTERFACE_EXPORT
 bool UNITY_INTERFACE_API SOb_ClearUnityReadbackBuffers(int32_t robot_id);
 
 UNITY_INTERFACE_EXPORT
-bool UNITY_INTERFACE_API SOb_LaunchVisionPipeline(int32_t robot_id, SObModel model);
+bool UNITY_INTERFACE_API SOb_LaunchVisionPipeline(
+    int32_t robot_id,
+    int32_t cam_stream_id,
+    SObModel model
+);
 UNITY_INTERFACE_EXPORT
-bool UNITY_INTERFACE_API SOb_StopVisionPipeline(int32_t robot_id);
+bool UNITY_INTERFACE_API SOb_StopVisionPipeline(int32_t robot_id, int32_t cam_stream_id);
 UNITY_INTERFACE_EXPORT
 bool UNITY_INTERFACE_API SOb_GetNextVisionPipelineImageSet(
     int32_t robot_id,
+    int32_t cam_stream_id,
     int32_t n_images_requested,
     uint8_t** images,
     float** depths
 );
 
 UNITY_INTERFACE_EXPORT
-bool UNITY_INTERFACE_API SOb_PushNextImageSetToUnityBuffers(int32_t robot_id);
+bool UNITY_INTERFACE_API SOb_PushNextImageSetToUnityBuffers(int32_t robot_id, int32_t cam_stream_id);
 
 UNITY_INTERFACE_EXPORT
-bool UNITY_INTERFACE_API SOb_PushNextVisionPipelineImageSetToUnityBuffers(int32_t robot_id);
+bool UNITY_INTERFACE_API SOb_PushNextVisionPipelineImageSetToUnityBuffers(int32_t robot_id, int32_t cam_stream_id);
 
 // Model stuff
 UNITY_INTERFACE_EXPORT
@@ -103,6 +116,8 @@ bool UNITY_INTERFACE_API SOb_SetUnityLogCallback(const SOb_LogCallback callback)
 UNITY_INTERFACE_EXPORT
 void UNITY_INTERFACE_API SOb_ToggleDebugDumps(const char* dump_path);
 
-}
+#ifdef __cplusplus
+} // extern "C"
+#endif
 
 #endif //SPOT_OBSERVER_H

src/cuda_kernels.cu

@@ -864,6 +864,11 @@ cudaError_t preprocess_depth_image2(
         err = cudaGetLastError();
         if (err != cudaSuccess) return err;
     }
+    // } else {
+    //     // Just copy input to output buffer
+    //     err = cudaMemcpyAsync(d_out, depth_image_in, N * sizeof(float), cudaMemcpyDeviceToDevice, stream);
+    //     if (err != cudaSuccess) return err;
+    // }
 
     if (downscale_factor > 1) {
         if ((downscale_factor & (downscale_factor - 1)) != 0) {
@@ -1159,7 +1164,7 @@ __global__ void nhwc_to_nchw_rgb_float_kernel(
 
             int out_p = out_y * out_W + out_x;
             int base = n * 3 * out_pixels_per_img + out_p;   // position of R plane element
-            out[base]                        = px.x * scale;
+            out[base]                          = px.x * scale;
             out[base +     out_pixels_per_img] = px.y * scale;
             out[base + 2 * out_pixels_per_img] = px.z * scale;
         }
@@ -1244,13 +1249,13 @@ __global__ void prefill_depth_from_cache(
     if (x >= width || y >= height) return;
 
     int idx = y * width + x;
-    float new_val = new_depth[idx];
-    float old_val = cached_depth[idx];
+    float new_val    = new_depth[idx];
+    float cached_val = cached_depth[idx];
 
     // Check if new depth value is valid
     bool new_valid = (new_val >= min_valid_depth && new_val <= max_valid_depth);
     // Update the new_depth buffer in-place to have gaps filled
-    output_depth[idx] = new_valid ? old_val : new_val;
+    output_depth[idx] = new_valid ? new_val : cached_val;
 }
 
 cudaError_t prefill_invalid_depth(
@@ -1305,9 +1310,13 @@ __global__ void update_depth_cache_kernel(
     bool old_valid = (old_val >= min_valid_depth && old_val <= max_valid_depth);
 
     if (new_valid) {
-        cached_depth[idx] = new_val;
+        if (old_valid) {
+            cached_depth[idx] = old_val * 0.5f + new_val * 0.5f;
+        } else {
+            cached_depth[idx] = new_val;
+        }
     } else if (old_valid) {
-        cached_depth[idx] = old_val * 0.5f + generated_val * 0.5f;
+        cached_depth[idx] = old_val * 0.8f + generated_val * 0.2f;
     } else {
         cached_depth[idx] = generated_val;
     }

src/dumper.cpp

@@ -34,6 +34,17 @@ bool create_directory_if_not_exists(const std::string& path) {
     return true; // Directory created successfully
 }
 
+static std::string get_file_path(
+    const std::string& path,
+    const std::string& subdir,
+    const std::string& base_name,
+    int32_t dump_id_major,
+    int32_t dump_id_minor,
+    const std::string& extension
+) {
+    return path + "/" + subdir + "/" + base_name + std::to_string(dump_id_major) + "_" + std::to_string(dump_id_minor) + extension;
+}
+
 bool ToggleDumping(const std::string& dump_path) {
     if (dump_path.empty()) {
         m_dumps_enabled = false;
@@ -70,15 +81,23 @@ static void _dump_RGB_image(
     int32_t height,
     int32_t channels,
     const std::string& subdir,
-    int32_t dump_id
+    int32_t dump_id_major,
+    int32_t dump_id_minor
 ) {
     // Generate filename
     std::string subdir_path = m_dump_path + "/" + subdir;
     if (!create_directory_if_not_exists(subdir_path)) {
         LogMessage("_dump_RGB_image: Failed to create directory: {}", subdir_path);
         return;
     }
-    std::string file_path = m_dump_path + "/" + subdir + "/rgb_" + std::to_string(dump_id) + ".png";
+    std::string file_path = get_file_path(
+        m_dump_path,
+        subdir,
+        "rgb_",
+        dump_id_major,
+        dump_id_minor,
+        ".png"
+    );
 
     // Write image to file
     const int stride_in_bytes = width * channels;
@@ -88,14 +107,14 @@ static void _dump_RGB_image(
 }
 
 
-void DumpRGBImageFromCudaHWC(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id) {
+void DumpRGBImageFromCudaHWC(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor) {
     if (!m_dumps_enabled) {
         return;
     }
 
     const size_t num_pixels = width * height;
     const size_t float_size = num_pixels * 3 * sizeof(float);
-    
+
     // Allocate host memory for float image
     std::vector<float> h_image(num_pixels * 3);
 
@@ -115,10 +134,10 @@ void DumpRGBImageFromCudaHWC(const float* image, int32_t width, int32_t height,
         h_image_u8[i * 3 + 2] = static_cast<uint8_t>(std::max(0.0f, std::min(1.0f, h_image[i * 3 + 2])) * 255.0f);
     }
 
-   return _dump_RGB_image(h_image_u8, width, height, 3, subdir, dump_id);
+   return _dump_RGB_image(h_image_u8, width, height, 3, subdir, dump_id_major, dump_id_minor);
 }
 
-void DumpRGBImageFromCudaCHW(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id) {
+void DumpRGBImageFromCudaCHW(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor) {
     if (!m_dumps_enabled) {
         return;
     }
@@ -145,7 +164,8 @@ void DumpRGBImageFromCudaCHW(const float* image, int32_t width, int32_t height,
         h_image_u8[i * 3 + 2] = static_cast<uint8_t>(std::max(0.0f, std::min(1.0f, h_image[i + 2 * num_pixels])) * 255.0f);
     }
 
-    return _dump_RGB_image(h_image_u8, width, height, 3, subdir, dump_id);
+    LogMessage("Dumping float RGB image of size {}x{} to {}/rgb_{}_{}.png", width, height, subdir, dump_id_major, dump_id_minor);
+    return _dump_RGB_image(h_image_u8, width, height, 3, subdir, dump_id_major, dump_id_minor);
 }
 
 void DumpRGBImageFromCuda(
@@ -154,7 +174,8 @@ void DumpRGBImageFromCuda(
     int32_t height,
     int32_t num_channels,
     const std::string& subdir,
-    int32_t dump_id
+    int32_t dump_id_major,
+    int32_t dump_id_minor
 ) {
     if (!m_dumps_enabled) {
         return;
@@ -173,11 +194,19 @@ void DumpRGBImageFromCuda(
         LogMessage("Failed to copy RGB image from device to host: {}", cudaGetErrorString(err));
         return;
     }
-    return _dump_RGB_image(h_image, width, height, num_channels, subdir, dump_id);
+    LogMessage("Dumping uint8 RGB image of size {}x{} to {}/rgb_{}_{}.png", width, height, subdir, dump_id_major, dump_id_minor);
 
+    return _dump_RGB_image(h_image, width, height, num_channels, subdir, dump_id_major, dump_id_minor);
 }
 
-void DumpDepthImageFromCuda(const float* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id) {
+void DumpDepthImageFromCuda(
+    const float* depth,
+    int32_t width,
+    int32_t height,
+    const std::string& subdir,
+    int32_t dump_id_major,
+    int32_t dump_id_minor
+) {
     if (!m_dumps_enabled) {
         return;
     }
@@ -223,7 +252,15 @@ void DumpDepthImageFromCuda(const float* depth, int32_t width, int32_t height, c
         LogMessage("Failed to create directory: {}", subdir_path);
         return;
     }
-    std::string file_path = m_dump_path + "/" + subdir + "/depth_" + std::to_string(dump_id) + ".png";
+
+    std::string file_path = get_file_path(
+        m_dump_path,
+        subdir,
+        "depth_",
+        dump_id_major,
+        dump_id_minor,
+        ".png"
+    );
 
     // Write image to file
     const int channels = 1; // Grayscale

src/include/dumper.h

@@ -5,11 +5,11 @@
 namespace SOb {
 
 bool ToggleDumping(const std::string& dump_path);
-void DumpRGBImageFromCudaCHW(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id);
-void DumpRGBImageFromCudaHWC(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id);
-void DumpRGBImageFromCuda(const uint8_t* image, int32_t width, int32_t height, int32_t num_channels, const std::string& subdir, int32_t dump_id);
-void DumpDepthImageFromCuda(const float* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id);
-void DumpDepthImageFromCuda(const uint16_t* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id);
+void DumpRGBImageFromCudaCHW(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor = 0);
+void DumpRGBImageFromCudaHWC(const float* image, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor = 0);
+void DumpRGBImageFromCuda(const uint8_t* image, int32_t width, int32_t height, int32_t num_channels, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor = 0);
+void DumpDepthImageFromCuda(const float* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor = 0);
+void DumpDepthImageFromCuda(const uint16_t* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id_major, int32_t dump_id_minor = 0);
 
 
 }

src/include/model.h

@@ -4,18 +4,15 @@
 
 #pragma once
 
+#include "utils.h"
+
 #include <string>
 #include <torch/script.h>
 #include <torch/torch.h>
 #include <onnxruntime_cxx_api.h>
 
 namespace SOb {
 
-struct TensorShape {
-    size_t N, C, H, W;
-    TensorShape(size_t n, size_t c, size_t h, size_t w) : N(n), C(c), H(h), W(w) {}
-};
-
 class MLModel {
 public:
     virtual ~MLModel() = default;