diff --git a/include/spot-observer.h b/include/spot-observer.h index 0311358..f619474 100644 --- a/include/spot-observer.h +++ b/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 diff --git a/src/cuda_kernels.cu b/src/cuda_kernels.cu index 6c559f8..3740a73 100644 --- a/src/cuda_kernels.cu +++ b/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; } diff --git a/src/dumper.cpp b/src/dumper.cpp index ba94516..5f4b869 100644 --- a/src/dumper.cpp +++ b/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,7 +81,8 @@ 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; @@ -78,7 +90,14 @@ static void _dump_RGB_image( 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 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(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(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 diff --git a/src/include/dumper.h b/src/include/dumper.h index 753b0a0..43d7476 100644 --- a/src/include/dumper.h +++ b/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); } \ No newline at end of file diff --git a/src/include/model.h b/src/include/model.h index 0029fc3..282335c 100644 --- a/src/include/model.h +++ b/src/include/model.h @@ -4,6 +4,8 @@ #pragma once +#include "utils.h" + #include #include #include @@ -11,11 +13,6 @@ 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; diff --git a/src/include/spot-connection.h b/src/include/spot-connection.h index 6750c1a..06dbad4 100644 --- a/src/include/spot-connection.h +++ b/src/include/spot-connection.h @@ -5,6 +5,7 @@ #pragma once #include "spot-observer.h" +#include "model.h" #include #include @@ -67,15 +68,18 @@ 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 robot_; - std::unique_ptr sdk_; +class SpotConnection; +class VisionPipeline; + +class SpotCamStream { + SpotCamStream(const SpotCamStream&) = delete; + SpotCamStream& operator=(const SpotCamStream&) = delete; + SpotCamStream() = delete; bosdyn::client::ImageClient* image_client_; @@ -83,7 +87,6 @@ class SpotConnection { ReaderWriterCBuf image_lifo_; std::atomic quit_requested_{false}; std::atomic 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& rgb_sources, const std::vector& 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, + 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 robot_; + std::unique_ptr sdk_; + bosdyn::client::ImageClient* image_client_; + + int32_t image_lifo_max_size_; + + std::unordered_map> cam_streams_; + std::unordered_map> 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 \ No newline at end of file diff --git a/src/include/unity-cuda-interop.h b/src/include/unity-cuda-interop.h index 355a830..9c21e1d 100644 --- a/src/include/unity-cuda-interop.h +++ b/src/include/unity-cuda-interop.h @@ -32,6 +32,7 @@ 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) @@ -39,8 +40,8 @@ bool registerOutputTextures( 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); } \ No newline at end of file diff --git a/src/include/utils.h b/src/include/utils.h index f84d9bb..3bae154 100644 --- a/src/include/utils.h +++ b/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)); diff --git a/src/include/vision-pipeline.h b/src/include/vision-pipeline.h index b2d802d..1ac2c5e 100644 --- a/src/include/vision-pipeline.h +++ b/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 @@ -54,6 +54,9 @@ class VisionPipeline { float* d_output_buffer_{nullptr}; uint8_t* d_rgb_data_{nullptr}; + int32_t dump_id = 0; + int32_t thread_num; + void pipelineWorker(std::stop_token stop_token); bool allocateCudaBuffers(); void deallocateCudaBuffers(); @@ -61,7 +64,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, diff --git a/src/spot-connection.cpp b/src/spot-connection.cpp index e395735..6b2d20b 100644 --- a/src/spot-connection.cpp +++ b/src/spot-connection.cpp @@ -6,7 +6,7 @@ #include "logger.h" #include "utils.h" #include "dumper.h" -#include "cuda_kernels.cuh" +#include "vision-pipeline.h" #include @@ -184,8 +184,8 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrFieldhost copies on default stream, this ensures correctness: - // checkCudaError(cudaStreamSynchronize(cuda_stream_), "sync before debug dump"); DumpDepthImageFromCuda( depth_write_ptr, cv_img.cols, cv_img.rows, - "pre-depth-cache", - n_depths_written + write_idx * n_images_per_response_ - ); - - // checkCudaError(update_depth_cache(depth_write_ptr, depth_cache_ptr, depth_width, depth_height, cuda_stream_), "update_depth_cache"); - - DumpDepthImageFromCuda( - depth_write_ptr, - cv_img.cols, - cv_img.rows, - "post-depth-cache", + "depth", n_depths_written + write_idx * n_images_per_response_ ); @@ -359,33 +348,40 @@ static std::vector convert_bitmask_to_spot_cam_vector(uint32_t bitma /////////////////////////////////////////////////////////////////////////////// -SpotConnection::SpotConnection() - : robot_(nullptr) - , image_client_(nullptr) - , image_lifo_(5) - , connected_(false) +SpotCamStream::SpotCamStream( + SpotConnection& robot, + bosdyn::client::ImageClient* image_client, + int32_t image_lifo_max_size +) + : robot_(robot) + , image_client_(image_client) + , image_lifo_(image_lifo_max_size) , streaming_(false) { - // Create SDK instance - sdk_ = bosdyn::client::CreateStandardSDK("SpotObserverConnection"); - if (!sdk_) { - LogMessage("SpotConnection::SpotConnection: Failed to create SDK instance"); - throw std::runtime_error("Failed to create Spot SDK instance"); - } + // Create one CUDA stream per SpotConnection and attach it to the buffer. + checkCudaError( + cudaStreamCreate(&cuda_stream_), + "cudaStreamCreate for SpotConnection" + ); + image_lifo_.attachCudaStream(cuda_stream_); + LogMessage("SpotCamStream::connect: Created CUDA stream {:#x} and attached to buffer", + size_t(cuda_stream_)); + } -SpotConnection::~SpotConnection() { +SpotCamStream::~SpotCamStream() { + quit_requested_.store(true); _joinStreamingThread(); - // Destroy the per-connection CUDA stream after the thread is stopped. + if (cuda_stream_) { checkCudaError(cudaStreamDestroy(cuda_stream_), "cudaStreamDestroy for SpotConnection"); cuda_stream_ = nullptr; - LogMessage("Destroyed CUDA stream for SpotConnection"); + LogMessage("SpotCamStream::~SpotConnection: Destroyed CUDA stream"); } // TODO: figure out how to cleanup image_client_ } -bosdyn::api::GetImageRequest SpotConnection::_createImageRequest( +bosdyn::api::GetImageRequest SpotCamStream::_createImageRequest( const std::vector& rgb_sources, const std::vector& depth_sources ) { @@ -413,13 +409,15 @@ bosdyn::api::GetImageRequest SpotConnection::_createImageRequest( return request; } -// Image producer thread that requests images from the robot -void SpotConnection::_spotCamReaderThread(std::stop_token stop_token) { - if (!image_client_) { - LogMessage("Image client not initialized"); - return; - } +void SpotCamStream::_startStreamingThread() { + // Create and start thread + image_streamer_thread_ = std::make_unique([this](std::stop_token stop_token) { + _spotCamReaderThread(stop_token); + }); +} +// Image producer thread that requests images from the robot +void SpotCamStream::_spotCamReaderThread(std::stop_token stop_token) { LogMessage("Producer thread started"); while (!stop_token.stop_requested() && !quit_requested_.load()) { @@ -439,6 +437,9 @@ void SpotConnection::_spotCamReaderThread(std::stop_token stop_token) { // Process and queue images image_lifo_.push(response.response.image_responses()); + auto end = std::chrono::high_resolution_clock::now(); + LogMessage("Total time for GetImage and push: {:.4f} ms", + std::chrono::duration(end - start).count()); } catch (const std::exception& e) { LogMessage("Error in producer thread: {}", e.what()); std::this_thread::sleep_for(std::chrono::seconds(1)); @@ -447,19 +448,7 @@ void SpotConnection::_spotCamReaderThread(std::stop_token stop_token) { LogMessage("Exiting producer thread."); } -void SpotConnection::_startStreamingThread() { - if (!image_client_) { - std::cerr << "Image client not initialized" << std::endl; - return; - } - - // Create and start thread - image_streamer_thread_ = std::make_unique([this](std::stop_token stop_token) { - _spotCamReaderThread(stop_token); - }); -} - -void SpotConnection::_joinStreamingThread() { +void SpotCamStream::_joinStreamingThread() { if (image_streamer_thread_) { image_streamer_thread_->request_stop(); if (image_streamer_thread_->joinable()) { @@ -469,76 +458,19 @@ void SpotConnection::_joinStreamingThread() { image_streamer_thread_.reset(); LogMessage("Streaming thread stopped"); } else { - LogMessage("SpotConnection::_joinStreamingThread: No streaming thread to join"); + LogMessage("SpotCamStream::_joinStreamingThread: No streaming thread to join"); return; } } -bool SpotConnection::connect( - const std::string& robot_ip, - const std::string& username, - const std::string& password -) { - try { - // Create robot using ClientSDK - bosdyn::client::Result> robot_result = sdk_->CreateRobot(robot_ip); - if (!robot_result.status) { - LogMessage("SpotConnection::connect: Failed to connect to robot: {}", - robot_result.status.message()); - return false; - } - - robot_ = std::move(robot_result.response); - - // Authenticate - bosdyn::common::Status auth_status = robot_->Authenticate(username, password); - if (!auth_status) { - LogMessage("SpotConnection::connect: Failed to authenticate: {}", auth_status.message()); - return false; - } - - // Create image client - bosdyn::client::Result image_client_result = - robot_->EnsureServiceClient(); - - if (!image_client_result.status) { - LogMessage("SpotConnection::connect: Failed to create image client: {}", - image_client_result.status.message()); - return false; - } - - image_client_ = image_client_result.response; - - LogMessage("SpotConnection::connect: Connected to Spot robot at {}", robot_ip); - - connected_ = true; - - // Create one CUDA stream per SpotConnection and attach it to the buffer. - checkCudaError( - cudaStreamCreate(&cuda_stream_), - "cudaStreamCreate for SpotConnection" - ); - image_lifo_.attachCudaStream(cuda_stream_); - LogMessage("SpotConnection::connect: Created CUDA stream {:#x} and attached to buffer", - size_t(cuda_stream_)); - - return true; - - } catch (const std::exception& e) { - LogMessage("SpotConnection::connect: Exception while connecting to robot {}: {}", - robot_ip, e.what()); - return false; - } -} - -bool SpotConnection::streamCameras(uint32_t cam_mask) { - if (!connected_) { - LogMessage("SpotConnection::streamCameras: Not connected to robot"); +bool SpotCamStream::streamCameras(uint32_t cam_mask) { + if (!robot_.connected_) { + LogMessage("SpotCamStream::streamCameras: Not connected to robot"); return false; } if (cam_mask == 0 || cam_mask >= SpotCamera::NUM_CAMERAS) { - LogMessage("SpotConnection::streamCameras: Invalid camera mask: {:#x}", cam_mask); + LogMessage("SpotCamStream::streamCameras: Invalid camera mask: {:#x}", cam_mask); return false; } @@ -574,21 +506,28 @@ bool SpotConnection::streamCameras(uint32_t cam_mask) { for (int32_t i = 0; i < max_connection_retries; i++) { bosdyn::client::GetImageResultType response = image_client_->GetImage(current_request_); if (!response.status) { - LogMessage("SpotConnection::streamCameras: Failed to get images: {}", + LogMessage("SpotCamStream::streamCameras: Failed to get images: {}", response.status.message()); - LogMessage("SpotConnection::streamCameras: Retrying... ({}/{})", + LogMessage("SpotCamStream::streamCameras: Retrying... ({}/{})", i + 1, max_connection_retries); std::this_thread::sleep_for(std::chrono::milliseconds(100)); continue; } - LogMessage("SpotConnection::streamCameras: Successfully retrieved images"); + LogMessage("SpotCamStream::streamCameras: Successfully retrieved images"); const auto& image_responses = response.response.image_responses(); if (image_responses.empty()) { - LogMessage("SpotConnection::streamCameras: No images received in response"); + LogMessage("SpotCamStream::streamCameras: No images received in response"); return false; } // Read image sizes - size_t rgb_ref_size = image_responses[0].shot().image().rows() * image_responses[0].shot().image().cols() * 4; + current_rgb_shape_ = TensorShape{ + size_t(num_cams_requested), + 4, //(image_responses[0].shot().image().pixel_format() == bosdyn::api::Image::PIXEL_FORMAT_RGBA_U8 ? 4 : 3) + size_t(image_responses[0].shot().image().rows()), + size_t(image_responses[0].shot().image().cols()) + }; + + size_t rgb_ref_size = current_rgb_shape_.H * current_rgb_shape_.W * 4; //(image_responses[0].shot().image().pixel_format() == bosdyn::api::Image::PIXEL_FORMAT_RGBA_U8 ? 4 : 3);; // For debugging purposes, ensure that all RGB images have the same size for (int32_t j = 1; j < num_cams_requested; j++) { @@ -596,19 +535,25 @@ bool SpotConnection::streamCameras(uint32_t cam_mask) { size_t rgb_size = img_response.shot().image().rows() * img_response.shot().image().cols() * 4; //(img_response.shot().image().pixel_format() == bosdyn::api::Image::PIXEL_FORMAT_RGBA_U8 ? 4 : 3); if (rgb_ref_size != rgb_size) { - LogMessage("SpotConnection::streamCameras: Inconsistent RGB image sizes" + LogMessage("SpotCamStream::streamCameras: Inconsistent RGB image sizes" "(expected {}, got {})", rgb_ref_size, rgb_size); return false; } } // Same thing for depth images - size_t depth_ref_size = image_responses[num_cams_requested].shot().image().rows() * - image_responses[num_cams_requested].shot().image().cols(); + current_depth_shape_ = TensorShape{ + size_t(num_cams_requested), + 1, + size_t(image_responses[num_cams_requested].shot().image().rows()), + size_t(image_responses[num_cams_requested].shot().image().cols()) + }; + + size_t depth_ref_size = current_depth_shape_.H * current_depth_shape_.W; for (int32_t j = num_cams_requested + 1; j < image_responses.size(); j++) { const auto& img_response = image_responses[j]; size_t depth_size = img_response.shot().image().rows() * img_response.shot().image().cols(); if (depth_ref_size != depth_size) { - LogMessage("SpotConnection::streamCameras: Inconsistent depth image sizes" + LogMessage("SpotCamStream::streamCameras: Inconsistent depth image sizes" "(expected {}, got {})", depth_ref_size, depth_size); return false; } @@ -628,7 +573,7 @@ bool SpotConnection::streamCameras(uint32_t cam_mask) { streaming_ = true; } catch (const std::exception& e) { - LogMessage("SpotConnection::streamCameras: Exception while getting images: {}", e.what()); + LogMessage("SpotCamStream::streamCameras: Exception while getting images: {}", e.what()); streaming_ = false; return false; } @@ -638,14 +583,13 @@ bool SpotConnection::streamCameras(uint32_t cam_mask) { return true; } -bool SpotConnection::getCurrentImages( +bool SpotCamStream::getCurrentImages( int32_t n_images_requested, uint8_t** images, float** depths ) const { auto [ret_images, ret_depths] = image_lifo_.pop(n_images_requested); if (ret_images == nullptr || ret_depths == nullptr) { - LogMessage("SpotConnection::getCurrentImages: No images available in the buffer"); return false; } @@ -657,4 +601,223 @@ bool SpotConnection::getCurrentImages( return true; } +/////////////////////////////////////////////////////////////////////////////// + +SpotConnection::SpotConnection( + const std::string& robot_ip, + const std::string& username, + const std::string& password +) + : robot_(nullptr) + , image_client_(nullptr) + , image_lifo_max_size_(5) + , connected_(false) +{ + // Create SDK instance + sdk_ = bosdyn::client::CreateStandardSDK("SpotObserverConnection"); + if (!sdk_) { + LogMessage("SpotConnection::SpotConnection: Failed to create SDK instance"); + throw std::runtime_error("Failed to create Spot SDK instance"); + } + + try { + // Create robot using ClientSDK + bosdyn::client::Result> robot_result = sdk_->CreateRobot(robot_ip); + if (!robot_result.status) { + throw std::runtime_error(std::format("SpotConnection::connect: Failed to connect to robot: {}", + robot_result.status.message())); + } + + robot_ = std::move(robot_result.response); + + // Authenticate + bosdyn::common::Status auth_status = robot_->Authenticate(username, password); + if (!auth_status) { + throw std::runtime_error(std::format("SpotConnection::connect: Failed to authenticate: {}", + auth_status.message())); + } + + // Create image client + bosdyn::client::Result image_client_result = + robot_->EnsureServiceClient(); + + if (!image_client_result.status) { + throw std::runtime_error(std::format("SpotConnection::connect: Failed to create image client: {}", + image_client_result.status.message())); + } + + image_client_ = image_client_result.response; + + LogMessage("SpotConnection::connect: Connected to Spot robot at {}", robot_ip); + + connected_ = true; + + } catch (const std::exception& e) { + throw std::runtime_error(std::format("SpotConnection::connect: Exception while connecting to robot {}: {}", + robot_ip, e.what())); + } +} + +SpotConnection::~SpotConnection() { + LogMessage("SpotConnection::~SpotConnection: Disconnecting from robot"); + + vision_pipelines_.clear(); + cam_streams_.clear(); + + if (connected_) { + robot_.reset(); + sdk_.reset(); + // TODO: figure out how to cleanup image_client_ + connected_ = false; + } + LogMessage("SpotConnection::~SpotConnection: Robot state fully cleaned up"); +} + +int32_t SpotConnection::createCamStream(uint32_t cam_mask) { + if (!connected_) { + LogMessage("SpotConnection::createCamStream: Not connected to robot"); + return -1; + } + + try { + int32_t stream_id = next_stream_id_++; + auto [it, inserted] = cam_streams_.try_emplace( + stream_id, + std::make_unique(*this, image_client_, image_lifo_max_size_) + ); + + LogMessage("SpotConnection::createCamStream: Created camera stream with mask {:#x}", + cam_mask); + + if (!it->second->streamCameras(cam_mask)) { + LogMessage("SpotConnection::createCamStream: Failed to start streaming cameras with mask {:#x}", + cam_mask); + cam_streams_.erase(it); + return -1; + } + + return stream_id; + + } catch (const std::exception& e) { + LogMessage("SpotConnection::createCamStream: Exception while creating camera stream: {}", + e.what()); + return -1; + } +} + +bool SpotConnection::removeCamStream(int32_t stream_id) { + auto cam_stream_it = cam_streams_.find(stream_id); + if (cam_stream_it == cam_streams_.end()) { + LogMessage("SpotConnection::removeCamStream: Camera stream {} doesn't exist", + stream_id); + return false; + } + try { + // Remove the associated vision pipeline if any + removeVisionPipeline(stream_id); + + cam_streams_.erase(cam_stream_it); + LogMessage("SpotConnection::removeCamStream: Removed camera stream {}", stream_id); + return true; + } catch (const std::exception& e) { + LogMessage("SpotConnection::removeCamStream: Exception while removing camera: {}", e.what()); + return false; + } +} + +SpotCamStream* SpotConnection::getCamStream(int32_t stream_id) { + auto cam_stream_it = cam_streams_.find(stream_id); + if (cam_stream_it == cam_streams_.end()) { + LogMessage("SpotConnection::getCamStream: Camera stream {} doesn't exist", + stream_id); + return nullptr; + } + return cam_stream_it->second.get(); +} + +bool SpotConnection::createVisionPipeline(MLModel& model, int32_t stream_id) { + SpotCamStream* cam_stream = getCamStream(stream_id); + if (cam_stream == nullptr || !cam_stream->isStreaming()) { + LogMessage("SpotConnection::createVisionPipeline: Camera stream {} doesn't exist. " + "Please create and start streaming before creating a vision pipeline.", + stream_id); + return false; + } + if (getVisionPipeline(stream_id) != nullptr) { + LogMessage("SpotConnection::createVisionPipeline: Vision pipeline for stream {} already exists", + stream_id); + return false; + } + // Get the expected tensor shapes from the camera stream + TensorShape rgb_shape = cam_stream->getCurrentRGBTensorShape(); + TensorShape depth_shape = cam_stream->getCurrentDepthTensorShape(); + + try { + auto [it, inserted] = vision_pipelines_.try_emplace( + stream_id, + std::make_unique( + model, + *cam_stream, + rgb_shape, + depth_shape, + depth_shape, + image_lifo_max_size_ + ) + ); + if (!inserted) { + LogMessage("SpotConnection::createVisionPipeline: Failed to insert vision pipeline for stream {}", + stream_id); + return false; + } + LogMessage("SpotConnection::createVisionPipeline: Created vision pipeline for stream {}", + stream_id); + + // Start the vision pipeline processing thread + if (!it->second->start()) { + LogMessage("SOb_LaunchVisionPipeline: Failed to start vision pipeline for stream {}", + stream_id); + return false; + } + + LogMessage("SOb_LaunchVisionPipeline: Successfully launched vision pipeline for stream {}", + stream_id); + + return true; + } catch (const std::exception& e) { + LogMessage("SpotConnection::createVisionPipeline: Exception while creating vision pipeline: {}", + e.what()); + return false; + } +} +bool SpotConnection::removeVisionPipeline(int32_t stream_id) { + auto vision_pipeline_it = vision_pipelines_.find(stream_id); + if (vision_pipeline_it == vision_pipelines_.end()) { + LogMessage("SpotConnection::removeVisionPipeline: Vision pipeline for stream {} doesn't exist", + stream_id); + return false; + } + + try { + vision_pipelines_.erase(vision_pipeline_it); + LogMessage("SpotConnection::removeVisionPipeline: Removed vision pipeline for stream {}", + stream_id); + return true; + } catch (const std::exception& e) { + LogMessage("SpotConnection::removeVisionPipeline: Exception while removing vision pipeline: {}", + e.what()); + return false; + } +} + +VisionPipeline* SpotConnection::getVisionPipeline(int32_t stream_id) { + auto vision_pipeline_it = vision_pipelines_.find(stream_id); + if (vision_pipeline_it == vision_pipelines_.end()) { + LogMessage("SpotConnection::getVisionPipeline: Vision pipeline for stream {} doesn't exist", + stream_id); + return nullptr; + } + return vision_pipeline_it->second.get(); +} + + } // namespace SOb diff --git a/src/spot-observer.cpp b/src/spot-observer.cpp index 609d6aa..c0fed3c 100644 --- a/src/spot-observer.cpp +++ b/src/spot-observer.cpp @@ -24,15 +24,12 @@ bool logging_enabled = true; // Map to hold robot connections by ID static int32_t __next_robot_id = 0; // Incremental ID for each robot connection -static std::unordered_map __robot_connections; +static std::unordered_map> __robot_connections; // Keeping track of loaded models std::unordered_map s_path_to_model_map; std::unordered_map s_model_to_path_map; -// Vision pipeline management -static std::unordered_map> __vision_pipelines; - // Needed to get Unity to load our DLL dependencies from the same directory as the plugin static bool SetDLLDirectory() { // Get the directory where our plugin is located @@ -81,16 +78,12 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserv break; case DLL_PROCESS_DETACH: LogMessage("DLL_PROCESS_DETACH"); - // Clean up all vision pipelines before DLL unload + // Clean up all robot connections try { - for (auto& [robot_id, pipeline] : __vision_pipelines) { - LogMessage("DLL_PROCESS_DETACH: Stopping vision pipeline for robot ID {}", robot_id); - pipeline->stop(); - } - __vision_pipelines.clear(); - LogMessage("DLL_PROCESS_DETACH: All vision pipelines stopped and cleared"); + __robot_connections.clear(); + LogMessage("DLL_PROCESS_DETACH: All Spot connections cleared"); } catch (const std::exception& e) { - LogMessage("DLL_PROCESS_DETACH: Exception while cleaning up vision pipelines: {}", e.what()); + LogMessage("DLL_PROCESS_DETACH: Exception while cleaning up Spot connections: {}", e.what()); } break; } @@ -101,14 +94,11 @@ static int32_t ConnectToSpot(const std::string& robot_ip, const std::string& use try { int32_t robot_id = __next_robot_id; __next_robot_id++; - auto [it, inserted] = __robot_connections.try_emplace(robot_id); - if (inserted) { - bool success = it->second.connect(robot_ip, username, password); - if (!success) { - __robot_connections.erase(it); - LogMessage("SOb::ConnectToSpot: Failed to connect to robot {}", robot_ip); - return -1; - } + auto [it, inserted] = + __robot_connections.try_emplace(robot_id, std::make_unique(robot_ip, username, password)); + if (!inserted) { + LogMessage("SOb::ConnectToSpot: Failed to connect to robot {}", robot_ip); + return -1; } return robot_id; @@ -120,6 +110,7 @@ static int32_t ConnectToSpot(const std::string& robot_ip, const std::string& use static bool getNextImageSet( int32_t robot_id, + int32_t cam_stream_id, int32_t n_images_requested, uint8_t** images, float** depths @@ -129,23 +120,33 @@ static bool getNextImageSet( LogMessage("SOb_GetNextImageSet: Robot ID {} not found", robot_id); return false; } - SpotConnection& robot = it->second; + SpotConnection& robot = *it->second; if (!robot.isConnected()) { LogMessage("SOb_GetNextImageSet: Robot ID {} is not connected", robot_id); return false; } - if (!robot.isStreaming()) { - LogMessage("SOb_GetNextImageSet: Robot ID {} is not streaming", robot_id); + + SpotCamStream* stream = robot.getCamStream(cam_stream_id); + if (!stream) { + LogMessage("SOb_GetNextImageSet: Camera stream ID {} not found for robot ID {}", + cam_stream_id, robot_id); + return false; + } + if (!stream->isStreaming()) { + LogMessage("SOb_GetNextImageSet: Camera stream ID {} not streaming for robot ID {}", + cam_stream_id, robot_id); return false; } - if (n_images_requested <= 0 || n_images_requested > robot.getCurrentNumCams()) { + + if (n_images_requested <= 0 || n_images_requested > stream->getCurrentNumCams()) { LogMessage("SOb_GetNextImageSet: Invalid number of images requested: {}", n_images_requested); return false; } // Pop images from the circular buffer - if (robot.getCurrentImages(n_images_requested, images, depths)) { - LogMessage("SOb_GetNextImageSet: Successfully retrieved {} images for robot ID {}", n_images_requested, robot_id); + if (stream->getCurrentImages(n_images_requested, images, depths)) { + LogMessage("SOb_GetNextImageSet: Successfully retrieved {} images for robot ID {} @ cam-stream ID {}", + n_images_requested, robot_id, cam_stream_id); return true; } else { // LogMessage("SOb_GetNextImageSet: Failed to retrieve images for robot ID {}", robot_id); @@ -157,17 +158,28 @@ static bool getNextImageSet( static bool getNextImageSetFromVisionPipeline( int32_t robot_id, + int32_t cam_stream_id, int32_t n_images_requested, uint8_t** images, float** depths ) { - auto pipeline_it = __vision_pipelines.find(robot_id); - if (pipeline_it == __vision_pipelines.end()) { - LogMessage("SOb_GetNextVisionPipelineImageSet: Vision pipeline not found for robot ID {}", robot_id); + auto it = __robot_connections.find(robot_id); + if (it == __robot_connections.end()) { + LogMessage("SOb_GetNextImageSet: Robot ID {} not found", robot_id); + return false; + } + SpotConnection& robot = *it->second; + if (!robot.isConnected()) { + LogMessage("SOb_GetNextImageSet: Robot ID {} is not connected", robot_id); + return false; + } + + VisionPipeline* pipeline = robot.getVisionPipeline(cam_stream_id); + if (!pipeline) { + LogMessage("SOb_GetNextVisionPipelineImageSet: Vision pipeline for camera stream ID {} not found for robot ID {}", + cam_stream_id, robot_id); return false; } - - auto& pipeline = pipeline_it->second; if (!pipeline->isRunning()) { LogMessage("SOb_GetNextVisionPipelineImageSet: Vision pipeline is not running for robot ID {}", robot_id); return false; @@ -326,20 +338,6 @@ bool UNITY_INTERFACE_API SOb_DisconnectFromSpot(int32_t robot_id) { return false; // Robot ID not found } try { - // First, clean up any associated vision pipeline - auto pipeline_it = SOb::__vision_pipelines.find(robot_id); - if (pipeline_it != SOb::__vision_pipelines.end()) { - SOb::LogMessage("SOb_DisconnectFromSpot: Stopping vision pipeline for robot ID {}", robot_id); - - // Take ownership of the pipeline to prevent access during cleanup - auto pipeline = std::move(pipeline_it->second); - SOb::__vision_pipelines.erase(pipeline_it); - - // Now safely stop the pipeline - pipeline->stop(); - SOb::LogMessage("SOb_DisconnectFromSpot: Vision pipeline stopped and removed for robot ID {}", robot_id); - } - // Then remove the robot connection SOb::__robot_connections.erase(it); // Remove from the map SOb::LogMessage("SOb_DisconnectFromSpot: Successfully disconnected robot ID {}", robot_id); @@ -352,27 +350,50 @@ bool UNITY_INTERFACE_API SOb_DisconnectFromSpot(int32_t robot_id) { // Start reading spot camera feeds. Runs in a separate thread. 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) { auto it = SOb::__robot_connections.find(robot_id); if (it == SOb::__robot_connections.end()) { - SOb::LogMessage("SOb_ReadCameraFeeds: Robot ID {} not found", robot_id); + SOb::LogMessage("SOb_CreateCameraStream: Robot ID {} not found", robot_id); + return -1; // Robot ID not found + } + + try { + int32_t stream_id = it->second->createCamStream(cam_bitmask); + if (stream_id < 0) { + SOb::LogMessage("SOb_CreateCameraStream: Failed to start reading cameras for robot ID {}", robot_id); + return -1; + } + SOb::LogMessage("SOb_CreateCameraStream: Successfully started reading cameras for robot ID {}", robot_id); + return stream_id; // Success + } catch (const std::exception& e) { + SOb::LogMessage("SOb_CreateCameraStream: Exception while reading cameras for robot ID {}: {}", robot_id, e.what()); + return -1; // Error during camera reading + } +} + +UNITY_INTERFACE_EXPORT +bool UNITY_INTERFACE_API SOb_DestroyCameraStream(int32_t robot_id, int32_t cam_stream_id) { + auto it = SOb::__robot_connections.find(robot_id); + if (it == SOb::__robot_connections.end()) { + SOb::LogMessage("SOb_DestroyCameraStream: Robot ID {} not found", robot_id); return false; // Robot ID not found } try { - bool success = it->second.streamCameras(cam_bitmask); - if (!success) { - SOb::LogMessage("SOb_ReadCameraFeeds: Failed to start reading cameras for robot ID {}", robot_id); - return false; + if (it->second->removeCamStream(cam_stream_id)) { + SOb::LogMessage("SOb_DestroyCameraStream: Successfully destroyed camera stream {} for robot ID {}", cam_stream_id, robot_id); + return true; // Success + } else { + SOb::LogMessage("SOb_DestroyCameraStream: Failed to destroy camera stream {} for robot ID {}", cam_stream_id, robot_id); + return false; // Failure } - SOb::LogMessage("SOb_ReadCameraFeeds: Successfully started reading cameras for robot ID {}", robot_id); - return true; // Success } catch (const std::exception& e) { - SOb::LogMessage("SOb_ReadCameraFeeds: Exception while reading cameras for robot ID {}: {}", robot_id, e.what()); - return false; // Error during camera reading + SOb::LogMessage("SOb_DestroyCameraStream: Exception while destroying camera stream {} for robot ID {}: {}", cam_stream_id, robot_id, e.what()); + return false; // Error during destruction } } + UNITY_INTERFACE_EXPORT SObModel UNITY_INTERFACE_API SOb_LoadModel(const char* modelPath, const char* backend) { if (!modelPath || !backend) { @@ -398,11 +419,15 @@ SObModel UNITY_INTERFACE_API SOb_LoadModel(const char* modelPath, const char* ba UNITY_INTERFACE_EXPORT void UNITY_INTERFACE_API SOb_UnloadModel(SObModel model) { - SOb::unloadModel(model); + try { + SOb::unloadModel(model); + } catch (const std::exception& e) { + SOb::LogMessage("SOb_UnloadModel: Exception while unloading model: {}", e.what()); + } } 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) { using namespace SOb; try { auto robot_it = __robot_connections.find(robot_id); @@ -411,8 +436,8 @@ bool UNITY_INTERFACE_API SOb_LaunchVisionPipeline(int32_t robot_id, SObModel mod return false; } - if (!robot_it->second.isConnected() || !robot_it->second.isStreaming()) { - LogMessage("SOb_LaunchVisionPipeline: Robot ID {} must be connected and streaming", robot_id); + if (!robot_it->second->isConnected()) { + LogMessage("SOb_LaunchVisionPipeline: Robot ID {} must be connected", robot_id); return false; } @@ -421,46 +446,16 @@ bool UNITY_INTERFACE_API SOb_LaunchVisionPipeline(int32_t robot_id, SObModel mod return false; } - // Check if pipeline already exists for this robot - if (__vision_pipelines.find(robot_id) != __vision_pipelines.end()) { - LogMessage("SOb_LaunchVisionPipeline: Vision pipeline already exists for robot ID {}", robot_id); - return false; - } - // Cast the model to the proper type auto ml_model = reinterpret_cast(model); // Create shared pointer to robot connection - const SpotConnection& spot_connection = robot_it->second; - - // Get batch size from the number of cameras being processed - int32_t batch_size = spot_connection.getCurrentNumCams(); - LogMessage("SOb_LaunchVisionPipeline: Using batch size {} based on number of cameras", batch_size); - - // Define tensor shapes using the actual batch size from spot connection - TensorShape input_shape(batch_size, 4, 480, 640); // NCHW format with dynamic batch size - TensorShape depth_shape(batch_size, 1, 480, 640); // Single channel depth with dynamic batch size - TensorShape output_shape(batch_size, 1, 480, 640); // Output shape with dynamic batch size - - // Create vision pipeline - auto pipeline = std::make_unique( - *ml_model, - spot_connection, - input_shape, - depth_shape, - output_shape, - 5 - ); - - // Start the pipeline - if (!pipeline->start()) { + SpotConnection& spot_connection = *robot_it->second; + if (!spot_connection.createVisionPipeline(*ml_model, cam_stream_id)) { LogMessage("SOb_LaunchVisionPipeline: Failed to start vision pipeline for robot ID {}", robot_id); return false; } - // Store the pipeline - __vision_pipelines[robot_id] = std::move(pipeline); - LogMessage("SOb_LaunchVisionPipeline: Successfully launched vision pipeline for robot ID {}", robot_id); return true; @@ -471,24 +466,30 @@ bool UNITY_INTERFACE_API SOb_LaunchVisionPipeline(int32_t robot_id, SObModel mod } 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) { using namespace SOb; try { - auto pipeline_it = __vision_pipelines.find(robot_id); - if (pipeline_it == __vision_pipelines.end()) { - LogMessage("SOb_StopVisionPipeline: Vision pipeline not found for robot ID {}", robot_id); + auto robot_it = __robot_connections.find(robot_id); + if (robot_it == __robot_connections.end()) { + LogMessage("SOb_StopVisionPipeline: Robot ID {} not found", robot_id); return false; } - auto& pipeline = pipeline_it->second; - pipeline->stop(); - __vision_pipelines.erase(pipeline_it); - - LogMessage("SOb_StopVisionPipeline: Successfully stopped vision pipeline for robot ID {}", robot_id); + // Create shared pointer to robot connection + SpotConnection& spot_connection = *robot_it->second; + if (!spot_connection.removeVisionPipeline(cam_stream_id)) { + LogMessage("SOb_StopVisionPipeline: Failed to stop vision pipeline for robot ID {} @ stream-ID {}", + robot_id, cam_stream_id); + return false; + } + + LogMessage("SOb_StopVisionPipeline: Successfully stopped vision pipeline for robot ID {} @ stream-ID {}", + robot_id, cam_stream_id); return true; - + } catch (const std::exception& e) { - LogMessage("SOb_StopVisionPipeline: Exception while stopping vision pipeline for robot ID {}: {}", robot_id, e.what()); + LogMessage("SOb_StopVisionPipeline: Exception while stopping vision pipeline for robot ID {} @ stream-ID: {}", + robot_id, cam_stream_id, e.what()); return false; } } @@ -496,6 +497,7 @@ bool UNITY_INTERFACE_API SOb_StopVisionPipeline(int32_t robot_id) { 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_tex, // ID3D12Resource* (aka texture) void* out_depth_tex, // ID3D12Resource* (aka texture) @@ -505,6 +507,7 @@ bool UNITY_INTERFACE_API SOb_RegisterUnityReadbackBuffers( try { bool ret = SOb::registerOutputTextures( robot_id, + cam_stream_id, cam_bit, out_img_tex, out_depth_tex, @@ -512,13 +515,16 @@ bool UNITY_INTERFACE_API SOb_RegisterUnityReadbackBuffers( depth_buffer_size ); if (!ret) { - SOb::LogMessage("SOb_RegisterOutputTextures: Failed to register output textures for robot ID {}", robot_id); + SOb::LogMessage("SOb_RegisterOutputTextures: Failed to register output textures for robot ID {} @ cam-stream ID {}", + robot_id, cam_stream_id); return false; } - SOb::LogMessage("SOb_RegisterOutputTextures: Successfully registered output textures for robot ID {}", robot_id); + SOb::LogMessage("SOb_RegisterOutputTextures: Successfully registered output textures for robot ID {} @ cam-stream ID {}", + robot_id, cam_stream_id); return true; } catch (const std::exception& e) { - SOb::LogMessage("SOb_RegisterOutputTextures: Exception while registering output textures for robot ID {}: {}", robot_id, e.what()); + SOb::LogMessage("SOb_RegisterOutputTextures: Exception while registering output textures for robot ID {} @ cam-stream ID {}: {}", + robot_id, cam_stream_id, e.what()); return false; } } @@ -542,6 +548,7 @@ bool UNITY_INTERFACE_API SOb_ClearUnityReadbackBuffers(int32_t robot_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 @@ -552,14 +559,16 @@ bool UNITY_INTERFACE_API SOb_GetNextImageSet( return false; } - bool ret = SOb::getNextImageSet(robot_id, n_images_requested, images, depths); + bool ret = SOb::getNextImageSet(robot_id, cam_stream_id, n_images_requested, images, depths); if (!ret) { - SOb::LogMessage("SOb_GetNextImageSet: Failed to get next image set for robot ID {}", robot_id); + // SOb::LogMessage("SOb_GetNextImageSet: Failed to get next image set for robot ID {} @ cam-stream ID {}", + // robot_id, cam_stream_id); return false; // Failed to get images } return ret; } catch (const std::exception& e) { - SOb::LogMessage("SOb::GetNextImageSet: Exception while getting next image set for robot ID {}: {}", robot_id, e.what()); + SOb::LogMessage("SOb::GetNextImageSet: Exception while getting next image set for robot ID {} @ cam-stream ID {}: {}", + robot_id, cam_stream_id, e.what()); return false; } } @@ -567,6 +576,7 @@ bool UNITY_INTERFACE_API SOb_GetNextImageSet( 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 @@ -577,7 +587,7 @@ bool UNITY_INTERFACE_API SOb_GetNextVisionPipelineImageSet( return false; } - bool ret = SOb::getNextImageSetFromVisionPipeline(robot_id, n_images_requested, images, depths); + bool ret = SOb::getNextImageSetFromVisionPipeline(robot_id, cam_stream_id, n_images_requested, images, depths); if (!ret) { SOb::LogMessage("SOb_GetNextVisionPipelineImageSet: Failed to get next image set from vision pipeline for robot ID {}", robot_id); return false; @@ -590,9 +600,9 @@ bool UNITY_INTERFACE_API SOb_GetNextVisionPipelineImageSet( } 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) { try { - bool ret = SOb::uploadNextImageSetToUnity(robot_id); + bool ret = SOb::uploadNextImageSetToUnity(robot_id, cam_stream_id); if (!ret) { SOb::LogMessage("SOb_PushNextImageSetToUnityBuffers: Failed to upload next image set for robot ID {}", robot_id); return false; // Failed to get images @@ -605,9 +615,9 @@ bool UNITY_INTERFACE_API SOb_PushNextImageSetToUnityBuffers(int32_t robot_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) { try { - bool ret = SOb::uploadNextVisionPipelineImageSetToUnity(robot_id); + bool ret = SOb::uploadNextVisionPipelineImageSetToUnity(robot_id, cam_stream_id); if (!ret) { SOb::LogMessage("SOb_PushNextVisionPipelineImageSetToUnityBuffers: Failed to upload next vision pipeline image set for robot ID {}", robot_id); return false; // Failed to get images @@ -665,5 +675,4 @@ void UNITY_INTERFACE_API SOb_ToggleDebugDumps(const char* dump_path) { SOb::LogMessage("Debug dumps enabled successfully"); } -} // extern "C" - +} // extern "C" \ No newline at end of file diff --git a/src/unity-cuda-interop.cpp b/src/unity-cuda-interop.cpp index 91683ad..1ed9e53 100644 --- a/src/unity-cuda-interop.cpp +++ b/src/unity-cuda-interop.cpp @@ -36,9 +36,9 @@ struct DX12InteropCacheEntry { static std::unordered_map s_InteropCache = {}; -// robot ID -> -static std::unordered_map> s_OutputRGBTextures = {}; -static std::unordered_map> s_OutputDepthTextures = {}; +// robot ID -> cam_stream_id -> +static std::unordered_map>> s_OutputRGBTextures = {}; +static std::unordered_map>> s_OutputDepthTextures = {}; // Store texture footprints for proper copying struct TextureFootprintInfo { @@ -247,6 +247,7 @@ 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* (texture or buffer) void* out_depth_tex, // ID3D12Resource* (texture or buffer) @@ -258,18 +259,26 @@ bool registerOutputTextures( LogMessage("Invalid output textures: out_img_tex = {}, out_depth_tex = {}", (void*)out_img_tex, (void*)out_depth_tex); return false; } - if (robot_id < 0 || cam_bit == 0 || cam_bit >= NUM_CAMERAS) { - LogMessage("Invalid robot ID or camera bitmask: robot_id = {}, cam_bit = {}", robot_id, cam_bit); + if (robot_id < 0) { + LogMessage("registerOutputTextures: Invalid robot ID {}", robot_id); return false; } - if (img_buffer_size <= 0 || depth_buffer_size <= 0) { - LogMessage("Invalid buffer sizes: img_buffer_size = {}, depth_buffer_size = {}", img_buffer_size, depth_buffer_size); + if (cam_stream_id < 0) { + LogMessage("registerOutputTextures: Invalid cam stream ID: cam_stream_id = {}", cam_stream_id); + return false; + } + if (cam_bit == 0 || cam_bit >= NUM_CAMERAS) { + LogMessage("registerOutputTextures: Invalid camera bitmask: cam_bit = {}", cam_bit); return false; } if (__num_set_bits(cam_bit) != 1) { LogMessage("Invalid camera bitmask: cam_bit = {:#x}. Must be a single bit set.", cam_bit); return false; } + if (img_buffer_size <= 0 || depth_buffer_size <= 0) { + LogMessage("Invalid buffer sizes: img_buffer_size = {}, depth_buffer_size = {}", img_buffer_size, depth_buffer_size); + return false; + } // Grab Unity's D3D12 resource ID3D12Resource* rgb_resource = reinterpret_cast(out_img_tex); @@ -342,10 +351,11 @@ bool registerOutputTextures( } // Register the output textures in the global maps - s_OutputRGBTextures[robot_id].push_back(rgb_resource); - s_OutputDepthTextures[robot_id].push_back(depth_resource); + s_OutputRGBTextures[robot_id][cam_stream_id].push_back(rgb_resource); + s_OutputDepthTextures[robot_id][cam_stream_id].push_back(depth_resource); - LogMessage("Registered output textures for robot ID {}: RGB = {}, Depth = {}", robot_id, (void*)rgb_resource, (void*)depth_resource); + LogMessage("Registered output textures for robot ID {} @ cam-stream ID {}: RGB = {}, Depth = {}", + robot_id, cam_stream_id, (void*)rgb_resource, (void*)depth_resource); return true; } @@ -357,23 +367,26 @@ bool clearOutputTextures(int32_t robot_id) { } auto erase_interop_entries = [&](auto& tex_map) { - for (auto& tex : tex_map) { - if (s_InteropCache.contains(tex)) { - DX12InteropCacheEntry& entry = s_InteropCache[tex]; - if (entry.extMem) { - cudaDestroyExternalMemory(entry.extMem); - entry.extMem = nullptr; + for (auto& cam_stream_map_to_tex : tex_map) { + for (auto& tex : cam_stream_map_to_tex.second) { + // Release interop entry + if (s_InteropCache.contains(tex)) { + DX12InteropCacheEntry& entry = s_InteropCache[tex]; + if (entry.extMem) { + cudaDestroyExternalMemory(entry.extMem); + entry.extMem = nullptr; + } + entry.cudaPtr = 0; + if (entry.sharedBuf) { + entry.sharedBuf->Release(); + entry.sharedBuf = nullptr; + } + entry.bufSize = 0; + s_InteropCache.erase(tex); } - entry.cudaPtr = 0; - if (entry.sharedBuf) { - entry.sharedBuf->Release(); - entry.sharedBuf = nullptr; + if (s_TextureFootprints.contains(tex)) { + s_TextureFootprints.erase(tex); } - entry.bufSize = 0; - s_InteropCache.erase(tex); - } - if (s_TextureFootprints.contains(tex)) { - s_TextureFootprints.erase(tex); } } }; @@ -390,10 +403,15 @@ bool clearOutputTextures(int32_t robot_id) { } // Helper function type for getting images -typedef bool (*GetImageSetFunc)(int32_t robot_id, int32_t n_images_requested, uint8_t** images, float** depths); +typedef bool (*GetImageSetFunc)(int32_t robot_id, int32_t cam_stream_id, int32_t n_images_requested, uint8_t** images, float** depths); // Common implementation for uploading image sets to Unity buffers -static bool uploadImageSetToUnityCommon(int32_t robot_id, GetImageSetFunc getImageSetFunc, const char* operation_name) { +static bool uploadImageSetToUnityCommon( + int32_t robot_id, + int32_t cam_stream_id, + GetImageSetFunc getImageSetFunc, + const char* operation_name +) { // Sanity checks if (robot_id < 0) { LogMessage("Invalid robot ID: {}", robot_id); @@ -434,28 +452,36 @@ static bool uploadImageSetToUnityCommon(int32_t robot_id, GetImageSetFunc getIma } // Check if we have registered textures for this robot ID - auto rgb_it = s_OutputRGBTextures.find(robot_id); - auto depth_it = s_OutputDepthTextures.find(robot_id); - if (rgb_it == s_OutputRGBTextures.end() || depth_it == s_OutputDepthTextures.end()) { + auto rgb_cam_stream_it = s_OutputRGBTextures.find(robot_id); + auto depth_cam_stream_it = s_OutputDepthTextures.find(robot_id); + if (rgb_cam_stream_it == s_OutputRGBTextures.end() || depth_cam_stream_it == s_OutputDepthTextures.end()) { LogMessage("No registered output textures for robot ID: {}", robot_id); return false; } + auto rgb_it = rgb_cam_stream_it->second.find(cam_stream_id); + auto depth_it = depth_cam_stream_it->second.find(cam_stream_id); + if (rgb_it == rgb_cam_stream_it->second.end() || depth_it == depth_cam_stream_it->second.end()) { + LogMessage("No registered output textures for cam-stream ID: {}", cam_stream_id); + return false; + } + std::vector& rgb_textures = rgb_it->second; std::vector& depth_textures = depth_it->second; size_t num_textures = rgb_textures.size(); if (num_textures != depth_textures.size()) { - LogMessage("Mismatch in number of RGB and depth textures for robot ID: {}", robot_id); + LogMessage("Mismatch in number of RGB and depth textures for robot ID: {} @ cam stream {}", + robot_id, cam_stream_id); return false; } // Grab the latest image set using the provided function uint8_t* rgb_images[NUM_CAMERAS]; float* depth_images[NUM_CAMERAS]; - bool ret = getImageSetFunc(robot_id, num_textures, rgb_images, depth_images); + bool ret = getImageSetFunc(robot_id, cam_stream_id, num_textures, rgb_images, depth_images); if (!ret) { - LogMessage("No new images ready for {}: {}", operation_name, robot_id); + LogMessage("No new images ready for {}: robot {}, cam stream {}", operation_name, robot_id, cam_stream_id); return false; } @@ -538,7 +564,8 @@ static bool uploadImageSetToUnityCommon(int32_t robot_id, GetImageSetFunc getIma } } - LogMessage("Copied images to shared buffers for {} robot ID: {}", operation_name, robot_id); + LogMessage("Copied images to shared buffers for {} robot ID: {}, cam stream {}", + operation_name, robot_id, cam_stream_id); // Copy to Unity resources s_CmdAlloc->Reset(); @@ -644,12 +671,12 @@ static bool uploadImageSetToUnityCommon(int32_t robot_id, GetImageSetFunc getIma return true; } -bool uploadNextImageSetToUnity(int32_t robot_id) { - return uploadImageSetToUnityCommon(robot_id, SOb_GetNextImageSet, "robot"); +bool uploadNextImageSetToUnity(int32_t robot_id, int32_t cam_stream_id) { + return uploadImageSetToUnityCommon(robot_id, cam_stream_id, SOb_GetNextImageSet, "robot"); } -bool uploadNextVisionPipelineImageSetToUnity(int32_t robot_id) { - return uploadImageSetToUnityCommon(robot_id, SOb_GetNextVisionPipelineImageSet, "vision pipeline"); +bool uploadNextVisionPipelineImageSetToUnity(int32_t robot_id, int32_t cam_stream_id) { + return uploadImageSetToUnityCommon(robot_id, cam_stream_id, SOb_GetNextVisionPipelineImageSet, "vision pipeline"); } } diff --git a/src/vision-pipeline.cpp b/src/vision-pipeline.cpp index 030379a..0ca3257 100644 --- a/src/vision-pipeline.cpp +++ b/src/vision-pipeline.cpp @@ -12,27 +12,29 @@ namespace SOb { -static int32_t dump_id = 900; +static int32_t thread_id_ = 0; VisionPipeline::VisionPipeline( MLModel& model, - const SpotConnection& spot_connection, + const SpotCamStream& spot_cam_stream_, const TensorShape& input_shape, const TensorShape& depth_shape, const TensorShape& output_shape, size_t max_results ) : model_(model) - , spot_connection_(std::move(spot_connection)) + , spot_cam_stream_(std::move(spot_cam_stream_)) , input_shape_(input_shape) , depth_shape_(depth_shape) , output_shape_(output_shape) , max_size_(max_results) - , cuda_stream_(spot_connection_.getCudaStream()) + , cuda_stream_(spot_cam_stream_.getCudaStream()) + , thread_num(thread_id_++) { } VisionPipeline::~VisionPipeline() { stop(); deallocateCudaBuffers(); + } bool VisionPipeline::start() { @@ -40,20 +42,26 @@ bool VisionPipeline::start() { return false; // Already running } - if (!spot_connection_.isConnected() || !spot_connection_.isStreaming()) { - LogMessage("SpotConnection must be connected and streaming before starting pipeline"); + if (!spot_cam_stream_.isStreaming()) { + LogMessage("SpotCamStream must be connected and streaming before starting pipeline"); return false; } - if (input_shape_.N != depth_shape_.N) { - LogMessage("Incompatible shapes: input_shape_.N != depth_shape_.N"); - } - if (input_shape_.N != output_shape_.N) { - LogMessage("Incompatible shapes: input_shape_.N != output_shape_.N"); + TensorShape stream_rgb_shape = spot_cam_stream_.getCurrentRGBTensorShape(); + if (input_shape_ != stream_rgb_shape) { + LogMessage("Incompatible shapes: input_shape {} != stream_rgb_shape {}", + input_shape_.to_string(), + stream_rgb_shape.to_string() + ); + return false; } - size_t num_cams = spot_connection_.getCurrentNumCams(); - if (input_shape_.N != num_cams) { - LogMessage("Incompatible shapes: input_shape_.N != spot_connection_.getCurrentNumCams()"); + TensorShape stream_depth_shape = spot_cam_stream_.getCurrentDepthTensorShape(); + if (depth_shape_ != stream_depth_shape) { + LogMessage("Incompatible shapes: depth_shape_ {} != stream_depth_shape {}", + depth_shape_.to_string(), + stream_depth_shape.to_string() + ); + return false; } if (!allocateCudaBuffers()) { @@ -169,8 +177,8 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { auto start_time = std::chrono::high_resolution_clock::now(); try { // Get current images from SpotConnection - if (!spot_connection_.getCurrentImages( - spot_connection_.getCurrentNumCams(), + if (!spot_cam_stream_.getCurrentImages( + input_shape_.N, rgb_images, depth_images) ) { @@ -224,29 +232,29 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { input_shape_float.W = input_shape_.H; } - TensorShape depth_shape_single = depth_shape_; + TensorShape depth_shape = depth_shape_; if (do_rotate_90_cw) { - depth_shape_single.H = depth_shape_.W / depth_scale_factor; - depth_shape_single.W = depth_shape_.H / depth_scale_factor; + depth_shape.H = depth_shape_.W / depth_scale_factor; + depth_shape.W = depth_shape_.H / depth_scale_factor; } else { - depth_shape_single.H = depth_shape_.H / depth_scale_factor; - depth_shape_single.W = depth_shape_.W / depth_scale_factor; + depth_shape.H = depth_shape_.H / depth_scale_factor; + depth_shape.W = depth_shape_.W / depth_scale_factor; } - TensorShape output_shape_single = output_shape_; + TensorShape output_shape = output_shape_; if (do_rotate_90_cw) { - output_shape_single.H = output_shape_.W; - output_shape_single.W = output_shape_.H; + output_shape.H = output_shape_.W; + output_shape.W = output_shape_.H; } else { - output_shape_single.H = output_shape_.H; - output_shape_single.W = output_shape_.W; + output_shape.H = output_shape_.H; + output_shape.W = output_shape_.W; } LogMessage("num_images_per_iter = {}", num_images_per_iter); for (size_t i = 0; i < num_images_per_iter; i++) { float* cur_rgb_input_ptr = cuda_ws_.d_rgb_float_data_ + i * 3 * input_shape_.H * input_shape_.W; float* cur_depth_input_ptr = cuda_ws_.d_depth_data_ + i * depth_shape_.C * depth_shape_.H * depth_shape_.W; - float* cur_preprocessed_depth_ptr = cuda_ws_.d_preprocessed_depth_data_ + i * depth_shape_single.C * depth_shape_single.H * depth_shape_single.W; + float* cur_preprocessed_depth_ptr = cuda_ws_.d_preprocessed_depth_data_ + i * depth_shape.C * depth_shape.H * depth_shape.W; float* cur_depth_output_ptr = d_depth_output_ptr + i * output_shape_.C * output_shape_.H * output_shape_.W; float* depth_cache_ptr = cuda_ws_.d_depth_cached_ + i * depth_shape_.C * depth_shape_.H * depth_shape_.W; @@ -263,18 +271,35 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { 100.0f, cuda_stream_ ), "prefill_invalid_depth"); - // - // checkCudaError(preprocess_depth_image2( - // cur_preprocessed_depth_ptr, - // cur_preprocessed_depth_ptr, - // depth_shape_.W, - // depth_shape_.H, - // false, - // depth_scale_factor, - // cuda_ws_.d_depth_preprocessor_workspace_, - // do_rotate_90_cw, - // cuda_stream_ // NEW - // ), "preprocess_depth_image"); + + DumpDepthImageFromCuda( + cur_preprocessed_depth_ptr, + depth_shape_.W, + depth_shape_.H, + "depth-post-prefill", + dump_id+i, + thread_num + ); + // Downscale + checkCudaError(preprocess_depth_image2( + cur_preprocessed_depth_ptr, + cur_preprocessed_depth_ptr, + depth_shape_.W, + depth_shape_.H, + false, + depth_scale_factor, + cuda_ws_.d_depth_preprocessor_workspace_, + do_rotate_90_cw, + cuda_stream_ + ), "preprocess_depth_image"); + DumpDepthImageFromCuda( + cur_preprocessed_depth_ptr, + depth_shape.W, + depth_shape.H, + "depth-post-downscale", + dump_id+i, + thread_num + ); } else { checkCudaError(preprocess_depth_image2( @@ -304,8 +329,8 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { cuda_ws_.d_preprocessed_depth_data_, d_depth_output_ptr, input_shape_float, - depth_shape_single, - output_shape_single + depth_shape, + output_shape ); if (!inference_success) { @@ -321,7 +346,6 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { for (size_t i = 0; i < num_images_per_iter; i++) { float* cur_rgb_input_ptr = cuda_ws_.d_rgb_float_data_ + i * 3 * input_shape_.H * input_shape_.W; float* cur_depth_input_ptr = cuda_ws_.d_depth_data_ + i * depth_shape_.C * depth_shape_.H * depth_shape_.W; - float* cur_preprocessed_depth_ptr = cuda_ws_.d_preprocessed_depth_data_ + i * depth_shape_.C * depth_shape_.H * depth_shape_.W; float* cur_depth_output_ptr = d_depth_output_ptr + i * output_shape_.C * output_shape_.H * output_shape_.W; float* depth_cache_ptr = cuda_ws_.d_depth_cached_ + i * depth_shape_.C * depth_shape_.H * depth_shape_.W; @@ -329,23 +353,23 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { float* temp_output_ptr = reinterpret_cast(cuda_ws_.d_depth_preprocessor_workspace_); checkCudaError(postprocess_depth_image( cur_depth_output_ptr, - output_shape_single.W, - output_shape_single.H, + output_shape_.W, + output_shape_.H, temp_output_ptr, do_rotate_90_cw, cuda_stream_ ), "postprocess_depth_image"); - - checkCudaError(update_depth_cache( - cur_depth_output_ptr, - cur_preprocessed_depth_ptr, - depth_cache_ptr, - output_shape_single.W, - output_shape_single.H, - 0.01f, - 100.0f, - cuda_stream_ - ), "update_depth_cache"); + // + // checkCudaError(update_depth_cache( + // cur_depth_output_ptr, + // cur_depth_input_ptr, + // depth_cache_ptr, + // output_shape_.W, + // output_shape_.H, + // 0.01f, + // 100.0f, + // cuda_stream_ + // ), "update_depth_cache"); // Ensure dumps see completed work (dumpers likely use default stream) checkCudaError(cudaStreamSynchronize(cuda_stream_), "sync before dumps"); @@ -355,23 +379,25 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { input_shape_.W, input_shape_.H, "input-rgb", - dump_id + dump_id+i, + thread_num ); DumpDepthImageFromCuda( cur_depth_input_ptr, depth_shape_.W, depth_shape_.H, "input-depth", - dump_id + dump_id+i, + thread_num ); DumpDepthImageFromCuda( - d_depth_output_ptr, + cur_depth_output_ptr, depth_shape_.W, depth_shape_.H, "output-depth", - dump_id + dump_id+i, + thread_num ); - dump_id++; } checkCudaError(cudaStreamSynchronize(cuda_stream_), "cudaStreamSynchronize after postprocess"); @@ -386,6 +412,8 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { LogMessage("VisionPipeline: Updating write index from {} to {}", write_idx_, (write_idx_ + 1) % max_size_); write_idx_ = (write_idx_ + 1) % max_size_; + // first_run_ = false; + dump_id += num_images_per_iter; } catch (const std::exception& e) { LogMessage("Exception in pipeline worker: {}", e.what()); @@ -423,9 +451,6 @@ bool VisionPipeline::getCurrentImages( uint8_t* rgb_data_out = d_rgb_data_ + read_idx * n_elems_rgb_total; float* depth_data_out = d_output_buffer_ + read_idx * n_elems_depth_total; - LogMessage("VisionPipeline::getCurrentImages, popping {} images from index {}", - input_shape_.N, read_idx); - for (int32_t i = 0; i < n_images_requested; i++) { images[i] = rgb_data_out + i * n_elems_per_rgb; depths[i] = depth_data_out + i * n_elems_per_depth; @@ -434,4 +459,4 @@ bool VisionPipeline::getCurrentImages( return true; } -} // namespace SOb \ No newline at end of file +} // namespace \ No newline at end of file diff --git a/tests/integ-test-dx12/integ-test-dx12.cpp b/tests/integ-test-dx12/integ-test-dx12.cpp index 1629a9c..f14ce6d 100644 --- a/tests/integ-test-dx12/integ-test-dx12.cpp +++ b/tests/integ-test-dx12/integ-test-dx12.cpp @@ -17,10 +17,12 @@ #include #include #include +#include #include "spot-observer.h" #include "d3dx12.h" #include +#include using Microsoft::WRL::ComPtr; using namespace std::chrono; @@ -54,26 +56,11 @@ static uint32_t __num_set_bits(uint32_t bitmask) { return __popcnt(bitmask); } -static std::vector convert_bitmask_to_spot_cam_vector(uint32_t bitmask) { - std::vector cams; - cams.reserve(__num_set_bits(bitmask)); - - if (bitmask & SpotCamera::BACK) cams.emplace_back(SpotCamera::BACK); - if (bitmask & SpotCamera::FRONTLEFT) cams.emplace_back(SpotCamera::FRONTLEFT); - if (bitmask & SpotCamera::FRONTRIGHT) cams.emplace_back(SpotCamera::FRONTRIGHT); - if (bitmask & SpotCamera::LEFT) cams.emplace_back(SpotCamera::LEFT); - if (bitmask & SpotCamera::RIGHT) cams.emplace_back(SpotCamera::RIGHT); - if (bitmask & SpotCamera::HAND) cams.emplace_back(SpotCamera::HAND); - - return cams; -} - -static int32_t connect_to_spot_and_start_cam_feed( +static int32_t connect_to_spot( const std::string& robot_ip, const std::string& username, - const std::string& password, - uint32_t cam_bitmask + const std::string& password ) { std::cout << "Connecting to Spot robot at " << robot_ip << " with user " << username << std::endl; int32_t spot_id = SOb_ConnectToSpot(robot_ip.c_str(), username.c_str(), password.c_str()); @@ -81,15 +68,32 @@ static int32_t connect_to_spot_and_start_cam_feed( std::cerr << "Failed to connect to Spot robot" << std::endl; return -1; } + std::cout << "Connected to Spot robot with ID: " << spot_id << std::endl; + return spot_id; +} - bool ret = SOb_ReadCameraFeeds(spot_id, cam_bitmask); - if (!ret) { +static int32_t start_cam_stream(int32_t spot_id, uint32_t cam_bitmask) { + int32_t cam_stream_id = SOb_CreateCameraStream(spot_id, cam_bitmask); + if (cam_stream_id < 0) { std::cerr << "Failed to start reading camera feeds" << std::endl; SOb_DisconnectFromSpot(spot_id); return -1; } + std::cout << "Started camera stream with ID: " << cam_stream_id << std::endl; + return cam_stream_id; +} - return spot_id; +static int32_t disconnect_from_spots(const int32_t spot_ids[], size_t num_spots) { + for (size_t i = 0; i < num_spots; i++) { + if (spot_ids[i] >= 0) { + if (SOb_DisconnectFromSpot(spot_ids[i])) { + std::cout << "Disconnected from Spot robot with ID: " << spot_ids[i] << std::endl; + } else { + std::cerr << "Failed to disconnect from Spot robot with ID: " << spot_ids[i] << std::endl; + } + } + } + return 0; } // =========================================================================================== @@ -112,7 +116,7 @@ int main(int argc, char* argv[]) { SOb_SetUnityLogCallback(TestLogCallback); SOb_ToggleLogging(true); - //SOb_ToggleDebugDumps("./spot-dump-dx12"); + //SOb_ToggleDebugDumps("./spot_dump_dx12"); // --------------------------------------------------------------------------------------- // 2. Initialize D3D12 @@ -174,36 +178,55 @@ int main(int argc, char* argv[]) { // --------------------------------------------------------------------------------------- // 3. Connect to Spot robots // --------------------------------------------------------------------------------------- - int32_t spot_ids[2]; - uint32_t cam_bitmask = FRONTLEFT | FRONTRIGHT; //| HAND | LEFT | RIGHT | BACK; - std::vector cams = convert_bitmask_to_spot_cam_vector(cam_bitmask); + int32_t spot_ids[2] = {-1, -1}; + std::unordered_map> cam_stream_ids; + std::vector cam_bitmasks = {FRONTRIGHT | FRONTLEFT, HAND}; for (size_t i = 0; i < 2; i++) { if (robot_ips[i] == "0") { spot_ids[i] = -1; } else { - spot_ids[i] = connect_to_spot_and_start_cam_feed(robot_ips[i], username, password, cam_bitmask); + spot_ids[i] = connect_to_spot(robot_ips[i], username, password); if (spot_ids[i] < 0) { - std::cerr << "Failed to connect to Spot robot " << i << std::endl; + disconnect_from_spots(spot_ids, 2); return -1; } } } std::cout << "Connected to Spot robots with IDs: " << spot_ids[0] << ", " << spot_ids[1] << std::endl; - uint32_t num_cameras = __num_set_bits(cam_bitmask); + for (size_t i = 0; i < 2; i++) { + if (spot_ids[i] < 0) continue; + for (uint32_t cam_bitmask : cam_bitmasks) { + int32_t cam_stream_id = start_cam_stream(spot_ids[i], cam_bitmask); + + if (cam_stream_id < 0) { + disconnect_from_spots(spot_ids, 2); + return -1; + } + cam_stream_ids[spot_ids[i]].push_back(cam_stream_id); + } + } + + // Calculate total cameras across all streams + uint32_t total_cameras = 0; + for (uint32_t cam_bitmask : cam_bitmasks) { + total_cameras += __num_set_bits(cam_bitmask); + } + uint32_t num_cameras = total_cameras; // --------------------------------------------------------------------------------------- - // 4. Create D3D12 resources for each camera + // 4. Create D3D12 resources for each camera stream // --------------------------------------------------------------------------------------- - std::cout << "Creating D3D12 textures for " << num_cameras << " cameras per robot...\n"; + std::cout << "Creating D3D12 textures for camera streams...\n"; struct CameraResources { ComPtr image_texture; ComPtr depth_texture; }; - std::vector> robot_resources(2, std::vector(num_cameras)); + // Structure: robot -> stream -> camera resources + std::unordered_map>> robot_resources; D3D12_HEAP_PROPERTIES heap_default = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_READBACK); D3D12_RESOURCE_DESC rgb_desc = CD3DX12_RESOURCE_DESC::Buffer(IMAGE_BUFSIZE); @@ -212,32 +235,39 @@ int main(int argc, char* argv[]) { for (size_t robot = 0; robot < 2; robot++) { if (spot_ids[robot] < 0) continue; - for (uint32_t cam = 0; cam < num_cameras; cam++) { - // Create image texture (RGB) - - ThrowIfFailed(device->CreateCommittedResource( - &heap_default, - D3D12_HEAP_FLAG_NONE, - &rgb_desc, - D3D12_RESOURCE_STATE_COMMON, - nullptr, - IID_PPV_ARGS(&robot_resources[robot][cam].image_texture) - ), "CreateCommittedResource for rgb texture"); - - // Create depth texture - ThrowIfFailed(device->CreateCommittedResource( - &heap_default, - D3D12_HEAP_FLAG_NONE, - &depth_desc, - D3D12_RESOURCE_STATE_COMMON, - nullptr, - IID_PPV_ARGS(&robot_resources[robot][cam].depth_texture) - ), "CreateCommittedResource for depth texture"); - - std::wstring img_name = L"Robot" + std::to_wstring(robot) + L"_Cam" + std::to_wstring(cam) + L"_Image"; - std::wstring depth_name = L"Robot" + std::to_wstring(robot) + L"_Cam" + std::to_wstring(cam) + L"_Depth"; - robot_resources[robot][cam].image_texture->SetName(img_name.c_str()); - robot_resources[robot][cam].depth_texture->SetName(depth_name.c_str()); + int32_t spot_id = spot_ids[robot]; + robot_resources[spot_id].resize(cam_bitmasks.size()); + + for (size_t stream = 0; stream < cam_bitmasks.size(); stream++) { + uint32_t num_cams_in_stream = __num_set_bits(cam_bitmasks[stream]); + robot_resources[spot_id][stream].resize(num_cams_in_stream); + + for (uint32_t cam = 0; cam < num_cams_in_stream; cam++) { + // Create image texture (RGB) + ThrowIfFailed(device->CreateCommittedResource( + &heap_default, + D3D12_HEAP_FLAG_NONE, + &rgb_desc, + D3D12_RESOURCE_STATE_COMMON, + nullptr, + IID_PPV_ARGS(&robot_resources[spot_id][stream][cam].image_texture) + ), "CreateCommittedResource for rgb texture"); + + // Create depth texture + ThrowIfFailed(device->CreateCommittedResource( + &heap_default, + D3D12_HEAP_FLAG_NONE, + &depth_desc, + D3D12_RESOURCE_STATE_COMMON, + nullptr, + IID_PPV_ARGS(&robot_resources[spot_id][stream][cam].depth_texture) + ), "CreateCommittedResource for depth texture"); + + std::wstring img_name = L"Robot" + std::to_wstring(robot) + L"_Stream" + std::to_wstring(stream) + L"_Cam" + std::to_wstring(cam) + L"_Image"; + std::wstring depth_name = L"Robot" + std::to_wstring(robot) + L"_Stream" + std::to_wstring(stream) + L"_Cam" + std::to_wstring(cam) + L"_Depth"; + robot_resources[spot_id][stream][cam].image_texture->SetName(img_name.c_str()); + robot_resources[spot_id][stream][cam].depth_texture->SetName(depth_name.c_str()); + } } } @@ -246,25 +276,33 @@ int main(int argc, char* argv[]) { // --------------------------------------------------------------------------------------- std::cout << "Registering D3D12 resources with SpotObserver...\n"; - int32_t robots_active = 0; + int32_t total_streams_active = 0; for (size_t robot = 0; robot < 2; robot++) { if (spot_ids[robot] < 0) continue; - for (uint32_t cam = 0; cam < num_cameras; cam++) { - bool reg_result = SOb_RegisterUnityReadbackBuffers( - spot_ids[robot], - 1 << cam, // Single bit for camera - robot_resources[robot][cam].image_texture.Get(), - robot_resources[robot][cam].depth_texture.Get(), - static_cast(IMAGE_BUFSIZE), - static_cast(DEPTH_BUFSIZE) - ); - - if (!reg_result) { - std::cerr << "Failed to register textures for robot " << robot << " camera " << cam << std::endl; - return -1; + int32_t spot_id = spot_ids[robot]; + for (size_t stream = 0; stream < cam_bitmasks.size(); stream++) { + uint32_t num_cams_in_stream = __num_set_bits(cam_bitmasks[stream]); + int32_t cam_stream_id = cam_stream_ids[spot_id][stream]; + + for (uint32_t cam = 0; cam < num_cams_in_stream; cam++) { + bool reg_result = SOb_RegisterUnityReadbackBuffers( + spot_id, + cam_stream_id, + 1 << cam, // Single bit for camera within the stream + robot_resources[spot_id][stream][cam].image_texture.Get(), + robot_resources[spot_id][stream][cam].depth_texture.Get(), + static_cast(IMAGE_BUFSIZE), + static_cast(DEPTH_BUFSIZE) + ); + + if (!reg_result) { + std::cerr << "Failed to register textures for robot " << robot << " stream " << stream << " camera " << cam << std::endl; + disconnect_from_spots(spot_ids, 2); + return -1; + } } - robots_active++; + total_streams_active++; } } @@ -284,17 +322,17 @@ int main(int argc, char* argv[]) { } std::cout << "Model loaded successfully!" << std::endl; - // Launch vision pipeline on both robots + // Launch vision pipeline on both robots (first stream only) for (size_t i = 0; i < 2; i++) { if (spot_ids[i] < 0) continue; - bool ret = SOb_LaunchVisionPipeline(spot_ids[i], model); + int32_t spot_id = spot_ids[i]; + int32_t cam_stream_id = cam_stream_ids[spot_id][0]; // Use first stream + bool ret = SOb_LaunchVisionPipeline(spot_id, cam_stream_id, model); if (!ret) { std::cerr << "Failed to launch vision pipeline on robot " << i << std::endl; - for (int32_t spot = 0; spot < 2; spot++) { - cv::destroyAllWindows(); - SOb_DisconnectFromSpot(spot_ids[spot]); - } + disconnect_from_spots(spot_ids, 2); SOb_UnloadModel(model); + cv::destroyAllWindows(); return -1; } std::cout << "Vision pipeline launched on robot " << i << std::endl; @@ -307,141 +345,138 @@ int main(int argc, char* argv[]) { std::cout << "Starting camera feed reading loop with OpenCV display...\n"; std::cout << "Press 'q' in any OpenCV window to quit...\n"; - int32_t frame_count = 0; - bool should_quit = false; - int32_t n_images_pending = num_cameras * robots_active; + // Initialized output pointers for GPU device pointers + std::vector num_images_requested_per_stream; + std::vector images_gpu; + std::vector depths_gpu; + for (uint32_t cam_bitmask : cam_bitmasks) { + uint32_t num_bits_set = __num_set_bits(cam_bitmask); + num_images_requested_per_stream.push_back(num_bits_set); + images_gpu.push_back(new uint8_t*[num_bits_set]); + depths_gpu.push_back(new float*[num_bits_set]); + } + + std::vector image_cpu_buffer(640 * 480 * 4); + std::vector depth_cpu_buffer(640 * 480); + + bool new_images = false; auto start_time = high_resolution_clock::now(); - auto frame_start = start_time; + bool exit_requested = false; + while (!exit_requested) { + if (new_images) { + auto end_time = high_resolution_clock::now(); + auto duration = duration_cast(end_time - start_time); + double latency_ms = double(duration.count()) / 1000.0; - for (int32_t iteration = 0; !should_quit && iteration < 10000; iteration++) { - for (size_t robot = 0; robot < 2; robot++) { - if (spot_ids[robot] < 0) continue; + std::cout << std::format("integ-test-dx12: total latency: {:.4f} ms\n", latency_ms); + start_time = end_time; // Reset start time for next push + new_images = false; + } - // Upload next image batch to Unity textures - bool upload_result; - if (using_vision_pipeline) { - upload_result = SOb_PushNextVisionPipelineImageSetToUnityBuffers(spot_ids[robot]); - } else { - upload_result = SOb_PushNextImageSetToUnityBuffers(spot_ids[robot]); - } - if (!upload_result) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); + for (int32_t spot = 0; spot < 2; spot++) { + int32_t spot_id = spot_ids[spot]; + if (spot_id < 0) { continue; } - std::cout << "Successfully uploaded image batch for robot " << robot << " at frame " << frame_count << std::endl; - - // Read back and display all camera data - for (uint32_t cam = 0; cam < num_cameras; cam++) { - uint8_t* image_data = nullptr; - float* depth_data = nullptr; - - D3D12_RANGE image_read_range = {0, IMAGE_BUFSIZE}; - D3D12_RANGE depth_read_range = {0, DEPTH_BUFSIZE}; - D3D12_RANGE no_write = {0, 0}; - - ThrowIfFailed(robot_resources[robot][cam].image_texture->Map( - 0, &image_read_range, reinterpret_cast(&image_data) - ), "Map image readback"); - - ThrowIfFailed(robot_resources[robot][cam].depth_texture->Map( - 0, &depth_read_range, reinterpret_cast(&depth_data) - ), "Map depth readback"); - - // Create OpenCV matrices from D3D12 readback data - size_t actual_width = WIDTH; - size_t actual_height = HEIGHT; - // switch (cams[cam]) { - // case SpotCamera::FRONTLEFT: - // case SpotCamera::FRONTRIGHT: - // actual_width = HEIGHT; - // actual_height = WIDTH; - // break; - // } - - cv::Mat image(actual_height, actual_width, CV_8UC4, image_data); - cv::Mat depth(actual_height, actual_width, CV_32FC1, depth_data); - - // Convert RGB to BGR for OpenCV display - cv::cvtColor(image, image, cv::COLOR_RGBA2BGR); - cv::normalize(depth, depth, 0, 1, cv::NORM_MINMAX); - - // Create window names - std::string image_window = "SPOT " + std::to_string(robot) + " RGB" + std::to_string(cam); - std::string depth_window = "SPOT " + std::to_string(robot) + " Depth" + std::to_string(cam); - - // Display the images - cv::imshow(image_window, image); - cv::imshow(depth_window, depth); - - n_images_pending -= 1; - - // Unmap the readback buffers - robot_resources[robot][cam].image_texture->Unmap(0, &no_write); - robot_resources[robot][cam].depth_texture->Unmap(0, &no_write); - - // Analyze data for stats (only for first camera for brevity) - if (cam == 0 && frame_count % 10 == 0) { - int img_nonzero = 0, depth_nonzero = 0; - - for (int i = 0; i < WIDTH * HEIGHT * CHANNELS; i++) { - if (std::abs(image_data[i]) > 1e-6) img_nonzero++; - } - for (int i = 0; i < WIDTH * HEIGHT; i++) { - if (std::abs(depth_data[i]) > 1e-6) depth_nonzero++; - } - - std::cout << std::format("Robot {}, Frame {}: Image nonzero={}, Depth nonzero={}\n", - robot, frame_count, img_nonzero, depth_nonzero); + for (int32_t stream = 0; stream < cam_stream_ids[spot_id].size(); stream++) { + int32_t cam_stream_id = cam_stream_ids[spot_id][stream]; + if (cam_stream_id < 0) { + std::cout << "Skipping Spot " << spot << " Stream " << stream << " as it is not valid." << std::endl; + continue; } - } - - // Check for quit key - int key = cv::waitKey(1) & 0xFF; - if (key == 'q' || key == 27) { // 'q' or ESC - should_quit = true; - break; - } - } - if (n_images_pending <= 0) { - if (n_images_pending < 0) { - std::cerr << "n_images_pending went negative!" << std::endl; - } - auto frame_end = high_resolution_clock::now(); - auto frame_duration = duration_cast(frame_end - frame_start); - - std::cout << std::format("Frame {} processing time: {:.2f} ms\n", - frame_count, double(frame_duration.count()) / 1000.0); + uint32_t num_images_requested = num_images_requested_per_stream[stream]; + uint8_t** images_set = images_gpu[stream]; + float** depths_set = depths_gpu[stream]; - frame_start = frame_end; + // Get GPU pointers from SpotObserver + if (using_vision_pipeline && stream == 0) { + if (!SOb_GetNextVisionPipelineImageSet(spot_id, cam_stream_id, int32_t(num_images_requested), images_set, depths_set)) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + continue; + } + } else { + if (!SOb_GetNextImageSet(spot_id, cam_stream_id, int32_t(num_images_requested), images_set, depths_set)) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + continue; + } + } + new_images = true; + + // Upload GPU data to D3D12 textures, then read back and display + for (uint32_t i = 0; i < num_images_requested; i++) { + // Copy from GPU to CPU buffers + cudaMemcpyAsync( + image_cpu_buffer.data(), + images_set[i], + 4 * 640 * 480, + cudaMemcpyDeviceToHost + ); + cudaMemcpyAsync( + depth_cpu_buffer.data(), + depths_set[i], + 640 * 480 * sizeof(float), + cudaMemcpyDeviceToHost + ); + cudaStreamSynchronize(0); // Wait for the copy to complete + + // Now write to D3D12 readback buffers + uint8_t* image_data = nullptr; + float* depth_data = nullptr; + + D3D12_RANGE image_read_range = {0, IMAGE_BUFSIZE}; + D3D12_RANGE depth_read_range = {0, DEPTH_BUFSIZE}; + D3D12_RANGE no_write = {0, 0}; + + ThrowIfFailed(robot_resources[spot_id][stream][i].image_texture->Map( + 0, &image_read_range, reinterpret_cast(&image_data) + ), "Map image readback"); + + ThrowIfFailed(robot_resources[spot_id][stream][i].depth_texture->Map( + 0, &depth_read_range, reinterpret_cast(&depth_data) + ), "Map depth readback"); + + // Copy from CPU buffers to D3D12 mapped memory + std::memcpy(image_data, image_cpu_buffer.data(), IMAGE_BUFSIZE); + std::memcpy(depth_data, depth_cpu_buffer.data(), DEPTH_BUFSIZE); + + // Create OpenCV matrices from D3D12 readback data + cv::Mat image(480, 640, CV_8UC4, image_data); + cv::Mat depth(480, 640, CV_32FC1, depth_data); + + cv::cvtColor(image, image, cv::COLOR_RGBA2BGR); + cv::normalize(depth, depth, 0, 1, cv::NORM_MINMAX); + + cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + " RGB" + std::to_string(i), image); + cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + " Depth" + std::to_string(i), depth); + + // Unmap the readback buffers + robot_resources[spot_id][stream][i].image_texture->Unmap(0, &no_write); + robot_resources[spot_id][stream][i].depth_texture->Unmap(0, &no_write); + } - n_images_pending = num_cameras * robots_active; + if (cv::waitKey(1) == 'q') { + exit_requested = true; + break; + } + } } - - frame_count++; } - auto end_time = high_resolution_clock::now(); - auto total_duration = duration_cast(end_time - start_time); - - std::cout << "Processed " << frame_count << " frames in " << total_duration.count() << " ms\n"; - std::cout << "Average frame time: " << (float(total_duration.count()) / float(frame_count)) << " ms\n"; - // --------------------------------------------------------------------------------------- // 7. Cleanup // --------------------------------------------------------------------------------------- std::cout << "Cleaning up...\n"; + disconnect_from_spots(spot_ids, 2); cv::destroyAllWindows(); - for (int robot = 0; robot < 2; robot++) { - if (spot_ids[robot] >= 0) { - SOb_DisconnectFromSpot(spot_ids[robot]); - } - } + if (model) SOb_UnloadModel(model); + for (auto image_set : images_gpu) delete[] image_set; + for (auto depth_set : depths_gpu) delete[] depth_set; CloseHandle(fenceEvt); diff --git a/tests/integ-test/integ-test.cpp b/tests/integ-test/integ-test.cpp index 1df7e9f..942ceae 100644 --- a/tests/integ-test/integ-test.cpp +++ b/tests/integ-test/integ-test.cpp @@ -14,11 +14,10 @@ static uint32_t __num_set_bits(uint32_t bitmask) { return __popcnt(bitmask); } -static int32_t connect_to_spot_and_start_cam_feed( +static int32_t connect_to_spot( const std::string& robot_ip, const std::string& username, - const std::string& password, - uint32_t cam_bitmask + const std::string& password ) { std::cout << "Connecting to Spot robot at " << robot_ip << " with user " << username << std::endl; int32_t spot_id = SOb_ConnectToSpot(robot_ip.c_str(), username.c_str(), password.c_str()); @@ -26,14 +25,32 @@ static int32_t connect_to_spot_and_start_cam_feed( std::cerr << "Failed to connect to Spot robot" << std::endl; return -1; } - bool ret = SOb_ReadCameraFeeds(spot_id, cam_bitmask); - if (!ret) { + std::cout << "Connected to Spot robot with ID: " << spot_id << std::endl; + return spot_id; +} + +static int32_t start_cam_stream(int32_t spot_id, uint32_t cam_bitmask) { + int32_t cam_stream_id = SOb_CreateCameraStream(spot_id, cam_bitmask); + if (cam_stream_id < 0) { std::cerr << "Failed to start reading camera feeds" << std::endl; SOb_DisconnectFromSpot(spot_id); return -1; } + std::cout << "Started camera stream with ID: " << cam_stream_id << std::endl; + return cam_stream_id; +} - return spot_id; +static int32_t disconnect_from_spots(const int32_t spot_ids[], size_t num_spots) { + for (size_t i = 0; i < num_spots; i++) { + if (spot_ids[i] >= 0) { + if (SOb_DisconnectFromSpot(spot_ids[i])) { + std::cout << "Disconnected from Spot robot with ID: " << spot_ids[i] << std::endl; + } else { + std::cerr << "Failed to disconnect from Spot robot with ID: " << spot_ids[i] << std::endl; + } + } + } + return 0; } int main(int argc, char* argv[]) { @@ -50,22 +67,37 @@ int main(int argc, char* argv[]) { //SOb_ToggleDebugDumps("./spot_dump"); - int32_t spot_ids[2]; + int32_t spot_ids[2] = {-1, -1}; + + std::unordered_map> cam_stream_ids; - uint32_t cam_bitmask = FRONTRIGHT | FRONTLEFT | HAND; + std::vector cam_bitmasks = {FRONTRIGHT | FRONTLEFT }; for (size_t i = 0; i < 2; i++) { if (robot_ips[i] == "0") { spot_ids[i] = -1; } else { - spot_ids[i] = connect_to_spot_and_start_cam_feed(robot_ips[i], username, password, cam_bitmask); + spot_ids[i] = connect_to_spot(robot_ips[i], username, password); if (spot_ids[i] < 0) { - std::cerr << "Failed to connect to Spot robot " << i << std::endl; + disconnect_from_spots(spot_ids, 2); return -1; } } } std::cout << "Connected to Spot robots with IDs: " << spot_ids[0] << ", " << spot_ids[1] << std::endl; + for (size_t i = 0; i < 2; i++) { + if (spot_ids[i] < 0) continue; + for (uint32_t cam_bitmask : cam_bitmasks) { + int32_t cam_stream_id = start_cam_stream(spot_ids[i], cam_bitmask); + + if (cam_stream_id < 0) { + disconnect_from_spots(spot_ids, 2); + return -1; + } + cam_stream_ids[spot_ids[i]].push_back(cam_stream_id); + } + } + // TODO: Setup a listener for ctrl-c to gracefully stop the connection // std::cout << "Press Ctrl-C to stop reading camera feeds..." << std::endl; @@ -77,10 +109,9 @@ int main(int argc, char* argv[]) { model = SOb_LoadModel(model_path, "cuda"); if (!model) { std::cerr << "Failed to load model from: " << argv[5] << std::endl; - for (int32_t spot = 0; spot < 2; spot++) { - cv::destroyAllWindows(); - SOb_DisconnectFromSpot(spot_ids[spot]); - } + disconnect_from_spots(spot_ids, 2); + cv::destroyAllWindows(); + return -1; } std::cout << "Model loaded successfully!" << std::endl; @@ -88,101 +119,115 @@ int main(int argc, char* argv[]) { // Launch vision pipeline on both robots for (size_t i = 0; i < 2; i++) { if (spot_ids[i] < 0) continue; - bool ret = SOb_LaunchVisionPipeline(spot_ids[i], model); + // Launch vision pipeline only on the first camera stream + bool ret = SOb_LaunchVisionPipeline(spot_ids[i], cam_stream_ids[spot_ids[i]][0], model); if (!ret) { std::cerr << "Failed to launch vision pipeline on robot " << i << std::endl; - for (int32_t spot = 0; spot < 2; spot++) { - cv::destroyAllWindows(); - SOb_DisconnectFromSpot(spot_ids[spot]); - } + disconnect_from_spots(spot_ids, 2); SOb_UnloadModel(model); + cv::destroyAllWindows(); return -1; } std::cout << "Vision pipeline launched on robot " << i << std::endl; } } - uint32_t num_images_requested = __num_set_bits(cam_bitmask); - uint8_t** images = new uint8_t*[num_images_requested]; - float** depths = new float*[num_images_requested]; + // Initialized output pointers + std::vector num_images_requested_per_stream; + std::vector images; + std::vector depths; - std::vector image_cpu_buffer(640 * 480 * 3); + for (uint32_t cam_bitmask : cam_bitmasks) { + uint32_t num_bits_set = __num_set_bits(cam_bitmask); + num_images_requested_per_stream.push_back(num_bits_set); + images.push_back(new uint8_t*[num_bits_set]); + depths.push_back(new float*[num_bits_set]); + } + + std::vector image_cpu_buffer(640 * 480 * 4); std::vector depth_cpu_buffer(640 * 480); - int32_t img_batch_id = 0; bool new_images = false; - while (true) { - static time_point start_time = high_resolution_clock::now(); - + time_point start_time = high_resolution_clock::now(); + bool exit_requested = false; + while (!exit_requested) { if (new_images) { time_point end_time = high_resolution_clock::now(); auto duration = duration_cast(end_time - start_time); double latency_ms = double(duration.count()) / 1000.0; - std::cout << std::format("integ-test: {}-rgbd read latency: {:.4f} ms\n", num_images_requested, latency_ms); + std::cout << std::format("integ-test: total latency: {:.4f} ms\n", latency_ms); start_time = end_time; // Reset start time for next push new_images = false; } for (int32_t spot = 0; spot < 2; spot++) { - if (spot_ids[spot] < 0) { + int32_t spot_id = spot_ids[spot]; + if (spot_id < 0) { std::cout << "Skipping Spot " << spot << " as it is not connected." << std::endl; continue; } - - if (using_vision_pipeline) { - if (!SOb_GetNextVisionPipelineImageSet(spot_ids[spot], int32_t(num_images_requested), images, depths)) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); + for (int32_t stream = 0; stream < cam_stream_ids[spot_id].size(); stream++) { + int32_t cam_stream_id = cam_stream_ids[spot_id][stream]; + if (cam_stream_id < 0) { + std::cout << "Skipping Spot " << spot << " Stream " << stream << " as it is not valid." << std::endl; continue; } - } else { - if (!SOb_GetNextImageSet(spot_ids[spot], int32_t(num_images_requested), images, depths)) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - continue; + + uint32_t num_images_requested = num_images_requested_per_stream[stream]; + uint8_t** images_set = images[stream]; + float** depths_set = depths[stream]; + + if (using_vision_pipeline && stream == 0) { + if (!SOb_GetNextVisionPipelineImageSet(spot_id, cam_stream_id, int32_t(num_images_requested), images_set, depths_set)) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + continue; + } + } else { + if (!SOb_GetNextImageSet(spot_id, cam_stream_id, int32_t(num_images_requested), images_set, depths_set)) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + continue; + } + } + new_images = true; + for (uint32_t i = 0; i < num_images_requested; i++) { + cudaMemcpyAsync( + image_cpu_buffer.data(), + images_set[i], + 4 * 640 * 480, + cudaMemcpyDeviceToHost + ); + cudaMemcpyAsync( + depth_cpu_buffer.data(), + depths_set[i], + 640 * 480 * sizeof(float), + cudaMemcpyDeviceToHost + ); + cudaStreamSynchronize(0); // Wait for the copy to complete + + cv::Mat image(480, 640, CV_8UC4, image_cpu_buffer.data()); + cv::Mat depth(480, 640, CV_32FC1, depth_cpu_buffer.data()); + + cv::cvtColor(image, image, cv::COLOR_RGBA2BGR); + cv::normalize(depth, depth, 0, 1, cv::NORM_MINMAX); + + cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + " RGB" + std::to_string(i), image); + cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + " Depth" + std::to_string(i), depth); + } + if (cv::waitKey(1) == 'q') { + exit_requested = true; + break; } - } - new_images = true; - for (uint32_t i = 0; i < num_images_requested; i++) { - // std::cout << "SPOT " << spot << " Image batch " << img_batch_id <<" Image " << i << ": "; - // std::cout << images[i] << " "; - // std::cout << depths[i] << std::endl; - cudaMemcpyAsync( - image_cpu_buffer.data(), - images[i], - 4 * 640 * 480, - cudaMemcpyDeviceToHost - ); - cudaMemcpyAsync( - depth_cpu_buffer.data(), - depths[i], - 640 * 480 * sizeof(float), - cudaMemcpyDeviceToHost - ); - cudaStreamSynchronize(0); // Wait for the copy to complete - - cv::Mat image(480, 640, CV_8UC4, image_cpu_buffer.data()); - cv::Mat depth(480, 640, CV_32FC1, depth_cpu_buffer.data()); - - cv::cvtColor(image, image, cv::COLOR_RGB2BGR); - cv::normalize(depth, depth, 0, 1, cv::NORM_MINMAX); - - cv::imshow("SPOT " + std::to_string(spot) + " RGB" + std::to_string(i), image); - cv::imshow("SPOT " + std::to_string(spot) + " Depth" + std::to_string(i), depth); - } - if (cv::waitKey(1) == 'q') { - break; } } - img_batch_id++; } - for (int32_t spot = 0; spot < 2; spot++) { - cv::destroyAllWindows(); - SOb_DisconnectFromSpot(spot_ids[spot]); - } + disconnect_from_spots(spot_ids, 2); + cv::destroyAllWindows(); + if (model) SOb_UnloadModel(model); - if (images) delete[] images; - if (depths) delete[] depths; + for (auto image_set : images) delete[] image_set; + for (auto depth_set : depths) delete[] depth_set; return 0; }