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/dumper.cpp

@@ -145,6 +145,7 @@ 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);
     }
 
+    LogMessage("Dumping float RGB image of size {}x{} to {}/rgb_{}.png", width, height, subdir, dump_id);
     return _dump_RGB_image(h_image_u8, width, height, 3, subdir, dump_id);
 }
 
@@ -173,6 +174,8 @@ void DumpRGBImageFromCuda(
         LogMessage("Failed to copy RGB image from device to host: {}", cudaGetErrorString(err));
         return;
     }
+    LogMessage("Dumping uint8 RGB image of size {}x{} to {}/rgb_{}.png", width, height, subdir, dump_id);
+
     return _dump_RGB_image(h_image, width, height, num_channels, subdir, dump_id);
 
 }

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;

src/include/spot-connection.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include "spot-observer.h"
+#include "model.h"
 
 #include <bosdyn/client/sdk/client_sdk.h>
 #include <bosdyn/client/robot/robot.h>
@@ -67,23 +68,25 @@ class ReaderWriterCBuf {
     // Attach a stream created/owned by SpotConnection.
     inline void attachCudaStream(cudaStream_t stream) { cuda_stream_ = stream; }
 
-    friend class SpotConnection;
+    friend class SpotCamStream;
 };
 
 ///////////////////////////////////////////////////////////////////////////////
 
-class SpotConnection {
-private:
-    std::unique_ptr<bosdyn::client::Robot> robot_;
-    std::unique_ptr<bosdyn::client::ClientSdk> sdk_;
+class SpotConnection;
+class VisionPipeline;
+
+class SpotCamStream {
+    SpotCamStream(const SpotCamStream&) = delete;
+    SpotCamStream& operator=(const SpotCamStream&) = delete;
+    SpotCamStream() = delete;
 
-    bosdyn::client::ImageClient* image_client_;
+    std::shared_ptr<bosdyn::client::ImageClient> image_client_;
 
     // Thread data
     ReaderWriterCBuf image_lifo_;
     std::atomic<bool> quit_requested_{false};
     std::atomic<int> num_samples_{0};
-    bool connected_;
     bool streaming_;
     // Camera feed params
     uint32_t current_cam_mask_;
@@ -96,7 +99,11 @@ class SpotConnection {
     // One CUDA stream per connection (owned here).
     cudaStream_t cuda_stream_{nullptr};
 
-private:
+    SpotConnection& robot_;
+
+    TensorShape current_rgb_shape_{0, 0, 0, 0};
+    TensorShape current_depth_shape_{0, 0, 0, 0};
+
     bosdyn::api::GetImageRequest _createImageRequest(
         const std::vector<std::string>& rgb_sources,
         const std::vector<std::string>& depth_sources
@@ -110,15 +117,12 @@ class SpotConnection {
     void _joinStreamingThread();
 
 public:
-
-    SpotConnection();
-    ~SpotConnection();
-
-    bool connect(
-        const std::string& robot_ip,
-        const std::string& username,
-        const std::string& password
+    SpotCamStream(
+        SpotConnection& robot,
+        std::shared_ptr<bosdyn::client::ImageClient> image_client,
+        int32_t image_lifo_max_size
     );
+    ~SpotCamStream();
 
     bool streamCameras(uint32_t cam_mask);
     bool getCurrentImages(
@@ -127,13 +131,51 @@ class SpotConnection {
         float** depths
     ) const;
 
-    bool     isConnected()       const { return connected_; }
     bool     isStreaming()       const { return streaming_; }
     uint32_t getCurrentCamMask() const { return current_cam_mask_; }
     int32_t  getCurrentNumCams() const { return current_num_cams_; }
 
+    TensorShape getCurrentRGBTensorShape() const { return current_rgb_shape_; }
+    TensorShape getCurrentDepthTensorShape() const { return current_depth_shape_; }
+
     // Return the connection's CUDA stream.
     const cudaStream_t getCudaStream() const { return cuda_stream_; }
 };
 
+class SpotConnection {
+    std::unique_ptr<bosdyn::client::Robot> robot_;
+    std::unique_ptr<bosdyn::client::ClientSdk> sdk_;
+    std::shared_ptr<bosdyn::client::ImageClient> image_client_;
+
+    int32_t image_lifo_max_size_;
+
+    std::unordered_map<int32_t, std::unique_ptr<SpotCamStream>> cam_streams_;
+    std::unordered_map<int32_t, std::unique_ptr<VisionPipeline>> vision_pipelines_;
+
+    bool connected_;
+
+    int32_t next_stream_id_{0xee1};
+
+public:
+
+    SpotConnection(
+        const std::string& robot_ip,
+        const std::string& username,
+        const std::string& password
+    );
+    ~SpotConnection();
+
+    int32_t createCamStream(uint32_t cam_mask);
+    bool removeCamStream(int32_t stream_id);
+    SpotCamStream* getCamStream(int32_t stream_id);
+
+    bool createVisionPipeline(MLModel& model, int32_t stream_id);
+    bool removeVisionPipeline(int32_t stream_id);
+    VisionPipeline* getVisionPipeline(int32_t stream_id);
+
+    bool isConnected() const { return connected_; }
+
+    friend class SpotCamStream;
+};
+
 } // namespace SOb

src/include/unity-cuda-interop.h

@@ -32,15 +32,16 @@ void shutdownUnityInterop();
 
 bool registerOutputTextures(
     int32_t robot_id,
+    int32_t cam_stream_id,
     uint32_t cam_bit,         // Single bit only
     void* out_img_tex,        // ID3D12Resource* (aka texture)
     void* out_depth_tex,      // ID3D12Resource* (aka texture)
     int32_t img_buffer_size,  // In bytes
     int32_t depth_buffer_size // In bytes
 );
 
-bool uploadNextImageSetToUnity(int32_t robot_id);
-bool uploadNextVisionPipelineImageSetToUnity(int32_t robot_id);
+bool uploadNextImageSetToUnity(int32_t robot_id, int32_t cam_stream_id);
+bool uploadNextVisionPipelineImageSetToUnity(int32_t robot_id, int32_t cam_stream_id);
 bool clearOutputTextures(int32_t robot_id);
 
 }

src/include/utils.h

@@ -14,6 +14,24 @@
 
 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) {}
+    bool operator==(const TensorShape& other) const {
+        return N == other.N && C == other.C && H == other.H && W == other.W;
+    }
+    bool operator!=(const TensorShape& other) const {
+        return !(*this == other);
+    }
+    size_t total_size() const {
+        return N * C * H * W;
+    }
+    std::string to_string() const {
+        return std::format("N={}, C={}, H={}, W={}", N, C, H, W);
+    }
+};
+
+
 inline void checkCudaError(cudaError_t error, const std::string& operation) {
     if (error != cudaSuccess) {
         throw std::runtime_error(operation + " failed: " + cudaGetErrorString(error));

src/include/vision-pipeline.h

@@ -30,7 +30,7 @@ class VisionPipeline {
     };
 
     MLModel& model_;
-    const SpotConnection& spot_connection_;
+    const SpotCamStream& spot_cam_stream_;
     bool first_run_{true};
 
     // Threading
@@ -61,7 +61,7 @@ class VisionPipeline {
 public:
     VisionPipeline(
         MLModel& model,
-        const SpotConnection& spot_connection,
+        const SpotCamStream& spot_cam_stream_,
         const TensorShape& input_shape,
         const TensorShape& depth_shape,
         const TensorShape& output_shape,