diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..beae08f --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,36 @@ +name: docs + +on: + push: + branches: ["main"] + pull_request: + +permissions: + contents: write + +jobs: + build-docs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install dependencies + run: pip install -r docs/requirements.txt + + - name: Ensure local mermaid JS exists + run: | + test -f docs/_static/mermaid.min.js || curl -sL https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js -o docs/_static/mermaid.min.js + + - name: Build docs + run: sphinx-build -W -b html docs docs/_build/html + + - name: Upload docs artifact + uses: actions/upload-artifact@v4 + with: + name: docs-html + path: docs/_build/html diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..258968d --- /dev/null +++ b/.gitignore @@ -0,0 +1,101 @@ +# C++ +**/cmake-build*/ +**/build*/ +**/*.so +**/*.log* + +# Omniverse +**/*.dmp +**/.thumbs + +# No USD files allowed in the repo +**/*.usd +**/*.usda +**/*.usdc +**/*.usdz + +# Python +.DS_Store +**/*.egg-info/ +**/__pycache__/ +**/.pytest_cache/ +**/*.pyc +**/*.pb + +# Docker/Singularity +**/*.sif +docker/cluster/exports/ +docker/.container.cfg + +# IDE +**/.idea/ +**/.vscode +# Don't ignore the top-level .vscode directory as it is +# used to configure VS Code settings +#!.vscode + +# Outputs +**/output/* +**/outputs/* +**/videos/* +**/wandb/* +**/.neptune/* +docker/artifacts/ +*.tmp +**/*.pt +**/*.pth + +# Doc Outputs +**/docs/_build/* +**/generated/* +**/docs/technical_reference + +# Isaac-Sim packman +_isaac_sim* +_repo +_build +.lastformat + +# RL-Games +**/runs/* +**/logs/* +**/recordings/* + +# Pre-Trained Checkpoints +/.pretrained_checkpoints/ + +# Teleop Recorded Dataset +/datasets/ +pipeline_out/ + +# Docker history +.isaac-lab-docker-history + +tex file extensions +**/*.aux* +**/*.fls* +**/*.fdb* +**/*.out +**/*.toc +**/*.lot +**/*.lof + +# backups +**/*.bak + +#data +data/* + +#env +**/.env +**/.venv +\.*_env +\.venv +**/venv + +.jepa_config.json + +# Vi files +**/*.swp +**/*.swo +docs/_static/mermaid.min.js diff --git a/CHANGELOG.md b/CHANGELOG.md index 4371aa1..440e5d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ Initial pre-release targeting the Humble ROS distribution. --- -## [pre-v0.1.0-m] — current +## [pre-v0.1.0-m] **Branch:** `vggt-no-arl-split` -> `main` MVP for decoder / encoder split. @@ -33,3 +33,14 @@ First tagged release targeting the GQ robot platforms. ### Added - Initial release for GQ platforms + +--- + +## [v1.0.0-main] — current + +First release support ROS2 Humble after RAL acceptance. + +### Added +- Refactored inference engine to support both tensorRT and ONNX runtimes using pluggable backends. +- Removed dependency on specific autonomous robot platforms and stacks using a generic adapter interface. +- Provide full documentation for installation, configuration, and usage of the Neuromesh library. diff --git a/VGGT_README.md b/VGGT_README.md deleted file mode 100644 index 8f5c55b..0000000 --- a/VGGT_README.md +++ /dev/null @@ -1,149 +0,0 @@ -# VGGT Model Implementation for 2-Robot Setup - -This implementation provides a ROS2 node for running VGGT (Visual Geometry Grounded Transformation) models in a multi-robot environment. - -## Overview - -The VGGT implementation includes: -- **Encoder**: Processes RGB images (392x518 resolution) to extract patch tokens (1x1036x1024) -- **Decoder**: Aggregates features from 2 robots to produce depth maps, 3D points, and confidence maps -- **Multi-robot communication**: Feature sharing between robots via ROS2 topics -- **Output publishers**: Depth images, pointclouds (with/without RGB colors) - -## Model Specifications - -### Encoder (vggt_image_encoder_2x.engine) -- **Input**: Images with dimensions 1x3x392x518 -- **Output**: patch_tokens with dimensions 1x1036x1024 - -### Decoder (vggt_aggregator_2x.engine) -- **Input**: patch_tokens with dimensions 2x1036x1024 (from 2 robots) -- **Outputs**: - - pose_enc: 1x2x9 - - depth: 1x2x392x518x1 - - depth_conf: 1x2x392x518 - - world_points: 1x2x392x518x3 - - world_points_conf: 1x2x392x518 - -## Files Created - -### Core Implementation -- `include/neuromesh_platform_r2/vggt_neuromesh_node.h` - Base VGGT node header -- `src/vggt_neuromesh_node.cpp` - Base VGGT node implementation -- `include/neuromesh_platform_r2/vggt_toy_implementation.h` - TensorRT integration header -- `src/vggt_toy_implementation.cpp` - TensorRT integration implementation - -### Launch Files -- `launch/vggt_model_neuromesh_launch.py` - Main launch file for VGGT setup -- `scripts/vggt_model_neuromesh_launch.sh` - Launch script for easy execution - -## Usage - -### Launch Single Robot -```bash -# Terminal 1 - Robot 1 (khonsu) -./scripts/vggt_model_neuromesh_launch.sh khonsu 1 - -# Terminal 2 - Robot 2 (anubis) -./scripts/vggt_model_neuromesh_launch.sh anubis 2 -``` - -### Launch with Custom Parameters -```bash -ros2 launch neuromesh_platform_r2 vggt_model_neuromesh_launch.py \ - name:=khonsu \ - agent_num:=1 \ - agent_list:="khonsu,anubis" \ - color_raw_topic:=/khonsu/sensors/camera_0/camera/color/image_raw -``` - -## Published Topics - -Each robot publishes the following topics: - -### Feature Sharing -- `/{robot_name}/features_{robot_name}` - Feature messages for multi-robot coordination - -### Depth Outputs -- `/{robot_name}/depth_robot1` - Depth image for current robot -- `/{robot_name}/depth_robot2` - Depth image for neighbor robot - -### Point Clouds -- `/{robot_name}/pointcloud_current` - 3D pointcloud for current robot -- `/{robot_name}/pointcloud_neighbor` - 3D pointcloud for neighbor robot -- `/{robot_name}/pointcloud_current_rgb` - RGB pointcloud for current robot -- `/{robot_name}/pointcloud_neighbor_rgb` - RGB pointcloud for neighbor robot - -## Subscribed Topics - -- `/{robot_name}/sensors/camera_0/camera/color/image_raw` - Input RGB camera images -- `/other_robot/features_other_robot` - Features from neighbor robots - -## Key Features - -### Image Preprocessing -- Automatic resizing to 392x518 resolution (VGGT input requirement) -- Normalization to [-1, 1] range -- HWC to CHW format conversion for neural network input - -### Multi-Robot Coordination -- Feature sharing between robots via ROS2 topics -- Temporal synchronization of features -- Automatic neighbor selection for decoder input - -### Output Processing -- Depth image generation from decoder outputs -- 3D pointcloud creation with confidence filtering -- RGB pointcloud generation using original camera colors -- Separate outputs for current robot vs neighbor robot data - -### Performance Monitoring -- Built-in timing measurements for encoder/decoder inference -- Debug logging for feature processing and tensor operations - -## Configuration - -### Default Agent Setup -- **Robot 1**: khonsu (agent_num: 1) -- **Robot 2**: anubis (agent_num: 2) -- **Agent List**: "khonsu,anubis" - -### Model Paths -- **Encoder**: `tensorrt_engine/models/vggt_onnx_2x/vggt_image_encoder_2x.engine` -- **Decoder**: `tensorrt_engine/models/vggt_onnx_2x/vggt_aggregator_2x.engine` - -### Cycle Timing -- **Encoder Cycle**: 3000ms (configurable) -- **Decoder Cycle**: 3000ms (configurable) - -## Coordinate Frames - -- **Depth Images**: Published in camera optical frame (`cam1_color_optical_frame`) -- **Point Clouds**: Published in robot map frame (`{robot_name}/map`) -- **Transform Broadcasting**: Automatic TF2 transform broadcasting between frames - -## Dependencies - -- ROS2 (tested with appropriate ROS2 distribution) -- TensorRT Engine service for model inference -- OpenCV for image processing -- cv_bridge for ROS-OpenCV conversion -- sensor_msgs for point cloud generation -- neuromesh_interfaces for custom message types - -## Build Instructions - -The VGGT implementation is integrated into the existing neuromesh_platform_r2 package: - -```bash -cd /path/to/workspace -colcon build --packages-select neuromesh_platform_r2 -source install/setup.bash -``` - -## Notes - -- Requires exactly 2 robots for proper operation -- TensorRT engine files must be present in the specified model directory -- Camera topics should follow the expected naming convention -- Point cloud confidence threshold is set to 0.5 (configurable in code) \ No newline at end of file diff --git a/VGGT_REFACTORING_DESIGN.md b/VGGT_REFACTORING_DESIGN.md deleted file mode 100644 index 82c2fbb..0000000 --- a/VGGT_REFACTORING_DESIGN.md +++ /dev/null @@ -1,372 +0,0 @@ -# VGGT Encoder-Decoder Separation Design Document - -## Executive Summary - -This document outlines the design for refactoring the Visual Geometry Grounded Transformation (VGGT) system from a single monolithic node into two separate ROS2 nodes: `vggt_encoder_node` and `vggt_decoder_node`. This separation aims to eliminate async scheduling conflicts, improve modularity, and enable better scalability for multi-robot systems. - -## Architecture Overview - -### Current Architecture (Monolithic) -``` -┌─────────────────────────────────────────┐ -│ vggt_neuromesh_node │ -│ ┌─────────────┐ ┌─────────────┐ │ -│ │ Encoder │───►│ Decoder │ │ -│ └─────────────┘ └─────────────┘ │ -│ │ │ │ -│ └───────┬───────────┘ │ -│ ▼ │ -│ TensorRT Service │ -└─────────────────────────────────────────┘ -``` - -### Proposed Architecture (Separated) -``` -┌─────────────────────┐ ┌─────────────────────┐ -│ vggt_encoder_node │ │ vggt_decoder_node │ -│ ┌─────────────┐ │ │ ┌─────────────┐ │ -│ │ Encoder │ │◄────────┤ │ Decoder │ │ -│ └─────────────┘ │Features │ └─────────────┘ │ -│ │ │ Topic │ │ │ -│ ▼ │ │ ▼ │ -│ TensorRT Service │ │ TensorRT Service │ -└─────────────────────┘ └─────────────────────┘ -``` - -## Node Design Details - -### 1. VGGT Encoder Node - -#### Purpose -Processes camera images and generates feature representations for multi-robot perception. - -#### Inputs -- **Camera Image**: Subscribes to `color_raw_topic` (sensor_msgs/Image) -- **Configuration**: ROS parameters for encoder settings - -#### Outputs -- **Features**: Publishes to `/robot_name/features_robot_name` (neuromesh_msgs/NeuroMeshFeatures) - -#### Core Components -```cpp -class VggtEncoderNode : public rclcpp::Node { -private: - // Subscriptions - rclcpp::Subscription::SharedPtr camera_sub_; - - // Publishers - rclcpp::Publisher::SharedPtr feature_pub_; - - // TensorRT client - rclcpp::Client::SharedPtr tensorrt_client_; - - // Timer for periodic processing - rclcpp::TimerBase::SharedPtr encoder_timer_; - - // Configuration - double encoder_cycle_interval_; // Configurable processing interval - std::string encoder_model_path_; - - // State management - std::shared_ptr latest_image_; - std::mutex image_mutex_; - bool processing_in_progress_; -}; -``` - -#### Key Methods -- `camera_callback()`: Stores latest image (non-blocking) -- `encoder_timer_callback()`: Triggers encoding at configured intervals -- `process_image()`: Preprocesses image and calls TensorRT -- `publish_features()`: Publishes encoded features with timestamp - -### 2. VGGT Decoder Node - -#### Purpose -Aggregates features from N robots and generates depth maps, point clouds, and pose information. - -#### Inputs -- **Self Features**: From local encoder node -- **Neighbor Features**: From other robots' encoder nodes -- **Configuration**: ROS parameters for decoder settings - -#### Outputs -- **Depth Images**: `/robot_name/depth_robotX` (sensor_msgs/Image) -- **Point Clouds**: `/robot_name/pointcloud_current`, `/robot_name/pointcloud_neighbor` -- **RGB Point Clouds**: With color information -- **Pose Information**: If needed - -#### Core Components -```cpp -class VggtDecoderNode : public rclcpp::Node { -private: - // Subscriptions - std::map::SharedPtr> feature_subs_; - - // Publishers - std::map::SharedPtr> depth_pubs_; - std::map::SharedPtr> pointcloud_pubs_; - - // TensorRT client - rclcpp::Client::SharedPtr tensorrt_client_; - - // Feature buffer - std::map feature_buffer_; - std::map feature_timestamps_; - std::mutex feature_mutex_; - - // Configuration - double decoder_cycle_interval_; - double feature_age_threshold_; // Max age for neighbor features (default: 10s) - std::vector robot_names_; - int num_robots_; // Configurable N - - // Timer for processing - rclcpp::TimerBase::SharedPtr decoder_timer_; -}; -``` - -#### Key Methods -- `feature_callback()`: Updates feature buffer for each robot -- `decoder_timer_callback()`: Triggers decoding at intervals -- `aggregate_features()`: Builds decoder tensor from N robot features -- `check_feature_freshness()`: Validates feature timestamps -- `process_decoder_output()`: Converts tensor outputs to ROS messages - -### 3. Feature Aggregation Strategy - -```cpp -std::vector VggtDecoderNode::aggregate_features() { - std::lock_guard lock(feature_mutex_); - std::vector aggregated_tensor; - - // Get current robot's features (always index 0) - auto self_features = feature_buffer_[robot_name_]; - - // Aggregate features from all robots - for (int i = 0; i < num_robots_; i++) { - if (i == 0) { - // Always use self features for index 0 - aggregated_tensor.insert(aggregated_tensor.end(), - self_features.data.begin(), - self_features.data.end()); - } else { - // Use neighbor features or fallback to self - std::string neighbor_name = robot_names_[i]; - - if (feature_buffer_.count(neighbor_name) > 0) { - auto age = this->now() - feature_timestamps_[neighbor_name]; - - if (age.seconds() < feature_age_threshold_) { - // Use neighbor features - aggregated_tensor.insert(aggregated_tensor.end(), - feature_buffer_[neighbor_name].data.begin(), - feature_buffer_[neighbor_name].data.end()); - } else { - // Features too old, use self features - RCLCPP_WARN(this->get_logger(), - "Features from %s are %.2f seconds old (threshold: %.2f). Using self features.", - neighbor_name.c_str(), age.seconds(), feature_age_threshold_); - aggregated_tensor.insert(aggregated_tensor.end(), - self_features.data.begin(), - self_features.data.end()); - } - } else { - // No features available, use self features - RCLCPP_WARN(this->get_logger(), - "No features available from %s. Using self features.", - neighbor_name.c_str()); - aggregated_tensor.insert(aggregated_tensor.end(), - self_features.data.begin(), - self_features.data.end()); - } - } - } - - return aggregated_tensor; -} -``` - -## Configuration System - -### Unified Configuration File -Create a YAML configuration file that both nodes can reference: - -```yaml -# config/vggt_config.yaml -vggt: - # Encoder settings - encoder: - cycle_interval: 3.0 # seconds - image_width: 518 - image_height: 392 - model_path: "$(find neuromesh_platform_r2)/../../tensorrt_engine/models/vggt_onnx_2x/vggt_image_encoder_2x.engine" - - # Decoder settings - decoder: - cycle_interval: 3.0 # seconds - feature_age_threshold: 10.0 # seconds - num_robots: 2 # Configurable N - model_path: "$(find neuromesh_platform_r2)/../../tensorrt_engine/models/vggt_onnx_2x/vggt_aggregator_2x.engine" - - # Common settings - robot_names: ["khonsu", "anubis"] # List of all robots - - # TensorRT service settings - tensorrt: - service_name: "tensorrt_request" - timeout: 30.0 # seconds -``` - -### Launch File Structure - -```python -# launch/vggt_separated_launch.py -import os -from ament_index_python.packages import get_package_share_directory -from launch import LaunchDescription -from launch.actions import DeclareLaunchArgument -from launch.substitutions import LaunchConfiguration -from launch_ros.actions import Node, LoadComposableNodes -from launch_ros.descriptions import ComposableNode - -def generate_launch_description(): - # Get package directory - pkg_dir = get_package_share_directory('neuromesh_platform_r2') - config_file = os.path.join(pkg_dir, 'config', 'vggt_config.yaml') - - # Declare launch arguments - robot_name = LaunchConfiguration('robot_name') - color_raw_topic = LaunchConfiguration('color_raw_topic') - - # Component container - container = Node( - package='rclcpp_components', - executable='component_container_mt', - name='vggt_container', - output='screen', - parameters=[config_file] - ) - - # Load encoder node - encoder_node = ComposableNode( - package='neuromesh_platform_r2', - plugin='neuromesh::VggtEncoderNode', - name='vggt_encoder', - parameters=[ - config_file, - {'robot_name': robot_name, - 'color_raw_topic': color_raw_topic} - ], - extra_arguments=[{'use_intra_process_comms': True}] - ) - - # Load decoder node - decoder_node = ComposableNode( - package='neuromesh_platform_r2', - plugin='neuromesh::VggtDecoderNode', - name='vggt_decoder', - parameters=[ - config_file, - {'robot_name': robot_name} - ], - extra_arguments=[{'use_intra_process_comms': True}] - ) - - # Load components - load_components = LoadComposableNodes( - target_container='vggt_container', - composable_node_descriptions=[encoder_node, decoder_node] - ) - - return LaunchDescription([ - DeclareLaunchArgument('robot_name', default_value='khonsu'), - DeclareLaunchArgument('color_raw_topic', default_value='/khonsu/color_raw'), - container, - load_components - ]) -``` - -## Implementation Roadmap - -### Phase 1: Infrastructure Setup (Week 1) -1. Create new header files: - - `include/neuromesh_platform_r2/vggt_encoder_node.h` - - `include/neuromesh_platform_r2/vggt_decoder_node.h` -2. Create base implementations with minimal functionality -3. Set up configuration file structure -4. Create unit tests framework - -### Phase 2: Encoder Node Implementation (Week 2) -1. Implement camera subscription and buffering -2. Port image preprocessing from current implementation -3. Integrate TensorRT service client for encoder -4. Implement feature publishing with proper timestamps -5. Add configurable processing intervals -6. Test encoder node independently - -### Phase 3: Decoder Node Implementation (Week 3) -1. Implement multi-robot feature subscription -2. Port feature aggregation logic with N-robot support -3. Implement feature freshness checking with warnings -4. Integrate TensorRT service client for decoder -5. Port output processing (depth, pointcloud generation) -6. Test decoder node with simulated features - -### Phase 4: Integration and Testing (Week 4) -1. Create integrated launch files -2. Test end-to-end system with 2 robots -3. Test with N robots (N > 2) -4. Verify backward compatibility with existing topics -5. Performance benchmarking vs monolithic node -6. Documentation and code cleanup - -### Phase 5: Optional Enhancements -1. Implement health monitoring/watchdog system -2. Add dynamic reconfiguration support -3. Create diagnostic publishers -4. Add feature quality metrics -5. Implement feature compression for bandwidth optimization - -## Migration Strategy - -### Backward Compatibility -- Maintain existing topic names and message formats -- Keep same output formats for depth images and point clouds -- Ensure launch file parameters are compatible - -### Gradual Migration -1. Deploy new nodes alongside existing monolithic node -2. Compare outputs to ensure consistency -3. Switch over once validated -4. Deprecate monolithic node - -## Testing Strategy - -### Unit Tests -- Test encoder preprocessing independently -- Test decoder tensor building with various feature combinations -- Test feature age validation logic - -### Integration Tests -- Test encoder-decoder communication -- Test multi-robot feature aggregation -- Test failure scenarios (missing features, old features) - -### System Tests -- End-to-end testing with real camera data -- Multi-robot coordination testing -- Performance comparison with monolithic approach - -## Benefits of This Design - -1. **Modularity**: Clean separation of encoder and decoder logic -2. **Scalability**: Easy to extend to N robots -3. **Maintainability**: Simpler codebase with focused responsibilities -4. **Flexibility**: Independent configuration and deployment -5. **Performance**: No async scheduling conflicts -6. **Debugging**: Easier to isolate issues to specific nodes - -## Conclusion - -This refactoring will transform the VGGT system into a more modular, scalable architecture while maintaining backward compatibility. The separation of encoder and decoder nodes eliminates the current async scheduling issues and provides a cleaner foundation for future enhancements. \ No newline at end of file diff --git a/docs/VGGT_README.md b/docs/VGGT_README.md new file mode 100644 index 0000000..ffcdb4f --- /dev/null +++ b/docs/VGGT_README.md @@ -0,0 +1,4 @@ +# VGGT README + +```{include} ../VGGT_README.md +``` diff --git a/docs/VGGT_REFACTORING_DESIGN.md b/docs/VGGT_REFACTORING_DESIGN.md new file mode 100644 index 0000000..6799961 --- /dev/null +++ b/docs/VGGT_REFACTORING_DESIGN.md @@ -0,0 +1,4 @@ +# VGGT Refactoring Design + +```{include} ../VGGT_REFACTORING_DESIGN.md +``` diff --git a/docs/_static/mermaid-init.js b/docs/_static/mermaid-init.js new file mode 100644 index 0000000..a384e38 --- /dev/null +++ b/docs/_static/mermaid-init.js @@ -0,0 +1,20 @@ +window.addEventListener("load", function () { + if (typeof mermaid === "undefined") { + return; + } + + mermaid.initialize({ startOnLoad: false, theme: "neutral" }); + + document.querySelectorAll("div.highlight-mermaid pre").forEach(function (pre) { + const source = pre.textContent; + const wrapper = document.createElement("div"); + wrapper.className = "mermaid"; + wrapper.textContent = source; + const outer = pre.closest("div.highlight-mermaid"); + if (outer) { + outer.parentElement.replaceChild(wrapper, outer); + } + }); + + mermaid.run(); +}); diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..2a4521a --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,17 @@ +# Architecture + +```mermaid +graph TD + NI[neuromesh_interfaces
ROS 2 msgs/srvs] --> EI[engine_interface] + EI --> TRT[tensorrt_engine] + EI --> ONNX[onnx_engine] + NP[neuromesh_platform_r2] --> EI + NP --> NI + NP --> MB[mission_maestro_bridge optional] +``` + +NeuroMesh separates model execution from application logic. + +- `engine_interface`: runtime backend abstraction. +- `neuromesh_platform_r2`: application nodes and orchestration. +- `mission_maestro_bridge`: optional integration bridge. diff --git a/docs/build.md b/docs/build.md new file mode 100644 index 0000000..b202391 --- /dev/null +++ b/docs/build.md @@ -0,0 +1,24 @@ +# Build Guide + +## System dependencies + +Install ROS 2 Humble base tools and common development dependencies: + +- `python3-colcon-common-extensions` +- `python3-rosdep` +- CMake and compiler toolchain + +## Build steps + +1. Run `rosdep install` for workspace dependencies. +2. Build with `colcon build --symlink-install`. +3. Source `install/setup.bash`. + +## Optional engines + +- `onnx_engine` for ONNX Runtime backend +- `tensorrt_engine` for TensorRT backend + +## Optional navigation bridge + +`mission_maestro_bridge` is optional. It only activates forwarding when `arl_mission_maestro` is present. diff --git a/docs/changelog.md b/docs/changelog.md new file mode 100644 index 0000000..63ae71b --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1,4 @@ +# Changelog + +```{include} ../CHANGELOG.md +``` diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..cac4d9b --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,56 @@ +"""NeuroMesh documentation""" + +from __future__ import annotations + +from datetime import datetime + +from pygments.lexers.special import TextLexer +from sphinx.highlighting import lexers + +project = "NeuroMesh" +author = "ARPL" +project_copyright = f"{datetime.now():%Y}, {author}" +release = "humble" + +extensions = [ + "myst_parser", + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", +] + +myst_enable_extensions = [ + "colon_fence", + "deflist", + "fieldlist", + "tasklist", +] + +html_theme = "sphinx_book_theme" +html_title = project +html_theme_options = { + "repository_url": "https://github.com/arplaboratory/neuromesh", + "use_repository_button": True, +} + +html_static_path = ["_static"] +html_js_files = [ + "mermaid.min.js", + "mermaid-init.js", +] + +templates_path = ["_templates"] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +source_suffix = { + ".rst": "restructuredtext", + ".md": "markdown", +} + +suppress_warnings = [ + "docutils", + "toc.not_included", + "ref.ref", + "misc.highlighting_failure", +] + +lexers["mermaid"] = TextLexer() diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 0000000..401c311 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,4 @@ +# Contributing + +```{include} ../CONTRIBUTING.md +``` diff --git a/docs/deployment/real_world.md b/docs/deployment/real_world.md new file mode 100644 index 0000000..845f57c --- /dev/null +++ b/docs/deployment/real_world.md @@ -0,0 +1,9 @@ +# Real-World Deployment + +Use namespaced launches for each robot and ensure synchronized clocks and transforms. + +Key recommendations: + +- Stable network transport +- Consistent frame naming (`map`, `/odom`) +- Health checks on inference and communication topics diff --git a/docs/deployment/vggt_setup.md b/docs/deployment/vggt_setup.md new file mode 100644 index 0000000..01cd9fc --- /dev/null +++ b/docs/deployment/vggt_setup.md @@ -0,0 +1,7 @@ +# VGGT Setup + +This branch includes VGGT-related execution paths in `neuromesh_platform_r2`. + +- Configure model paths in launch parameters +- Validate tensor dimensions in `engine_interface` +- Verify QoS settings across robots diff --git a/docs/design/vggt_refactoring.md b/docs/design/vggt_refactoring.md new file mode 100644 index 0000000..5e44ef3 --- /dev/null +++ b/docs/design/vggt_refactoring.md @@ -0,0 +1,7 @@ +# VGGT Refactoring Design + +The refactoring focuses on: + +- separating encoder/decoder responsibilities, +- decoupling engine backends, +- and introducing optional bridges for external navigation interfaces. diff --git a/docs/getting_started.md b/docs/getting_started.md new file mode 100644 index 0000000..84f8a96 --- /dev/null +++ b/docs/getting_started.md @@ -0,0 +1,19 @@ +# Getting Started + +## Requirements + +- Ubuntu 22.04 +- ROS 2 Humble +- Python 3.10+ +- Colcon and ROS build tools + +## Quick start + +1. Build your workspace with `colcon`. +2. Source ROS 2 and workspace overlays. +3. Launch one of the `neuromesh_platform_r2` launch files. + +## Notes + +- This branch targets ROS 2 Humble. +- Navigation integrations are optional and should be connected through adapters/bridges. diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..00ebf14 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,49 @@ +# NeuroMesh Documentation + +NeuroMesh is a modular and decentralized framework for multi-robot collaborative learning and inference. + +```{toctree} +:maxdepth: 2 +:caption: Getting Started + +getting_started +build +``` + +```{toctree} +:maxdepth: 2 +:caption: Architecture + +architecture +packages/neuromesh_interfaces +packages/engine_interface +packages/tensorrt_engine +packages/onnx_engine +packages/neuromesh_platform_r2 +``` + +```{toctree} +:maxdepth: 2 +:caption: Deployment + +deployment/real_world +deployment/vggt_setup +``` + +```{toctree} +:maxdepth: 2 +:caption: Design + +design/vggt_refactoring +VGGT_README +VGGT_REFACTORING_DESIGN +``` + +```{toctree} +:maxdepth: 1 +:caption: Project + +changelog +contributing +license +``` diff --git a/docs/license.md b/docs/license.md new file mode 100644 index 0000000..127fd93 --- /dev/null +++ b/docs/license.md @@ -0,0 +1,5 @@ +# License + +```{include} ../LICENSE +:literal: +``` diff --git a/docs/packages/engine_interface.md b/docs/packages/engine_interface.md new file mode 100644 index 0000000..092de82 --- /dev/null +++ b/docs/packages/engine_interface.md @@ -0,0 +1,6 @@ +# engine_interface + +Provides the backend-agnostic execution interface used by application nodes. + +- Loads selected engine plugin +- Runs model inference with configured tensors diff --git a/docs/packages/neuromesh_interfaces.md b/docs/packages/neuromesh_interfaces.md new file mode 100644 index 0000000..9c583d4 --- /dev/null +++ b/docs/packages/neuromesh_interfaces.md @@ -0,0 +1,6 @@ +# neuromesh_interfaces + +Defines ROS 2 interfaces used across NeuroMesh packages. + +- Feature/tensor transport messages +- State and communication messages diff --git a/docs/packages/neuromesh_platform_r2.md b/docs/packages/neuromesh_platform_r2.md new file mode 100644 index 0000000..29589e4 --- /dev/null +++ b/docs/packages/neuromesh_platform_r2.md @@ -0,0 +1,7 @@ +# neuromesh_platform_r2 + +Application-level ROS 2 nodes for NeuroMesh deployments. + +- VGGT and GAT node pipelines +- Topic-based goal publishing +- Adapter/factory integration for optional navigation bridges diff --git a/docs/packages/onnx_engine.md b/docs/packages/onnx_engine.md new file mode 100644 index 0000000..55ce5dd --- /dev/null +++ b/docs/packages/onnx_engine.md @@ -0,0 +1,6 @@ +# onnx_engine + +ONNX Runtime backend plugin. + +- Portable inference path +- Compatible with `engine_interface` diff --git a/docs/packages/tensorrt_engine.md b/docs/packages/tensorrt_engine.md new file mode 100644 index 0000000..0ffb59d --- /dev/null +++ b/docs/packages/tensorrt_engine.md @@ -0,0 +1,6 @@ +# tensorrt_engine + +TensorRT backend plugin for accelerated inference. + +- Optimized runtime for NVIDIA GPUs +- Compatible with `engine_interface` diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..2de76c5 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,3 @@ +sphinx>=7.2 +myst-parser>=3.0 +sphinx-book-theme>=1.1 diff --git a/engine_interface/.gitignore b/engine_interface/.gitignore new file mode 100644 index 0000000..ca29f7c --- /dev/null +++ b/engine_interface/.gitignore @@ -0,0 +1,7 @@ +third_party/* +build/* +.vscode/* +resources/* +**/*.trt +**/*.engine +**/*.data diff --git a/engine_interface/.gitmodules b/engine_interface/.gitmodules new file mode 100644 index 0000000..e69de29 diff --git a/engine_interface/CMakeLists.txt b/engine_interface/CMakeLists.txt new file mode 100644 index 0000000..98321bd --- /dev/null +++ b/engine_interface/CMakeLists.txt @@ -0,0 +1,77 @@ +cmake_minimum_required(VERSION 3.5) +project(engine_interface) +set(CMAKE_BUILD_TYPE Release) + +if(BUILD_TESTING) + message("Building tests for engine_interface") + find_package(ament_cmake_gmock REQUIRED) + ament_add_gmock(test_engine_interface test/test_engine_interface.cpp) + target_link_libraries(test_engine_interface engine_interface) +endif() + +# Default to C++14 +if(NOT CMAKE_CXX_STANDARD) + set(CMAKE_CXX_STANDARD 14) +endif() + +if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") + add_compile_options(-Wall -Wextra -Wpedantic) +endif() + +set (dependencies + "std_msgs" + "rclcpp" + "neuromesh_interfaces" + "rclcpp_components" + "pluginlib" +) + +find_package(ament_cmake REQUIRED) +foreach(dep ${dependencies}) + find_package(${dep} REQUIRED) +endforeach() + +include_directories(include) + +add_library(${PROJECT_NAME} SHARED + src/engine_interface_node.cpp + src/inference_engine_base.cpp +) + +target_include_directories(engine_interface PUBLIC + $ + $ +) + +ament_target_dependencies( + engine_interface + ${dependencies} +) + +rclcpp_components_register_nodes(engine_interface "engine_interface::EngineInterfaceNode") + + install(TARGETS + engine_interface + ARCHIVE DESTINATION lib + LIBRARY DESTINATION lib + RUNTIME DESTINATION bin +) + +install(DIRECTORY launch + DESTINATION share/${PROJECT_NAME}) + +install(DIRECTORY models + DESTINATION share/${PROJECT_NAME} + OPTIONAL + ) + +install(DIRECTORY include/ + DESTINATION include) + + +ament_export_include_directories(include) +ament_export_libraries(engine_interface) +ament_export_dependencies( + ${dependencies} +) +ament_package() diff --git a/engine_interface/include/engine_interface/engine_interface_node.h b/engine_interface/include/engine_interface/engine_interface_node.h new file mode 100644 index 0000000..13f923b --- /dev/null +++ b/engine_interface/include/engine_interface/engine_interface_node.h @@ -0,0 +1,83 @@ +#ifndef ENGINE_INTERFACE_NODE_H +#define ENGINE_INTERFACE_NODE_H + +#define INFERENCE_HELPER_ENABLE_TENSORRT + +#include +#include + +#include "rclcpp/rclcpp.hpp" +#include + +#include "neuromesh_interfaces/msg/tensor.hpp" +#include "neuromesh_interfaces/srv/tensor_request.hpp" +#include "std_msgs/msg/header.hpp" + +#include "engine_interface/inference_engine_base.hpp" + + +namespace engine_interface { + +class EngineInterfaceNode : public rclcpp::Node { +public: + EngineInterfaceNode( + rclcpp::NodeOptions options, + const std::string& plugin_class = "engine_interface::BaseEngine" + ); + ~EngineInterfaceNode() override = default; + +private: + pluginlib::ClassLoader engine_loader_; + std::shared_ptr engine_; + + rclcpp::Publisher::SharedPtr + tensor_publisher_; + rclcpp::Subscription::SharedPtr + tensor_subscription_; + rclcpp::Service::SharedPtr service_; + + // params + std::string models_param; + std::vector model_names; + + std::unordered_map model_paths; + std::unordered_map>> + input_dimensions; + std::unordered_map>> + output_dimensions; + std::unordered_map input_dimensions_strings; + std::unordered_map output_dimensions_strings; + std::unordered_map tensor_type_params; + std::string tensor_qos_param; + + // engine + std::unordered_map> engines; + std::unordered_map> input_lengths; + std::unordered_map> output_lengths; + std::unordered_map tensor_typelengths; + + std::vector failed_models; + + // callbacks + void tensor_request_callback( + const std::shared_ptr + request, + const std::shared_ptr + response); + + // execution + std::vector + execute(const std::string &model, + const std::vector &tensor_msgs); + + // helper functions + int tensor_string_to_typelength(std::string input); + std::vector string_to_vector(std::string in); + std::vector string_to_dims_single(std::string in); + std::vector> string_to_dims(std::string in); + + // Convert string to ROS2 QoS profile + rmw_qos_profile_t parseQoSString(const std::string &str); +}; +} // namespace engine_interface +#endif diff --git a/engine_interface/include/engine_interface/inference_engine_base.hpp b/engine_interface/include/engine_interface/inference_engine_base.hpp new file mode 100644 index 0000000..c8948a8 --- /dev/null +++ b/engine_interface/include/engine_interface/inference_engine_base.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include +#include + +namespace engine_interface +{ + +class BaseEngine +{ +public: + virtual ~BaseEngine(); + + // Load model (path, input/output dims, etc.) + virtual bool loadModel(const std::string& model_path, + const std::vector>& input_dims, + int type_length) = 0; + + virtual void runInference(const std::vector &inputTensors, + const std::vector &inputSizes, + std::vector &outputTensors, + const std::vector &outputSizes) = 0; +}; + +} // namespace engine_interface diff --git a/engine_interface/models/README.md b/engine_interface/models/README.md new file mode 100644 index 0000000..53904e3 --- /dev/null +++ b/engine_interface/models/README.md @@ -0,0 +1,2 @@ +# Prepare your models + diff --git a/engine_interface/package.xml b/engine_interface/package.xml new file mode 100644 index 0000000..d42521c --- /dev/null +++ b/engine_interface/package.xml @@ -0,0 +1,32 @@ + + + + engine_interface + 0.0.0 + Engine agnostic interface for distributed inference + Long Quang + TODO: License declaration + + ament_cmake + + neuromesh_interfaces + rclcpp + std_msgs + pluginlib + + rclcpp_components + rclcpp_components + + ament_lint_auto + ament_lint_common + ament_cmake_gmock + + + ament_cmake + + + Engine Interface Node + + + + diff --git a/tensorrt_engine/src/engine_modified.cpp b/engine_interface/src/engine_interface_node.cpp similarity index 81% rename from tensorrt_engine/src/engine_modified.cpp rename to engine_interface/src/engine_interface_node.cpp index 7b0aba1..1e78adc 100644 --- a/tensorrt_engine/src/engine_modified.cpp +++ b/engine_interface/src/engine_interface_node.cpp @@ -1,83 +1,27 @@ -#include "tensorrt_engine/engine_modified.h" -#include "tensorrt_engine/trt_engine_modified.h" - -#include "cv_bridge/cv_bridge.h" +#include "engine_interface/engine_interface_node.h" #include -namespace tensorrt_engine_node { - -std::vector string_to_dims_single(std::string in) { - std::stringstream stream(in); - std::string element; - std::vector out; - - while (getline(stream, element, ',')) { - out.push_back(std::stoi(element)); - } - return out; -} - -std::vector> string_to_dims(std::string in) { - std::stringstream stream(in); - std::string element; - std::vector> out; - - while (getline(stream, element, ';')) { - std::vector dims = string_to_dims_single(element); - out.push_back(dims); - } - return out; -} - -std::vector string_to_vector(std::string in) { - std::stringstream stream(in); - std::string element; - - std::vector out; - - while (getline(stream, element, ',')) { - out.push_back(element); - } - return out; -} - -void TensorRTEngineNode::tensor_request_callback( - const std::shared_ptr - request, - const std::shared_ptr - response) { - auto now = this->get_clock()->now(); - double timestamp = now.seconds() + now.nanoseconds() / 1e9; - RCLCPP_DEBUG(this->get_logger(), "Received service request."); - RCLCPP_DEBUG(this->get_logger(), "Time of receiving service call %.9f", - timestamp); - RCLCPP_DEBUG(this->get_logger(), "Number of input tensors: %zu", - request->tensor1.size()); - for (size_t i = 0; i < request->tensor1.size(); i++) { - RCLCPP_DEBUG(this->get_logger(), "Tensor size (%d): %ld", i, (request->tensor1[i]).data.size()); - } - - std::vector input_tensors = - request->tensor1; - response->tensor2 = execute(request->model_name, input_tensors); -} - -TensorRTEngineNode::TensorRTEngineNode(rclcpp::NodeOptions options) - : Node("TensorRTEngineNode", - options.allow_undeclared_parameters(true) - .automatically_declare_parameters_from_overrides(true)) { - // params - // this->declare_parameter("model_names", ""); //only declared - // parameters +namespace engine_interface { + +EngineInterfaceNode::EngineInterfaceNode( + rclcpp::NodeOptions options, + const std::string& plugin_class +) + : Node("EngineInterfaceNode", + options.allow_undeclared_parameters(true) + .automatically_declare_parameters_from_overrides(true)), + engine_loader_("engine_interface", plugin_class) +{ + // set param vars this->declare_parameter("tensor_qos_profile", "default"); - // set param vars - this->get_parameter("model_names", models_param); // values separated by comma + this->get_parameter("model_names", models_param); model_names = string_to_vector(models_param); + std::string plugin_cls = this->get_parameter("engine_type").as_string(); for (std::vector::iterator it = model_names.begin(); - it != model_names.end(); it++) { + it != model_names.end(); it++) { std::string m = *it; this->get_parameter(m + ".model_path", model_paths[m]); this->get_parameter(m + ".input_dimensions", input_dimensions_strings[m]); @@ -87,11 +31,11 @@ TensorRTEngineNode::TensorRTEngineNode(rclcpp::NodeOptions options) RCLCPP_DEBUG(this->get_logger(), "Loading parameters."); RCLCPP_DEBUG(this->get_logger(), "Model Name: %s", m.c_str()); RCLCPP_DEBUG(this->get_logger(), "Input Dimensions: %s", - input_dimensions_strings[m].c_str()); + input_dimensions_strings[m].c_str()); RCLCPP_DEBUG(this->get_logger(), "Output Dimensions: %s", - output_dimensions_strings[m].c_str()); + output_dimensions_strings[m].c_str()); RCLCPP_DEBUG(this->get_logger(), "Tensor Type: %s", - tensor_type_params[m].c_str()); + tensor_type_params[m].c_str()); // here we set default values if (!tensor_type_params.count(m)) { @@ -119,7 +63,6 @@ TensorRTEngineNode::TensorRTEngineNode(rclcpp::NodeOptions options) RCLCPP_DEBUG(this->get_logger(), "Getting output dimensions"); output_dimensions[m] = string_to_dims(output_dimensions_strings[m]); RCLCPP_DEBUG(this->get_logger(), "Got output dimensions"); - RCLCPP_DEBUG(this->get_logger(), "Getting input lengths"); input_lengths[m].resize(input_dimensions[m].size()); for (size_t i = 0; i < input_dimensions[m].size(); ++i) { @@ -158,42 +101,42 @@ TensorRTEngineNode::TensorRTEngineNode(rclcpp::NodeOptions options) // // Publisher and subscriber setup tensor_publisher_ = this->create_publisher( "tensorrt_output", tensor_qos); - - // tensor_subscription_ = - // this->create_subscription( - // "tensorrt_input", tensor_qos, - // std::bind(&TensorRTEngineNode::tensor_callback, this, - // std::placeholders::_1)); - - RCLCPP_DEBUG(this->get_logger(), "Creating service"); + RCLCPP_INFO(this->get_logger(), "Creating service"); service_ = this->create_service( "tensorrt_request", - std::bind(&TensorRTEngineNode::tensor_request_callback, this, + std::bind(&EngineInterfaceNode::tensor_request_callback, this, std::placeholders::_1, std::placeholders::_2)); for (std::vector::iterator it = model_names.begin(); it != model_names.end(); it++) { std::string m = *it; - // initialize engine - // TODO try/catch engine creation failure - RCLCPP_DEBUG(this->get_logger(), "Resetting engine"); - engines[m].reset(new TRTEngine(model_paths[m], input_dimensions[m], - tensor_typelengths[m])); - RCLCPP_DEBUG(this->get_logger(), "Engine reset"); - - if (engines[m].get() == NULL) { + try { + engines[m] = engine_loader_.createSharedInstance(plugin_cls); + RCLCPP_INFO(this->get_logger(), "Loading model"); + engines[m]->loadModel(model_paths[m], + input_dimensions[m], + tensor_typelengths[m]); + } catch (const pluginlib::PluginlibException& ex) { + RCLCPP_ERROR(this->get_logger(), "Failed to load engine plugin: %s", ex.what()); + } + + if (engines[m].get() == NULL) { RCLCPP_ERROR(this->get_logger(), "Failed to initialize engine %s.", m.c_str()); } + } -} -std::vector TensorRTEngineNode::execute( + RCLCPP_DEBUG(this->get_logger(), "Done loading parameters."); + +} +std::vector EngineInterfaceNode::execute( const std::string &model, - const std::vector &tensor_msgs) { + const std::vector &tensor_msgs) +{ RCLCPP_DEBUG(this->get_logger(), "Execute function for model %s", model.c_str()); - + if (tensor_msgs.empty()) { RCLCPP_ERROR(this->get_logger(), "Expected at least 1 input tensor, got 0"); return {[]() { @@ -202,7 +145,6 @@ std::vector TensorRTEngineNode::execute( return t; }()}; } - // Prepare input tensors and sizes std::vector inputTensors; std::vector inputSizes; @@ -275,9 +217,9 @@ std::vector TensorRTEngineNode::execute( output_msg.result = 0; output_msg.data_type = 9; // float32 output_msgs.push_back(std::move(output_msg)); - RCLCPP_DEBUG(this->get_logger(), "I: %d", i); + RCLCPP_DEBUG(this->get_logger(), "I: %ld", i); RCLCPP_DEBUG(this->get_logger(), "MODEL: %s", model.c_str()); - RCLCPP_DEBUG(this->get_logger(), "OUTPUT TENSOR: %ld", output_dimensions[model][i]); + RCLCPP_DEBUG(this->get_logger(), "OUTPUT TENSOR: %d", output_dimensions[model][i][0]); } RCLCPP_DEBUG(this->get_logger(), "Returning output messages"); @@ -285,7 +227,63 @@ std::vector TensorRTEngineNode::execute( return output_msgs; } -int TensorRTEngineNode::tensor_string_to_typelength(std::string input) { +void EngineInterfaceNode::tensor_request_callback( + const std::shared_ptr + request, + const std::shared_ptr + response) { + auto now = this->get_clock()->now(); + double timestamp = now.seconds() + now.nanoseconds() / 1e9; + RCLCPP_DEBUG(this->get_logger(), "Received service request."); + RCLCPP_DEBUG(this->get_logger(), "Time of receiving service call %.9f", + timestamp); + RCLCPP_DEBUG(this->get_logger(), "Number of input tensors: %zu", + request->tensor1.size()); + for (size_t i = 0; i < request->tensor1.size(); i++) { + RCLCPP_DEBUG(this->get_logger(), "Tensor size (%ld): %ld", i, (request->tensor1[i]).data.size()); + } + + std::vector input_tensors = + request->tensor1; + response->tensor2 = execute(request->model_name, input_tensors); +} + +std::vector EngineInterfaceNode::string_to_dims_single(std::string in) { + std::stringstream stream(in); + std::string element; + std::vector out; + + while (getline(stream, element, ',')) { + out.push_back(std::stoi(element)); + } + return out; +} + +std::vector> EngineInterfaceNode::string_to_dims(std::string in) { + std::stringstream stream(in); + std::string element; + std::vector> out; + + while (getline(stream, element, ';')) { + std::vector dims = string_to_dims_single(element); + out.push_back(dims); + } + return out; +} + +std::vector EngineInterfaceNode::string_to_vector(std::string in) { + std::stringstream stream(in); + std::string element; + + std::vector out; + + while (getline(stream, element, ',')) { + out.push_back(element); + } + return out; +} + +int EngineInterfaceNode::tensor_string_to_typelength(std::string input) { if (input == "fp32") return 4; @@ -304,7 +302,7 @@ int TensorRTEngineNode::tensor_string_to_typelength(std::string input) { // Convert string to ROS2 QoS profile // from // https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_nvblox/nvblox_ros_common/src/qos.cpp#L26 -rmw_qos_profile_t TensorRTEngineNode::parseQoSString(const std::string &str) { +rmw_qos_profile_t EngineInterfaceNode::parseQoSString(const std::string &str) { std::string profile = str; // Convert to upper case. std::transform(profile.begin(), profile.end(), profile.begin(), ::toupper); @@ -332,7 +330,7 @@ rmw_qos_profile_t TensorRTEngineNode::parseQoSString(const std::string &str) { << ". Returning profile: DEFAULT"); return rmw_qos_profile_default; } -} // namespace tensorrt_engine_node +} // namespace engine_interface #include "rclcpp_components/register_node_macro.hpp" -RCLCPP_COMPONENTS_REGISTER_NODE(tensorrt_engine_node::TensorRTEngineNode) +RCLCPP_COMPONENTS_REGISTER_NODE(engine_interface::EngineInterfaceNode) diff --git a/engine_interface/src/inference_engine_base.cpp b/engine_interface/src/inference_engine_base.cpp new file mode 100644 index 0000000..014e2ec --- /dev/null +++ b/engine_interface/src/inference_engine_base.cpp @@ -0,0 +1,4 @@ +#include "engine_interface/inference_engine_base.hpp" +namespace engine_interface { +BaseEngine::~BaseEngine() = default; +} diff --git a/engine_interface/test/engine_interface_test_fixture.hpp b/engine_interface/test/engine_interface_test_fixture.hpp new file mode 100644 index 0000000..053253b --- /dev/null +++ b/engine_interface/test/engine_interface_test_fixture.hpp @@ -0,0 +1,11 @@ +#ifndef ENGINE_INTERFACE_TEST_FIXTURE_HPP +#define ENGINE_INTERFACE_TEST_FIXTURE_HPP + +#include "mock_inference_engine.hpp" + +class EngineInterfaceTest : public ::testing::Test { + protected: + engine_interface::MockInferenceEngine mock_engine; + }; + +#endif // ENGINE_INTERFACE_TEST_FIXTURE_HPP \ No newline at end of file diff --git a/engine_interface/test/mock_inference_engine.hpp b/engine_interface/test/mock_inference_engine.hpp new file mode 100644 index 0000000..264a397 --- /dev/null +++ b/engine_interface/test/mock_inference_engine.hpp @@ -0,0 +1,30 @@ +#ifndef MOCK_INFERENCE_ENGINE_HPP +#define MOCK_INFERENCE_ENGINE_HPP + +#include +#include "engine_interface/inference_engine_base.hpp" + +using ::testing::Return; +using ::testing::_; + +namespace engine_interface { + +// Mock class for BaseEngine +class MockInferenceEngine : public BaseEngine { +public: + MOCK_METHOD(bool, loadModel, + (const std::string& model_path, + const std::vector>& input_dims, + int type_length), + (override)); + MOCK_METHOD(void, runInference, + (const std::vector& inputTensors, + const std::vector& inputSizes, + std::vector& outputTensors, + const std::vector& outputSizes), + (override)); +}; + +} // namespace engine_interface + +#endif // MOCK_INFERENCE_ENGINE_HPP \ No newline at end of file diff --git a/engine_interface/test/test_engine_interface.cpp b/engine_interface/test/test_engine_interface.cpp new file mode 100644 index 0000000..7b3279b --- /dev/null +++ b/engine_interface/test/test_engine_interface.cpp @@ -0,0 +1,47 @@ +#include +#include "engine_interface_test_fixture.hpp" +#include "mock_inference_engine.hpp" + +TEST_F(EngineInterfaceTest, LoadModelSuccess) { + std::vector> dims = {{1, 3, 224, 224}}; + EXPECT_CALL(mock_engine, loadModel("model.onnx", dims, 4)) + .Times(1) + .WillOnce(Return(true)); + + bool result = mock_engine.loadModel("model.onnx", dims, 4); + EXPECT_TRUE(result); +} + +TEST_F(EngineInterfaceTest, LoadModelFailure) { + std::vector> dims = {{1, 3, 224, 224}}; + EXPECT_CALL(mock_engine, loadModel(_, _, _)) + .Times(1) + .WillOnce(Return(false)); + + bool result = mock_engine.loadModel("bad_model.onnx", dims, 4); + EXPECT_FALSE(result); +} + +TEST_F(EngineInterfaceTest, RunInference) { + std::vector input_tensors(1, nullptr); + std::vector input_sizes(1, 1024); + std::vector output_tensors(1, nullptr); + std::vector output_sizes(1, 1024); + + EXPECT_CALL(mock_engine, runInference(_, _, _, _)) + .Times(1) + .WillOnce([](const std::vector& inputs, const std::vector& input_sizes, + std::vector& outputs, const std::vector& output_sizes) { + // Simulate inference by populating output tensors + for (size_t i = 0; i < outputs.size(); ++i) { + outputs[i] = new int[output_sizes[i]]; // Allocate memory for output + } + }); + + mock_engine.runInference(input_tensors, input_sizes, output_tensors, output_sizes); + EXPECT_EQ(output_tensors.size(), 1); + EXPECT_NE(output_tensors[0], nullptr); + // Clean up allocated memory + delete[] static_cast(output_tensors[0]); + output_tensors[0] = nullptr; +} \ No newline at end of file diff --git a/mission_maestro_bridge/launch/mission_maestro_bridge.launch.py b/mission_maestro_bridge/launch/mission_maestro_bridge.launch.py new file mode 100644 index 0000000..0d2a7b9 --- /dev/null +++ b/mission_maestro_bridge/launch/mission_maestro_bridge.launch.py @@ -0,0 +1,25 @@ +from launch import LaunchDescription +from launch_ros.actions import Node + + +def generate_launch_description() -> LaunchDescription: + return LaunchDescription( + [ + Node( + package="mission_maestro_bridge", + executable="mission_maestro_bridge", + name="mission_maestro_bridge", + output="screen", + parameters=[ + { + "mission_goal_topic": "mission_goals", + "mission_yaml_topic": "mission_yaml", + "mission_command_topic": "mission_command", + "mission_yaml_service": "maestro_yaml", + "mission_command_service": "maestro_command", + "radius": 2.0, + } + ], + ) + ] + ) diff --git a/mission_maestro_bridge/mission_maestro_bridge/__init__.py b/mission_maestro_bridge/mission_maestro_bridge/__init__.py new file mode 100644 index 0000000..ac384db --- /dev/null +++ b/mission_maestro_bridge/mission_maestro_bridge/__init__.py @@ -0,0 +1 @@ +"""Mission Maestro bridge package.""" diff --git a/mission_maestro_bridge/mission_maestro_bridge/bridge_node.py b/mission_maestro_bridge/mission_maestro_bridge/bridge_node.py new file mode 100644 index 0000000..0a2db99 --- /dev/null +++ b/mission_maestro_bridge/mission_maestro_bridge/bridge_node.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from typing import Optional + +import rclpy +from geometry_msgs.msg import PoseArray +from rclpy.node import Node +from std_msgs.msg import String, UInt8 +import yaml + + +class MissionMaestroBridge(Node): + def __init__(self) -> None: + super().__init__("mission_maestro_bridge") + + self.declare_parameter("mission_goal_topic", "mission_goals") + self.declare_parameter("mission_yaml_topic", "mission_yaml") + self.declare_parameter("mission_command_topic", "mission_command") + self.declare_parameter("mission_yaml_service", "maestro_yaml") + self.declare_parameter("mission_command_service", "maestro_command") + self.declare_parameter("radius", 2.0) + + mission_goal_topic = self.get_parameter("mission_goal_topic").get_parameter_value().string_value + mission_yaml_topic = self.get_parameter("mission_yaml_topic").get_parameter_value().string_value + mission_command_topic = self.get_parameter("mission_command_topic").get_parameter_value().string_value + self._yaml_service_name = self.get_parameter("mission_yaml_service").get_parameter_value().string_value + self._command_service_name = self.get_parameter("mission_command_service").get_parameter_value().string_value + self._radius = self.get_parameter("radius").get_parameter_value().double_value + + self._maestro_available = False + self._maestro_yaml_srv = None + self._maestro_cmd_srv = None + self._yaml_client = None + self._command_client = None + + try: + from arl_mission_maestro.srv import MaestroCommand, MaestroMissionYaml + + self._maestro_yaml_srv = MaestroMissionYaml + self._maestro_cmd_srv = MaestroCommand + self._yaml_client = self.create_client(MaestroMissionYaml, self._yaml_service_name) + self._command_client = self.create_client(MaestroCommand, self._command_service_name) + self._maestro_available = True + self.get_logger().info( + f"Mission Maestro bridge enabled: services '{self._yaml_service_name}' and " + f"'{self._command_service_name}'" + ) + except (ImportError, ModuleNotFoundError) as exc: + self.get_logger().warn( + "arl_mission_maestro not available in this environment. " + f"Bridge will stay passive. Details: {exc}" + ) + + self.create_subscription(PoseArray, mission_goal_topic, self._on_pose_array, 10) + self.create_subscription(String, mission_yaml_topic, self._on_yaml, 10) + self.create_subscription(UInt8, mission_command_topic, self._on_command, 10) + + self.get_logger().info( + f"Listening on topics: goals='{mission_goal_topic}', yaml='{mission_yaml_topic}', " + f"command='{mission_command_topic}'" + ) + + def _wait_for_service(self, service_type: str) -> bool: + if not self._maestro_available: + return False + client = self._yaml_client if service_type == "yaml" else self._command_client + if client is None: + return False + if client.service_is_ready(): + return True + ok = client.wait_for_service(timeout_sec=0.5) + if not ok: + name = self._yaml_service_name if service_type == "yaml" else self._command_service_name + self.get_logger().warn(f"Mission Maestro service not ready: {name}") + return ok + + def _call_yaml_service(self, yaml_payload: str) -> None: + if not self._wait_for_service("yaml"): + return + req = self._maestro_yaml_srv.Request() + req.yaml_as_string = yaml_payload + self._yaml_client.call_async(req) + self.get_logger().info("Forwarded mission YAML to maestro_yaml service") + + def _call_command_service(self, command: int) -> None: + if not self._wait_for_service("command"): + return + req = self._maestro_cmd_srv.Request() + req.command = int(command) + self._command_client.call_async(req) + self.get_logger().info(f"Forwarded mission command to maestro_command service: {command}") + + def _build_yaml_from_goals(self, msg: PoseArray) -> str: + frame_id = msg.header.frame_id if msg.header.frame_id else "map" + waypoints = [] + for idx, pose in enumerate(msg.poses): + waypoints.append( + { + "name": f"waypoint{idx + 1}", + "pose": [float(pose.position.x), float(pose.position.y), float(pose.position.z)], + "radius": float(self._radius), + } + ) + + payload = { + "version": 2.0, + "frameid": frame_id, + "waypoints": waypoints, + } + return yaml.safe_dump(payload, sort_keys=False) + + def _on_pose_array(self, msg: PoseArray) -> None: + if not self._maestro_available: + return + yaml_payload = self._build_yaml_from_goals(msg) + self._call_yaml_service(yaml_payload) + + def _on_yaml(self, msg: String) -> None: + if not self._maestro_available: + return + if not msg.data: + return + self._call_yaml_service(msg.data) + + def _on_command(self, msg: UInt8) -> None: + if not self._maestro_available: + return + self._call_command_service(int(msg.data)) + + +def main(args: Optional[list[str]] = None) -> None: + rclpy.init(args=args) + node = MissionMaestroBridge() + try: + rclpy.spin(node) + finally: + node.destroy_node() + rclpy.shutdown() diff --git a/mission_maestro_bridge/package.xml b/mission_maestro_bridge/package.xml new file mode 100644 index 0000000..90c640a --- /dev/null +++ b/mission_maestro_bridge/package.xml @@ -0,0 +1,23 @@ + + + + mission_maestro_bridge + 0.1.0 + Bridge for NeuroMesh mission topics to arl_mission_maestro services. + Long Quang + MIT + + ament_python + + rclpy + geometry_msgs + std_msgs + python3-yaml + + ament_lint_auto + ament_lint_common + + + ament_python + + diff --git a/mission_maestro_bridge/resource/mission_maestro_bridge b/mission_maestro_bridge/resource/mission_maestro_bridge new file mode 100644 index 0000000..867981b --- /dev/null +++ b/mission_maestro_bridge/resource/mission_maestro_bridge @@ -0,0 +1 @@ +mission_maestro_bridge diff --git a/mission_maestro_bridge/setup.cfg b/mission_maestro_bridge/setup.cfg new file mode 100644 index 0000000..a56de40 --- /dev/null +++ b/mission_maestro_bridge/setup.cfg @@ -0,0 +1,5 @@ +[develop] +script_dir=$base/lib/mission_maestro_bridge + +[install] +install_scripts=$base/lib/mission_maestro_bridge diff --git a/mission_maestro_bridge/setup.py b/mission_maestro_bridge/setup.py new file mode 100644 index 0000000..6d52ede --- /dev/null +++ b/mission_maestro_bridge/setup.py @@ -0,0 +1,26 @@ +from setuptools import find_packages, setup + +package_name = "mission_maestro_bridge" + +setup( + name=package_name, + version="0.1.0", + packages=find_packages(exclude=["test"]), + data_files=[ + ("share/ament_index/resource_index/packages", [f"resource/{package_name}"]), + (f"share/{package_name}", ["package.xml"]), + (f"share/{package_name}/launch", ["launch/mission_maestro_bridge.launch.py"]), + ], + install_requires=["setuptools", "PyYAML"], + zip_safe=True, + maintainer="Long Quang", + maintainer_email="longquang@nyu.edu", + description="Bridge generic NeuroMesh mission topics to arl_mission_maestro services.", + license="MIT", + tests_require=["pytest"], + entry_points={ + "console_scripts": [ + "mission_maestro_bridge = mission_maestro_bridge.bridge_node:main", + ], + }, +) diff --git a/neuromesh_platform_r2/CMakeLists.txt b/neuromesh_platform_r2/CMakeLists.txt index 4e874d6..2558a90 100755 --- a/neuromesh_platform_r2/CMakeLists.txt +++ b/neuromesh_platform_r2/CMakeLists.txt @@ -10,56 +10,10 @@ if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") add_compile_options(-Wall -Wextra -Wpedantic) endif() -# find dependencies -find_package(ament_cmake REQUIRED) -find_package(rosidl_default_generators REQUIRED) -find_package(std_msgs REQUIRED) -find_package(geometry_msgs REQUIRED) -find_package(nav_msgs REQUIRED) -find_package(rclcpp REQUIRED) -find_package(sensor_msgs REQUIRED) -find_package(neuromesh_interfaces REQUIRED) -find_package(ZLIB) -find_package(OpenCV 4 REQUIRED) -find_package(cv_bridge REQUIRED) -find_package(rclcpp_components REQUIRED) -find_package(yaml-cpp REQUIRED) -find_package(tf2 REQUIRED) -find_package(tf2_ros REQUIRED) -find_package(tf2_sensor_msgs REQUIRED) -find_package(Eigen3 REQUIRED) -find_package(tf2_geometry_msgs REQUIRED) -find_package(geometry_msgs REQUIRED) -find_package(visualization_msgs REQUIRED) -find_package(interactive_markers REQUIRED) -find_package(rviz_rendering REQUIRED) -find_package(rviz_common REQUIRED) -find_package(std_srvs REQUIRED) -find_package(rclcpp_components REQUIRED) -find_package(PCL REQUIRED) -find_package(pcl_conversions REQUIRED) -find_package(message_filters REQUIRED) -find_package(realsense2_camera_msgs REQUIRED) - -include_directories(include) -include_directories( - include ${EIGEN3_INCLUDE_DIR} ${PCL_INCLUDE_DIRS}) - -add_executable(visualization_node - src/visualization_node.cpp -) - -add_executable(decompression_node - src/decompression_node.cpp -) - -add_executable(compression_node - src/compression_node.cpp -) - set (dependencies "std_msgs" "rclcpp" + "ament_index_cpp" "sensor_msgs" "neuromesh_interfaces" "rclcpp_components" @@ -72,6 +26,7 @@ set (dependencies "nav_msgs" "Eigen3" "tf2_geometry_msgs" + "tf2_sensor_msgs" "visualization_msgs" "interactive_markers" "rviz_rendering" @@ -79,8 +34,36 @@ set (dependencies "std_srvs" "pcl_conversions" "tf2" +"message_filters" ) + +# find dependencies +find_package(ament_cmake REQUIRED) +find_package(rosidl_default_generators REQUIRED) +find_package(ZLIB) +find_package(OpenCV 4 REQUIRED) +find_package(PCL REQUIRED) +foreach(dep ${dependencies}) + find_package(${dep} REQUIRED) +endforeach() + +include_directories(include) +include_directories( + include ${EIGEN3_INCLUDE_DIR} ${PCL_INCLUDE_DIRS}) + +add_executable(visualization_node + src/visualization_node.cpp +) + +add_executable(decompression_node + src/decompression_node.cpp +) + +add_executable(compression_node + src/compression_node.cpp +) + # Modified Codes add_library(dust3r_example SHARED src/dust3r_toy_implementation_batch.cpp @@ -97,22 +80,6 @@ ament_target_dependencies(dust3r_example target_link_libraries(dust3r_example ${YAML_CPP_LIBRARIES}) rclcpp_components_register_nodes(dust3r_example "neuromeshNode::ToyImplementation") -# VGGT Codes -add_library(vggt_example SHARED - src/vggt_toy_implementation.cpp - src/vggt_neuromesh_node.cpp) - -set_target_properties(vggt_example PROPERTIES - COMPILE_DEFINITIONS "COMPOSITION_BUILDING_DLL" -) - - ament_target_dependencies(vggt_example - ${dependencies} - ) - -target_link_libraries(vggt_example ${YAML_CPP_LIBRARIES}) -rclcpp_components_register_nodes(vggt_example "vggtNode::VggtToyImplementation") - # VGGT Separated Nodes add_library(vggt_separated SHARED src/vggt_encoder_node.cpp @@ -148,6 +115,23 @@ target_link_libraries(control_implementation ${YAML_CPP_LIBRARIES}) rclcpp_components_register_nodes(control_implementation "ControlneuromeshNode::ControlImplementation") +# GAT Codes +add_library(gat_example SHARED + src/gat_model_implementation.cpp + src/gat_model_neuromesh_node.cpp + src/gat_planner_model_implementation.cpp + src/gat_planner_model_neuromesh_node.cpp) +set_target_properties(gat_example PROPERTIES + COMPILE_DEFINITIONS "COMPOSITION_BUILDING_DLL" +) +ament_target_dependencies(gat_example + ${dependencies} + ) +target_link_libraries(gat_example ${YAML_CPP_LIBRARIES}) +rclcpp_components_register_nodes(gat_example + "GATneuromeshNode::GATImplementation" + "GATPlannerNeuromeshNode::GATPlannerImplementation") + # odom_republisher add_library(odom_republisher SHARED src/odom_republisher.cpp) @@ -165,6 +149,24 @@ ament_target_dependencies(odom_republisher rclcpp_components_register_nodes(odom_republisher "odom_republisher::OdomRepublisher") +add_library(starting_poses_sender SHARED + src/starting_poses_sender.cpp) +set_target_properties(starting_poses_sender PROPERTIES + COMPILE_DEFINITIONS "COMPOSITION_BUILDING_DLL" +) +ament_target_dependencies(starting_poses_sender + "ament_index_cpp" + "geometry_msgs" + "nav_msgs" + "rclcpp" + "rclcpp_components" + "std_msgs" + "tf2" + "tf2_ros" + "yaml-cpp" +) +rclcpp_components_register_nodes(starting_poses_sender "goal_sender::GoalSenderNode") + # Add include folder target_include_directories(visualization_node PUBLIC "include/") target_include_directories(decompression_node PUBLIC "include/") @@ -211,7 +213,7 @@ ament_target_dependencies(depth_completion_node "PCL" "pcl_conversions" "message_filters" - "realsense2_camera_msgs" + # "realsense2_camera_msgs" ) target_link_libraries(depth_completion_node @@ -234,7 +236,8 @@ endif() install(TARGETS odom_republisher dust3r_example - vggt_example + gat_example + starting_poses_sender vggt_separated #control_implementation ARCHIVE DESTINATION lib @@ -253,7 +256,8 @@ install(FILES launch/full_except_bridge.py launch/dust3r_model_neuromesh_launch.py launch/vggt_model_neuromesh_launch.py - launch/vggt_separated_launch.py + launch/gat_model_neuromesh_launch.py + launch/gat_planner_model_neuromesh_launch.py launch/depth_completion_launch.py DESTINATION share/${PROJECT_NAME}/launch ) @@ -264,9 +268,12 @@ install(DIRECTORY config install(PROGRAMS scripts/dust3r_model_neuromesh_launch.sh scripts/vggt_model_neuromesh_launch.sh - scripts/vggt_separated_launch.sh DESTINATION share/${PROJECT_NAME}) +install(DIRECTORY include/ + DESTINATION include) + ament_export_dependencies(rosidl_default_runtime) -ament_export_libraries(vggt_separated dust3r_example vggt_example odom_republisher) +ament_export_include_directories(include) +ament_export_libraries(vggt_separated dust3r_example odom_republisher starting_poses_sender) ament_package() diff --git a/neuromesh_platform_r2/config/gat_planner_goal1_pos.yaml b/neuromesh_platform_r2/config/gat_planner_goal1_pos.yaml new file mode 100644 index 0000000..ae3620f --- /dev/null +++ b/neuromesh_platform_r2/config/gat_planner_goal1_pos.yaml @@ -0,0 +1,23 @@ +voxl2_1: + goals: + - position: + x: -0.8 + y: -1.35 + +voxl2_2: + goals: + - position: + x: -0.8 + y: 0.8 + +voxl2_5: + goals: + - position: + x: 1.6 + y: -1.35 + +voxl2_9: + goals: + - position: + x: 1.6 + y: 0.8 diff --git a/neuromesh_platform_r2/config/gat_planner_start_pos.yaml b/neuromesh_platform_r2/config/gat_planner_start_pos.yaml new file mode 100644 index 0000000..89722ee --- /dev/null +++ b/neuromesh_platform_r2/config/gat_planner_start_pos.yaml @@ -0,0 +1,23 @@ +voxl2_1: + goals: + - position: + x: -1.0 + y: -0.65 + +voxl2_2: + goals: + - position: + x: 0.12 + y: 0.65 + +voxl2_5: + goals: + - position: + x: 1.25 + y: -0.65 + +voxl2_9: + goals: + - position: + x: 2.32 + y: 0.65 diff --git a/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_implementation.h b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_implementation.h new file mode 100644 index 0000000..9d03fb0 --- /dev/null +++ b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_implementation.h @@ -0,0 +1,42 @@ +#ifndef GAT_IMPLEMENTATION_HEADER_H +#define GAT_IMPLEMENTATION_HEADER_H + +#include "neuromesh_interfaces/srv/tensor_request.hpp" +#include "neuromesh_platform_r2/gat_neuromesh_node.h" +#include "rclcpp/rclcpp.hpp" +#include +#include +#include + +namespace GATneuromeshNode { +class GATImplementation : public GATneuromeshNode { +public: + GATImplementation(const rclcpp::NodeOptions &options); + +protected: + // Perform inference + std::future>> + performInference( + const std::string &model_name, + const std::vector &tensors); + + // Build decoder + bool buildDecoderTensor( + std::map + &agent_features, + neuromesh_interfaces::msg::Tensor &own_feature, + neuromesh_interfaces::msg::Tensor &aggregated_tensor); + + rclcpp::Subscription::SharedPtr + feature_subscription_; + rclcpp::Publisher::SharedPtr + feature_publisher_; + rclcpp::Client::SharedPtr + tensor_client_; + rclcpp::Subscription::SharedPtr pos_sub_; + + std::map goal_poses_; + std::map current_states_; +}; +} // namespace GATneuromeshNode +#endif diff --git a/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_neuromesh_node.h b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_neuromesh_node.h new file mode 100644 index 0000000..6c4c199 --- /dev/null +++ b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_neuromesh_node.h @@ -0,0 +1,210 @@ +#ifndef GAT_neuromesh_NODE_HEADER_H +#define GAT_neuromesh_NODE_HEADER_H + +#include "rclcpp/rclcpp.hpp" +#include "geometry_msgs/msg/pose.hpp" +#include "nav_msgs/msg/odometry.hpp" +#include "neuromesh_interfaces/msg/comm_message.hpp" +#include "neuromesh_interfaces/msg/feature.hpp" +#include "neuromesh_interfaces/msg/state_vector.hpp" +#include "neuromesh_interfaces/msg/tensor.hpp" +#include "sensor_msgs/image_encodings.hpp" +#include "sensor_msgs/msg/image.hpp" +#include "std_msgs/msg/string.hpp" +#include "yaml-cpp/yaml.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "service_adapters/adapter_factory.hpp" +#include "types/navigation_goal.hpp" + +namespace GATneuromeshNode { +class GATneuromeshNode : public rclcpp::Node { + // FUNCTIONS +public: + // Constructor + GATneuromeshNode(const rclcpp::NodeOptions &options); + +protected: + void + feature_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg); + + virtual std::future< + std::vector>> + performInference( + const std::string &model_name, + const std::vector &tensors); + + // Convert tensor of features to a feature message + neuromesh_interfaces::msg::Feature + buildFeatureMessage(const neuromesh_interfaces::msg::Tensor &tensor); + + // Aggregates features from available agents (and itself) into a single tensor + // PLACEHOLDER + virtual bool buildDecoderTensor( + std::map + &agent_features, + neuromesh_interfaces::msg::Tensor &own_feature, + neuromesh_interfaces::msg::Tensor &aggregated_tensor); + + // Adds a subscription to the feature_subscriptions_ map + void createSubscription( + std::map::SharedPtr> + &subscription_map, + std::string id, rclcpp::QoS qos); + + // Remove subscription form the feature_subscriptions_ map + void removeSubscription( + std::map::SharedPtr> + subscription_map, + std::string id); + + // Adds a subscription to the gnn_subscriptions_ map + void createGNNSubscription( + std::map::SharedPtr> + &gnn_subscription_map, + std::string id, rclcpp::QoS qos); + + // Remove subscription form the gnn_subscriptions_ map + void removeGNNSubscription( + std::map::SharedPtr> + gnn_subscription_map, + std::string id); + + // Convert string to ROS2 QoS profile + rmw_qos_profile_t parseQoSString(const std::string &str); + + // Split agent string parameter into vector of agent ids + std::set splitAgentString(std::string str); + + // Parameter handling + void load_goal_poses_from_yaml(); + + // Calculate distances to goal poses + neuromesh_interfaces::msg::Tensor + calculate_goal_distances(const std::vector &state_vector); + + // Set encoder status as to whether or not it's already been run this cycle + void run_encoder_cycle(); + + // Methods for GNN result handling + void run_decoder_cycle(); + neuromesh_interfaces::msg::Feature + build_gnn_msg(const neuromesh_interfaces::msg::Tensor &gnn_result); + void + gnn_result_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg); + void prepare_second_stage_decoding(); + + // transpose tensor + // works for ints + neuromesh_interfaces::msg::Tensor + convert_to_nchw(const neuromesh_interfaces::msg::Tensor &input); + + // measure time + void startClock(std::string phase); + void stopClock(std::string phase); + int64_t checkClock(std::string phase); + std::map> + times; // int = milliseconds, bool = has timer stopped. if bool=false, + // then int=starting time + + // VARIABLES + + // Lists of agent ids that describe 1) all agents 2) the available ones + std::set all_agents; + std::set available_agents; + + // publishers and subscriptions + rclcpp::Publisher::SharedPtr + feature_publisher_; + + std::map::SharedPtr> + feature_subscriptions_; + std::map::SharedPtr> + gnn_subscriptions_; + + // repeating function to keep track of cycles + rclcpp::TimerBase::SharedPtr decoder_timer_; // process features directly + rclcpp::TimerBase::SharedPtr + encoder_timer_; // update bool to be ready for encoder to run + + bool fresh_encoder_cycle; // if true encoder is ready to run. + bool waypoint_cmd_sent_; // tracking sending waypoints only once + double goals_sending_delay_; // Delay sending goals to robots + std::future>> + encoder_result; + std::future>> + gnn_result_future; + + // Publishers and subscribers for GNN results + rclcpp::Publisher::SharedPtr + gnn_result_publisher_; + rclcpp::Subscription::SharedPtr + gnn_result_subscriber_; + rclcpp::Publisher::SharedPtr + second_decoder_result_publisher_; + + // variables for features + std::map + feature_buffer_; // To store all features + std::map + feature_buffer_timestamp_; // To store timestamps of all features + + void pos_callback(const nav_msgs::msg::Odometry::SharedPtr msg); + + Eigen::Vector3f quaternion_to_euler(const geometry_msgs::msg::Quaternion &q); + + // parameters + std::string encoder_model_name_; + std::string decoder_model1_name_; + std::string decoder_model2_name_; + std::string topic_prefix_; + std::string output_topic_; + int decoder_cycle_length_; + int encoder_cycle_length_; + int encoder_await_length_; + std::string id_; + std::string image_qos_profile_; + std::string features_qos_profile_; + std::string output_qos_profile_; + std::string agents_; + bool to_nchw_; + bool pos_callback_complete = false; + neuromesh_interfaces::msg::Tensor encoder_output_tensor; + std::string goal_poses_yaml_file; + std::string planning_frame_; + std::string service_adapter_type_; + + std::map current_states_; + std::map + received_features_; + + std::vector goal_poses_; + + std::map + received_gnn_results_; + + std::future>> + second_decoder_result_future; + + std::shared_ptr goal_adapter_; + + std::shared_ptr tf_buffer_; + std::shared_ptr tf_listener_; +}; +} // namespace GATneuromeshNode + +#endif diff --git a/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_planner_implementation.h b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_planner_implementation.h new file mode 100644 index 0000000..da07236 --- /dev/null +++ b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_planner_implementation.h @@ -0,0 +1,42 @@ +#ifndef GAT_PLANNER_IMPLEMENTATION_HEADER_H +#define GAT_PLANNER_IMPLEMENTATION_HEADER_H + +#include "neuromesh_interfaces/srv/tensor_request.hpp" +#include "neuromesh_platform_r2/gat_planner_neuromesh_node.h" +#include "rclcpp/rclcpp.hpp" +#include +#include +#include + +namespace GATPlannerNeuromeshNode { +class GATPlannerImplementation : public GATPlannerNeuromeshNode { +public: + GATPlannerImplementation(const rclcpp::NodeOptions &options); + +protected: + // Perform inference + std::future>> + performInference( + const std::string &model_name, + const std::vector &tensors); + + // Build decoder + bool buildDecoderTensor( + std::map + &agent_features, + neuromesh_interfaces::msg::Tensor &own_feature, + neuromesh_interfaces::msg::Tensor &aggregated_tensor); + + rclcpp::Subscription::SharedPtr + feature_subscription_; + rclcpp::Publisher::SharedPtr + feature_publisher_; + rclcpp::Client::SharedPtr + tensor_client_; + rclcpp::Subscription::SharedPtr pos_sub_; + + std::map goal_poses_; + std::map current_states_; +}; +} // namespace GATPlannerNeuromeshNode +#endif diff --git a/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_planner_neuromesh_node.h b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_planner_neuromesh_node.h new file mode 100644 index 0000000..8c94cee --- /dev/null +++ b/neuromesh_platform_r2/include/neuromesh_platform_r2/gat_planner_neuromesh_node.h @@ -0,0 +1,224 @@ +#ifndef GAT_planner_neuromesh_NODE_HEADER_H +#define GAT_planner_neuromesh_NODE_HEADER_H + +#include "rclcpp/rclcpp.hpp" + +#include "geometry_msgs/msg/pose.hpp" +#include "nav_msgs/msg/odometry.hpp" +#include "neuromesh_interfaces/msg/comm_message.hpp" +#include "neuromesh_interfaces/msg/feature.hpp" +#include "neuromesh_interfaces/msg/state_vector.hpp" +#include "neuromesh_interfaces/msg/tensor.hpp" +#include "sensor_msgs/image_encodings.hpp" +#include "sensor_msgs/msg/image.hpp" +#include "std_msgs/msg/string.hpp" +#include "yaml-cpp/yaml.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "service_adapters/adapter_factory.hpp" +#include "types/navigation_goal.hpp" + +namespace GATPlannerNeuromeshNode { +class GATPlannerNeuromeshNode : public rclcpp::Node { + // FUNCTIONS +public: + // Constructor + GATPlannerNeuromeshNode(const rclcpp::NodeOptions &options); + +protected: + void + feature_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg); + + virtual std::future< + std::vector>> + performInference( + const std::string &model_name, + const std::vector &tensors); + + // Convert tensor of features to a feature message + neuromesh_interfaces::msg::Feature + buildFeatureMessage(const neuromesh_interfaces::msg::Tensor &tensor); + + // Aggregates features from available agents (and itself) into a single tensor + // PLACEHOLDER + virtual bool buildDecoderTensor( + std::map + &agent_features, + neuromesh_interfaces::msg::Tensor &own_feature, + neuromesh_interfaces::msg::Tensor &aggregated_tensor); + + // Adds a subscription to the feature_subscriptions_ map + void createSubscription( + std::map::SharedPtr> + &subscription_map, + std::string id, rclcpp::QoS qos); + + void createPosSubscription( + std::map::SharedPtr> + &pos_subscription_map, + std::string id, rclcpp::QoS qos); + + // Remove subscription form the feature_subscriptions_ map + void removeSubscription( + std::map::SharedPtr> + subscription_map, + std::string id); + + // Adds a subscription to the gnn_subscriptions_ map + void createGNNSubscription( + std::map::SharedPtr> + &gnn_subscription_map, + std::string id, rclcpp::QoS qos); + + // Remove subscription form the gnn_subscriptions_ map + void removeGNNSubscription( + std::map::SharedPtr> + gnn_subscription_map, + std::string id); + + // Convert string to ROS2 QoS profile + rmw_qos_profile_t parseQoSString(const std::string &str); + + // Split agent string parameter into vector of agent ids + std::set splitAgentString(std::string str); + + // Parameter handling + void load_goal_poses_from_yaml(); + + // Calculate input features + neuromesh_interfaces::msg::Tensor + calculate_input_features(const std::vector &state_vector, const std::vector neighbors_ids); + std::vector get_closest_neighbors(); + + // Set encoder status as to whether or not it's already been run this cycle + void run_encoder_cycle(); + + // Methods for GNN result handling + void run_decoder_cycle(); + neuromesh_interfaces::msg::Feature + build_gnn_msg(const neuromesh_interfaces::msg::Tensor &gnn_result); + void + gnn_result_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg); + void prepare_second_stage_decoding(); + + // transpose tensor + // works for ints + neuromesh_interfaces::msg::Tensor + convert_to_nchw(const neuromesh_interfaces::msg::Tensor &input); + + // measure time + void startClock(std::string phase); + void stopClock(std::string phase); + int64_t checkClock(std::string phase); + std::map> + times; // int = milliseconds, bool = has timer stopped. if bool=false, + // then int=starting time + + // VARIABLES + + // Lists of agent ids that describe 1) all agents 2) the available ones + std::set all_agents; + std::set available_agents; + + // publishers and subscriptions + rclcpp::Publisher::SharedPtr + feature_publisher_; + + std::map::SharedPtr> + feature_subscriptions_; + std::map::SharedPtr> + pos_subscriptions_; + std::map::SharedPtr> + gnn_subscriptions_; + + // repeating function to keep track of cycles + rclcpp::TimerBase::SharedPtr decoder_timer_; // process features directly + rclcpp::TimerBase::SharedPtr + encoder_timer_; // update bool to be ready for encoder to run + + bool fresh_encoder_cycle; // if true encoder is ready to run. + bool waypoint_cmd_sent_; // tracking sending waypoints only once + double goals_sending_delay_; // Delay sending goals to robots + std::future>> + encoder_result; + std::future>> + gnn_result_future; + + // Publishers and subscribers for GNN results + rclcpp::Publisher::SharedPtr + gnn_result_publisher_; + rclcpp::Subscription::SharedPtr + gnn_result_subscriber_; + rclcpp::Publisher::SharedPtr + second_decoder_result_publisher_; + + std::string service_adapter_type_; + std::shared_ptr goal_adapter_; + + // variables for features + std::map + feature_buffer_; // To store all features + std::map + feature_buffer_timestamp_; // To store timestamps of all features + + void pos_callback(const nav_msgs::msg::Odometry::SharedPtr msg); + + void neighbor_pos_callback(const nav_msgs::msg::Odometry::SharedPtr msg, const std::string &id); + + Eigen::Vector3f quaternion_to_euler(const geometry_msgs::msg::Quaternion &q); + + // parameters + std::string encoder_model_name_; + std::string decoder_model1_name_; + std::string decoder_model2_name_; + std::string topic_prefix_; + std::string pos_topic_prefix_; + std::string output_topic_; + int decoder_cycle_length_; + int encoder_cycle_length_; + int encoder_await_length_; + std::string id_; + std::string image_qos_profile_; + std::string features_qos_profile_; + std::string output_qos_profile_; + std::string agents_; + bool to_nchw_; + bool pos_callback_complete = false; + neuromesh_interfaces::msg::Tensor encoder_output_tensor; + std::string goal_poses_yaml_file; + std::string planning_frame_; + + std::map current_states_; + std::map + received_features_; + + std::vector goal_poses_; + + std::map + received_gnn_results_; + + std::future>> + second_decoder_result_future; + + std::shared_ptr tf_buffer_; + std::shared_ptr tf_listener_; +}; +} // namespace GATPlannerNeuromeshNode + +#endif diff --git a/neuromesh_platform_r2/include/service_adapters/adapter_factory.hpp b/neuromesh_platform_r2/include/service_adapters/adapter_factory.hpp new file mode 100644 index 0000000..f4f356e --- /dev/null +++ b/neuromesh_platform_r2/include/service_adapters/adapter_factory.hpp @@ -0,0 +1,124 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include "types/navigation_goal.hpp" + +class ServiceInterface { +public: + virtual ~ServiceInterface() = default; + virtual bool sendGoal(const NavigationGoal & goal) = 0; + virtual std::string type() const = 0; +}; + +class NullServiceAdapter : public ServiceInterface { +public: + explicit NullServiceAdapter(rclcpp::Logger logger) + : logger_(logger) + { + } + + bool sendGoal(const NavigationGoal & goal) override + { + RCLCPP_WARN( + logger_, + "Navigation adapter is disabled; dropping goal in frame '%s' at (%.3f, %.3f, %.3f)", + goal.planning_frame.c_str(), goal.x, goal.y, goal.z); + return false; + } + + std::string type() const override + { + return "none"; + } + +private: + rclcpp::Logger logger_; +}; + +class PoseStampedTopicAdapter : public ServiceInterface { +public: + explicit PoseStampedTopicAdapter(const rclcpp::Node::SharedPtr & node) + : node_(node) + { + if (!node_->has_parameter("navigation_goal_topic")) { + node_->declare_parameter("navigation_goal_topic", "goal_pose"); + } + node_->get_parameter("navigation_goal_topic", goal_topic_); + publisher_ = node_->create_publisher(goal_topic_, 10); + } + + bool sendGoal(const NavigationGoal & goal) override + { + geometry_msgs::msg::PoseStamped msg; + msg.header.stamp = node_->now(); + msg.header.frame_id = goal.planning_frame; + msg.pose.position.x = goal.x; + msg.pose.position.y = goal.y; + msg.pose.position.z = goal.z; + msg.pose.orientation.w = 1.0; + + publisher_->publish(msg); + RCLCPP_INFO( + node_->get_logger(), + "Published navigation goal on '%s' in frame '%s' at (%.3f, %.3f, %.3f)", + goal_topic_.c_str(), goal.planning_frame.c_str(), goal.x, goal.y, goal.z); + return true; + } + + std::string type() const override + { + return "topic"; + } + +private: + rclcpp::Node::SharedPtr node_; + std::string goal_topic_; + rclcpp::Publisher::SharedPtr publisher_; +}; + +class AdapterFactory { +public: + static std::shared_ptr + create(const std::string & adapter_type, const rclcpp::Node::SharedPtr & node) + { + std::string normalized = adapter_type; + std::transform( + normalized.begin(), normalized.end(), normalized.begin(), + [](unsigned char c) {return static_cast(std::tolower(c));}); + + if ( + normalized.empty() || normalized == "none" || normalized == "noop" || + normalized == "null") + { + return std::make_shared(node->get_logger()); + } + + if ( + normalized == "topic" || normalized == "pose_topic" || + normalized == "posestamped" || normalized == "missionmaestro" || + normalized == "mavmanager") + { + if (normalized == "missionmaestro" || normalized == "mavmanager") { + RCLCPP_WARN( + node->get_logger(), + "Adapter type '%s' is treated as the generic topic adapter. " + "Use 'topic' and bridge it externally if you need a specific navigation backend.", + adapter_type.c_str()); + } + return std::make_shared(node); + } + + RCLCPP_WARN( + node->get_logger(), + "Unknown navigation adapter '%s'; using no-op adapter instead.", + adapter_type.c_str()); + return std::make_shared(node->get_logger()); + } +}; \ No newline at end of file diff --git a/neuromesh_platform_r2/include/types/navigation_goal.hpp b/neuromesh_platform_r2/include/types/navigation_goal.hpp new file mode 100644 index 0000000..4919cdf --- /dev/null +++ b/neuromesh_platform_r2/include/types/navigation_goal.hpp @@ -0,0 +1,11 @@ +#pragma once + +#include + +struct NavigationGoal { + double x{0.0}; + double y{0.0}; + double z{0.0}; + std::string planning_frame{"map"}; + std::string robot_id; +}; \ No newline at end of file diff --git a/neuromesh_platform_r2/launch/gat_model_neuromesh_launch.py b/neuromesh_platform_r2/launch/gat_model_neuromesh_launch.py new file mode 100755 index 0000000..f985ba9 --- /dev/null +++ b/neuromesh_platform_r2/launch/gat_model_neuromesh_launch.py @@ -0,0 +1,388 @@ +import os +from ast import literal_eval + +import yaml +from ament_index_python.packages import get_package_share_directory +from launch import LaunchDescription +from launch.actions import DeclareLaunchArgument, OpaqueFunction +from launch.conditions import IfCondition +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import ComposableNodeContainer, Node +from launch_ros.descriptions import ComposableNode + + +def launch_setup(context): + agent_list = LaunchConfiguration("agent_list").perform(context) + name = LaunchConfiguration("name").perform(context) + goal_file = LaunchConfiguration("goal_file").perform(context) + start_file = LaunchConfiguration("start_file").perform(context) + start_position = LaunchConfiguration("start_position") + use_mission_maestro_bridge = LaunchConfiguration("use_mission_maestro_bridge") + odom_republisher = LaunchConfiguration("odom_republisher") + publish_static_map = LaunchConfiguration("publish_static_map").perform(context) + planning_frame = LaunchConfiguration("planning_frame") + engine_plugin_package = LaunchConfiguration("engine_plugin_package").perform(context) + engine_type = LaunchConfiguration("engine_type").perform(context) + + launch_list = [] + + # If using static map read in start location from yaml, publish map transform + start_x = 0.0 + start_y = 0.0 + complete_agent_list = [] + with open(start_file) as f: + start_positions = yaml.safe_load(f) + start_x = start_positions[name]["goals"][0]["position"]["x"] + start_y = start_positions[name]["goals"][0]["position"]["y"] + for k, v in start_positions.items(): + complete_agent_list.append(k) + + print(f"{name} start positions, x: {start_x}, y: {start_y}") + + if (engine_plugin_package == "tensorrt_engine") and (engine_type == "engine_interface::TRTEngine"): + model = "trt" + elif (engine_plugin_package == "onnx_engine") and ( + engine_type == "engine_interface::ONNXEngine" + ): + model = "onnx" + else: + raise ValueError( + f"Invalid engine_plugin_package {engine_plugin_package} or engine_type {engine_type}" + ) + + tf2_static_pub = Node( + package="tf2_ros", + executable="static_transform_publisher", + name="map_tf2_static_pub", + arguments=[ + "--x", + str(start_x), + "--y", + str(start_y), + "--z", + "0", + "--yaw", + "0", + "--pitch", + "0", + "--roll", + "0", + "--frame-id", + "map", + "--child-frame-id", + name + "/map", + ], + condition=IfCondition(publish_static_map), + ) + launch_list.append(tf2_static_pub) + + launch_list.append( + Node( + package="mission_maestro_bridge", + executable="mission_maestro_bridge", + namespace=name, + name="mission_maestro_bridge", + parameters=[ + { + "mission_goal_topic": "mission_goals", + "mission_yaml_topic": "mission_yaml", + "mission_command_topic": "mission_command", + "mission_yaml_service": "maestro_yaml", + "mission_command_service": "maestro_command", + } + ], + condition=IfCondition(use_mission_maestro_bridge), + output="screen", + ) + ) + + # LaunchConfiguration('agent_list') was serialized into string + if isinstance(agent_list, str): + try: + agent_list = literal_eval(agent_list) + except (ValueError, SyntaxError): + agent_list = agent_list.split(",") + + remappings = [] + for i, agent in enumerate(complete_agent_list): + remappings.append((f"/{name}/features_{agent}", f"/{agent}/features_{agent}")) + remappings.append( + (f"/{name}/gnn_output_{agent}", f"/{agent}/gnn_output_{agent}") + ) + + remappings.append(("position_topic", f"/{name}/odometry/local")) + + composable_nodes = [] + composable_nodes.append( + ComposableNode( + package="neuromesh_platform_r2", + plugin="goal_sender::GoalSenderNode", + name="starting_poses_sender", + namespace=name, + parameters=[ + { + "robot_id": name, + "start_poses_yaml_file": start_file, + } + ], + condition=IfCondition(start_position), + ) + ) + composable_nodes.append( + ComposableNode( + package="engine_interface", + namespace=name, + name=["engine", LaunchConfiguration("agent_num")], + plugin="engine_interface::EngineInterfaceNode", + parameters=[ + { + "engine_type": engine_type, + "model_names": "encoder,decoder1,decoder2", + "encoder.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat/encoder." + model, + "encoder.input_dimensions": "1,1,5", + "encoder.output_dimensions": "1,1,16", + "encoder.tensor_type": "fp32", + "decoder1.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat/gat_layer1." + model, + "decoder1.input_dimensions": "1,1,16;1,4,16;1,1,16", + "decoder1.output_dimensions": "1,1,16", + "decoder1.tensor_type": "fp32", + "decoder2.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat/gat_layer2." + model, + "decoder2.input_dimensions": "1,16;1,4,16;1,1,16", + "decoder2.output_dimensions": "1,1,5", + "decoder2.tensor_type": "fp32", + } + ], + ) + ) + composable_nodes.append( + ComposableNode( + package="neuromesh_platform_r2", + namespace=name, + name=["neuromesh"], + plugin="GATneuromeshNode::GATImplementation", + parameters=[ + { + "id": name, + "encoder_model_name": "encoder", + "decoder_model1_name": "decoder1", + "decoder_model2_name": "decoder2", + "encoder_cycle_length": 3000, + "decoder_cycle_length": 3000, + "output_topic": "gnn_output_", + "agents": ",".join(complete_agent_list), + "goal_poses_yaml_file": goal_file, + "planning_frame": planning_frame, + "goals_sending_delay": 2.0, + } + ], + remappings=remappings, + ) + ) + + # If the number of agents is less than complete list, add the odom_republisher nodes + if len(complete_agent_list) > len(agent_list) and LaunchConfiguration("odom_republisher").perform(context) == "True": + print(f"Number of agents is less than {len(complete_agent_list)}. Adding odom_republisher nodes!") + missing_agents = set(complete_agent_list) - set(agent_list) + for missing_agent in missing_agents: + remappings = [] + for i, agent in enumerate(complete_agent_list): + # print(f"remapping {agent} to {missing_agent}") + remappings.append((f"/{missing_agent}/features_{agent}", f"/{agent}/features_{agent}")) + remappings.append( + (f"/{missing_agent}/gnn_output_{agent}", f"/{agent}/gnn_output_{agent}") + ) + + remappings.append(("position_topic", f"/{name}/odometry/local")) + composable_nodes.append( + ComposableNode( + package="neuromesh_platform_r2", + plugin="odom_republisher::OdomRepublisher", + name=f"odom_republisher_{missing_agent}", + remappings=[ + ("original/odom", f"/{name}/odometry/local"), + ("republished/odom", f"/{missing_agent}/odometry/local"), + ], + condition=IfCondition(odom_republisher), + ) + ) + print( + f"Added odom_republisher for {missing_agent} using {name}'s odometry" + ) + composable_nodes.append( + ComposableNode( + package="engine_interface", + namespace=missing_agent, + name=["engine", LaunchConfiguration("agent_num")], + plugin="engine_interface::EngineInterfaceNode", + parameters=[ + { + "engine_type": engine_type, + "model_names": "encoder,decoder1,decoder2", + "encoder.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat/encoder." + model, + "encoder.input_dimensions": "1,1,5", + "encoder.output_dimensions": "1,1,16", + "encoder.tensor_type": "fp32", + "decoder1.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat/gat_layer1." + model, + "decoder1.input_dimensions": "1,1,16;1,4,16;1,1,16", + "decoder1.output_dimensions": "1,1,16", + "decoder1.tensor_type": "fp32", + "decoder2.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat/gat_layer2." + model, + "decoder2.input_dimensions": "1,16;1,4,16;1,1,16", + "decoder2.output_dimensions": "1,1,5", + "decoder2.tensor_type": "fp32", + } + ], + condition=IfCondition(odom_republisher), + ) + ) + composable_nodes.append( + ComposableNode( + package="neuromesh_platform_r2", + namespace=missing_agent, + name=["neuromesh"], + plugin="GATneuromeshNode::GATImplementation", + parameters=[ + { + "id": missing_agent, + "encoder_model_name": "encoder", + "decoder_model1_name": "decoder1", + "decoder_model2_name": "decoder2", + "encoder_cycle_length": 3000, + "decoder_cycle_length": 3000, + "output_topic": "gnn_output_", + "agents": ",".join(complete_agent_list), + "goal_poses_yaml_file": goal_file, + "planning_frame": planning_frame, + "goals_sending_delay": 2.0, + } + ], + remappings=remappings, + condition=IfCondition(odom_republisher), + ) + ) + + + launch_list.append( + ComposableNodeContainer( + name="neuromesh_container", + namespace=name, + package="rclcpp_components", + executable="component_container", + composable_node_descriptions=composable_nodes, + # prefix="xterm -e gdb --args", + # arguments=[ + # "--ros-args", + # "--log-level", + # "DEBUG"], + output="screen", + ) + ) + return launch_list + + +def generate_launch_description(): + name_arg = DeclareLaunchArgument( + name="name", default_value="warty", description=("Which agent we are running") + ) + + planning_frame_arg = DeclareLaunchArgument( + name="planning_frame", + default_value="map", + description=("Which frame to send goals in"), + ) + + agent_num_arg = DeclareLaunchArgument( + name="agent_num", default_value="1", description=("Which agent we are running") + ) + + agent_list_arg = DeclareLaunchArgument( + name="agent_list", + default_value="khonsu,osiris,ra,sekhmet,anubis", + description=("List of all agents present (including self)"), + ) + + publish_static_map_arg = DeclareLaunchArgument( + name="publish_static_map", + default_value="False", + description="run publish static map", + ) + + goal_file_arg = DeclareLaunchArgument( + name="goal_file", + default_value=os.path.join( + get_package_share_directory("neuromesh_platform_r2"), + "config", + "goal_poses.yaml", + ), + description=("Location of config containing starting positions"), + ) + start_file_arg = DeclareLaunchArgument( + name="start_file", + default_value=os.path.join( + get_package_share_directory("neuromesh_platform_r2"), + "config", + "start_poses.yaml", + ), + description=("Location of config containing starting positions"), + ) + + start_position_arg = DeclareLaunchArgument( + name="start_position", + default_value="False", + description=("Whether or not to run start goals node"), + ) + use_mission_maestro_bridge_arg = DeclareLaunchArgument( + name="use_mission_maestro_bridge", + default_value="False", + description=( + "Whether or not to run optional bridge from generic mission topics " + "to arl_mission_maestro services" + ), + ) + odom_republisher_arg = DeclareLaunchArgument( + name="odom_republisher", + default_value="False", + description=( + "Whether or not to start the sekhmet republisher needed for < 5 robots" + ), + ) + engine_plugin_package_arg = DeclareLaunchArgument( + name="engine_plugin_package", + default_value="onnx_engine", + description=( + "The package containing the engine plugin to use, e.g. tensorrt_engine or onnx_engine" + ), + ) + engine_type_arg = DeclareLaunchArgument( + name="engine_type", + default_value="engine_interface::ONNXEngine", + description=( + "The type of engine to use, e.g. engine_interface::TRTEngine or engine_interface::ONNXEngine" + ), + ) + + opaque_function_action = OpaqueFunction(function=launch_setup) + + return LaunchDescription( + [ + name_arg, + agent_num_arg, + agent_list_arg, + goal_file_arg, + start_file_arg, + start_position_arg, + use_mission_maestro_bridge_arg, + publish_static_map_arg, + odom_republisher_arg, + planning_frame_arg, + engine_plugin_package_arg, + engine_type_arg, + opaque_function_action, + ] + ) diff --git a/neuromesh_platform_r2/launch/gat_planner_model_neuromesh_launch.py b/neuromesh_platform_r2/launch/gat_planner_model_neuromesh_launch.py new file mode 100755 index 0000000..dd3aa6f --- /dev/null +++ b/neuromesh_platform_r2/launch/gat_planner_model_neuromesh_launch.py @@ -0,0 +1,346 @@ +import os +from ast import literal_eval + +import yaml +from ament_index_python.packages import get_package_share_directory +from launch import LaunchDescription +from launch.actions import DeclareLaunchArgument, OpaqueFunction +from launch.conditions import IfCondition +from launch.substitutions import LaunchConfiguration +from launch_ros.actions import ComposableNodeContainer, Node +from launch_ros.descriptions import ComposableNode + + +def launch_setup(context): + agent_list = LaunchConfiguration("agent_list").perform(context) + name = LaunchConfiguration("name").perform(context) + goal_file = LaunchConfiguration("goal_file").perform(context) + start_file = LaunchConfiguration("start_file").perform(context) + start_position = LaunchConfiguration("start_position") + # odom_republisher = LaunchConfiguration("odom_republisher") + publish_static_map = LaunchConfiguration("publish_static_map").perform(context) + planning_frame = LaunchConfiguration("planning_frame") + engine_plugin_package = LaunchConfiguration("engine_plugin_package").perform(context) + engine_type = LaunchConfiguration("engine_type").perform(context) + + launch_list = [] + + # If using static map read in start location from yaml, publish map transform + start_x = 0.0 + start_y = 0.0 + complete_agent_list = [] + with open(start_file) as f: + start_positions = yaml.safe_load(f) + start_x = start_positions[name]["goals"][0]["position"]["x"] + start_y = start_positions[name]["goals"][0]["position"]["y"] + for k, v in start_positions.items(): + complete_agent_list.append(k) + + print(f"{name} start positions, x: {start_x}, y: {start_y}") + + if (engine_plugin_package == "tensorrt_engine") and (engine_type == "engine_interface::TRTEngine"): + model = "trt" + elif (engine_plugin_package == "onnx_engine") and ( + engine_type == "engine_interface::ONNXEngine" + ): + model = "onnx" + else: + raise ValueError( + f"Invalid engine_plugin_package {engine_plugin_package} or engine_type {engine_type}" + ) + + tf2_static_pub = Node( + package="tf2_ros", + executable="static_transform_publisher", + name="map_tf2_static_pub", + arguments=[ + "--x", + str(start_x), + "--y", + str(start_y), + "--z", + "0", + "--yaw", + "0", + "--pitch", + "0", + "--roll", + "0", + "--frame-id", + "map", + "--child-frame-id", + name + "/map", + ], + condition=IfCondition(publish_static_map), + ) + launch_list.append(tf2_static_pub) + + + # LaunchConfiguration('agent_list') was serialized into string + if isinstance(agent_list, str): + try: + agent_list = literal_eval(agent_list) + except (ValueError, SyntaxError): + agent_list = agent_list.split(",") + + remappings = [] + for i, agent in enumerate(complete_agent_list): + remappings.append((f"/{name}/features_{agent}", f"/{agent}/features_{agent}")) + remappings.append((f"/{name}/pos_{agent}", f"/{agent}/odom")) + remappings.append( + (f"/{name}/gnn_output_{agent}", f"/{agent}/gnn_output_{agent}") + ) + + remappings.append(("position_topic", f"/{name}/odom")) + + composable_nodes = [] + # composable_nodes.append( + # ComposableNode( + # package="neuromesh_platform_r2", + # plugin="goal_sender::GoalSenderNode", + # name="starting_poses_sender", + # namespace=name, + # parameters=[ + # { + # "robot_id": name, + # "start_poses_yaml_file": start_file, + # } + # ], + # condition=IfCondition(start_position), + # ) + # ) + composable_nodes.append( + ComposableNode( + package="engine_interface", + namespace=name, + name=["engine", LaunchConfiguration("agent_num")], + plugin="engine_interface::EngineInterfaceNode", + parameters=[ + { + "engine_type": engine_type, + "model_names": "encoder,decoder1,decoder2", + "encoder.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat_planner/encoder." + model, + "encoder.input_dimensions": "1,1,14", + "encoder.output_dimensions": "1,1,64", + "encoder.tensor_type": "fp32", + "decoder1.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat_planner/gat_layer1." + model, + "decoder1.input_dimensions": "1,1,64;1,2,64;1,1,64", ## ooops, here: size of messages = second input is not fixed, is it a pb?? + "decoder1.output_dimensions": "1,1,64", + "decoder1.tensor_type": "fp32", + "decoder2.model_path": get_package_share_directory(engine_plugin_package) + + "/models/gat_planner/gat_layer2." + model, + "decoder2.input_dimensions": "1,1,64;1,2,64;1,1,64", ## ooops, here: size of messages = second input is not fixed, is it a pb?? + "decoder2.output_dimensions": "1,1,2", + "decoder2.tensor_type": "fp32", + } + ], + ) + ) + + composable_nodes.append( + ComposableNode( + package="neuromesh_platform_r2", + namespace=name, + name=["neuromesh"], + plugin="GATPlannerNeuromeshNode::GATPlannerImplementation", + parameters=[ + { + "id": name, + "encoder_model_name": "encoder", + "decoder_model1_name": "decoder1", + "decoder_model2_name": "decoder2", + "encoder_cycle_length": 3000, + "decoder_cycle_length": 3000, + "output_topic": "gnn_output_", + "agents": ",".join(complete_agent_list), + "goal_poses_yaml_file": goal_file, + "planning_frame": planning_frame, + "goals_sending_delay": 2.0, + } + ], + remappings=remappings, + ) + ) + + # If the number of agents is less than complete list, add the odom_republisher nodes + # if len(complete_agent_list) > len(agent_list) and LaunchConfiguration("odom_republisher").perform(context) == "True": + # print(f"Number of agents is less than {len(complete_agent_list)}. Adding odom_republisher nodes!") + # missing_agents = set(complete_agent_list) - set(agent_list) + # for missing_agent in missing_agents: + # remappings = [] + # for i, agent in enumerate(complete_agent_list): + # # print(f"remapping {agent} to {missing_agent}") + # remappings.append((f"/{missing_agent}/features_{agent}", f"/{agent}/features_{agent}")) + # remappings.append( + # (f"/{missing_agent}/gnn_output_{agent}", f"/{agent}/gnn_output_{agent}") + # ) + + # remappings.append(("position_topic", f"/{name}/odometry/local")) + # composable_nodes.append( + # ComposableNode( + # package="neuromesh_platform_r2", + # plugin="odom_republisher::OdomRepublisher", + # name=f"odom_republisher_{missing_agent}", + # remappings=[ + # ("original/odom", f"/{name}/odometry/local"), + # ("republished/odom", f"/{missing_agent}/odometry/local"), + # ], + # condition=IfCondition(odom_republisher), + # ) + # ) + # print( + # f"Added odom_republisher for {missing_agent} using {name}'s odometry" + # ) + # composable_nodes.append( + # ComposableNode( + # package="engine_interface", + # namespace=missing_agent, + # name=["engine", LaunchConfiguration("agent_num")], + # plugin="engine_interface::EngineInterfaceNode", + # parameters=[ + # { + # "engine_type": engine_type, + # "model_names": "encoder,decoder1,decoder2", + # "encoder.model_path": get_package_share_directory(engine_plugin_package) + # + "/models/gat/encoder." + model, + # "encoder.input_dimensions": "1,1,5", + # "encoder.output_dimensions": "1,1,16", + # "encoder.tensor_type": "fp32", + # "decoder1.model_path": get_package_share_directory(engine_plugin_package) + # + "/models/gat/gat_layer1." + model, + # "decoder1.input_dimensions": "1,1,16;1,4,16;1,1,16", + # "decoder1.output_dimensions": "1,1,16", + # "decoder1.tensor_type": "fp32", + # "decoder2.model_path": get_package_share_directory(engine_plugin_package) + # + "/models/gat/gat_layer2." + model, + # "decoder2.input_dimensions": "1,16;1,4,16;1,1,16", + # "decoder2.output_dimensions": "1,1,5", + # "decoder2.tensor_type": "fp32", + # } + # ], + # condition=IfCondition(odom_republisher), + # ) + # ) + # composable_nodes.append( + # ComposableNode( + # package="neuromesh_platform_r2", + # namespace=missing_agent, + # name=["neuromesh"], + # plugin="GATneuromeshNode::GATImplementation", + # parameters=[ + # { + # "id": missing_agent, + # "encoder_model_name": "encoder", + # "decoder_model1_name": "decoder1", + # "decoder_model2_name": "decoder2", + # "encoder_cycle_length": 3000, + # "decoder_cycle_length": 3000, + # "output_topic": "gnn_output_", + # "agents": ",".join(complete_agent_list), + # "goal_poses_yaml_file": goal_file, + # "planning_frame": planning_frame, + # "goals_sending_delay": 2.0, + # } + # ], + # remappings=remappings, + # condition=IfCondition(odom_republisher), + # ) + # ) + + + launch_list.append( + ComposableNodeContainer( + name="neuromesh_container", + namespace=name, + package="rclcpp_components", + executable="component_container", + composable_node_descriptions=composable_nodes, + # prefix="xterm -e gdb --args", + # arguments=[ + # "--ros-args", + # "--log-level", + # "DEBUG"], + output="screen", + ) + ) + return launch_list + + +def generate_launch_description(): + name_arg = DeclareLaunchArgument( + name="name", default_value="voxl2_1", description=("Which agent we are running") + ) + + planning_frame_arg = DeclareLaunchArgument( + name="planning_frame", + default_value="map", + description=("Which frame to send goals in"), + ) + + agent_num_arg = DeclareLaunchArgument( + name="agent_num", default_value="4", description=("Which agent we are running") + ) + + agent_list_arg = DeclareLaunchArgument( + name="agent_list", + default_value="voxl2_1,voxl2_2,voxl2_5,voxl2_9", + description=("List of all agents present (including self)"), + ) + + publish_static_map_arg = DeclareLaunchArgument( + name="publish_static_map", + default_value="False", + description="run publish static map", + ) + + goal_file_arg = DeclareLaunchArgument( + name="goal_file", + default_value=os.path.join( + get_package_share_directory("neuromesh_platform_r2"), + "config", + "gat_planner_goal1_pos.yaml", + ), + description=("Location of config containing starting positions"), + ) + start_file_arg = DeclareLaunchArgument( + name="start_file", + default_value=os.path.join( + get_package_share_directory("neuromesh_platform_r2"), + "config", + "gat_planner_start_pos.yaml", + ), + description=("Location of config containing starting positions"), + ) + engine_plugin_package_arg = DeclareLaunchArgument( + name="engine_plugin_package", + default_value="onnx_engine", + description=( + "The package containing the engine plugin to use, e.g. tensorrt_engine or onnx_engine" + ), + ) + engine_type_arg = DeclareLaunchArgument( + name="engine_type", + default_value="engine_interface::ONNXEngine", + description=( + "The type of engine to use, e.g. engine_interface::TRTEngine or engine_interface::ONNXEngine" + ), + ) + + opaque_function_action = OpaqueFunction(function=launch_setup) + + return LaunchDescription( + [ + name_arg, + agent_num_arg, + agent_list_arg, + goal_file_arg, + start_file_arg, + publish_static_map_arg, + planning_frame_arg, + engine_plugin_package_arg, + engine_type_arg, + opaque_function_action, + ] + ) diff --git a/neuromesh_platform_r2/launch/vggt_model_neuromesh_launch.py b/neuromesh_platform_r2/launch/vggt_model_neuromesh_launch.py index 1e60ce8..9f29673 100644 --- a/neuromesh_platform_r2/launch/vggt_model_neuromesh_launch.py +++ b/neuromesh_platform_r2/launch/vggt_model_neuromesh_launch.py @@ -1,116 +1,165 @@ import os from ament_index_python.packages import get_package_share_directory from launch import LaunchDescription -from launch.actions import LogInfo, DeclareLaunchArgument, TimerAction, OpaqueFunction -from launch.substitutions import EnvironmentVariable, LaunchConfiguration, TextSubstitution +from launch.actions import DeclareLaunchArgument, OpaqueFunction +from launch.substitutions import LaunchConfiguration, TextSubstitution from launch_ros.actions import ComposableNodeContainer from launch_ros.descriptions import ComposableNode def launch_setup(context): - name = LaunchConfiguration('name').perform(context) + name = LaunchConfiguration('robot_name').perform(context) agent_list = LaunchConfiguration('agent_list').perform(context) agent_num = LaunchConfiguration('agent_num').perform(context) color_raw_topic = LaunchConfiguration('color_raw_topic').perform(context) log_level = LaunchConfiguration('log_level').perform(context) - + + # Path to config file + config_file = os.path.join( + get_package_share_directory('neuromesh_platform_r2'), + 'config', + 'vggt_config.yaml' + ) + + # Build remappings for inter-robot feature topics remappings = [] for agent in agent_list.split(','): if agent != name: remappings.append((f'features_{agent}', f'/{agent}/features_{agent}')) - + composable_nodes = [ - ComposableNode( - package= 'neuromesh_platform_r2', - namespace= name, - name= ['vggt_neuromesh'], - plugin='vggtNode::VggtToyImplementation', - parameters=[{'id': name, - 'encoder_model_name' : 'vggt_encoder', - 'decoder_model_name' : 'vggt_decoder', - 'encoder_cycle_length' : 3000, - 'decoder_cycle_length' : 3000, - 'agents': agent_list, - 'ints_to_floats': True, - 'vggt_decoder_output_dimensions' : "1,2,9;1,2,392,518,1;1,2,392,518;1,2,392,518,3;1,2,392,518", - }], - - remappings=[('camera', color_raw_topic),] + remappings, + # TensorRT Engine Node for Encoder + ComposableNode( + package='tensorrt_engine', + namespace=name, + name="tensorrt_encoder", + plugin="tensorrt_engine_node::TensorRTEngineNode", + parameters=[{ + 'model_names': 'vggt_image_encoder_2x.engine', + 'vggt_image_encoder_2x.engine.model_path': os.path.join( + get_package_share_directory('tensorrt_engine'), + 'models/vggt_onnx_2x/vggt_image_encoder_2x.engine' + ), + 'vggt_image_encoder_2x.engine.input_dimensions': "1,3,392,518", + 'vggt_image_encoder_2x.engine.output_dimensions': "1,1036,1024", + 'vggt_image_encoder_2x.engine.tensor_type': "fp32", + }], + remappings=[ + ('tensorrt_request', 'tensorrt_request_encoder'), + ('tensorrt_output', 'tensorrt_output_encoder') + ], ), - - ComposableNode( - package= 'tensorrt_engine', - namespace= name, - name=["engine", agent_num], + + # TensorRT Engine Node for Decoder + ComposableNode( + package='tensorrt_engine', + namespace=name, + name="tensorrt_decoder", plugin="tensorrt_engine_node::TensorRTEngineNode", - parameters=[{'model_names' : 'vggt_encoder,vggt_decoder', - 'vggt_encoder.model_path': get_package_share_directory('tensorrt_engine') + "/models/vggt_onnx_2x/vggt_image_encoder_2x.engine", - 'vggt_encoder.input_dimensions' : "1,3,392,518", - 'vggt_encoder.output_dimensions' : "1,1036,1024", - 'vggt_encoder.tensor_type': "fp32", - 'vggt_decoder.model_path': get_package_share_directory('tensorrt_engine') + "/models/vggt_onnx_2x/vggt_aggregator_2x.engine", - 'vggt_decoder.input_dimensions' : "2,1036,1024", - 'vggt_decoder.output_dimensions' : "1,2,9;1,2,392,518,1;1,2,392,518;1,2,392,518,3;1,2,392,518", - 'vggt_decoder.tensor_type': "fp32", - }], + parameters=[{ + 'model_names': 'vggt_aggregator_2x.engine', + 'vggt_aggregator_2x.engine.model_path': os.path.join( + get_package_share_directory('tensorrt_engine'), + 'models/vggt_onnx_2x/vggt_aggregator_2x.engine' + ), + 'vggt_aggregator_2x.engine.input_dimensions': "2,1036,1024", + 'vggt_aggregator_2x.engine.output_dimensions': "1,2,9;1,2,392,518,1;1,2,392,518;1,2,392,518,3;1,2,392,518", + 'vggt_aggregator_2x.engine.tensor_type': "fp32", + }], + remappings=[ + ('tensorrt_request', 'tensorrt_request_decoder'), + ('tensorrt_output', 'tensorrt_output_decoder') + ], ), + # VGGT Encoder Node + ComposableNode( + package='neuromesh_platform_r2', + namespace=name, + name='vggt_encoder', + plugin='neuromesh::VggtEncoderNode', + parameters=[ + config_file, + { + 'robot_name': name, + 'color_raw_topic': color_raw_topic, + 'vggt.encoder.model_path': os.path.join( + get_package_share_directory('tensorrt_engine'), + 'models/vggt_onnx_2x/vggt_image_encoder_2x.engine' + ) + } + ], + remappings=[('camera', color_raw_topic)], + ), + # VGGT Decoder Node + ComposableNode( + package='neuromesh_platform_r2', + namespace=name, + name='vggt_decoder', + plugin='neuromesh::VggtDecoderNode', + parameters=[ + config_file, + { + 'robot_name': name, + 'vggt.robot_names': agent_list.split(','), + 'vggt.decoder.model_path': os.path.join( + get_package_share_directory('tensorrt_engine'), + 'models/vggt_onnx_2x/vggt_aggregator_2x.engine' + ) + } + ], + remappings=remappings, + ), ] - + return [ComposableNodeContainer( - name='vggt_container', + name='vggt_separated_container', namespace=name, package='rclcpp_components', executable='component_container_mt', # Use multi-threaded executor composable_node_descriptions=composable_nodes, output='screen', - arguments=['--ros-args', '--log-level', log_level], - additional_env={'ROS_DOMAIN_ID': EnvironmentVariable('ROS_DOMAIN_ID', default_value='0')}, - # Use multi-threaded executor with 4 threads - ros_arguments=['--ros-args', '--log-level', log_level, '-p', 'use_intra_process_comms:=true'], + # Use multi-threaded executor with intra-process communication + arguments=['--ros-args', '--log-level', log_level, '-p', 'use_intra_process_comms:=true'], )] def generate_launch_description(): - name_arg = DeclareLaunchArgument( - name='name', default_value='khonsu', - description=( - 'Which robot we are running' + robot_name_arg = DeclareLaunchArgument( + name='robot_name', + default_value='khonsu', + description='Which robot we are running' ) - ) - + color_raw_topic_arg = DeclareLaunchArgument( name='color_raw_topic', default_value=[ TextSubstitution(text='/'), - LaunchConfiguration('name'), + LaunchConfiguration('robot_name'), TextSubstitution(text='/sensors/camera_0/camera/color/image_raw') ], description='Camera topic to be remapped', ) - + agent_num_arg = DeclareLaunchArgument( - name='agent_num', default_value='1', - description=( - 'Which agent we are running' - ) + name='agent_num', + default_value='1', + description='Which agent we are running' ) - + agent_list_arg = DeclareLaunchArgument( - name='agent_list', default_value='khonsu,anubis', - description=( - 'List of all agents present (including self) - default 2 robots for VGGT' - ) + name='agent_list', + default_value='khonsu,anubis', + description='List of all agents present (including self) - default 2 robots for VGGT' ) - + log_level_arg = DeclareLaunchArgument( - name='log_level', default_value='INFO', - description=( - 'Log level for all nodes (DEBUG, INFO, WARN, ERROR, FATAL)' - ) + name='log_level', + default_value='INFO', + description='Log level for all nodes (DEBUG, INFO, WARN, ERROR, FATAL)' ) - + opaque_function_action = OpaqueFunction(function=launch_setup) - + return LaunchDescription([ - name_arg, + robot_name_arg, agent_num_arg, agent_list_arg, color_raw_topic_arg, diff --git a/neuromesh_platform_r2/package.xml b/neuromesh_platform_r2/package.xml index 43b9c6b..d630537 100755 --- a/neuromesh_platform_r2/package.xml +++ b/neuromesh_platform_r2/package.xml @@ -12,19 +12,24 @@ rosidl_default_runtime + ament_index_cpp rclcpp std_msgs sensor_msgs geometry_msgs builtin_interfaces neuromesh_interfaces - tensorrt_engine + + onnx_engine + engine_interface cv_bridge tf2 tf2_ros pcl_conversions nav_msgs visualization_msgs + tensorrt + yaml-cpp python3-opencv rclcpp_components rclcpp_components diff --git a/neuromesh_platform_r2/scripts/gat_model_neuromesh_launch.sh b/neuromesh_platform_r2/scripts/gat_model_neuromesh_launch.sh new file mode 100755 index 0000000..4677f0b --- /dev/null +++ b/neuromesh_platform_r2/scripts/gat_model_neuromesh_launch.sh @@ -0,0 +1,31 @@ +#!/bin/bash + +# Check if a robot name is provided +if [ $# -eq 0 ]; then + echo "Usage: $0 " + exit 1 +fi + +# Use the provided robot name +ROBOT_NAME=$1 + +# Start a new tmux session +SESSION="gat_neuromesh_launch" +tmux new-session -d -s $SESSION + +# Rename the first window +tmux rename-window -t $SESSION:0 'gat_neuromesh_launch' + +# Create a layout with 7 panes +tmux split-window -h -t $SESSION:0.0 # Split horizontally (creates pane 1) +tmux split-window -v -t $SESSION:0.1 # Split right pane vertically (creates pane 2) +tmux split-window -v -t $SESSION:0.2 # Split bottom-right pane vertically (creates pane 3) +tmux split-window -v -t $SESSION:0.3 # Split bottom-right pane vertically again (creates pane 4) +tmux split-window -v -t $SESSION:0.4 # Split left pane vertically (creates pane 5) +tmux split-window -v -t $SESSION:0.5 # Split bottom-left pane vertically (creates pane 6) + +# Pane 2 (middle-right top) +tmux send-keys -t $SESSION:0.2 'sleep 30; ./ros2_bag_record.sh' C-m + +# Attach to the tmux session +tmux attach -t $SESSION diff --git a/neuromesh_platform_r2/src/gat_model_implementation.cpp b/neuromesh_platform_r2/src/gat_model_implementation.cpp new file mode 100644 index 0000000..631faa0 --- /dev/null +++ b/neuromesh_platform_r2/src/gat_model_implementation.cpp @@ -0,0 +1,110 @@ +#include "neuromesh_platform_r2/gat_implementation.h" + +namespace GATneuromeshNode { +GATImplementation::GATImplementation(const rclcpp::NodeOptions &options) + : GATneuromeshNode(options) { + + // Subscribe to odom topic to get current pose of robot + pos_sub_ = this->create_subscription( + "position_topic", 10, + std::bind(&GATImplementation::pos_callback, this, std::placeholders::_1)); + + // Subscribe to gnn_features topic received from the other neighbour robots + gnn_result_subscriber_ = + this->create_subscription( + "gnn_result_topic", 10, + std::bind(&GATImplementation::gnn_result_callback, this, + std::placeholders::_1)); + + this->tensor_client_ = + create_client( + "tensorrt_request"); +} + +std::future>> +GATImplementation::performInference( + const std::string &model_name, + const std::vector &tensors) { + if (!tensor_client_->wait_for_service(std::chrono::seconds(1))) { + RCLCPP_ERROR(this->get_logger(), "Engine not reachable via service."); + // Complex stuff just to return future that resolves to empty tensor + std::promise< + std::vector>> + prom; + std::future>> + r = prom.get_future(); + std::vector> t( + 1, std::make_shared()); + t[0]->result = 3; // Cannot reach engine error code + prom.set_value(t); + return r; + } + + // Create a request to send to the service server + auto request = + std::make_shared(); + request->model_name = model_name; + request->tensor1 = tensors; + + // Call the service and wait for the response + std::shared_future< + std::shared_ptr> + future = tensor_client_->async_send_request(request); + + std::future>> + return_tensors = std::async(std::launch::async, [future]() { + std::vector> + output_tensors; + for (const auto &tensor : future.get()->tensor2) { + output_tensors.emplace_back( + std::make_shared(tensor)); + } + return output_tensors; + }); + + return return_tensors; +} + +bool GATImplementation::buildDecoderTensor( + std::map + &agent_features, + neuromesh_interfaces::msg::Tensor &own_feature, + neuromesh_interfaces::msg::Tensor &aggregated_tensor) { + if (agent_features.size() != 5) { + RCLCPP_INFO(this->get_logger(), "Expected 5 feature messages, got %zu", + agent_features.size()); + return false; + } + + std::vector sorted_ids; + for (const auto &pair : agent_features) { + sorted_ids.push_back(pair.first); + } + std::sort(sorted_ids.begin(), sorted_ids.end()); + + own_feature.name = "own_features"; + aggregated_tensor.name = "other_features"; + own_feature.data_type = 9; + aggregated_tensor.data_type = 9; + + for (const std::string &id : sorted_ids) { + const auto &feature_msg = agent_features.at(id); + if (id == this->id_) { + own_feature = feature_msg->tensor; + } else { + if (aggregated_tensor.data.empty()) { + aggregated_tensor = feature_msg->tensor; + } else { + aggregated_tensor.data.insert(aggregated_tensor.data.end(), + feature_msg->tensor.data.begin(), + feature_msg->tensor.data.end()); + } + } + } + + return true; +} +} // namespace GATneuromeshNode + +#include "rclcpp_components/register_node_macro.hpp" +RCLCPP_COMPONENTS_REGISTER_NODE(GATneuromeshNode::GATImplementation) diff --git a/neuromesh_platform_r2/src/gat_model_neuromesh_node.cpp b/neuromesh_platform_r2/src/gat_model_neuromesh_node.cpp new file mode 100644 index 0000000..8306e10 --- /dev/null +++ b/neuromesh_platform_r2/src/gat_model_neuromesh_node.cpp @@ -0,0 +1,682 @@ +#include "neuromesh_platform_r2/gat_neuromesh_node.h" +#include "service_adapters/adapter_factory.hpp" +#include "rclcpp/rclcpp.hpp" +#include "types/navigation_goal.hpp" +#include "chrono" + +namespace GATneuromeshNode { +GATneuromeshNode :: GATneuromeshNode(const rclcpp::NodeOptions &options): Node("GAT_neuromesh_node", options) +{ + // Declare node parameters + this->declare_parameter("encoder_model_name", "default_encoder_model"); + this->declare_parameter("decoder_model1_name", "default_decoder_model"); + this->declare_parameter("decoder_model2_name", "default_decoder_model"); + this->declare_parameter("topic_prefix", "features_"); + this->declare_parameter("output_topic", "gnn_output_"); + this->declare_parameter("planning_frame", "map"); + this->declare_parameter("decoder_cycle_length", 1000); + this->declare_parameter("encoder_cycle_length", 1000); + this->declare_parameter("encoder_await_length", 10000); + this->declare_parameter("id", "default_id"); + this->declare_parameter("image_qos_profile", "default"); + this->declare_parameter("features_qos_profile", "default"); + this->declare_parameter("output_qos_profile", "default"); + this->declare_parameter("agents", ""); + this->declare_parameter("to_nchw", true); + this->declare_parameter("ints_to_floats", true); + this->declare_parameter("goal_poses_yaml_file", "goal_poses.yaml"); + this->declare_parameter("goals_sending_delay", 10.0); + this->declare_parameter("service_adapter_type", "topic"); + + // Get node parameters + this->get_parameter("encoder_model_name", encoder_model_name_); + this->get_parameter("decoder_model1_name", decoder_model1_name_); + this->get_parameter("decoder_model2_name", decoder_model2_name_); + this->get_parameter("topic_prefix", topic_prefix_); + this->get_parameter("output_topic", output_topic_); + this->get_parameter("planning_frame", planning_frame_); + this->get_parameter("decoder_cycle_length", decoder_cycle_length_); + this->get_parameter("encoder_cycle_length", encoder_cycle_length_); + this->get_parameter("encoder_await_length", encoder_await_length_); + this->get_parameter("id", id_); + this->get_parameter("image_qos_profile", image_qos_profile_); + this->get_parameter("features_qos_profile", features_qos_profile_); //for both input and output + this->get_parameter("output_qos_profile", output_qos_profile_); + this->get_parameter("agents", agents_); + this->get_parameter("to_nchw", to_nchw_); + this->get_parameter("goal_poses_yaml_file", goal_poses_yaml_file); + this->get_parameter("goals_sending_delay", goals_sending_delay_); + this->get_parameter("service_adapter_type", service_adapter_type_); + + auto feature_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(features_qos_profile_)); + auto output_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(output_qos_profile_)); + feature_publisher_ = this->create_publisher(topic_prefix_ + id_, feature_qos); + + // Initialize GNN result publisher and subscriber + gnn_result_publisher_ = this->create_publisher( + output_topic_ + id_, feature_qos + ); + + second_decoder_result_publisher_ = this->create_publisher( + "second_decoder_result_topic", 10 // topic name and queue size + ); + + //PLACEHOLDER: update available_agents + + all_agents = splitAgentString(agents_); + for (const auto& agent : all_agents) { + RCLCPP_DEBUG(this->get_logger(), "%s", agent.c_str()); + } + all_agents.erase(id_); //remove self from list + + // Print out state vector elements + neuromesh_interfaces::msg::StateVector tmp; + for (uint i = 0; i < 6; ++i) + { + tmp.state_vector.push_back(0.); + } + current_states_.insert(std::pair(id_, tmp)); + + available_agents = all_agents; + + // Print out all agent names in the agent list + for (std::string id : all_agents){ + RCLCPP_DEBUG(this->get_logger(), "Going through all agents to create subscriptions"); + RCLCPP_DEBUG(this->get_logger(), "Id: %s", id.c_str()); + this->createSubscription(feature_subscriptions_, id, feature_qos); + this->createGNNSubscription(gnn_subscriptions_, id, feature_qos); + } + + // Extra print statements for debugging + for(const auto& [key, value] : current_states_) { + RCLCPP_DEBUG( + get_logger(), + "Key: %s, State Vector Size: %zu", + key.c_str(), + value.state_vector.size() + ); + } + + // Load the goal poses for the robots from the yaml file + load_goal_poses_from_yaml(); + + // Create timers for encoder and decoder cycles + decoder_timer_ = this->create_wall_timer(std::chrono::duration(decoder_cycle_length_), std::bind(&GATneuromeshNode::run_decoder_cycle, this)); + encoder_timer_ = this->create_wall_timer(std::chrono::duration(encoder_cycle_length_), std::bind(&GATneuromeshNode::run_encoder_cycle, this)); + fresh_encoder_cycle = true; + waypoint_cmd_sent_ = false; + + // Initialize tf buffers + tf_buffer_ = std::make_shared(this->get_clock()); + tf_listener_ = std::make_shared(*tf_buffer_); +} + +void GATneuromeshNode::load_goal_poses_from_yaml() { + try { + // Clear existing goal poses + goal_poses_.clear(); + + // Load the YAML file + YAML::Node config = YAML::LoadFile(goal_poses_yaml_file); + + RCLCPP_DEBUG(this->get_logger(), "Loading goals for all robots"); + + // Iterate through all top-level nodes (robot names) + for (const auto& robot_entry : config) { + std::string robot_name = robot_entry.first.as(); + + // Check if this robot has a goals section + if (!robot_entry.second["goals"] || !robot_entry.second["goals"].IsSequence()) { + RCLCPP_WARN(this->get_logger(), "No valid goals found for robot '%s'", robot_name.c_str()); + continue; + } + + // Process each goal for this robot + for (const auto& goal : robot_entry.second["goals"]) { + if (!goal["position"] || !goal["position"]["x"] || !goal["position"]["y"]) { + RCLCPP_WARN(this->get_logger(), "Skipping malformed goal entry for robot '%s'", robot_name.c_str()); + continue; + } + + geometry_msgs::msg::Pose pose; + pose.position.x = goal["position"]["x"].as(); + pose.position.y = goal["position"]["y"].as(); + pose.position.z = 0.0; // Set to 0 if not needed + + goal_poses_.push_back(pose); + + RCLCPP_DEBUG( + this->get_logger(), + "Loaded goal for robot '%s': Position (%.2f, %.2f)", + robot_name.c_str(), + pose.position.x, + pose.position.y + ); + } + } + + RCLCPP_DEBUG( + this->get_logger(), + "Successfully loaded %zu goal poses for all robots", + goal_poses_.size() + ); + + } catch (const YAML::Exception& e) { + RCLCPP_ERROR( + this->get_logger(), + "Failed to load goals from YAML file: %s", + e.what() + ); + } +} + +void GATneuromeshNode::pos_callback(const nav_msgs::msg::Odometry::SharedPtr msg) { + + // Parse position data + pos_callback_complete = true; + + nav_msgs::msg::Odometry rel_odom = *msg; + geometry_msgs::msg::TransformStamped transformStamped; + try { + transformStamped = tf_buffer_->lookupTransform(planning_frame_, rel_odom.header.frame_id, rclcpp::Time(0)); + } catch (tf2::TransformException &ex) { + RCLCPP_ERROR(this->get_logger(), "Transform error: %s", ex.what()); + return; + } + tf2::doTransform(rel_odom.pose.pose, rel_odom.pose.pose, transformStamped); + + current_states_[id_].state_vector[0] = rel_odom.pose.pose.position.x; + current_states_[id_].state_vector[1] = rel_odom.pose.pose.position.y; + + // Set orientation to [0, 0, 0, 1] + current_states_[id_].state_vector[2] = 0.0; + current_states_[id_].state_vector[3] = 0.0; + current_states_[id_].state_vector[4] = 0.0; + current_states_[id_].state_vector[5] = 1.0; + + // Keep spamming others until goal is sent + if(!waypoint_cmd_sent_) { + run_encoder_cycle(); + } + +} + +neuromesh_interfaces::msg::Tensor GATneuromeshNode::calculate_goal_distances(const std::vector& state_vector) { + neuromesh_interfaces::msg::Tensor tensor_msg; + + // Set tensor dimensions + tensor_msg.shape.dims = {1, static_cast(goal_poses_.size())}; + + tensor_msg.data_type = 9; // float32 + + std::vector distance_values; + for (size_t j = 0; j < goal_poses_.size(); ++j) { + // Calculate Euclidean distance + float dx = state_vector[0] - goal_poses_[j].position.x; + float dy = state_vector[1] - goal_poses_[j].position.y; + + // Store the actual Euclidean distance + distance_values.push_back(dx * dx + dy * dy); + } + + // Directly copy float values to tensor data + tensor_msg.data.resize(distance_values.size() * sizeof(float)); + std::copy( + reinterpret_cast(distance_values.data()), + reinterpret_cast(distance_values.data() + distance_values.size()), + tensor_msg.data.begin() + ); + + // Set data type explicitly + tensor_msg.data_type = 9; + + // Set shape explicitly + tensor_msg.shape.dims = {1, static_cast(distance_values.size())}; + + // Optionally set strides + tensor_msg.strides = {static_cast(distance_values.size()), 1}; + + return tensor_msg; +} + +neuromesh_interfaces::msg::Feature GATneuromeshNode::buildFeatureMessage(const neuromesh_interfaces::msg::Tensor& tensor) +{ + neuromesh_interfaces::msg::Feature feature_msg = neuromesh_interfaces::msg::Feature(); + + feature_msg.tensor = tensor; + feature_msg.id = id_; + feature_msg.timestamp = this->get_clock()->now(); + return feature_msg; +} + +void GATneuromeshNode::feature_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg) { + std::string uuid = msg->id; + received_features_[uuid] = msg; + feature_buffer_timestamp_[uuid] = msg->timestamp.sec + msg->timestamp.nanosec*1e-9; + auto feature_now = this->get_clock()->now(); + double feature_timestamp = feature_now.seconds() + feature_now.nanoseconds() / 1e9; + RCLCPP_DEBUG(this->get_logger(), "Delay between receiving and publishing features %.9f", feature_timestamp - feature_buffer_timestamp_[uuid]); + + // std::cout << "Delay between receiving and publishing features " << feature_now.seconds() - feature_buffer_timestamp_[uuid] << " seconds" << std::endl + // << "feature_buff_timestamp_[uuid]: " << feature_buffer_timestamp_[uuid] << std::endl + // << "msg->timestamp.sec: " << msg->timestamp.sec << std::endl + // << "msg->timestamp.nanosec: " << msg->timestamp.nanosec << std::endl + // << "msg->timestamp.nanosec*1e-9: " << msg->timestamp.nanosec*1e-9 << std::endl + // // << "feature_timestamp: " << feature_timestamp << std::endl + // << "feature_now.seconds(): " << feature_now.seconds() << std::endl + // << "feature_now.nanoseconds(): " << feature_now.nanoseconds() << std::endl + // << "feature_now.nanoseconds() / 1e9: " << feature_now.nanoseconds() / 1e9 << std::endl + // << "Clock type: " << this->get_clock()->get_clock_type() << std::endl; +} + +//perform inference on tensor using model called model_name +//PLACEHOLDER +std::future>> GATneuromeshNode::performInference(const std::string& model_name, const std::vector& tensors) +{ + std::promise>> prom; + std::future>> r = prom.get_future(); + std::vector> t(1, std::make_shared()); + t[0]->result = 1; // Cannot reach engine error code + prom.set_value(std::move(t)); + return r; +} + +//Aggregates features from available agents (and itself) into a single tensor +//PLACEHOLDER +bool GATneuromeshNode::buildDecoderTensor( + std::map& agent_features, + neuromesh_interfaces::msg::Tensor& own_feature, + neuromesh_interfaces::msg::Tensor& aggregated_tensor) +{ + return true; +} + +//Adds a subscription to the feature_subscriptions_ map +void GATneuromeshNode::createSubscription(std::map::SharedPtr>& subscription_map, std::string id, rclcpp::QoS qos) +{ + std::string topic = topic_prefix_ + id; + + RCLCPP_DEBUG(this->get_logger(), "creating subscription for topic %s", topic.c_str()); + + rclcpp::Subscription::SharedPtr feature_subscription_ = + this->create_subscription( + topic, + qos, + std::bind(&GATneuromeshNode::feature_callback, this, std::placeholders::_1)); + + subscription_map.insert( {id, feature_subscription_} ); +} + +//Remove subscription form the feature_subscriptions_ map +void GATneuromeshNode::removeSubscription(std::map::SharedPtr> subscription_map, std::string id) +{ + subscription_map.erase(id); +} + +//Adds a subscription to the gnn_subscriptions_ map +void GATneuromeshNode::createGNNSubscription(std::map::SharedPtr>& gnn_subscription_map, std::string id, rclcpp::QoS qos) +{ + std::string topic = output_topic_ + id; + + RCLCPP_DEBUG(this->get_logger(), "creating gnn subscription for topic %s", topic.c_str()); + + rclcpp::Subscription::SharedPtr gnn_subscription_ = + this->create_subscription( + topic, + qos, + std::bind(&GATneuromeshNode::gnn_result_callback, this, std::placeholders::_1)); + + gnn_subscription_map.insert( {id, gnn_subscription_} ); +} + +//Remove subscription form the gnn_subscriptions_ map +void GATneuromeshNode::removeGNNSubscription(std::map::SharedPtr> gnn_subscription_map, std::string id) +{ + gnn_subscription_map.erase(id); +} + +//Convert string to ROS2 QoS profile +//from https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_nvblox/nvblox_ros_common/src/qos.cpp#L26 +rmw_qos_profile_t GATneuromeshNode::parseQoSString(const std::string& str) +{ + std::string profile = str; + // Convert to upper case. + std::transform(profile.begin(), profile.end(), profile.begin(), ::toupper); + + if (profile == "SYSTEM_DEFAULT") { + return rmw_qos_profile_system_default; + } + if (profile == "DEFAULT") { + return rmw_qos_profile_default; + } + if (profile == "PARAMETER_EVENTS") { + return rmw_qos_profile_parameter_events; + } + if (profile == "SERVICES_DEFAULT") { + return rmw_qos_profile_services_default; + } + if (profile == "PARAMETERS") { + return rmw_qos_profile_parameters; + } + if (profile == "SENSOR_DATA") { + return rmw_qos_profile_sensor_data; + } + RCLCPP_WARN_STREAM( + rclcpp::get_logger("parseQosString"), + "Unknown QoS profile: " << profile << ". Returning profile: DEFAULT"); + return rmw_qos_profile_default; +} + +//Split agent string parameter into vector of agent ids +std::set GATneuromeshNode::splitAgentString(std::string str) +{ + std::set agents; + const std::string delimiter = ","; + + size_t pos = 0; + std::string token; + while ((pos = str.find(delimiter)) != std::string::npos) { + token = str.substr(0, pos); + agents.insert(token); + str.erase(0, pos + delimiter.length()); + } + agents.insert(str); + return agents; +} + +void GATneuromeshNode::run_encoder_cycle() { + // Run the encoder on all feature vectors + if (fresh_encoder_cycle && !current_states_.empty() && pos_callback_complete && !goal_poses_.empty() && !waypoint_cmd_sent_) { + fresh_encoder_cycle = false; + startClock("encoder_inference"); + + neuromesh_interfaces::msg::Tensor goal_distances = calculate_goal_distances( + std::vector(current_states_[id_].state_vector.begin(), + current_states_[id_].state_vector.end()) + ); + + auto encoder_now = this->get_clock()->now(); + RCLCPP_DEBUG(this->get_logger(), "Performing inference"); + encoder_result = performInference(encoder_model_name_, {goal_distances}); + auto encoder_end = this->get_clock()->now(); + + RCLCPP_DEBUG(this->get_logger(), "Finished performing Inference on Encoder"); + } + + // Keep publishing features to others + if (!fresh_encoder_cycle && encoder_result.valid()) { + auto encoder_status = encoder_result.wait_for(std::chrono::milliseconds(0)); + if (encoder_status == std::future_status::ready) { + stopClock("encoder_inference"); + std::vector> encoded_features = encoder_result.get(); + + std::shared_ptr feature_tensor = encoded_features[0]; + + // Print out encoder model output + float max_prob, next; + size_t size = feature_tensor->data.size(); + std::memcpy(&max_prob, &feature_tensor->data[0], sizeof(float)); + std::stringstream ss; + ss << "[" << std::fixed << std::setprecision(2) << max_prob; + for (size_t i = 4; i < size; i +=4) { + if (i + 3 < size) + { + std::memcpy(&next, &feature_tensor->data[i], sizeof(float)); + ss << " " << next << " "; + if (((i/4) % 5) == 0 ){ + ss << "\n"; + } + } + } + ss << "]"; + + RCLCPP_DEBUG_STREAM(this->get_logger(), "Encoder features output: " << "\n Size: " << size + << "\n Output: " << ss.str()); + + neuromesh_interfaces::msg::Feature feature_msg = buildFeatureMessage(*feature_tensor.get()); + + // Publish features + RCLCPP_DEBUG_THROTTLE(this->get_logger(), + *this->get_clock(), + 1000, + "Publishing features."); + feature_publisher_->publish(feature_msg); + + // Store features + received_features_[this->id_] = std::make_shared(feature_msg); + + + feature_buffer_timestamp_[this->id_] = feature_msg.timestamp.sec + feature_msg.timestamp.nanosec*1e-9; + + fresh_encoder_cycle = true; + RCLCPP_DEBUG(this->get_logger(), "Encoder cycle completed"); + + } else if (checkClock("encoder_inference") >= encoder_await_length_) { + RCLCPP_WARN(this->get_logger(), "Encoder inference timed out"); + fresh_encoder_cycle = true; + } + } +} + +void GATneuromeshNode::run_decoder_cycle() { + // Handle previous decoder inference result + if (gnn_result_future.valid()) { + auto decoder_status = gnn_result_future.wait_for(std::chrono::milliseconds(0)); + if (decoder_status == std::future_status::ready) { + stopClock("decoder_inference"); + std::vector> gnn_result = gnn_result_future.get(); + + std::shared_ptr gnn_result_tensor = gnn_result[0]; + + // Print first decoder output values + float max_prob, next; + size_t size = gnn_result_tensor->data.size(); + std::memcpy(&max_prob, &gnn_result_tensor->data[0], sizeof(float)); + std::stringstream ss; + ss << "[" << std::fixed << std::setprecision(2) << max_prob; + for (size_t i = 4; i < size; i +=4) { + if (i + 3 < size) + { + std::memcpy(&next, &gnn_result_tensor->data[i], sizeof(float)); + ss << " " << next << " "; + if (((i/4) % 5) == 0 ){ + ss << "\n"; + } + } + } + ss << "]"; + + const auto& data = gnn_result_tensor->data; + + // Publish the GNN result to other robots + neuromesh_interfaces::msg::Feature gnn_msg = build_gnn_msg(*gnn_result_tensor.get()); + + // Publish the result + gnn_result_publisher_->publish(gnn_msg); + + RCLCPP_DEBUG(this->get_logger(), "First Decoder inference completed and published"); + + // Store local GNN result + received_gnn_results_[this->id_] = std::make_shared(gnn_msg); + + // Prepare for potential second-stage decoding + prepare_second_stage_decoding(); + + RCLCPP_DEBUG(this->get_logger(), "Second Decoder inference completed and published"); + } + } + + // Check if we have enough feature messages (5 total) + if (!received_features_.empty() && !waypoint_cmd_sent_) { + startClock("decoder_inference"); + + // Step 4: Combine features of other robots + // Prepare encoder output tensor (from previous encoder cycle) + encoder_output_tensor; + neuromesh_interfaces::msg::Tensor aggregated_tensor; + if (buildDecoderTensor(received_features_, encoder_output_tensor, aggregated_tensor)) { + std::vector decoder_tensors = {encoder_output_tensor, aggregated_tensor, encoder_output_tensor}; + auto first_gnn_start = this->get_clock()->now(); + gnn_result_future = performInference(decoder_model1_name_, decoder_tensors); + auto first_gnn_end = this->get_clock()->now(); + //std::cout <<"Computation time first gnn round" <get_logger(), "Failed to build decoder tensor"); + } + + // Clear received features for the next cycle + received_features_.clear(); + feature_buffer_timestamp_.clear(); + } +} + +neuromesh_interfaces::msg::Feature GATneuromeshNode::build_gnn_msg(const neuromesh_interfaces::msg::Tensor& gnn_result) { + // Convert GNN result to Tensor message + neuromesh_interfaces::msg::Feature result_msg = neuromesh_interfaces::msg::Feature(); + result_msg.tensor = gnn_result; + + // Add metadata if needed + result_msg.id = this->id_; + result_msg.timestamp = this->get_clock()->now(); + + return result_msg; +} + +void GATneuromeshNode::gnn_result_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg) { + // Store received GNN result + std::string gnn_id = msg->id; + received_gnn_results_[gnn_id] = msg; + // Optional: Log received result + RCLCPP_DEBUG(this->get_logger(),"Received GNN result from robot %d", msg->id); +} + +void GATneuromeshNode::prepare_second_stage_decoding() { + + if (second_decoder_result_future.valid()) { + auto second_decoder_status = second_decoder_result_future.wait_for(std::chrono::milliseconds(0)); + if (second_decoder_status == std::future_status::ready) { + stopClock("decoder_inference"); + std::vector> second_decoder_result = second_decoder_result_future.get(); + + // Ensure the vector is not empty + if (!second_decoder_result.empty()) { + // Get the first tensor + std::shared_ptr second_decoder_result_tensor = second_decoder_result[0]; + + // Find the index of the maximum probability + int max_prob_index = 0; + + // Final step to get decoder output, print the values and send it as waypoint goals to robots + // Assuming the Tensor has a data field that is a vector of floats + if (!second_decoder_result_tensor->data.empty()) { + float max_prob, next; + size_t size = second_decoder_result_tensor->data.size(); + std::memcpy(&max_prob, &second_decoder_result_tensor->data[0], sizeof(float)); + + if (!goal_adapter_) { + goal_adapter_ = AdapterFactory::create(service_adapter_type_, this->shared_from_this()); + } + if (!goal_adapter_) { + RCLCPP_ERROR(this->get_logger(), "Failed to create service adapter of type '%s'", service_adapter_type_.c_str()); + return; + } + + NavigationGoal goal; + goal.x = goal_poses_[max_prob_index].position.x; + goal.y = goal_poses_[max_prob_index].position.y; + goal.planning_frame = planning_frame_; + goal_adapter_->sendGoal(goal); + } + } + } + } + + // Check if we have received GNN results from all robots + if (!received_gnn_results_.empty()) { + // Prepare aggregated GNN results for second-stage decoder + neuromesh_interfaces::msg::Tensor gnn_output_tensor; + neuromesh_interfaces::msg::Tensor aggregated_gnn_output_tensor; + if (buildDecoderTensor(received_gnn_results_, gnn_output_tensor, aggregated_gnn_output_tensor)) { + + second_decoder_result_future = performInference(decoder_model2_name_, {gnn_output_tensor, aggregated_gnn_output_tensor, encoder_output_tensor}); + + auto gnn_now = this->get_clock()->now(); + + for (const auto& entry : feature_buffer_timestamp_) { + const std::string& robot_name = entry.first; + double timestamp = entry.second; + } + + } else { + RCLCPP_ERROR(this->get_logger(), "Failed to build decoder tensor"); + } + + // Clear received GNN results for next cycle + received_gnn_results_.clear(); + } +} + +void GATneuromeshNode::startClock(std::string phase){ + int64_t now_time = std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); + + times[phase] = {now_time, false}; +} +void GATneuromeshNode::stopClock(std::string phase){ + if (times[phase].second){ + return; // clock already stopped + } + int64_t start_time = times[phase].first; + int64_t now_time = std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); + + times[phase] = {now_time - start_time, true}; +} + +int64_t GATneuromeshNode::checkClock(std::string phase){ + if (times[phase].second){ + return times[phase].first; // clock already stopped + } + + //stopclock calculations without saving + int64_t start_time = times[phase].first; + int64_t now_time = std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); + + return now_time - start_time; +} + +//TODO generalize to a transpose function +//tested for use with integers +neuromesh_interfaces::msg::Tensor GATneuromeshNode::convert_to_nchw(const neuromesh_interfaces::msg::Tensor& input){ + std::vector new_data; + const std::vector& old_data = input.data; + + std::vector strides = input.strides; + std::vector dims = input.shape.dims; + unsigned long bytedepth = input.strides[0] / input.shape.dims[0]; + + + for(int i = 0; i < dims[2]; i++){ + for(int j = 0; j < dims[0]; j++){ + for(int k = 0; k < dims[1]; k++){ + new_data.push_back(old_data[ (j * strides[0]) + (k * strides[1]) + (i * strides[2]) ] ); + } + } + } + + neuromesh_interfaces::msg::Tensor new_tensor = std::move(input); + new_tensor.data = new_data; + new_tensor.shape.dims = {input.shape.dims[2], input.shape.dims[0], input.shape.dims[1]}; + new_tensor.strides = {dims[1] * dims[2] * bytedepth, dims[2] * bytedepth, bytedepth}; + + return new_tensor; +} + +float int_to_scaled_float(int i){ return static_cast(i) / 255.0;} +} diff --git a/neuromesh_platform_r2/src/gat_planner_model_implementation.cpp b/neuromesh_platform_r2/src/gat_planner_model_implementation.cpp new file mode 100644 index 0000000..0b3f92f --- /dev/null +++ b/neuromesh_platform_r2/src/gat_planner_model_implementation.cpp @@ -0,0 +1,133 @@ +#include "neuromesh_platform_r2/gat_planner_implementation.h" + +namespace GATPlannerNeuromeshNode { +GATPlannerImplementation::GATPlannerImplementation(const rclcpp::NodeOptions &options) + : GATPlannerNeuromeshNode(options) { + + // Subscribe to odom topic to get current pose of robot + pos_sub_ = this->create_subscription( + "position_topic", 10, + std::bind(&GATPlannerImplementation::pos_callback, this, std::placeholders::_1)); + + // Subscribe to gnn_features topic received from the other neighbour robots + gnn_result_subscriber_ = + this->create_subscription( + "gnn_result_topic", 10, + std::bind(&GATPlannerImplementation::gnn_result_callback, this, + std::placeholders::_1)); + + this->tensor_client_ = + create_client( + "tensorrt_request"); +} + +std::future>> +GATPlannerImplementation::performInference( + const std::string &model_name, + const std::vector &tensors) { + + if (!tensor_client_->wait_for_service(std::chrono::seconds(1))) { + RCLCPP_ERROR(this->get_logger(), "Engine not reachable via service."); + // Complex stuff just to return future that resolves to empty tensor + std::promise< + std::vector>> + prom; + std::future>> + r = prom.get_future(); + std::vector> t( + 1, std::make_shared()); + t[0]->result = 3; // Cannot reach engine error code + prom.set_value(t); + return r; + } + + // Create a request to send to the service server + auto request = + std::make_shared(); + request->model_name = model_name; + request->tensor1 = tensors; + + // Call the service and wait for the response + std::shared_future< + std::shared_ptr> + future = tensor_client_->async_send_request(request); + + std::future>> + return_tensors = std::async(std::launch::async, [future]() { + std::vector> + output_tensors; + for (const auto &tensor : future.get()->tensor2) { + output_tensors.emplace_back( + std::make_shared(tensor)); + } + return output_tensors; + }); + + // RCLCPP_INFO(this->get_logger(), "Here all good"); + // std::vector> r_tensors = return_tensors.get(); + // auto tensor_ptr = r_tensors[0]; // Assume we're just printing the first tensor + // if (tensor_ptr->data_type != 9) { // 9 = float32 + // RCLCPP_WARN(this->get_logger(), "Tensor data_type is not float32 (got %d)", tensor_ptr->data_type); + // return; + // } + + // size_t num_floats = tensor_ptr->data.size() / sizeof(float); + // const float* float_data = reinterpret_cast(tensor_ptr->data.data()); + + // std::ostringstream oss; + // oss << "return_tensors: ["; + // for (size_t i = 0; i < num_floats; ++i) { + // oss << float_data[i]; + // if (i < num_floats - 1) + // oss << ", "; + // } + // oss << "]"; + + // RCLCPP_INFO(this->get_logger(), "return_tensors = %s", oss.str().c_str()); + + return return_tensors; +} + +bool GATPlannerImplementation::buildDecoderTensor( + std::map + &agent_features, + neuromesh_interfaces::msg::Tensor &own_feature, + neuromesh_interfaces::msg::Tensor &aggregated_tensor) { + if (agent_features.size() != 5) { + RCLCPP_INFO(this->get_logger(), "Expected 5 feature messages, got %zu", + agent_features.size()); + return false; + } + + std::vector sorted_ids; + for (const auto &pair : agent_features) { + sorted_ids.push_back(pair.first); + } + std::sort(sorted_ids.begin(), sorted_ids.end()); + + own_feature.name = "own_features"; + aggregated_tensor.name = "other_features"; + own_feature.data_type = 9; + aggregated_tensor.data_type = 9; + + for (const std::string &id : sorted_ids) { + const auto &feature_msg = agent_features.at(id); + if (id == this->id_) { + own_feature = feature_msg->tensor; + } else { + if (aggregated_tensor.data.empty()) { + aggregated_tensor = feature_msg->tensor; + } else { + aggregated_tensor.data.insert(aggregated_tensor.data.end(), + feature_msg->tensor.data.begin(), + feature_msg->tensor.data.end()); + } + } + } + + return true; +} +} // namespace GATPlannerNeuromeshNode + +#include "rclcpp_components/register_node_macro.hpp" +RCLCPP_COMPONENTS_REGISTER_NODE(GATPlannerNeuromeshNode::GATPlannerImplementation) diff --git a/neuromesh_platform_r2/src/gat_planner_model_neuromesh_node.cpp b/neuromesh_platform_r2/src/gat_planner_model_neuromesh_node.cpp new file mode 100644 index 0000000..475578d --- /dev/null +++ b/neuromesh_platform_r2/src/gat_planner_model_neuromesh_node.cpp @@ -0,0 +1,772 @@ +#include "neuromesh_platform_r2/gat_planner_neuromesh_node.h" +#include "rclcpp/rclcpp.hpp" +#include "chrono" + +namespace GATPlannerNeuromeshNode { +GATPlannerNeuromeshNode :: GATPlannerNeuromeshNode(const rclcpp::NodeOptions &options): Node("GAT_planner_neuromesh_node", options) +{ + // Declare node parameters + this->declare_parameter("encoder_model_name", "default_encoder_model"); + this->declare_parameter("decoder_model1_name", "default_decoder_model"); + this->declare_parameter("decoder_model2_name", "default_decoder_model"); + this->declare_parameter("topic_prefix", "features_"); + this->declare_parameter("pos_topic_prefix", "pos_"); + this->declare_parameter("output_topic", "gnn_output_"); + this->declare_parameter("planning_frame", "map"); + this->declare_parameter("decoder_cycle_length", 1000); + this->declare_parameter("encoder_cycle_length", 1000); + this->declare_parameter("encoder_await_length", 10000); + this->declare_parameter("id", "default_id"); + this->declare_parameter("image_qos_profile", "default"); + this->declare_parameter("features_qos_profile", "default"); + this->declare_parameter("output_qos_profile", "default"); + this->declare_parameter("agents", ""); + this->declare_parameter("to_nchw", true); + this->declare_parameter("ints_to_floats", true); + this->declare_parameter("goal_poses_yaml_file", "goal_poses.yaml"); + this->declare_parameter("goals_sending_delay", 10.0); + this->declare_parameter("service_adapter_type", "topic"); + + // Get node parameters + this->get_parameter("encoder_model_name", encoder_model_name_); + this->get_parameter("decoder_model1_name", decoder_model1_name_); + this->get_parameter("decoder_model2_name", decoder_model2_name_); + this->get_parameter("topic_prefix", topic_prefix_); + this->get_parameter("pos_topic_prefix", pos_topic_prefix_); + this->get_parameter("output_topic", output_topic_); + this->get_parameter("planning_frame", planning_frame_); + this->get_parameter("decoder_cycle_length", decoder_cycle_length_); + this->get_parameter("encoder_cycle_length", encoder_cycle_length_); + this->get_parameter("encoder_await_length", encoder_await_length_); + this->get_parameter("id", id_); + this->get_parameter("image_qos_profile", image_qos_profile_); + this->get_parameter("features_qos_profile", features_qos_profile_); //for both input and output + this->get_parameter("output_qos_profile", output_qos_profile_); + this->get_parameter("agents", agents_); + this->get_parameter("to_nchw", to_nchw_); + this->get_parameter("goal_poses_yaml_file", goal_poses_yaml_file); + this->get_parameter("goals_sending_delay", goals_sending_delay_); + this->get_parameter("service_adapter_type", service_adapter_type_); + + auto feature_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(features_qos_profile_)); + auto output_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(output_qos_profile_)); + feature_publisher_ = this->create_publisher(topic_prefix_ + id_, feature_qos); + + // Initialize GNN result publisher and + gnn_result_publisher_ = this->create_publisher( + output_topic_ + id_, feature_qos + ); + + second_decoder_result_publisher_ = this->create_publisher( + "second_decoder_result_topic", 10 // topic name and queue size + ); + + //PLACEHOLDER: update available_agents + + all_agents = splitAgentString(agents_); + for (const auto& agent : all_agents) { + RCLCPP_DEBUG(this->get_logger(), "%s", agent.c_str()); + } + all_agents.erase(id_); //remove self from list + + // Print out state vector elements + neuromesh_interfaces::msg::StateVector tmp; + for (uint i = 0; i < 6; ++i) + { + tmp.state_vector.push_back(0.); + } + current_states_.insert(std::pair(id_, tmp)); + + // for (const auto& id : all_agents) { + // current_states_.emplace(id, tmp); + // } + + available_agents = all_agents; + + // Print out all agent names in the agent list + for (std::string id : all_agents){ + RCLCPP_DEBUG(this->get_logger(), "Going through all agents to create subscriptions"); + RCLCPP_DEBUG(this->get_logger(), "Id: %s", id.c_str()); + this->createSubscription(feature_subscriptions_, id, feature_qos); + this->createPosSubscription(pos_subscriptions_, id, feature_qos); + this->createGNNSubscription(gnn_subscriptions_, id, feature_qos); + } + + // Extra print statements for debugging + for(const auto& [key, value] : current_states_) { + RCLCPP_DEBUG( + get_logger(), + "Key: %s, State Vector Size: %zu", + key.c_str(), + value.state_vector.size() + ); + } + + // Load the goal poses for the robots from the yaml file + load_goal_poses_from_yaml(); + + // Create timers for encoder and decoder cycles + decoder_timer_ = this->create_wall_timer(std::chrono::duration(decoder_cycle_length_), std::bind(&GATPlannerNeuromeshNode::run_decoder_cycle, this)); + encoder_timer_ = this->create_wall_timer(std::chrono::duration(encoder_cycle_length_), std::bind(&GATPlannerNeuromeshNode::run_encoder_cycle, this)); + fresh_encoder_cycle = true; + waypoint_cmd_sent_ = false; + + // Initialize tf buffers + tf_buffer_ = std::make_shared(this->get_clock()); + tf_listener_ = std::make_shared(*tf_buffer_); +} + +void GATPlannerNeuromeshNode::load_goal_poses_from_yaml() { + try { + // Clear existing goal poses + goal_poses_.clear(); + + // Load the YAML file + YAML::Node config = YAML::LoadFile(goal_poses_yaml_file); + + RCLCPP_DEBUG(this->get_logger(), "Loading goals for all robots"); + + // Iterate through all top-level nodes (robot names) + for (const auto& robot_entry : config) { + std::string robot_name = robot_entry.first.as(); + + // Check if this robot has a goals section + if (!robot_entry.second["goals"] || !robot_entry.second["goals"].IsSequence()) { + RCLCPP_WARN(this->get_logger(), "No valid goals found for robot '%s'", robot_name.c_str()); + continue; + } + + // Process each goal for this robot + for (const auto& goal : robot_entry.second["goals"]) { + if (!goal["position"] || !goal["position"]["x"] || !goal["position"]["y"]) { + RCLCPP_WARN(this->get_logger(), "Skipping malformed goal entry for robot '%s'", robot_name.c_str()); + continue; + } + + geometry_msgs::msg::Pose pose; + pose.position.x = goal["position"]["x"].as(); + pose.position.y = goal["position"]["y"].as(); + pose.position.z = 0.0; // Set to 0 if not needed + + goal_poses_.push_back(pose); + + RCLCPP_DEBUG( + this->get_logger(), + "Loaded goal for robot '%s': Position (%.2f, %.2f)", + robot_name.c_str(), + pose.position.x, + pose.position.y + ); + } + } + + RCLCPP_DEBUG( + this->get_logger(), + "Successfully loaded %zu goal poses for all robots", + goal_poses_.size() + ); + + } catch (const YAML::Exception& e) { + RCLCPP_ERROR( + this->get_logger(), + "Failed to load goals from YAML file: %s", + e.what() + ); + } +} + +void GATPlannerNeuromeshNode::pos_callback(const nav_msgs::msg::Odometry::SharedPtr msg) { + + // Parse position data + pos_callback_complete = true; + + current_states_[id_].state_vector[0] = msg->pose.pose.position.x; + current_states_[id_].state_vector[1] = msg->pose.pose.position.y; + + // Set orientation to [0, 0, 0, 1] + current_states_[id_].state_vector[2] = 0.0; + current_states_[id_].state_vector[3] = 0.0; + current_states_[id_].state_vector[4] = 0.0; + current_states_[id_].state_vector[5] = 1.0; + + // Keep spamming others until goal is sent + if(!waypoint_cmd_sent_) { + run_encoder_cycle(); + } + +} + +void GATPlannerNeuromeshNode::neighbor_pos_callback(const nav_msgs::msg::Odometry::SharedPtr msg, const std::string &id) { + // if empty key, initialize state_vector + auto& state = current_states_[id]; + if (state.state_vector.size() != 6) { + state.state_vector.resize(6, 0.0); + } + + current_states_[id].state_vector[0] = msg->pose.pose.position.x; + current_states_[id].state_vector[1] = msg->pose.pose.position.y; + + // Set orientation to [0, 0, 0, 1] + current_states_[id].state_vector[2] = 0.0; + current_states_[id].state_vector[3] = 0.0; + current_states_[id].state_vector[4] = 0.0; + current_states_[id].state_vector[5] = 1.0; +} + +neuromesh_interfaces::msg::Tensor GATPlannerNeuromeshNode::calculate_input_features(const std::vector& state_vector, const std::vector neighbors_ids) { + neuromesh_interfaces::msg::Tensor tensor_msg; + // Set tensor dimensions + tensor_msg.shape.dims = {1, static_cast((1 + goal_poses_.size() + 2)*2)}; // ego robot pos + relative pos to 4 goals + relative pos to 2 neighbors + + tensor_msg.data_type = 9; // float32 + + std::vector input_feature; + input_feature.push_back(state_vector[0] / 2.0f); + input_feature.push_back(state_vector[1] / 2.0f); + + // Relative pos to the goals + for (size_t j = 0; j < goal_poses_.size(); ++j) { + float goal_x = goal_poses_[j].position.x; + float goal_y = goal_poses_[j].position.y; + + float dx = (goal_x - state_vector[0]) / 4.0f; // hard coded normalization by 2*env_bound + float dy = (goal_y - state_vector[1]) / 4.0f; + + input_feature.push_back(dx); + input_feature.push_back(dy); + } + // Relative pos to the robots + for (size_t n = 0; n < neighbors_ids.size(); ++n) {\ + float neighbor_x = current_states_[neighbors_ids[n]].state_vector[0]; + float neighbor_y = current_states_[neighbors_ids[n]].state_vector[1]; + + float dx = (neighbor_x - state_vector[0]) / 4.0f; // hard coded normalization by 2*env_bound + float dy = (neighbor_y - state_vector[1]) / 4.0f; + + // Write into input_feat starting at index 2 + input_feature.push_back(dx); + input_feature.push_back(dy); + } + + // Directly copy float values to tensor data + tensor_msg.data.resize(input_feature.size() * sizeof(float)); + std::copy( + reinterpret_cast(input_feature.data()), + reinterpret_cast(input_feature.data() + input_feature.size()), + tensor_msg.data.begin() + ); + + // Set data type explicitly + tensor_msg.data_type = 9; + + // Set shape explicitly + tensor_msg.shape.dims = {1, static_cast(input_feature.size())}; + + // Optionally set strides + tensor_msg.strides = {static_cast(input_feature.size()), 1}; + + return tensor_msg; +} + +std::vector GATPlannerNeuromeshNode::get_closest_neighbors() { + const auto& ego_state = current_states_[id_]; + double ego_x = ego_state.state_vector[0]; + double ego_y = ego_state.state_vector[1]; + + std::vector> distances; + for (const auto& [agent_id, state] : current_states_) { + if (agent_id == id_) continue; // skip ego + + double dx = state.state_vector[0] - ego_x; + double dy = state.state_vector[1] - ego_y; + double dist = std::sqrt(dx * dx + dy * dy); + + distances.emplace_back(agent_id, dist); + } + std::sort(distances.begin(), distances.end(), + [](const auto& a, const auto& b) { + return a.second < b.second; + }); + std::vector closest_ids; + for (size_t i = 0; i < std::min(size_t(2), distances.size()); ++i) { + closest_ids.push_back(distances[i].first); + } + return closest_ids; +} + +neuromesh_interfaces::msg::Feature GATPlannerNeuromeshNode::buildFeatureMessage(const neuromesh_interfaces::msg::Tensor& tensor) +{ + neuromesh_interfaces::msg::Feature feature_msg = neuromesh_interfaces::msg::Feature(); + + feature_msg.tensor = tensor; + feature_msg.id = id_; + feature_msg.timestamp = this->get_clock()->now(); + return feature_msg; +} + +void GATPlannerNeuromeshNode::feature_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg) { + std::string uuid = msg->id; + received_features_[uuid] = msg; + feature_buffer_timestamp_[uuid] = msg->timestamp.sec + msg->timestamp.nanosec*1e-9; + auto feature_now = this->get_clock()->now(); + double feature_timestamp = feature_now.seconds() + feature_now.nanoseconds() / 1e9; + RCLCPP_DEBUG(this->get_logger(), "Delay between receiving and publishing features %.9f", feature_timestamp - feature_buffer_timestamp_[uuid]); + + // std::cout << "Delay between receiving and publishing features " << feature_now.seconds() - feature_buffer_timestamp_[uuid] << " seconds" << std::endl + // << "feature_buff_timestamp_[uuid]: " << feature_buffer_timestamp_[uuid] << std::endl + // << "msg->timestamp.sec: " << msg->timestamp.sec << std::endl + // << "msg->timestamp.nanosec: " << msg->timestamp.nanosec << std::endl + // << "msg->timestamp.nanosec*1e-9: " << msg->timestamp.nanosec*1e-9 << std::endl + // // << "feature_timestamp: " << feature_timestamp << std::endl + // << "feature_now.seconds(): " << feature_now.seconds() << std::endl + // << "feature_now.nanoseconds(): " << feature_now.nanoseconds() << std::endl + // << "feature_now.nanoseconds() / 1e9: " << feature_now.nanoseconds() / 1e9 << std::endl + // << "Clock type: " << this->get_clock()->get_clock_type() << std::endl; +} + +//perform inference on tensor using model called model_name +//PLACEHOLDER +std::future>> GATPlannerNeuromeshNode::performInference(const std::string& model_name, const std::vector& tensors) +{ + std::promise>> prom; + std::future>> r = prom.get_future(); + std::vector> t(1, std::make_shared()); + t[0]->result = 1; // Cannot reach engine error code + prom.set_value(std::move(t)); + return r; +} + +//Aggregates features from available agents (and itself) into a single tensor +//PLACEHOLDER +bool GATPlannerNeuromeshNode::buildDecoderTensor( + std::map& agent_features, + neuromesh_interfaces::msg::Tensor& own_feature, + neuromesh_interfaces::msg::Tensor& aggregated_tensor) +{ + return true; +} + +//Adds a subscription to the feature_subscriptions_ map +void GATPlannerNeuromeshNode::createSubscription(std::map::SharedPtr>& subscription_map, std::string id, rclcpp::QoS qos) +{ + std::string topic = topic_prefix_ + id; + + RCLCPP_DEBUG(this->get_logger(), "creating subscription for topic %s", topic.c_str()); + + rclcpp::Subscription::SharedPtr feature_subscription_ = + this->create_subscription( + topic, + qos, + std::bind(&GATPlannerNeuromeshNode::feature_callback, this, std::placeholders::_1)); + + subscription_map.insert( {id, feature_subscription_} ); +} + +void GATPlannerNeuromeshNode::createPosSubscription(std::map::SharedPtr>& pos_subscription_map, std::string id, rclcpp::QoS qos) +{ + std::string topic = pos_topic_prefix_ + id; + + RCLCPP_DEBUG(this->get_logger(), "creating subscription for topic %s", topic.c_str()); + + auto lambda_callback = [this, id](const nav_msgs::msg::Odometry::SharedPtr msg) { + this->neighbor_pos_callback(msg, id); + }; + + rclcpp::Subscription::SharedPtr pos_subscription_ = + this->create_subscription( + topic, + qos, + lambda_callback); + + pos_subscription_map.insert( {id, pos_subscription_} ); +} + +//Remove subscription form the feature_subscriptions_ map +void GATPlannerNeuromeshNode::removeSubscription(std::map::SharedPtr> subscription_map, std::string id) +{ + subscription_map.erase(id); +} + +//Adds a subscription to the gnn_subscriptions_ map +void GATPlannerNeuromeshNode::createGNNSubscription(std::map::SharedPtr>& gnn_subscription_map, std::string id, rclcpp::QoS qos) +{ + std::string topic = output_topic_ + id; + + RCLCPP_DEBUG(this->get_logger(), "creating gnn subscription for topic %s", topic.c_str()); + + rclcpp::Subscription::SharedPtr gnn_subscription_ = + this->create_subscription( + topic, + qos, + std::bind(&GATPlannerNeuromeshNode::gnn_result_callback, this, std::placeholders::_1)); + + gnn_subscription_map.insert( {id, gnn_subscription_} ); +} + +//Remove subscription form the gnn_subscriptions_ map +void GATPlannerNeuromeshNode::removeGNNSubscription(std::map::SharedPtr> gnn_subscription_map, std::string id) +{ + gnn_subscription_map.erase(id); +} + +//Convert string to ROS2 QoS profile +//from https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_nvblox/nvblox_ros_common/src/qos.cpp#L26 +rmw_qos_profile_t GATPlannerNeuromeshNode::parseQoSString(const std::string& str) +{ + std::string profile = str; + // Convert to upper case. + std::transform(profile.begin(), profile.end(), profile.begin(), ::toupper); + + if (profile == "SYSTEM_DEFAULT") { + return rmw_qos_profile_system_default; + } + if (profile == "DEFAULT") { + return rmw_qos_profile_default; + } + if (profile == "PARAMETER_EVENTS") { + return rmw_qos_profile_parameter_events; + } + if (profile == "SERVICES_DEFAULT") { + return rmw_qos_profile_services_default; + } + if (profile == "PARAMETERS") { + return rmw_qos_profile_parameters; + } + if (profile == "SENSOR_DATA") { + return rmw_qos_profile_sensor_data; + } + RCLCPP_WARN_STREAM( + rclcpp::get_logger("parseQosString"), + "Unknown QoS profile: " << profile << ". Returning profile: DEFAULT"); + return rmw_qos_profile_default; +} + +//Split agent string parameter into vector of agent ids +std::set GATPlannerNeuromeshNode::splitAgentString(std::string str) +{ + std::set agents; + const std::string delimiter = ","; + + size_t pos = 0; + std::string token; + while ((pos = str.find(delimiter)) != std::string::npos) { + token = str.substr(0, pos); + agents.insert(token); + str.erase(0, pos + delimiter.length()); + } + agents.insert(str); + return agents; +} + +void GATPlannerNeuromeshNode::run_encoder_cycle() { + // Run the encoder on all feature vectors + // RCLCPP_INFO(this->get_logger(), "current_states_ size = %d", current_states_.size()); + if (fresh_encoder_cycle && current_states_.size()==4 && pos_callback_complete && !goal_poses_.empty() && !waypoint_cmd_sent_) { + fresh_encoder_cycle = false; + startClock("encoder_inference"); + + // TO CHANGE + std::vector closest_neighbors_ids = get_closest_neighbors(); + neuromesh_interfaces::msg::Tensor input_features = calculate_input_features( + std::vector(current_states_[id_].state_vector.begin(), + current_states_[id_].state_vector.end()), closest_neighbors_ids + ); + size_t num_floats = input_features.data.size() / sizeof(float); + const float* float_data = reinterpret_cast(input_features.data.data()); + + std::ostringstream oss; + oss << "input_features: ["; + for (size_t i = 0; i < num_floats; ++i) { + oss << float_data[i]; + if (i < num_floats - 1) + oss << ", "; + } + oss << "]"; + + // RCLCPP_INFO(this->get_logger(), "input features = %s", oss.str().c_str()); + + auto encoder_now = this->get_clock()->now(); + // RCLCPP_INFO(this->get_logger(), "Performing inference"); + encoder_result = performInference(encoder_model_name_, {input_features}); + auto encoder_end = this->get_clock()->now(); + + RCLCPP_DEBUG(this->get_logger(), "Finished performing Inference on Encoder"); + } + + // Keep publishing features to others + if (!fresh_encoder_cycle && encoder_result.valid()) { + RCLCPP_INFO(this->get_logger(), "I am here"); + auto encoder_status = encoder_result.wait_for(std::chrono::milliseconds(0)); + if (encoder_status == std::future_status::ready) { + stopClock("encoder_inference"); + std::vector> encoded_features = encoder_result.get(); + + std::shared_ptr feature_tensor = encoded_features[0]; + + // Print out encoder model output + float max_prob, next; + size_t size = feature_tensor->data.size(); + std::memcpy(&max_prob, &feature_tensor->data[0], sizeof(float)); + std::stringstream ss; + ss << "[" << std::fixed << std::setprecision(2) << max_prob; + for (size_t i = 4; i < size; i +=4) { + if (i + 3 < size) + { + std::memcpy(&next, &feature_tensor->data[i], sizeof(float)); + ss << " " << next << " "; + if (((i/4) % 5) == 0 ){ + ss << "\n"; + } + } + } + ss << "]"; + + RCLCPP_INFO(this->get_logger(), "Encoder features output:\n Size: %zu\n Output: %s", size, ss.str().c_str()); + + neuromesh_interfaces::msg::Feature feature_msg = buildFeatureMessage(*feature_tensor.get()); + + // Publish features + RCLCPP_DEBUG_THROTTLE(this->get_logger(), + *this->get_clock(), + 1000, + "Publishing features."); + feature_publisher_->publish(feature_msg); + + // Store features + received_features_[this->id_] = std::make_shared(feature_msg); + + + feature_buffer_timestamp_[this->id_] = feature_msg.timestamp.sec + feature_msg.timestamp.nanosec*1e-9; + + fresh_encoder_cycle = true; + RCLCPP_DEBUG(this->get_logger(), "Encoder cycle completed"); + + } else if (checkClock("encoder_inference") >= encoder_await_length_) { + RCLCPP_WARN(this->get_logger(), "Encoder inference timed out"); + fresh_encoder_cycle = true; + } + } +} + +void GATPlannerNeuromeshNode::run_decoder_cycle() { + // Handle previous decoder inference result + if (gnn_result_future.valid()) { + auto decoder_status = gnn_result_future.wait_for(std::chrono::milliseconds(0)); + if (decoder_status == std::future_status::ready) { + stopClock("decoder_inference"); + std::vector> gnn_result = gnn_result_future.get(); + + std::shared_ptr gnn_result_tensor = gnn_result[0]; + + // Print first decoder output values + float max_prob, next; + size_t size = gnn_result_tensor->data.size(); + std::memcpy(&max_prob, &gnn_result_tensor->data[0], sizeof(float)); + std::stringstream ss; + ss << "[" << std::fixed << std::setprecision(2) << max_prob; + for (size_t i = 4; i < size; i +=4) { + if (i + 3 < size) + { + std::memcpy(&next, &gnn_result_tensor->data[i], sizeof(float)); + ss << " " << next << " "; + if (((i/4) % 5) == 0 ){ + ss << "\n"; + } + } + } + ss << "]"; + + const auto& data = gnn_result_tensor->data; + + // Publish the GNN result to other robots + neuromesh_interfaces::msg::Feature gnn_msg = build_gnn_msg(*gnn_result_tensor.get()); + + // Publish the result + gnn_result_publisher_->publish(gnn_msg); + + RCLCPP_DEBUG(this->get_logger(), "First Decoder inference completed and published"); + + // Store local GNN result + received_gnn_results_[this->id_] = std::make_shared(gnn_msg); + + // Prepare for potential second-stage decoding + prepare_second_stage_decoding(); + + RCLCPP_DEBUG(this->get_logger(), "Second Decoder inference completed and published"); + } + } + + // Check if we have enough feature messages (5 total) + if (!received_features_.empty() && !waypoint_cmd_sent_) { + startClock("decoder_inference"); + + // Step 4: Combine features of other robots + // Prepare encoder output tensor (from previous encoder cycle) + encoder_output_tensor; + neuromesh_interfaces::msg::Tensor aggregated_tensor; + if (buildDecoderTensor(received_features_, encoder_output_tensor, aggregated_tensor)) { + std::vector decoder_tensors = {encoder_output_tensor, aggregated_tensor, encoder_output_tensor}; + auto first_gnn_start = this->get_clock()->now(); + gnn_result_future = performInference(decoder_model1_name_, decoder_tensors); + auto first_gnn_end = this->get_clock()->now(); + //std::cout <<"Computation time first gnn round" <get_logger(), "Failed to build decoder tensor"); + } + + // Clear received features for the next cycle + received_features_.clear(); + feature_buffer_timestamp_.clear(); + } +} + +neuromesh_interfaces::msg::Feature GATPlannerNeuromeshNode::build_gnn_msg(const neuromesh_interfaces::msg::Tensor& gnn_result) { + // Convert GNN result to Tensor message + neuromesh_interfaces::msg::Feature result_msg = neuromesh_interfaces::msg::Feature(); + result_msg.tensor = gnn_result; + + // Add metadata if needed + result_msg.id = this->id_; + result_msg.timestamp = this->get_clock()->now(); + + return result_msg; +} + +void GATPlannerNeuromeshNode::gnn_result_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg) { + // Store received GNN result + std::string gnn_id = msg->id; + received_gnn_results_[gnn_id] = msg; + // Optional: Log received result + RCLCPP_DEBUG(this->get_logger(),"Received GNN result from robot %d", msg->id); +} + +void GATPlannerNeuromeshNode::prepare_second_stage_decoding() { + + if (second_decoder_result_future.valid()) { + auto second_decoder_status = second_decoder_result_future.wait_for(std::chrono::milliseconds(0)); + if (second_decoder_status == std::future_status::ready) { + stopClock("decoder_inference"); + std::vector> second_decoder_result = second_decoder_result_future.get(); + + // Ensure the vector is not empty + if (!second_decoder_result.empty()) { + // Get the first tensor + std::shared_ptr second_decoder_result_tensor = second_decoder_result[0]; + + // Find the index of the maximum probability + int max_prob_index = 0; + + // Final step to get decoder output, print the values and send it as waypoint goals to robots + // Assuming the Tensor has a data field that is a vector of floats + if (!second_decoder_result_tensor->data.empty()) { + float max_prob, next; + size_t size = second_decoder_result_tensor->data.size(); + std::memcpy(&max_prob, &second_decoder_result_tensor->data[0], sizeof(float)); + + if (!goal_adapter_) { + goal_adapter_ = AdapterFactory::create(service_adapter_type_, this->shared_from_this()); + } + if (!goal_adapter_) { + RCLCPP_ERROR(this->get_logger(), "Failed to create service adapter of type '%s'", service_adapter_type_.c_str()); + return; + } + + NavigationGoal goal; + goal.x = goal_poses_[max_prob_index].position.x; + goal.y = goal_poses_[max_prob_index].position.y; + goal.planning_frame = planning_frame_; + goal_adapter_->sendGoal(goal); + } + } + } + } + + // Check if we have received GNN results from all robots + if (!received_gnn_results_.empty()) { + // Prepare aggregated GNN results for second-stage decoder + neuromesh_interfaces::msg::Tensor gnn_output_tensor; + neuromesh_interfaces::msg::Tensor aggregated_gnn_output_tensor; + if (buildDecoderTensor(received_gnn_results_, gnn_output_tensor, aggregated_gnn_output_tensor)) { + + second_decoder_result_future = performInference(decoder_model2_name_, {gnn_output_tensor, aggregated_gnn_output_tensor, encoder_output_tensor}); + + auto gnn_now = this->get_clock()->now(); + + for (const auto& entry : feature_buffer_timestamp_) { + const std::string& robot_name = entry.first; + double timestamp = entry.second; + } + + } else { + RCLCPP_ERROR(this->get_logger(), "Failed to build decoder tensor"); + } + + // Clear received GNN results for next cycle + received_gnn_results_.clear(); + } +} + +void GATPlannerNeuromeshNode::startClock(std::string phase){ + int64_t now_time = std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); + + times[phase] = {now_time, false}; +} +void GATPlannerNeuromeshNode::stopClock(std::string phase){ + if (times[phase].second){ + return; // clock already stopped + } + int64_t start_time = times[phase].first; + int64_t now_time = std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); + + times[phase] = {now_time - start_time, true}; +} + +int64_t GATPlannerNeuromeshNode::checkClock(std::string phase){ + if (times[phase].second){ + return times[phase].first; // clock already stopped + } + + //stopclock calculations without saving + int64_t start_time = times[phase].first; + int64_t now_time = std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); + + return now_time - start_time; +} + +//TODO generalize to a transpose function +//tested for use with integers +neuromesh_interfaces::msg::Tensor GATPlannerNeuromeshNode::convert_to_nchw(const neuromesh_interfaces::msg::Tensor& input){ + std::vector new_data; + const std::vector& old_data = input.data; + + std::vector strides = input.strides; + std::vector dims = input.shape.dims; + unsigned long bytedepth = input.strides[0] / input.shape.dims[0]; + + + for(int i = 0; i < dims[2]; i++){ + for(int j = 0; j < dims[0]; j++){ + for(int k = 0; k < dims[1]; k++){ + new_data.push_back(old_data[ (j * strides[0]) + (k * strides[1]) + (i * strides[2]) ] ); + } + } + } + + neuromesh_interfaces::msg::Tensor new_tensor = std::move(input); + new_tensor.data = new_data; + new_tensor.shape.dims = {input.shape.dims[2], input.shape.dims[0], input.shape.dims[1]}; + new_tensor.strides = {dims[1] * dims[2] * bytedepth, dims[2] * bytedepth, bytedepth}; + + return new_tensor; +} + +float int_to_scaled_float(int i){ return static_cast(i) / 255.0;} +} diff --git a/neuromesh_platform_r2/src/starting_poses_sender.cpp b/neuromesh_platform_r2/src/starting_poses_sender.cpp new file mode 100644 index 0000000..9d9961c --- /dev/null +++ b/neuromesh_platform_r2/src/starting_poses_sender.cpp @@ -0,0 +1,288 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace goal_sender +{ +class GoalSenderNode : public rclcpp::Node +{ +public: + explicit GoalSenderNode(const rclcpp::NodeOptions & options) + : Node("goal_sender_node", options) + { + this->declare_parameter("start_poses_yaml_file", "start_poses.yaml"); + this->declare_parameter("id", ""); + this->declare_parameter("robot_id", ""); + this->declare_parameter("initialization_timeout", 10.0); + this->declare_parameter("mission_goal_topic", "mission_goals"); + this->declare_parameter("mission_yaml_topic", "mission_yaml"); + this->declare_parameter("mission_command_topic", "mission_command"); + this->declare_parameter("publish_mission_yaml", true); + this->declare_parameter("start_command_value", 0); + + std::string robot_id = this->get_parameter("robot_id").as_string(); + std::string id = this->get_parameter("id").as_string(); + id_ = !robot_id.empty() ? robot_id : id; + if (id_.empty()) { + id_ = this->get_namespace(); + if (!id_.empty() && id_[0] == '/') { + id_.erase(0, 1); + } + } + + std::string start_poses_yaml_file = this->get_parameter("start_poses_yaml_file").as_string(); + yaml_file_path_ = resolveYamlPath(start_poses_yaml_file); + publish_mission_yaml_ = this->get_parameter("publish_mission_yaml").as_bool(); + start_command_value_ = this->get_parameter("start_command_value").as_int(); + + std::string mission_goal_topic = this->get_parameter("mission_goal_topic").as_string(); + std::string mission_yaml_topic = this->get_parameter("mission_yaml_topic").as_string(); + std::string mission_command_topic = this->get_parameter("mission_command_topic").as_string(); + + goal_pub_ = this->create_publisher(mission_goal_topic, 10); + command_pub_ = this->create_publisher(mission_command_topic, 10); + if (publish_mission_yaml_) { + yaml_pub_ = this->create_publisher(mission_yaml_topic, 10); + } + + tf_buffer_ = std::make_unique(this->get_clock()); + tf_listener_ = std::make_shared(*tf_buffer_); + + std::string odom_topic = "/" + id_ + "/odom"; + odom_sub_ = this->create_subscription( + odom_topic, + 10, + std::bind(&GoalSenderNode::odom_callback, this, std::placeholders::_1)); + + // Set initialization timeout + double timeout = this->get_parameter("initialization_timeout").as_double(); + init_timeout_ = this->now() + rclcpp::Duration::from_seconds(timeout); + + timer_ = this->create_wall_timer( + std::chrono::seconds(1), + std::bind(&GoalSenderNode::timerCallback, this)); + + loadGoalsFromYaml(); + + RCLCPP_INFO(this->get_logger(), "Waiting for robot %s initialization...", id_.c_str()); + } + +private: + std::string resolveYamlPath(const std::string & yaml_file) const + { + std::ifstream direct_file(yaml_file); + if (direct_file.good()) { + return yaml_file; + } + + const std::string package_share = + ament_index_cpp::get_package_share_directory("neuromesh_platform_r2"); + const std::string config_candidate = package_share + "/config/" + yaml_file; + std::ifstream config_file(config_candidate); + if (config_file.good()) { + return config_candidate; + } + + return yaml_file; + } + + void odom_callback(const nav_msgs::msg::Odometry::SharedPtr msg) + { + if (!odom_received_) { + odom_received_ = true; + RCLCPP_INFO(this->get_logger(), "Received first odometry message for robot %s", id_.c_str()); + } + } + + bool check_transform_available() + { + try { + geometry_msgs::msg::TransformStamped transform = + tf_buffer_->lookupTransform("map", id_ + "/odom", tf2::TimePointZero); + + if (!tf_available_) { + RCLCPP_INFO(this->get_logger(), "Transform from odom to map is now available for robot %s", + id_.c_str()); + tf_available_ = true; + } + return true; + } + catch (const tf2::TransformException & ex) { + return false; + } + } + + void loadGoalsFromYaml() + { + try { + YAML::Node config = YAML::LoadFile(yaml_file_path_); + + if (!config[id_]) { + RCLCPP_ERROR( + this->get_logger(), + "Robot ID '%s' not found in YAML file", id_.c_str()); + return; + } + + const YAML::Node& robot_goals = config[id_]["goals"]; + + if (!robot_goals || !robot_goals.IsSequence()) { + RCLCPP_ERROR( + this->get_logger(), + "No valid goals found for robot '%s'", id_.c_str()); + return; + } + + for (const auto& goal : robot_goals) { + geometry_msgs::msg::Pose pose; + + if (goal["position"] && + goal["position"]["x"] && + goal["position"]["y"]) { + + pose.position.x = goal["position"]["x"].as(); + pose.position.y = goal["position"]["y"].as(); + + goal_poses_.push_back(pose); + } else { + RCLCPP_WARN( + this->get_logger(), + "Skipping malformed goal entry for robot '%s'", id_.c_str()); + } + } + + RCLCPP_INFO( + this->get_logger(), + "Loaded %zu goals for robot '%s'", goal_poses_.size(), id_.c_str()); + } catch (const YAML::Exception& e) { + RCLCPP_ERROR( + this->get_logger(), + "Failed to load YAML file: %s", e.what()); + } + } + + YAML::Node createMissionYaml() + { + YAML::Node yaml_string; + yaml_string["version"] = 2.0; + yaml_string["frameid"] = id_ + "/map"; + + yaml_string["waypoints"] = YAML::Node(YAML::NodeType::Sequence); + + for (size_t i = 0; i < goal_poses_.size(); ++i) { + YAML::Node wp_node; + const auto& pose = goal_poses_[i]; + + std::vector pose_data{ + static_cast(pose.position.x), + static_cast(pose.position.y), + static_cast(pose.position.z) + }; + + wp_node["name"] = "waypoint" + std::to_string(i + 1); + wp_node["pose"] = pose_data; + wp_node["pose"].SetStyle(YAML::EmitterStyle::Flow); + wp_node["radius"] = 2.0; + + yaml_string["waypoints"].push_back(wp_node); + } + + return yaml_string; + } + + void sendGoals() + { + if (goals_sent_ || goal_poses_.empty()) { + return; + } + + YAML::Node mission_yaml = createMissionYaml(); + + geometry_msgs::msg::PoseArray goal_array; + goal_array.header.stamp = this->now(); + goal_array.header.frame_id = id_ + "/map"; + goal_array.poses = goal_poses_; + goal_pub_->publish(goal_array); + + if (publish_mission_yaml_ && yaml_pub_) { + std_msgs::msg::String yaml_msg; + yaml_msg.data = YAML::Dump(mission_yaml); + yaml_pub_->publish(yaml_msg); + } + + std_msgs::msg::UInt8 command_msg; + command_msg.data = static_cast(start_command_value_); + command_pub_->publish(command_msg); + + RCLCPP_INFO( + this->get_logger(), + "Published %zu starting goals for robot %s using generic mission topics.", + goal_poses_.size(), id_.c_str()); + + goals_sent_ = true; + } + + void timerCallback() + { + if (goals_sent_) { + timer_->cancel(); + return; + } + + check_transform_available(); + + bool robot_ready = odom_received_ && tf_available_; + + if (robot_ready) + { + RCLCPP_INFO(this->get_logger(), "Robot %s is ready, sending goals...", id_.c_str()); + sendGoals(); + } + else if (this->now() > init_timeout_) { + RCLCPP_ERROR(this->get_logger(), + "Robot %s initialization timeout. Status:", id_.c_str()); + RCLCPP_ERROR(this->get_logger(), "Odometry received: %s", + odom_received_ ? "yes" : "no"); + RCLCPP_ERROR(this->get_logger(), "Transform available: %s", + tf_available_ ? "yes" : "no"); + timer_->cancel(); + } + } + + std::unique_ptr tf_buffer_; + std::shared_ptr tf_listener_; + + rclcpp::Subscription::SharedPtr odom_sub_; + rclcpp::Publisher::SharedPtr goal_pub_; + rclcpp::Publisher::SharedPtr yaml_pub_; + rclcpp::Publisher::SharedPtr command_pub_; + + rclcpp::TimerBase::SharedPtr timer_; + + std::string yaml_file_path_; + std::string id_; + double waypoint_radius_; + int start_command_value_{0}; + bool publish_mission_yaml_{true}; + bool goals_sent_{false}; + bool odom_received_{false}; + bool tf_available_{false}; + rclcpp::Time init_timeout_; + std::vector goal_poses_; +}; + +} // namespace goal_sender + +#include "rclcpp_components/register_node_macro.hpp" +RCLCPP_COMPONENTS_REGISTER_NODE(goal_sender::GoalSenderNode) diff --git a/neuromesh_platform_r2/src/vggt_neuromesh_node.cpp b/neuromesh_platform_r2/src/vggt_neuromesh_node.cpp deleted file mode 100644 index b9f8036..0000000 --- a/neuromesh_platform_r2/src/vggt_neuromesh_node.cpp +++ /dev/null @@ -1,762 +0,0 @@ -#include "neuromesh_platform_r2/vggt_neuromesh_node.h" -#include "rclcpp/rclcpp.hpp" -#include "cv_bridge/cv_bridge.h" -#include -#include -#include -#include "chrono" - -namespace vggtNode { -vggtNode::vggtNode(const rclcpp::NodeOptions &options): Node("vggt_node", options) -{ - // Declare node parameters - this->declare_parameter("encoder_model_name", "vggt_encoder"); - this->declare_parameter("decoder_model_name", "vggt_decoder"); - this->declare_parameter("topic_prefix", "features_"); - this->declare_parameter("output_topic", "vggt_output"); - this->declare_parameter("decoder_cycle_length", 3000); - this->declare_parameter("encoder_cycle_length", 3000); - this->declare_parameter("encoder_await_length", 10000); - this->declare_parameter("id", "default_id"); - this->declare_parameter("image_qos_profile", "default"); - this->declare_parameter("features_qos_profile", "default"); - this->declare_parameter("output_qos_profile", "default"); - this->declare_parameter("agents", ""); - this->declare_parameter("to_nchw", true); - this->declare_parameter("ints_to_floats", true); - - // Get node parameters - this->get_parameter("encoder_model_name", encoder_model_name_); - this->get_parameter("decoder_model_name", decoder_model_name_); - this->get_parameter("topic_prefix", topic_prefix_); - this->get_parameter("output_topic", output_topic_); - this->get_parameter("decoder_cycle_length", decoder_cycle_length_); - this->get_parameter("encoder_cycle_length", encoder_cycle_length_); - this->get_parameter("encoder_await_length", encoder_await_length_); - this->get_parameter("id", id_); - this->get_parameter("image_qos_profile", image_qos_profile_); - this->get_parameter("features_qos_profile", features_qos_profile_); - this->get_parameter("output_qos_profile", output_qos_profile_); - this->get_parameter("agents", agents_); - this->get_parameter("to_nchw", to_nchw_); - this->get_parameter("ints_to_floats", ints_to_floats_); - - // Declare VGGT decoder output dimensions parameter - this->declare_parameter("vggt_decoder_output_dimensions", "1,2,9;1,2,392,518,1;1,2,392,518;1,2,392,518,3;1,2,392,518"); - this->get_parameter("vggt_decoder_output_dimensions", decoder_output_dimensions_str); - - auto feature_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(features_qos_profile_)); - auto output_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(output_qos_profile_)); - - // Create publishers - feature_publisher_ = this->create_publisher(topic_prefix_ + id_, feature_qos); - output_publisher_ = this->create_publisher(output_topic_, output_qos); - - // VGGT-specific publishers - depth_robot1_publisher_ = this->create_publisher("depth_robot1", output_qos); - depth_robot2_publisher_ = this->create_publisher("depth_robot2", output_qos); - pointcloud_current_publisher_ = this->create_publisher("pointcloud_current", output_qos); - pointcloud_neighbor_publisher_ = this->create_publisher("pointcloud_neighbor", output_qos); - pointcloud_current_rgb_publisher_ = this->create_publisher("pointcloud_current_rgb", output_qos); - pointcloud_neighbor_rgb_publisher_ = this->create_publisher("pointcloud_neighbor_rgb", output_qos); - - // Add TransformBroadcaster - tf_broadcaster_ = std::make_unique(*this); - - // Add a timer to broadcast the transform periodically - transform_timer_ = this->create_wall_timer( - std::chrono::milliseconds(100), - std::bind(&vggtNode::broadcast_transform, this)); - - // Parse agent list - all_agents = splitAgentString(agents_); - RCLCPP_INFO(this->get_logger(), "VGGT Agents:"); - for (const auto& agent : all_agents) { - RCLCPP_INFO(this->get_logger(), "%s", agent.c_str()); - } - all_agents.erase(id_); // remove self from list - - // Parse decoder output dimensions - decoder_output_dims = string_to_dims(decoder_output_dimensions_str); - - // Create camera subscription - auto image_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(image_qos_profile_)); - camera_subscription_ = this->create_subscription( - "camera", image_qos, - std::bind(&vggtNode::camera_callback, this, std::placeholders::_1)); - - // Create feature subscriptions for other agents - for (const auto& agent : all_agents) { - createSubscription(feature_subscriptions_, agent, feature_qos); - } - - // Create timers - decoder_timer_ = this->create_wall_timer( - std::chrono::milliseconds(decoder_cycle_length_), - std::bind(&vggtNode::process_features, this)); - - encoder_timer_ = this->create_wall_timer( - std::chrono::milliseconds(encoder_cycle_length_), - std::bind(&vggtNode::run_encoder_cycle, this)); - - // Create timer to check encoder results more frequently (100ms) - encoder_result_timer_ = this->create_wall_timer( - std::chrono::milliseconds(100), - std::bind(&vggtNode::process_encoder_result, this)); - - fresh_encoder_cycle = true; - - RCLCPP_INFO(this->get_logger(), "VGGT Node initialized with encoder: %s, decoder: %s", - encoder_model_name_.c_str(), decoder_model_name_.c_str()); -} - -void vggtNode::camera_callback(const sensor_msgs::msg::Image::SharedPtr msg) { - RCLCPP_INFO(this->get_logger(), "=== Camera callback triggered ==="); - - { - std::lock_guard lock(camera_msg_mutex_); - latest_camera_msg_ = msg; - } - - RCLCPP_INFO(this->get_logger(), "fresh_encoder_cycle: %s, encoder_cycle_count: %d", - fresh_encoder_cycle ? "true" : "false", encoder_cycle_count_); - - if (fresh_encoder_cycle) { - RCLCPP_INFO(this->get_logger(), "Starting encoder inference..."); - startClock("encoder_inference"); - - // Convert image to tensor with VGGT preprocessing - auto tensor = imageToTensor(msg); - RCLCPP_INFO(this->get_logger(), "Image converted to tensor: dims=[%s], data_size=%zu", - tensor.shape.dims.empty() ? "empty" : - std::accumulate(tensor.shape.dims.begin(), tensor.shape.dims.end(), std::string(), - [](const std::string& a, uint32_t b) { return a.empty() ? std::to_string(b) : a + ", " + std::to_string(b); }).c_str(), - tensor.data.size()); - - // Perform encoder inference - std::vector input_tensors = {tensor}; - RCLCPP_INFO(this->get_logger(), "Calling performInference for encoder..."); - // Always create a new future for each encoder cycle - encoder_result = performInference(encoder_model_name_, input_tensors); - - fresh_encoder_cycle = false; - encoder_cycle_count_++; - RCLCPP_INFO(this->get_logger(), "Encoder inference started, cycle count: %d", encoder_cycle_count_); - } else { - RCLCPP_DEBUG(this->get_logger(), "Skipping encoder inference - not a fresh cycle"); - } -} - -neuromesh_interfaces::msg::Tensor -vggtNode::imageToTensor(const sensor_msgs::msg::Image::SharedPtr msg) { - neuromesh_interfaces::msg::Tensor tensor; - - try { - // Convert ROS image to OpenCV - cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::RGB8); - cv::Mat image = cv_ptr->image; - - // Resize to VGGT input dimensions: 392x518 - cv::Mat resized_image; - cv::resize(image, resized_image, cv::Size(tensor_width_, tensor_height_)); - - // Normalize to [-1, 1] range - cv::Mat normalized_image; - resized_image.convertTo(normalized_image, CV_32F, 2.0/255.0, -1.0); - - // Set tensor dimensions [1, 3, 392, 518] - tensor.shape.dims = {tensor_batch_size_, tensor_channels_, tensor_height_, tensor_width_}; - - // Convert HWC to CHW format for neural network input - std::vector channels(3); - cv::split(normalized_image, channels); - - // Prepare float data - std::vector float_data; - float_data.reserve(tensor_batch_size_ * tensor_channels_ * tensor_height_ * tensor_width_); - - // Add data in CHW order - for (int c = 0; c < 3; ++c) { - float* channel_data = reinterpret_cast(channels[c].data); - size_t channel_size = tensor_height_ * tensor_width_; - for (size_t i = 0; i < channel_size; ++i) { - float_data.push_back(channel_data[i]); - } - } - - // Convert float data to uint8 data - tensor.data_type = 9; // float32 - tensor.data.resize(float_data.size() * sizeof(float)); - std::memcpy(tensor.data.data(), float_data.data(), float_data.size() * sizeof(float)); - - // Set metadata - tensor.name = msg->header.frame_id + std::to_string(msg->header.stamp.sec) + - std::to_string(msg->header.stamp.nanosec); - - RCLCPP_DEBUG(this->get_logger(), "Image converted to tensor: [%u, %u, %u, %u]", - tensor.shape.dims[0], tensor.shape.dims[1], tensor.shape.dims[2], tensor.shape.dims[3]); - - } catch (cv_bridge::Exception& e) { - RCLCPP_ERROR(this->get_logger(), "cv_bridge exception: %s", e.what()); - } - - return tensor; -} - -void vggtNode::feature_callback(const neuromesh_interfaces::msg::Feature::SharedPtr msg) { - RCLCPP_INFO(this->get_logger(), "=== Feature callback from agent: %s ===", msg->id.c_str()); - RCLCPP_INFO(this->get_logger(), "Feature tensor dims: [%s], data_size: %zu", - msg->tensor.shape.dims.empty() ? "empty" : - std::accumulate(msg->tensor.shape.dims.begin(), msg->tensor.shape.dims.end(), std::string(), - [](const std::string& a, uint32_t b) { return a.empty() ? std::to_string(b) : a + ", " + std::to_string(b); }).c_str(), - msg->tensor.data.size()); - - feature_buffer_[msg->id] = msg; - feature_buffer_timestamp_[msg->id] = this->get_clock()->now().seconds(); - - RCLCPP_INFO(this->get_logger(), "Feature buffer updated, current size: %zu", feature_buffer_.size()); -} - -void vggtNode::process_encoder_result() { - // Check if we have our own encoder result - if (encoder_result.valid() && - encoder_result.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - - RCLCPP_INFO(this->get_logger(), "Encoder result is ready, retrieving..."); - auto encoder_results = encoder_result.get(); - RCLCPP_INFO(this->get_logger(), "Retrieved %zu encoder results", encoder_results.size()); - // Note: After get(), encoder_result becomes invalid and will remain so until next encoder cycle - - if (!encoder_results.empty() && encoder_results[0]) { - RCLCPP_INFO(this->get_logger(), "Encoder result tensor dims: [%s], data_size: %zu", - encoder_results[0]->shape.dims.empty() ? "empty" : - std::accumulate(encoder_results[0]->shape.dims.begin(), encoder_results[0]->shape.dims.end(), std::string(), - [](const std::string& a, uint32_t b) { return a.empty() ? std::to_string(b) : a + ", " + std::to_string(b); }).c_str(), - encoder_results[0]->data.size()); - - // Check if the tensor has valid data - if (encoder_results[0]->data.empty() || encoder_results[0]->shape.dims.empty()) { - RCLCPP_ERROR(this->get_logger(), "Encoder returned empty tensor data or dimensions!"); - } else { - // Store our own feature - auto own_feature = buildFeatureMessage(*encoder_results[0]); - feature_buffer_[id_] = std::make_shared(own_feature); - feature_buffer_timestamp_[id_] = this->get_clock()->now().seconds(); - - // Publish our feature for other agents - feature_publisher_->publish(own_feature); - - stopClock("encoder_inference"); - RCLCPP_INFO(this->get_logger(), "Encoder inference took: %ld ms", checkClock("encoder_inference")); - RCLCPP_INFO(this->get_logger(), "Published feature for agent: %s with tensor size: %zu bytes", - id_.c_str(), own_feature.tensor.data.size()); - } - } else { - RCLCPP_WARN(this->get_logger(), "Encoder results empty or null!"); - } - } -} - -void vggtNode::process_features() { - auto now = std::chrono::system_clock::now(); - auto time_since_epoch = now.time_since_epoch(); - auto seconds = std::chrono::duration_cast(time_since_epoch); - auto nanoseconds = std::chrono::duration_cast(time_since_epoch - seconds); - - RCLCPP_INFO(this->get_logger(), "=== START process_features at %ld.%09ld ===", - seconds.count(), nanoseconds.count()); - - try { - // Process encoder results if available - process_encoder_result(); - - // Check if we have features from neighbor agents for decoder - RCLCPP_INFO(this->get_logger(), "Current feature buffer size: %zu", feature_buffer_.size()); - if (feature_buffer_.size() >= 2) { // Need at least 2 agents (self + 1 neighbor) - RCLCPP_INFO(this->get_logger(), "Have enough features for decoder, building decoder tensor..."); - neuromesh_interfaces::msg::Tensor own_tensor, neighbor_tensor; - - if (buildDecoderTensor(feature_buffer_, feature_buffer_timestamp_, own_tensor, neighbor_tensor)) { - RCLCPP_INFO(this->get_logger(), "Decoder tensor built successfully, starting decoder inference..."); - startClock("decoder_inference"); - - // Perform decoder inference - VGGT expects single concatenated tensor [2, 1036, 1024] - std::vector decoder_inputs = {own_tensor}; - RCLCPP_INFO(this->get_logger(), "Calling performInference with decoder model: %s", decoder_model_name_.c_str()); - RCLCPP_INFO(this->get_logger(), "Decoder input tensor size: %zu bytes", decoder_inputs[0].data.size()); - - decoder_result_future = performInference(decoder_model_name_, decoder_inputs); - RCLCPP_INFO(this->get_logger(), "Decoder inference started"); - } else { - RCLCPP_WARN(this->get_logger(), "Failed to build decoder tensor"); - } - } else { - RCLCPP_DEBUG(this->get_logger(), "Not enough features yet (need 2, have %zu)", feature_buffer_.size()); - } - - // Check if decoder result is ready (non-blocking check) - if (decoder_result_future.valid() && - decoder_result_future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - - RCLCPP_INFO(this->get_logger(), "Decoder result is ready, retrieving..."); - auto decoder_results = decoder_result_future.get(); - RCLCPP_INFO(this->get_logger(), "Retrieved %zu decoder results", decoder_results.size()); - - if (decoder_results.size() >= 5) { // pose_enc, depth, depth_conf, world_points, world_points_conf - stopClock("decoder_inference"); - RCLCPP_INFO(this->get_logger(), "Decoder inference took: %ld ms", checkClock("decoder_inference")); - - // Extract decoder outputs - auto pose_enc = decoder_results[0]; - auto depth = decoder_results[1]; - auto depth_conf = decoder_results[2]; - auto world_points = decoder_results[3]; - auto world_points_conf = decoder_results[4]; - - // Create header for outputs - std_msgs::msg::Header header; - header.stamp = this->get_clock()->now(); - header.frame_id = id_ + "/map"; - - // Publish depth images - auto depth_robot1 = createDepthImage(*depth, 0, header); - auto depth_robot2 = createDepthImage(*depth, 1, header); - if (depth_robot1) depth_robot1_publisher_->publish(*depth_robot1); - if (depth_robot2) depth_robot2_publisher_->publish(*depth_robot2); - - // Publish pointclouds - auto pc_current = createPointCloud(*world_points, *world_points_conf, 0, header); - auto pc_neighbor = createPointCloud(*world_points, *world_points_conf, 1, header); - if (pc_current) pointcloud_current_publisher_->publish(*pc_current); - if (pc_neighbor) pointcloud_neighbor_publisher_->publish(*pc_neighbor); - - // Publish RGB pointclouds - sensor_msgs::msg::Image::SharedPtr rgb_image; - { - std::lock_guard lock(camera_msg_mutex_); - rgb_image = latest_camera_msg_; - } - - if (rgb_image) { - auto pc_current_rgb = createPointCloud(*world_points, *world_points_conf, 0, header, true, rgb_image); - auto pc_neighbor_rgb = createPointCloud(*world_points, *world_points_conf, 1, header, true, rgb_image); - if (pc_current_rgb) pointcloud_current_rgb_publisher_->publish(*pc_current_rgb); - if (pc_neighbor_rgb) pointcloud_neighbor_rgb_publisher_->publish(*pc_neighbor_rgb); - } - } - } else if (decoder_result_future.valid()) { - RCLCPP_DEBUG(this->get_logger(), "Decoder result not ready yet, will check again next cycle"); - } - - RCLCPP_INFO(this->get_logger(), "=== END process_features ==="); - - } catch (const std::exception& e) { - RCLCPP_ERROR(this->get_logger(), "Exception in process_features: %s", e.what()); - } catch (...) { - RCLCPP_ERROR(this->get_logger(), "Unknown exception in process_features"); - } -} - -bool vggtNode::buildDecoderTensor( - std::map buffer, - std::map buffer_timestamp, - neuromesh_interfaces::msg::Tensor &own_tensor, - neuromesh_interfaces::msg::Tensor &neighbour_tensor) { - - // Find our own feature and one neighbor feature - auto own_it = buffer.find(id_); - if (own_it == buffer.end()) { - RCLCPP_WARN(this->get_logger(), "Own feature not found in buffer"); - return false; - } - - // Find a neighbor feature (just take the first one that's not us) - neuromesh_interfaces::msg::Feature::SharedPtr neighbor_feature = nullptr; - for (const auto& [agent_id, feature] : buffer) { - if (agent_id != id_) { - neighbor_feature = feature; - break; - } - } - - if (!neighbor_feature) { - RCLCPP_WARN(this->get_logger(), "No neighbor feature found in buffer"); - return false; - } - - // Build tensors from features - own_tensor = own_it->second->tensor; - neighbour_tensor = neighbor_feature->tensor; - - // Ensure correct dimensions for VGGT decoder: 2x1036x1024 - if (own_tensor.shape.dims.size() >= 3 && own_tensor.shape.dims[1] == 1036 && own_tensor.shape.dims[2] == 1024) { - // Concatenate tensors for decoder input - neuromesh_interfaces::msg::Tensor combined_tensor; - combined_tensor.shape.dims = {2, 1036, 1024}; - combined_tensor.data_type = 9; // float32 - - // Extract float data from own tensor - std::vector own_float_data(own_tensor.data.size() / sizeof(float)); - std::memcpy(own_float_data.data(), own_tensor.data.data(), own_tensor.data.size()); - - // Extract float data from neighbor tensor - std::vector neighbor_float_data(neighbour_tensor.data.size() / sizeof(float)); - std::memcpy(neighbor_float_data.data(), neighbour_tensor.data.data(), neighbour_tensor.data.size()); - - // Combine float data - std::vector combined_float_data; - combined_float_data.insert(combined_float_data.end(), own_float_data.begin(), own_float_data.end()); - combined_float_data.insert(combined_float_data.end(), neighbor_float_data.begin(), neighbor_float_data.end()); - - // Convert back to uint8 data - combined_tensor.data.resize(combined_float_data.size() * sizeof(float)); - std::memcpy(combined_tensor.data.data(), combined_float_data.data(), combined_float_data.size() * sizeof(float)); - - own_tensor = combined_tensor; - neighbour_tensor = combined_tensor; // For decoder, we pass the same combined tensor - - RCLCPP_DEBUG(this->get_logger(), "Built decoder tensor with dimensions [%u, %u, %u]", - combined_tensor.shape.dims[0], combined_tensor.shape.dims[1], combined_tensor.shape.dims[2]); - return true; - } else { - RCLCPP_WARN(this->get_logger(), "Invalid feature tensor dimensions for VGGT decoder"); - return false; - } -} - -sensor_msgs::msg::Image::SharedPtr -vggtNode::createDepthImage(const neuromesh_interfaces::msg::Tensor &depth_tensor, - int robot_idx, const std_msgs::msg::Header &header) { - auto depth_image = std::make_shared(); - - // Set header - depth_image->header = header; - depth_image->header.frame_id = "cam1_color_optical_frame"; - - // Set image properties for depth (392x518x1) - depth_image->height = 392; - depth_image->width = 518; - depth_image->encoding = sensor_msgs::image_encodings::TYPE_32FC1; - depth_image->is_bigendian = false; - depth_image->step = depth_image->width * sizeof(float); - - // Extract depth data for specific robot (robot_idx: 0 or 1) - // Tensor shape: 1x2x392x518x1 - if (depth_tensor.shape.dims.size() >= 5 && - depth_tensor.shape.dims[0] == 1 && depth_tensor.shape.dims[1] == 2 && - depth_tensor.shape.dims[2] == 392 && depth_tensor.shape.dims[3] == 518) { - - size_t depth_slice_size = 392 * 518; - size_t start_idx = robot_idx * depth_slice_size * sizeof(float); - size_t data_size = depth_slice_size * sizeof(float); - - if (start_idx + data_size <= depth_tensor.data.size()) { - depth_image->data.resize(data_size); - std::memcpy(depth_image->data.data(), - &depth_tensor.data[start_idx], - data_size); - } - } - - return depth_image; -} - -sensor_msgs::msg::PointCloud2::SharedPtr -vggtNode::createPointCloud(const neuromesh_interfaces::msg::Tensor &world_points_tensor, - const neuromesh_interfaces::msg::Tensor &world_points_conf_tensor, - int robot_idx, const std_msgs::msg::Header &header, - bool use_rgb, const sensor_msgs::msg::Image::SharedPtr rgb_image) { - auto pointcloud = std::make_shared(); - - // Set header - pointcloud->header = header; - - // Tensor shape: 1x2x392x518x3 for world_points - if (world_points_tensor.shape.dims.size() >= 5 && - world_points_tensor.shape.dims[0] == 1 && world_points_tensor.shape.dims[1] == 2 && - world_points_tensor.shape.dims[2] == 392 && world_points_tensor.shape.dims[3] == 518 && - world_points_tensor.shape.dims[4] == 3) { - - size_t height = 392; - size_t width = 518; - size_t points_per_robot = height * width; - size_t start_idx = robot_idx * points_per_robot * 3; // 3 for XYZ - - // Set up point cloud fields - sensor_msgs::PointCloud2Modifier modifier(*pointcloud); - if (use_rgb && rgb_image) { - modifier.setPointCloud2Fields(4, - "x", 1, sensor_msgs::msg::PointField::FLOAT32, - "y", 1, sensor_msgs::msg::PointField::FLOAT32, - "z", 1, sensor_msgs::msg::PointField::FLOAT32, - "rgb", 1, sensor_msgs::msg::PointField::UINT32); - } else { - modifier.setPointCloud2Fields(3, - "x", 1, sensor_msgs::msg::PointField::FLOAT32, - "y", 1, sensor_msgs::msg::PointField::FLOAT32, - "z", 1, sensor_msgs::msg::PointField::FLOAT32); - } - - // Count valid points (where confidence > threshold) - size_t valid_points = 0; - float conf_threshold = 0.5; // Confidence threshold - - // Extract confidence data as float array - std::vector conf_data(world_points_conf_tensor.data.size() / sizeof(float)); - std::memcpy(conf_data.data(), world_points_conf_tensor.data.data(), world_points_conf_tensor.data.size()); - - for (size_t i = 0; i < points_per_robot; ++i) { - size_t conf_idx = robot_idx * points_per_robot + i; - if (conf_idx < conf_data.size() && - conf_data[conf_idx] > conf_threshold) { - valid_points++; - } - } - - modifier.resize(valid_points); - - // Fill point cloud data - sensor_msgs::PointCloud2Iterator iter_x(*pointcloud, "x"); - sensor_msgs::PointCloud2Iterator iter_y(*pointcloud, "y"); - sensor_msgs::PointCloud2Iterator iter_z(*pointcloud, "z"); - - std::unique_ptr> iter_rgb_ptr; - if (use_rgb && rgb_image) { - iter_rgb_ptr = std::make_unique>(*pointcloud, "rgb"); - } - - // Extract world points data as float array - std::vector points_data(world_points_tensor.data.size() / sizeof(float)); - std::memcpy(points_data.data(), world_points_tensor.data.data(), world_points_tensor.data.size()); - - // Convert RGB image to OpenCV for color extraction - cv::Mat rgb_cv; - if (use_rgb && rgb_image) { - try { - cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(rgb_image, sensor_msgs::image_encodings::RGB8); - cv::resize(cv_ptr->image, rgb_cv, cv::Size(518, 392)); // Resize to match point cloud resolution - } catch (cv_bridge::Exception& e) { - RCLCPP_WARN(this->get_logger(), "Failed to convert RGB image: %s", e.what()); - use_rgb = false; - } - } - - for (size_t i = 0; i < points_per_robot; ++i) { - size_t conf_idx = robot_idx * points_per_robot + i; - if (conf_idx < conf_data.size() && - conf_data[conf_idx] > conf_threshold) { - - size_t xyz_base_idx = start_idx + i * 3; - if (xyz_base_idx + 2 < points_data.size()) { - *iter_x = points_data[xyz_base_idx]; - *iter_y = points_data[xyz_base_idx + 1]; - *iter_z = points_data[xyz_base_idx + 2]; - - if (use_rgb && rgb_image && !rgb_cv.empty() && iter_rgb_ptr) { - // Map point index back to image coordinates - int img_y = i / width; - int img_x = i % width; - - if (img_y < rgb_cv.rows && img_x < rgb_cv.cols) { - cv::Vec3b color = rgb_cv.at(img_y, img_x); - uint32_t rgb_value = (color[0] << 16) | (color[1] << 8) | color[2]; - **iter_rgb_ptr = rgb_value; - ++(*iter_rgb_ptr); - } - } - - ++iter_x; - ++iter_y; - ++iter_z; - } - } - } - } - - return pointcloud; -} - -std::future>> -vggtNode::performInference(const std::string &model_name, - const std::vector &tensors) { - // Placeholder implementation - in real system this would call TensorRT service - std::promise>> prom; - std::future>> result = prom.get_future(); - std::vector> output_tensors(1, std::make_shared()); - output_tensors[0]->result = 1; // Cannot reach engine error code - prom.set_value(std::move(output_tensors)); - return result; -} - -neuromesh_interfaces::msg::Feature -vggtNode::buildFeatureMessage(const neuromesh_interfaces::msg::Tensor &tensor) { - RCLCPP_INFO(this->get_logger(), "=== buildFeatureMessage called ==="); - RCLCPP_INFO(this->get_logger(), "Input tensor dims: [%s], data_size: %zu", - tensor.shape.dims.empty() ? "empty" : - std::accumulate(tensor.shape.dims.begin(), tensor.shape.dims.end(), std::string(), - [](const std::string& a, uint32_t b) { return a.empty() ? std::to_string(b) : a + ", " + std::to_string(b); }).c_str(), - tensor.data.size()); - - neuromesh_interfaces::msg::Feature feature_msg; - feature_msg.tensor = tensor; - feature_msg.id = id_; - feature_msg.timestamp = this->get_clock()->now(); - - RCLCPP_INFO(this->get_logger(), "Built feature message for agent: %s", id_.c_str()); - return feature_msg; -} - -void vggtNode::createSubscription( - std::map::SharedPtr> &subscription_map, - std::string id, rclcpp::QoS qos) { - - auto callback = [this, id](const neuromesh_interfaces::msg::Feature::SharedPtr msg) { - this->feature_callback(msg); - }; - - subscription_map[id] = this->create_subscription( - topic_prefix_ + id, qos, callback); -} - -void vggtNode::removeSubscription( - std::map::SharedPtr> subscription_map, - std::string id) { - subscription_map.erase(id); -} - -rmw_qos_profile_t vggtNode::parseQoSString(const std::string &str) { - std::string profile = str; - if (profile == "SYSTEM_DEFAULT") { - return rmw_qos_profile_system_default; - } - if (profile == "DEFAULT") { - return rmw_qos_profile_default; - } - if (profile == "PARAMETER_EVENTS") { - return rmw_qos_profile_parameter_events; - } - if (profile == "SERVICES_DEFAULT") { - return rmw_qos_profile_services_default; - } - if (profile == "PARAMETERS") { - return rmw_qos_profile_parameters; - } - if (profile == "SENSOR_DATA") { - return rmw_qos_profile_sensor_data; - } - RCLCPP_WARN_STREAM(rclcpp::get_logger("parseQosString"), - "Unknown QoS profile: " << profile << ". Returning profile: DEFAULT"); - return rmw_qos_profile_default; -} - -std::set vggtNode::splitAgentString(std::string str) { - std::set agents; - const std::string delimiter = ","; - - size_t pos = 0; - std::string token; - while ((pos = str.find(delimiter)) != std::string::npos) { - token = str.substr(0, pos); - agents.insert(token); - str.erase(0, pos + delimiter.length()); - } - agents.insert(str); - return agents; -} - -void vggtNode::run_encoder_cycle() { - auto now = std::chrono::system_clock::now(); - auto time_since_epoch = now.time_since_epoch(); - auto seconds = std::chrono::duration_cast(time_since_epoch); - auto nanoseconds = std::chrono::duration_cast(time_since_epoch - seconds); - - RCLCPP_INFO(this->get_logger(), "=== run_encoder_cycle called at %ld.%09ld, setting fresh_encoder_cycle = true ===", - seconds.count(), nanoseconds.count()); - fresh_encoder_cycle = true; -} - -void vggtNode::broadcast_transform() { - geometry_msgs::msg::TransformStamped t; - t.header.stamp = this->get_clock()->now(); - t.header.frame_id = "cam1_color_optical_frame"; - t.child_frame_id = id_ + "/map"; - - t.transform.translation.x = 0.0; - t.transform.translation.y = 0.0; - t.transform.translation.z = 0.0; - t.transform.rotation.x = 0.0; - t.transform.rotation.y = 0.0; - t.transform.rotation.z = 0.0; - t.transform.rotation.w = 1.0; - - tf_broadcaster_->sendTransform(t); -} - -void vggtNode::revertTensorDimensions(neuromesh_interfaces::msg::Tensor &tensor) { - // Implementation for reverting tensor dimensions if needed -} - -neuromesh_interfaces::msg::Tensor -vggtNode::convert_to_nchw(const neuromesh_interfaces::msg::Tensor &input) { - // Implementation for converting tensor to NCHW format - return input; // Placeholder -} - -neuromesh_interfaces::msg::Tensor -vggtNode::tensor_ints_to_floats(neuromesh_interfaces::msg::Tensor &input) { - // Implementation for converting ints to floats - return input; // Placeholder -} - -std::vector vggtNode::string_to_dims_single(std::string in) { - std::stringstream stream(in); - std::string element; - std::vector out; - - while (getline(stream, element, ',')) { - out.push_back(std::stoi(element)); - } - return out; -} - -std::vector> vggtNode::string_to_dims(std::string in) { - std::stringstream stream(in); - std::string element; - std::vector> out; - - while (getline(stream, element, ';')) { - std::vector dims = string_to_dims_single(element); - out.push_back(dims); - } - return out; -} - -void vggtNode::startClock(std::string phase) { - int64_t now_time = std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()).count(); - times[phase] = {now_time, false}; -} - -void vggtNode::stopClock(std::string phase) { - if (times[phase].second) { - return; // clock already stopped - } - int64_t start_time = times[phase].first; - int64_t now_time = std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()).count(); - times[phase] = {now_time - start_time, true}; -} - -int64_t vggtNode::checkClock(std::string phase) { - if (times[phase].second) { - return times[phase].first; // clock already stopped - } - - // stopclock calculations without saving - int64_t start_time = times[phase].first; - int64_t now_time = std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()).count(); - return now_time - start_time; -} - -} // namespace vggtNode \ No newline at end of file diff --git a/neuromesh_platform_r2/src/vggt_toy_implementation.cpp b/neuromesh_platform_r2/src/vggt_toy_implementation.cpp deleted file mode 100644 index 09f3bb6..0000000 --- a/neuromesh_platform_r2/src/vggt_toy_implementation.cpp +++ /dev/null @@ -1,360 +0,0 @@ -#include "neuromesh_platform_r2/vggt_toy_implementation.h" -#include -#include - -namespace vggtNode { -VggtToyImplementation::VggtToyImplementation(const rclcpp::NodeOptions &options) - : vggtNode(options) { - - std::string tensor_qos_profile_; - - this->declare_parameter("tensor_qos_profile", "default"); - this->get_parameter("tensor_qos_profile", tensor_qos_profile_); - - // Create separate callback groups for encoder and decoder to prevent blocking - auto encoder_callback_group = this->create_callback_group(rclcpp::CallbackGroupType::MutuallyExclusive); - auto decoder_callback_group = this->create_callback_group(rclcpp::CallbackGroupType::MutuallyExclusive); - - auto main_opt = rclcpp::SubscriptionOptions(); - main_opt.callback_group = encoder_callback_group; - - // Redefine subscriptions with callback groups - auto image_qos = rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(output_qos_profile_)); - - camera_subscription_ = this->create_subscription( - "camera", - image_qos, - std::bind(&VggtToyImplementation::camera_callback, this, std::placeholders::_1), - main_opt); - - // Create decoder timer in its own callback group - decoder_timer_ = this->create_wall_timer( - std::chrono::duration(decoder_cycle_length_), - std::bind(&VggtToyImplementation::process_features, this), - decoder_callback_group); - - // Create encoder timer in its own callback group - encoder_timer_ = this->create_wall_timer( - std::chrono::duration(encoder_cycle_length_), - std::bind(&VggtToyImplementation::run_encoder_cycle, this), - encoder_callback_group); - - // Define TensorRT service client - this->tensor_client_ = create_client("tensorrt_request"); - - RCLCPP_INFO(this->get_logger(), "VGGT Toy Implementation initialized"); -} - -std::future>> -VggtToyImplementation::performInference( - const std::string& model_name, - const std::vector& tensors) { - - RCLCPP_INFO(this->get_logger(), "=== START performInference for model: %s ===", model_name.c_str()); - - try { - if (!tensor_client_->wait_for_service(std::chrono::milliseconds(100))) { - RCLCPP_ERROR(this->get_logger(), "TensorRT engine not reachable via service."); - - // Return future that resolves to error tensor - std::promise>> prom; - std::future>> result = prom.get_future(); - std::vector> error_tensors(1, - std::make_shared()); - error_tensors[0]->result = 3; // Cannot reach engine error code - prom.set_value(error_tensors); - return result; - } - - RCLCPP_INFO(this->get_logger(), "TensorRT service is available"); - - // Create service request - auto request = std::make_shared(); - request->model_name = model_name; - request->tensor1 = tensors; - - RCLCPP_INFO(this->get_logger(), "Sending TensorRT request for model: %s with %zu tensors", - model_name.c_str(), tensors.size()); - for (size_t i = 0; i < tensors.size(); ++i) { - RCLCPP_INFO(this->get_logger(), " Tensor %zu: dims=[%s], data_size=%zu bytes", i, - tensors[i].shape.dims.empty() ? "empty" : - std::accumulate(tensors[i].shape.dims.begin(), tensors[i].shape.dims.end(), std::string(), - [](const std::string& a, uint32_t b) { return a.empty() ? std::to_string(b) : a + ", " + std::to_string(b); }).c_str(), - tensors[i].data.size()); - } - - // Call TensorRT service - RCLCPP_INFO(this->get_logger(), "Calling async_send_request..."); - auto future_and_request_id = tensor_client_->async_send_request(request); - std::shared_future> future = - future_and_request_id.future.share(); - RCLCPP_INFO(this->get_logger(), "Service request sent successfully"); - - // Process response asynchronously - std::future>> return_tensors = - std::async(std::launch::async, [this, future, model_name]() { - RCLCPP_INFO(this->get_logger(), "Async task started for model: %s", model_name.c_str()); - std::vector> output_tensors; - - try { - // Wait for response with timeout - RCLCPP_INFO(this->get_logger(), "Waiting for TensorRT response with 30s timeout..."); - auto status = future.wait_for(std::chrono::seconds(30)); - if (status == std::future_status::timeout) { - RCLCPP_ERROR(this->get_logger(), "TensorRT service timeout for model: %s", model_name.c_str()); - auto error_tensor = std::make_shared(); - error_tensor->result = 4; // Timeout error - output_tensors.push_back(error_tensor); - return output_tensors; - } - - RCLCPP_INFO(this->get_logger(), "Getting response from future..."); - auto response = future.get(); - RCLCPP_INFO(this->get_logger(), "Received TensorRT response for model: %s with %zu output tensors", - model_name.c_str(), response->tensor2.size()); - - for (size_t i = 0; i < response->tensor2.size(); ++i) { - RCLCPP_INFO(this->get_logger(), "Processing output tensor %zu", i); - output_tensors.emplace_back(std::make_shared(response->tensor2[i])); - } - - RCLCPP_INFO(this->get_logger(), "Async task completed successfully for model: %s", model_name.c_str()); - return output_tensors; - - } catch (const std::exception& e) { - RCLCPP_ERROR(this->get_logger(), "Exception in async task for model %s: %s", - model_name.c_str(), e.what()); - auto error_tensor = std::make_shared(); - error_tensor->result = 5; // Exception error - output_tensors.push_back(error_tensor); - return output_tensors; - } catch (...) { - RCLCPP_ERROR(this->get_logger(), "Unknown exception in async task for model: %s", model_name.c_str()); - auto error_tensor = std::make_shared(); - error_tensor->result = 5; // Exception error - output_tensors.push_back(error_tensor); - return output_tensors; - } - }); - - RCLCPP_INFO(this->get_logger(), "=== END performInference (returning future) ==="); - return return_tensors; - - } catch (const std::exception& e) { - RCLCPP_ERROR(this->get_logger(), "Exception in performInference: %s", e.what()); - - // Return future that resolves to error tensor - std::promise>> prom; - std::future>> result = prom.get_future(); - std::vector> error_tensors(1, - std::make_shared()); - error_tensors[0]->result = 6; // Exception error - prom.set_value(error_tensors); - return result; - } -} - -bool VggtToyImplementation::buildDecoderTensor( - std::map buffer, - std::map buffer_timestamp, - neuromesh_interfaces::msg::Tensor& own_tensor, - neuromesh_interfaces::msg::Tensor& neighbour_tensor) { - - RCLCPP_INFO(this->get_logger(), "=== START buildDecoderTensor ==="); - RCLCPP_INFO(this->get_logger(), "Building VGGT decoder tensor with %zu features", buffer.size()); - - try { - startClock("vggt_decoder_tensor"); - - // VGGT requires exactly 2 agents (current + 1 neighbor) - if (buffer.size() < 2) { - RCLCPP_WARN(this->get_logger(), "Not enough features for VGGT decoder (need 2, have %zu)", buffer.size()); - return false; - } - - // Log all agents in buffer - RCLCPP_INFO(this->get_logger(), "Agents in buffer:"); - for (const auto& [agent_id, feature] : buffer) { - RCLCPP_INFO(this->get_logger(), " - Agent: %s, Feature ptr: %p", - agent_id.c_str(), feature.get()); - } - - // Find our own feature - RCLCPP_INFO(this->get_logger(), "Looking for own feature with id: %s", id_.c_str()); - auto own_feature_it = buffer.find(id_); - if (own_feature_it == buffer.end()) { - RCLCPP_WARN(this->get_logger(), "Own feature not found in buffer for agent: %s", id_.c_str()); - return false; - } - RCLCPP_INFO(this->get_logger(), "Found own feature"); - - // Validate own feature pointer - if (!own_feature_it->second) { - RCLCPP_ERROR(this->get_logger(), "Own feature pointer is null!"); - return false; - } - - // Find neighbor feature (select the most recent one if multiple neighbors) - neuromesh_interfaces::msg::Feature::SharedPtr neighbor_feature = nullptr; - std::string neighbor_id; - double most_recent_timestamp = 0.0; - - RCLCPP_INFO(this->get_logger(), "Searching for neighbor feature..."); - for (const auto& [agent_id, feature] : buffer) { - if (agent_id != id_) { - double timestamp = buffer_timestamp[agent_id]; - RCLCPP_INFO(this->get_logger(), " Checking neighbor: %s, timestamp: %f", - agent_id.c_str(), timestamp); - if (timestamp > most_recent_timestamp) { - most_recent_timestamp = timestamp; - neighbor_feature = feature; - neighbor_id = agent_id; - } - } - } - - if (!neighbor_feature) { - RCLCPP_WARN(this->get_logger(), "No neighbor feature found for VGGT decoder"); - return false; - } - RCLCPP_INFO(this->get_logger(), "Selected neighbor: %s with timestamp: %f", - neighbor_id.c_str(), most_recent_timestamp); - - // Extract feature tensors - RCLCPP_INFO(this->get_logger(), "Extracting feature tensors..."); - auto own_feature_tensor = own_feature_it->second->tensor; - auto neighbor_feature_tensor = neighbor_feature->tensor; - - // Log tensor sizes - RCLCPP_INFO(this->get_logger(), "Own tensor data size: %zu bytes", own_feature_tensor.data.size()); - RCLCPP_INFO(this->get_logger(), "Neighbor tensor data size: %zu bytes", neighbor_feature_tensor.data.size()); - - // Validate feature dimensions for VGGT: should be 1x1036x1024 - RCLCPP_INFO(this->get_logger(), "Validating tensor dimensions..."); - if (own_feature_tensor.shape.dims.size() < 3) { - RCLCPP_ERROR(this->get_logger(), "Own feature tensor has only %zu dimensions, expected at least 3", - own_feature_tensor.shape.dims.size()); - return false; - } - - RCLCPP_INFO(this->get_logger(), "Own feature dimensions: [%u, %u, %u]", - own_feature_tensor.shape.dims[0], own_feature_tensor.shape.dims[1], own_feature_tensor.shape.dims[2]); - - if (own_feature_tensor.shape.dims[0] != 1 || - own_feature_tensor.shape.dims[1] != 1036 || - own_feature_tensor.shape.dims[2] != 1024) { - RCLCPP_WARN(this->get_logger(), "Invalid own feature dimensions for VGGT: [%u, %u, %u]", - own_feature_tensor.shape.dims[0], own_feature_tensor.shape.dims[1], own_feature_tensor.shape.dims[2]); - return false; - } - - if (neighbor_feature_tensor.shape.dims.size() < 3) { - RCLCPP_ERROR(this->get_logger(), "Neighbor feature tensor has only %zu dimensions, expected at least 3", - neighbor_feature_tensor.shape.dims.size()); - return false; - } - - RCLCPP_INFO(this->get_logger(), "Neighbor feature dimensions: [%u, %u, %u]", - neighbor_feature_tensor.shape.dims[0], neighbor_feature_tensor.shape.dims[1], neighbor_feature_tensor.shape.dims[2]); - - if (neighbor_feature_tensor.shape.dims[0] != 1 || - neighbor_feature_tensor.shape.dims[1] != 1036 || - neighbor_feature_tensor.shape.dims[2] != 1024) { - RCLCPP_WARN(this->get_logger(), "Invalid neighbor feature dimensions for VGGT: [%u, %u, %u]", - neighbor_feature_tensor.shape.dims[0], neighbor_feature_tensor.shape.dims[1], neighbor_feature_tensor.shape.dims[2]); - return false; - } - - // Validate data sizes - size_t expected_size = 1 * 1036 * 1024 * sizeof(float); - RCLCPP_INFO(this->get_logger(), "Expected tensor data size: %zu bytes", expected_size); - - if (own_feature_tensor.data.size() != expected_size) { - RCLCPP_ERROR(this->get_logger(), "Own tensor data size mismatch! Expected: %zu, Got: %zu", - expected_size, own_feature_tensor.data.size()); - return false; - } - - if (neighbor_feature_tensor.data.size() != expected_size) { - RCLCPP_ERROR(this->get_logger(), "Neighbor tensor data size mismatch! Expected: %zu, Got: %zu", - expected_size, neighbor_feature_tensor.data.size()); - return false; - } - - // Create combined tensor for VGGT decoder input: 2x1036x1024 - RCLCPP_INFO(this->get_logger(), "Creating combined tensor..."); - neuromesh_interfaces::msg::Tensor combined_tensor; - combined_tensor.shape.dims = {2, 1036, 1024}; - combined_tensor.data_type = 9; // float32 - - // Extract float data from tensors - RCLCPP_INFO(this->get_logger(), "Extracting float data from own tensor..."); - std::vector own_float_data(own_feature_tensor.data.size() / sizeof(float)); - if (!own_feature_tensor.data.empty()) { - std::memcpy(own_float_data.data(), own_feature_tensor.data.data(), own_feature_tensor.data.size()); - RCLCPP_INFO(this->get_logger(), "Successfully extracted %zu floats from own tensor", own_float_data.size()); - } else { - RCLCPP_ERROR(this->get_logger(), "Own tensor data is empty!"); - return false; - } - - RCLCPP_INFO(this->get_logger(), "Extracting float data from neighbor tensor..."); - std::vector neighbor_float_data(neighbor_feature_tensor.data.size() / sizeof(float)); - if (!neighbor_feature_tensor.data.empty()) { - std::memcpy(neighbor_float_data.data(), neighbor_feature_tensor.data.data(), neighbor_feature_tensor.data.size()); - RCLCPP_INFO(this->get_logger(), "Successfully extracted %zu floats from neighbor tensor", neighbor_float_data.size()); - } else { - RCLCPP_ERROR(this->get_logger(), "Neighbor tensor data is empty!"); - return false; - } - - // Concatenate features: first our own, then neighbor's - RCLCPP_INFO(this->get_logger(), "Concatenating features..."); - std::vector combined_float_data; - combined_float_data.reserve(2 * 1036 * 1024); - combined_float_data.insert(combined_float_data.end(), own_float_data.begin(), own_float_data.end()); - combined_float_data.insert(combined_float_data.end(), neighbor_float_data.begin(), neighbor_float_data.end()); - RCLCPP_INFO(this->get_logger(), "Combined float data size: %zu elements", combined_float_data.size()); - - // Convert back to uint8 data - RCLCPP_INFO(this->get_logger(), "Converting back to uint8 data..."); - size_t combined_data_size = combined_float_data.size() * sizeof(float); - combined_tensor.data.resize(combined_data_size); - std::memcpy(combined_tensor.data.data(), combined_float_data.data(), combined_data_size); - RCLCPP_INFO(this->get_logger(), "Combined tensor data size: %zu bytes", combined_tensor.data.size()); - - // Set metadata - combined_tensor.name = "vggt_combined_features_" + id_ + "_" + neighbor_id; - RCLCPP_INFO(this->get_logger(), "Combined tensor name: %s", combined_tensor.name.c_str()); - - // For VGGT decoder, we pass the same combined tensor as both inputs - own_tensor = combined_tensor; - neighbour_tensor = combined_tensor; - - stopClock("vggt_decoder_tensor"); - RCLCPP_INFO(this->get_logger(), "VGGT decoder tensor built in %ld ms", checkClock("vggt_decoder_tensor")); - RCLCPP_INFO(this->get_logger(), "Built VGGT decoder tensor: [%u, %u, %u] with %zu elements", - combined_tensor.shape.dims[0], combined_tensor.shape.dims[1], combined_tensor.shape.dims[2], - combined_float_data.size()); - - RCLCPP_INFO(this->get_logger(), "=== END buildDecoderTensor SUCCESS ==="); - return true; - - } catch (const std::exception& e) { - RCLCPP_ERROR(this->get_logger(), "Exception in buildDecoderTensor: %s", e.what()); - return false; - } catch (...) { - RCLCPP_ERROR(this->get_logger(), "Unknown exception in buildDecoderTensor"); - return false; - } -} - -void VggtToyImplementation::tensor_callback(const neuromesh_interfaces::msg::Tensor::SharedPtr /* msg */) { - RCLCPP_DEBUG(this->get_logger(), "Received tensor callback"); - // Handle any additional tensor processing if needed -} - -} // namespace vggtNode - -#include "rclcpp_components/register_node_macro.hpp" -RCLCPP_COMPONENTS_REGISTER_NODE(vggtNode::VggtToyImplementation) \ No newline at end of file diff --git a/neuromesh_platform_r2/test/test_gat.cpp b/neuromesh_platform_r2/test/test_gat.cpp new file mode 100644 index 0000000..e69de29 diff --git a/onnx_engine/.gitignore b/onnx_engine/.gitignore new file mode 100644 index 0000000..ca29f7c --- /dev/null +++ b/onnx_engine/.gitignore @@ -0,0 +1,7 @@ +third_party/* +build/* +.vscode/* +resources/* +**/*.trt +**/*.engine +**/*.data diff --git a/onnx_engine/.gitmodules b/onnx_engine/.gitmodules new file mode 100644 index 0000000..e69de29 diff --git a/onnx_engine/CMakeLists.txt b/onnx_engine/CMakeLists.txt new file mode 100644 index 0000000..876980c --- /dev/null +++ b/onnx_engine/CMakeLists.txt @@ -0,0 +1,92 @@ +cmake_minimum_required(VERSION 3.5) +project(onnx_engine) +set(CMAKE_BUILD_TYPE Release) + +if(BUILD_TESTING) + message("Building tests for onnx_engine") + find_package(ament_cmake_gmock REQUIRED) + ament_add_gmock(test_onnx_engine test/test_onnx_engine.cpp) + target_link_libraries(test_onnx_engine onnx_engine) +endif() + +# Default to C++14 +if(NOT CMAKE_CXX_STANDARD) + set(CMAKE_CXX_STANDARD 14) +endif() + +if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") + add_compile_options(-Wall -Wextra -Wpedantic) +endif() + +set (dependencies + "engine_interface" + "pluginlib" +) + +find_package(ament_cmake REQUIRED) +find_package(PkgConfig REQUIRED) +foreach(dep ${dependencies}) + find_package(${dep} REQUIRED) +endforeach() + +include_directories(include) + +# Using Pkg Config to find ONNX Runtime +pkg_check_modules(ONNXRUNTIME REQUIRED libonnxruntime) +include_directories(${ONNXRUNTIME_INCLUDE_DIRS}) +link_directories(${ONNXRUNTIME_LIBRARY_DIRS}) + +# Using find_package to find ONNX Runtime +# message("Finding ONNX Runtime") +# message(onnxruntime_INCLUDE_DIRS: ${onnxruntime_INCLUDE_DIRS}) +# message(onnxruntime_LIBRARIES: ${onnxruntime_LIBRARIES}) +# find_package(onnxruntime REQUIRED) +# include_directories(${onnxruntime_INCLUDE_DIRS}) + +add_library(onnx_engine SHARED + src/onnx_engine.cpp) + +target_include_directories(onnx_engine PUBLIC + $ + $ +) + +target_link_libraries(onnx_engine ${ONNXRUNTIME_LIBRARIES}) +# target_link_libraries(onnx_engine +# ${onnxruntime_LIBRARIES} +# ) +# set_target_properties(onnx_engine PROPERTIES +# COMPILE_DEFINITIONS "COMPOSITION_BUILDING_DLL" +# ) +ament_target_dependencies( + onnx_engine + ${dependencies} +) + +install(FILES plugin_description.xml + DESTINATION share/${PROJECT_NAME}) +pluginlib_export_plugin_description_file(engine_interface plugin_description.xml) + +install(TARGETS + onnx_engine + EXPORT export_${PROJECT_NAME} + ARCHIVE DESTINATION lib + LIBRARY DESTINATION lib + RUNTIME DESTINATION bin +) + +#install(DIRECTORY launch +# DESTINATION share/${PROJECT_NAME}) +install(DIRECTORY models + DESTINATION share/${PROJECT_NAME}) + +install(FILES plugin_description.xml + DESTINATION share/${PROJECT_NAME}) + +ament_export_include_directories(include) +ament_export_libraries(onnx_engine) +ament_export_targets(export_${PROJECT_NAME}) +ament_export_dependencies( + ${dependencies} +) +ament_package() diff --git a/onnx_engine/include/onnx_engine/onnx_engine.h b/onnx_engine/include/onnx_engine/onnx_engine.h new file mode 100644 index 0000000..032f8f3 --- /dev/null +++ b/onnx_engine/include/onnx_engine/onnx_engine.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine_interface/inference_engine_base.hpp" +#include +#include +#include + +namespace engine_interface +{ + +class ONNXEngine : public BaseEngine +{ + public: + ONNXEngine(); + ~ONNXEngine() override; + + bool loadModel(const std::string& model_path, + const std::vector>& input_dims, + int type_length) override; + + void runInference(const std::vector &inputTensors, + const std::vector &inputSizes, + std::vector &outputTensors, + const std::vector &outputSizes) override; + + private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine_interface diff --git a/onnx_engine/models/README.md b/onnx_engine/models/README.md new file mode 100644 index 0000000..53904e3 --- /dev/null +++ b/onnx_engine/models/README.md @@ -0,0 +1,2 @@ +# Prepare your models + diff --git a/tensorrt_engine/models/dust3r_decoder_tensor_params.onnx b/onnx_engine/models/dust3r/dust3r_decoder_tensor_params.onnx similarity index 100% rename from tensorrt_engine/models/dust3r_decoder_tensor_params.onnx rename to onnx_engine/models/dust3r/dust3r_decoder_tensor_params.onnx diff --git a/tensorrt_engine/models/dust3r_encoder_single_mini_params.onnx b/onnx_engine/models/dust3r/dust3r_encoder_single_mini_params.onnx old mode 100755 new mode 100644 similarity index 100% rename from tensorrt_engine/models/dust3r_encoder_single_mini_params.onnx rename to onnx_engine/models/dust3r/dust3r_encoder_single_mini_params.onnx diff --git a/onnx_engine/models/gat/encoder.onnx b/onnx_engine/models/gat/encoder.onnx new file mode 100644 index 0000000..32372d9 --- /dev/null +++ b/onnx_engine/models/gat/encoder.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5e48423b1fc189be0918e5e4850fbcce39e29bf02b4df3edd4c70b9a7891ccce +size 2249 diff --git a/tensorrt_engine/models/encoder_local.onnx b/onnx_engine/models/gat/encoder_local.onnx similarity index 100% rename from tensorrt_engine/models/encoder_local.onnx rename to onnx_engine/models/gat/encoder_local.onnx diff --git a/onnx_engine/models/gat/gat_layer1.onnx b/onnx_engine/models/gat/gat_layer1.onnx new file mode 100644 index 0000000..2a67704 --- /dev/null +++ b/onnx_engine/models/gat/gat_layer1.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dc7b3301268c45f27bcdbd35b5ab9bce515230c79e7de59604f11459b1e26820 +size 13824 diff --git a/onnx_engine/models/gat/gat_layer2.onnx b/onnx_engine/models/gat/gat_layer2.onnx new file mode 100644 index 0000000..8ba27e8 --- /dev/null +++ b/onnx_engine/models/gat/gat_layer2.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4cd0967fcb22118211f73d055a8d2b0d2c84240a781522fd300ae79d5fcd5b44 +size 13076 diff --git a/tensorrt_engine/models/multi_head_gat_layer1.onnx b/onnx_engine/models/gat/multi_head_gat_layer1.onnx similarity index 100% rename from tensorrt_engine/models/multi_head_gat_layer1.onnx rename to onnx_engine/models/gat/multi_head_gat_layer1.onnx diff --git a/tensorrt_engine/models/multi_head_gat_layer2.onnx b/onnx_engine/models/gat/multi_head_gat_layer2.onnx similarity index 100% rename from tensorrt_engine/models/multi_head_gat_layer2.onnx rename to onnx_engine/models/gat/multi_head_gat_layer2.onnx diff --git a/tensorrt_engine/models/gnn_post_combined.onnx b/onnx_engine/models/gnn/gnn_post_combined.onnx similarity index 100% rename from tensorrt_engine/models/gnn_post_combined.onnx rename to onnx_engine/models/gnn/gnn_post_combined.onnx diff --git a/onnx_engine/models/vggt_2x/vggt_aggregator.onnx b/onnx_engine/models/vggt_2x/vggt_aggregator.onnx new file mode 100644 index 0000000..857c3ad --- /dev/null +++ b/onnx_engine/models/vggt_2x/vggt_aggregator.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2016573d8892a4c580dcfb20ab1befc9b5deb66d30f100879db90e629958c5e6 +size 10714079 diff --git a/onnx_engine/models/vggt_2x/vggt_image_encoder.onnx b/onnx_engine/models/vggt_2x/vggt_image_encoder.onnx new file mode 100644 index 0000000..a974ac1 --- /dev/null +++ b/onnx_engine/models/vggt_2x/vggt_image_encoder.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:10151ee3b220d93abc2c28d538d7bf8e37f90967986ea8be36adcfa1cb6730dd +size 1822389 diff --git a/onnx_engine/models/vggt_2x_8805/vggt_aggregator.onnx b/onnx_engine/models/vggt_2x_8805/vggt_aggregator.onnx new file mode 100644 index 0000000..857c3ad --- /dev/null +++ b/onnx_engine/models/vggt_2x_8805/vggt_aggregator.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2016573d8892a4c580dcfb20ab1befc9b5deb66d30f100879db90e629958c5e6 +size 10714079 diff --git a/onnx_engine/models/vggt_2x_8805/vggt_image_encoder.onnx b/onnx_engine/models/vggt_2x_8805/vggt_image_encoder.onnx new file mode 100644 index 0000000..a974ac1 --- /dev/null +++ b/onnx_engine/models/vggt_2x_8805/vggt_image_encoder.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:10151ee3b220d93abc2c28d538d7bf8e37f90967986ea8be36adcfa1cb6730dd +size 1822389 diff --git a/onnx_engine/package.xml b/onnx_engine/package.xml new file mode 100644 index 0000000..069b3c7 --- /dev/null +++ b/onnx_engine/package.xml @@ -0,0 +1,25 @@ + + + + onnx_engine + 0.0.0 + Onnx engine plugin implementation + Long Quang + TODO: License declaration + + ament_cmake + PkgConfig + + + engine_interface + pluginlib + + ament_lint_auto + ament_lint_common + ament_cmake_gmock + + + ament_cmake + + + diff --git a/onnx_engine/plugin_description.xml b/onnx_engine/plugin_description.xml new file mode 100644 index 0000000..2451d25 --- /dev/null +++ b/onnx_engine/plugin_description.xml @@ -0,0 +1,7 @@ + + + ONNX Runtime inference engine plugin + + diff --git a/onnx_engine/src/onnx_engine.cpp b/onnx_engine/src/onnx_engine.cpp new file mode 100644 index 0000000..b066770 --- /dev/null +++ b/onnx_engine/src/onnx_engine.cpp @@ -0,0 +1,221 @@ +#include "onnx_engine/onnx_engine.h" +#include +#include + +namespace engine_interface +{ + +class ONNXEngine::Impl +{ +public: + Ort::Env env_; + Ort::SessionOptions session_options_; + std::unique_ptr session_; + std::vector input_names_; + std::vector> input_shapes_; + std::vector output_names_; + std::vector> output_shapes_; + int type_length_; + + Impl() + : env_(ORT_LOGGING_LEVEL_WARNING, "onnx_engine"), + session_options_(), + type_length_(0) + { + session_options_.SetInterOpNumThreads(1); + session_options_.SetIntraOpNumThreads(1); + session_options_.SetGraphOptimizationLevel(ORT_DISABLE_ALL); + + // TODO: If using CUDA, set up CUDA options + // if (_UseCuda) { + // OrtCUDAProviderOptions cuda_options; + // cuda_options.device_id = 0; + // cuda_options.cudnn_conv_algo_search = OrtCudnnConvAlgoSearchExhaustive; + // cuda_options.arena_extend_strategy = 0; + // cuda_options.do_copy_in_default_stream = 0; + // session_options_.AppendExecutionProvider_CUDA(cuda_options); + // } + } +}; + +ONNXEngine::ONNXEngine() + : impl_(std::make_unique()) +{} + +ONNXEngine::~ONNXEngine() = default; + +bool ONNXEngine::loadModel(const std::string& model_path, + const std::vector>& input_dims, + int type_length) +{ + // std::cout << "Loading ONNX model from: " << model_path << std::endl; + impl_->type_length_ = type_length; + // std::cout << "Type length: " << impl_->type_length_ << std::endl; + try + { + impl_->session_ = std::make_unique(impl_->env_, + model_path.c_str(), + impl_->session_options_); + // std::cout << "Session created successfully." << std::endl; + } + catch (const Ort::Exception& e) + { + std::cerr << "Error creating ONNX session: " << e.what() << ". Code: " << e.GetOrtErrorCode() << std::endl; + return -1; + } + + Ort::AllocatorWithDefaultOptions allocator; + ONNXTensorElementDataType type; + Ort::TypeInfo* type_info; + + impl_->input_names_.clear(); + impl_->input_shapes_.clear(); + impl_->output_names_.clear(); + impl_->output_shapes_.clear(); + + size_t num_inputs = impl_->session_->GetInputCount(); + size_t num_outputs = impl_->session_->GetOutputCount(); + + // std::cout << "Number of inputs: " << num_inputs << std::endl; + // std::cout << "Number of outputs: " << num_outputs << std::endl; + + for (size_t i = 0; i < num_inputs; ++i) { + char* name = impl_->session_->GetInputName(i, allocator); + if (name == nullptr) { + throw std::runtime_error("Failed to get input name for index " + std::to_string(i)); + } + impl_->input_names_.emplace_back(name); + allocator.Free(name); + + type_info = new Ort::TypeInfo(impl_->session_->GetInputTypeInfo(i)); + auto tensor_info = type_info->GetTensorTypeAndShapeInfo(); + type = tensor_info.GetElementType(); + impl_->input_shapes_.push_back(tensor_info.GetShape()); + + std::cout << "Input " << i << ": name = " << impl_->input_names_.back() + << ", shape = ["; + for (const auto& dim : impl_->input_shapes_.back()) { + std::cout << dim << " "; + } + std::cout << "], type = " << type << std::endl; + + delete(type_info); + } + + for (size_t i = 0; i < num_outputs; ++i) { + char* name = impl_->session_->GetOutputName(i, allocator); + if (name == nullptr) { + throw std::runtime_error("Failed to get output name for index " + std::to_string(i)); + } + impl_->output_names_.emplace_back(name); + allocator.Free(name); + + type_info = new Ort::TypeInfo(impl_->session_->GetOutputTypeInfo(i)); + auto tensor_info = type_info->GetTensorTypeAndShapeInfo(); + type = tensor_info.GetElementType(); + impl_->input_shapes_.push_back(tensor_info.GetShape()); + + std::cout << "Output " << i << ": name = " << impl_->output_names_.back() + << ", shape = ["; + for (const auto& dim : impl_->input_shapes_.back()) { + std::cout << dim << " "; + } + std::cout << "], type = " << type << std::endl; + + delete(type_info); + } + + return true; +} + +void ONNXEngine::runInference(const std::vector &inputTensors, + const std::vector &inputSizes, + std::vector &outputTensors, + const std::vector &outputSizes) +{ + // std::cout << "Running inference..." << std::endl; + if (!impl_->session_) { + throw std::runtime_error("Session is not initialized. Call loadModel first."); + } + if (inputTensors.size() != impl_->input_names_.size()) { + throw std::runtime_error( + "Number of input tensors doesn't match number of input names"); + } + if (outputTensors.size() != impl_->output_names_.size()) { + throw std::runtime_error( + "Number of output tensors doesn't match number of output names"); + } + + std::vector ort_inputs, ort_outputs; + Ort::MemoryInfo mem_info{ nullptr }; + try + { + mem_info = std::move(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)); + } + catch (Ort::Exception &e) + { + std::cerr << "ONNX exception caught: " << e.what() << ". Code: " << e.GetOrtErrorCode() << std::endl; + } + + for (size_t i = 0; i < inputTensors.size(); ++i) { + const float* input_data = reinterpret_cast(inputTensors[i]); + const std::vector& input_shape = impl_->input_shapes_[i]; + size_t input_numel = 1; + for (auto d : input_shape) input_numel *= d; + + ort_inputs.emplace_back( + Ort::Value::CreateTensor( + mem_info, const_cast(input_data), + input_numel, + input_shape.data(), + input_shape.size() + ) + ); + } + + // std::cout << "Input tensors converted to Ort::Value." << std::endl; + + std::vector input_names_c; + std::vector output_names_c; + for (const auto& name : impl_->input_names_) input_names_c.push_back(name.c_str()); + for (const auto& name : impl_->output_names_) output_names_c.push_back(name.c_str()); + + // run the inference + // std::cout << "Running ONNX session..." << std::endl; + try + { + ort_outputs = impl_->session_->Run( + Ort::RunOptions{nullptr}, + input_names_c.data(), ort_inputs.data(), ort_inputs.size(), + output_names_c.data(), output_names_c.size() + ); + } + catch (const Ort::Exception& e) + { + std::cerr << "ONNX exception caught during inference: " << e.what() << ". Code: " << e.GetOrtErrorCode() << std::endl; + throw; + } + + for (size_t i = 0; i < ort_outputs.size(); ++i) + { + auto& ort_val = ort_outputs[i]; + auto tensor_info = ort_val.GetTensorTypeAndShapeInfo(); + auto shape = tensor_info.GetShape(); + size_t numel = 1; + for (auto d : shape) { + if (d <= 0) throw std::runtime_error("Invalid output shape"); + numel *= d; + } + if (tensor_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) + throw std::runtime_error("Output tensor type mismatch"); + if (outputSizes[i] < static_cast(numel * sizeof(float))) + throw std::runtime_error("Output buffer too small"); + float* output_data = ort_val.GetTensorMutableData(); + std::memcpy(outputTensors[i], output_data, numel * sizeof(float)); + } +} + +} // namespace engine_interface + +#include +PLUGINLIB_EXPORT_CLASS(engine_interface::ONNXEngine, engine_interface::BaseEngine) diff --git a/onnx_engine/test/test_onnx_engine.cpp b/onnx_engine/test/test_onnx_engine.cpp new file mode 100644 index 0000000..bae6092 --- /dev/null +++ b/onnx_engine/test/test_onnx_engine.cpp @@ -0,0 +1,30 @@ +#include +#include "onnx_engine/onnx_engine.h" + +TEST(ONNXEngineTest, LoadModelSuccess) { + engine_interface::ONNXEngine engine; + EXPECT_NO_THROW({ + }); + + EXPECT_NO_THROW({ + bool loaded = engine.loadModel("test_model.onnx", {{1, 3, 224, 224}}, 4); + EXPECT_TRUE(loaded); + }); +} + +TEST(ONNXEngineTest, InferenceProducesOutput) { + engine_interface::ONNXEngine engine; + engine.loadModel("test_model.onnx", {{1, 3, 224, 224}}, 4); + + std::vector input_tensors(1, nullptr); + std::vector input_sizes(1, 1024); + std::vector output_tensors(1, nullptr); + std::vector output_sizes(1, 1024); + + engine.runInference(input_tensors, input_sizes, output_tensors, output_sizes); + EXPECT_EQ(output_tensors.size(), 1); + EXPECT_NE(output_tensors[0], nullptr); + // Clean up allocated memory + delete[] static_cast(output_tensors[0]); + output_tensors[0] = nullptr; +} \ No newline at end of file diff --git a/tensorrt_engine/.gitignore b/tensorrt_engine/.gitignore index 06df182..ca29f7c 100644 --- a/tensorrt_engine/.gitignore +++ b/tensorrt_engine/.gitignore @@ -1,5 +1,7 @@ third_party/* -models/* build/* .vscode/* -resources/* \ No newline at end of file +resources/* +**/*.trt +**/*.engine +**/*.data diff --git a/tensorrt_engine/CMakeLists.txt b/tensorrt_engine/CMakeLists.txt index cfbe6f3..2814fe0 100644 --- a/tensorrt_engine/CMakeLists.txt +++ b/tensorrt_engine/CMakeLists.txt @@ -1,5 +1,13 @@ cmake_minimum_required(VERSION 3.5) project(tensorrt_engine) +set(CMAKE_BUILD_TYPE Release) + +if(BUILD_TESTING) + message("Building tests for tensorrt_engine") + find_package(ament_cmake_gmock REQUIRED) + ament_add_gmock(test_trt_engine test/test_trt_engine.cpp) + target_link_libraries(test_trt_engine tensorrt_engine) +endif() # Default to C++14 if(NOT CMAKE_CXX_STANDARD) @@ -10,54 +18,53 @@ if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") add_compile_options(-Wall -Wextra -Wpedantic) endif() - +set (dependencies + "rclcpp" + "std_msgs" + "sensor_msgs" + "cv_bridge" + "neuromesh_interfaces" + "rclcpp_components" + "engine_interface" + "pluginlib" +) # find dependencies find_package(ament_cmake REQUIRED) find_package(OpenCV REQUIRED) -find_package(std_msgs REQUIRED) -find_package(sensor_msgs REQUIRED) -find_package(cv_bridge REQUIRED) -find_package(rclcpp REQUIRED) -find_package(neuromesh_interfaces REQUIRED) -find_package(rclcpp_components REQUIRED) find_package(CUDA REQUIRED) -# uncomment the following section in order to fill in -# further dependencies manually. -# find_package( REQUIRED) - -#add_executable(tensorrt_engine src/engine_node.cpp src/main.cpp src/trt_engine.cpp ) -#target_include_directories(tensorrt_engine PUBLIC "include/") +foreach(dep ${dependencies}) + find_package(${dep} REQUIRED) +endforeach() include_directories(include) add_library(tensorrt_engine SHARED - src/engine_modified.cpp src/trt_engine_modified.cpp) + src/trt_engine_modified.cpp) + +target_include_directories(tensorrt_engine PUBLIC + $ + $ + ${OpenCV_INCLUDE_DIRS} ${CUDA_INCLUDE_DIRS}) -set_target_properties(tensorrt_engine PROPERTIES - COMPILE_DEFINITIONS "COMPOSITION_BUILDING_DLL" -) + +target_link_libraries(tensorrt_engine ${OpenCV_LIBS} ${CUDA_LIBRARIES} nvinfer nvonnxparser nvinfer_plugin cudnn) +# set_target_properties(tensorrt_engine PROPERTIES +# COMPILE_DEFINITIONS "COMPOSITION_BUILDING_DLL" +# ) ament_target_dependencies( tensorrt_engine - rclcpp - std_msgs - sensor_msgs - cv_bridge - neuromesh_interfaces - rclcpp_components + ${dependencies} ) - #install(TARGETS tensorrt_engine - #DESTINATION lib/${PROJECT_NAME}) - -target_include_directories(tensorrt_engine PUBLIC ${OpenCV_INCLUDE_DIRS} ${CUDA_INCLUDE_DIRS} tensorrt) -target_link_libraries(tensorrt_engine ${OpenCV_LIBS} ${CUDA_LIBRARIES} nvinfer nvonnxparser nvinfer_plugin cudnn) - -rclcpp_components_register_nodes(tensorrt_engine "tensorrt_engine_node::TensorRTEngineNode") +install(FILES plugin_description.xml + DESTINATION share/${PROJECT_NAME}) +pluginlib_export_plugin_description_file(engine_interface plugin_description.xml) - install(TARGETS +install(TARGETS tensorrt_engine + EXPORT export_${PROJECT_NAME} ARCHIVE DESTINATION lib LIBRARY DESTINATION lib RUNTIME DESTINATION bin @@ -69,4 +76,5 @@ install(DIRECTORY models DESTINATION share/${PROJECT_NAME}) ament_export_libraries(tensorrt_engine) +ament_export_targets(export_${PROJECT_NAME}) ament_package() diff --git a/tensorrt_engine/include/tensorrt_engine/engine_modified.h b/tensorrt_engine/include/tensorrt_engine/engine_modified.h deleted file mode 100644 index c8dd333..0000000 --- a/tensorrt_engine/include/tensorrt_engine/engine_modified.h +++ /dev/null @@ -1,81 +0,0 @@ -#ifndef TENSORRT_ENGINE_H -#define TENSORRT_ENGINE_H - -#define INFERENCE_HELPER_ENABLE_TENSORRT - -#include -#include - -#include "rclcpp/rclcpp.hpp" - -#include "neuromesh_interfaces/msg/tensor.hpp" -#include "neuromesh_interfaces/srv/tensor_request.hpp" -#include "sensor_msgs/msg/image.hpp" - -#include "std_msgs/msg/header.hpp" - -#include - -#include "tensorrt_engine/trt_engine_modified.h" - -namespace tensorrt_engine_node { -class TensorRTEngineNode : public rclcpp::Node { -public: - TensorRTEngineNode(rclcpp::NodeOptions options); - ~TensorRTEngineNode() {}; - -private: - rclcpp::Publisher::SharedPtr - tensor_publisher_; - rclcpp::Subscription::SharedPtr - tensor_subscription_; - rclcpp::Service::SharedPtr service_; - - // params - std::string models_param; - std::vector model_names; - - std::unordered_map model_paths; - std::unordered_map>> - input_dimensions; - std::unordered_map>> - output_dimensions; - std::unordered_map input_dimensions_strings; - std::unordered_map output_dimensions_strings; - std::unordered_map tensor_type_params; - std::string tensor_qos_param; - - // engine - std::unordered_map> engines; - - std::unordered_map> input_lengths; - std::unordered_map> output_lengths; - std::unordered_map tensor_typelengths; - - std::vector failed_models; - - // functions - // callbacks - void tensor_request_callback( - const std::shared_ptr - request, - const std::shared_ptr - response); - - void - tensor_callback(const std::shared_ptr msg); - - // execution - std::vector - execute(const std::string &model, - const std::vector &tensor_msgs); - - // helper functions - int tensor_string_to_typelength(std::string input); - std::vector construct_dims(int width, int height, bool rgb, bool nchw); - - // Convert string to ROS2 QoS profile - rmw_qos_profile_t parseQoSString(const std::string &str); -}; -} // namespace tensorrt_engine_node -#endif diff --git a/tensorrt_engine/include/tensorrt_engine/engine_node.h b/tensorrt_engine/include/tensorrt_engine/engine_node.h deleted file mode 100644 index cff2386..0000000 --- a/tensorrt_engine/include/tensorrt_engine/engine_node.h +++ /dev/null @@ -1,79 +0,0 @@ -#ifndef TENSORRT_ENGINE_H -#define TENSORRT_ENGINE_H - -#define INFERENCE_HELPER_ENABLE_TENSORRT - -#include -#include - -#include "rclcpp/rclcpp.hpp" - -#include "neuromesh_interfaces/msg/tensor.hpp" -#include "neuromesh_interfaces/srv/tensor_request.hpp" -#include "sensor_msgs/msg/image.hpp" - -#include "std_msgs/msg/header.hpp" - -#include - -#include "tensorrt_engine/trt_engine.h" - -namespace tensorrt_engine_node { -class TensorRTEngineNode : public rclcpp::Node { -public: - TensorRTEngineNode(rclcpp::NodeOptions options); - ~TensorRTEngineNode() {}; - -private: - rclcpp::Publisher::SharedPtr - tensor_publisher_; - rclcpp::Subscription::SharedPtr - tensor_subscription_; - rclcpp::Service::SharedPtr service_; - - // params - std::string models_param; - std::vector model_names; - - std::unordered_map model_paths; - std::unordered_map> input_dimensions; - std::unordered_map> output_dimensions; - std::unordered_map input_dimensions_strings; - std::unordered_map output_dimensions_strings; - std::unordered_map tensor_type_params; - std::string tensor_qos_param; - - // engine - std::unordered_map> engines; - - std::unordered_map input_lengths; - std::unordered_map output_lengths; - std::unordered_map tensor_typelengths; - - std::vector failed_models; - - // functions - // callbacks - void tensor_request_callback( - const std::shared_ptr - request, - const std::shared_ptr - response); - - void - tensor_callback(const std::shared_ptr msg); - - // execution - neuromesh_interfaces::msg::Tensor - execute(const std::string model, - const neuromesh_interfaces::msg::Tensor &tensor_msg); - - // helper functions - int tensor_string_to_typelength(std::string input); - std::vector construct_dims(int width, int height, bool rgb, bool nchw); - - // Convert string to ROS2 QoS profile - rmw_qos_profile_t parseQoSString(const std::string &str); -}; -} // namespace tensorrt_engine_node -#endif diff --git a/tensorrt_engine/include/tensorrt_engine/logger.h b/tensorrt_engine/include/tensorrt_engine/logger.h index 8836e22..e88b4f2 100644 --- a/tensorrt_engine/include/tensorrt_engine/logger.h +++ b/tensorrt_engine/include/tensorrt_engine/logger.h @@ -1,4 +1,7 @@ +#pragma once + #include +#include class Logger : public nvinfer1::ILogger { public: diff --git a/tensorrt_engine/include/tensorrt_engine/trt_engine.h b/tensorrt_engine/include/tensorrt_engine/trt_engine.h deleted file mode 100644 index 62202db..0000000 --- a/tensorrt_engine/include/tensorrt_engine/trt_engine.h +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef TRT_ENGINE_H -#define TRT_ENGINE_H - -#include -#include -#include - -#include -#include - -#include - -#include - -class TRTEngine { -public: - - TRTEngine(std::string model_filename, std::vector input_dims, int type_length, int batchSize=1); - ~TRTEngine(); - - void runInference( const void* inputTensor, int inputSize, void* outputTensor, int outputSize); - -protected: - Logger logger; - std::unique_ptr engine; - - std::unique_ptr context; - std::vector bindings; - -}; - - -#endif \ No newline at end of file diff --git a/tensorrt_engine/include/tensorrt_engine/trt_engine_modified.h b/tensorrt_engine/include/tensorrt_engine/trt_engine_modified.h index 7dc8a6b..3222e89 100644 --- a/tensorrt_engine/include/tensorrt_engine/trt_engine_modified.h +++ b/tensorrt_engine/include/tensorrt_engine/trt_engine_modified.h @@ -1,35 +1,32 @@ -#ifndef TRT_ENGINE_H -#define TRT_ENGINE_H +#pragma once +#include "engine_interface/inference_engine_base.hpp" #include #include #include -#include -#include - -#include - #include -class TRTEngine { -public: +namespace engine_interface { - TRTEngine(std::string model_filename, std::vector> input_dims, int type_length, int batchSize=1); - ~TRTEngine(); +class TRTEngine : public BaseEngine +{ + public: + TRTEngine(); + ~TRTEngine() override; - void runInference(const std::vector& inputTensors, const std::vector& inputSizes, - std::vector& outputTensors, const std::vector& outputSizes); + bool loadModel(const std::string& model_path, + const std::vector>& input_dims, + int type_length) override; -protected: - Logger logger; - std::unique_ptr engine; - std::unique_ptr runtime; - - std::unique_ptr context; - std::vector bindings; + void runInference(const std::vector& inputTensors, + const std::vector& inputSizes, + std::vector& outputTensors, + const std::vector& outputSizes) override; + private: + class Impl; + std::unique_ptr impl_; }; - -#endif +} // namespace engine_interface \ No newline at end of file diff --git a/tensorrt_engine/models/encoder.onnx b/tensorrt_engine/models/encoder.onnx deleted file mode 100644 index bdf89c5..0000000 --- a/tensorrt_engine/models/encoder.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f1dbfc57d5e491cbf7c1c9620176c1454a8698564aeeaeec92125729004dfe30 -size 19164 diff --git a/tensorrt_engine/package.xml b/tensorrt_engine/package.xml index 29a3a3c..b25cb12 100644 --- a/tensorrt_engine/package.xml +++ b/tensorrt_engine/package.xml @@ -3,28 +3,27 @@ tensorrt_engine 0.0.0 - TODO: Package description + TensorRT engine plugin implementation arpl TODO: License declaration ament_cmake - - neuromesh_interfaces + engine_interface + pluginlib + rclcpp_components - rclcpp_components tensorrt + rclcpp_components + ament_lint_auto ament_lint_common + ament_cmake_gmock ament_cmake - - - TensorRT Engine Node - - + diff --git a/tensorrt_engine/plugin_description.xml b/tensorrt_engine/plugin_description.xml new file mode 100644 index 0000000..ee25302 --- /dev/null +++ b/tensorrt_engine/plugin_description.xml @@ -0,0 +1,7 @@ + + + TensorRT inference engine plugin + + diff --git a/tensorrt_engine/src/engine_node.cpp b/tensorrt_engine/src/engine_node.cpp deleted file mode 100644 index 2316034..0000000 --- a/tensorrt_engine/src/engine_node.cpp +++ /dev/null @@ -1,265 +0,0 @@ -#include "tensorrt_engine/engine_node.h" -#include "tensorrt_engine/trt_engine.h" - -#include "cv_bridge/cv_bridge.h" - -#include - -namespace tensorrt_engine_node { -std::vector string_to_dims(std::string in) { - std::stringstream stream(in); - std::string element; - - std::vector out; - - while (getline(stream, element, ',')) { - out.push_back(std::stoi(element)); - } - return out; -} -std::vector string_to_vector(std::string in) { - std::stringstream stream(in); - std::string element; - - std::vector out; - - while (getline(stream, element, ',')) { - out.push_back(element); - } - return out; -} - -void TensorRTEngineNode::tensor_request_callback( - const std::shared_ptr - request, - const std::shared_ptr - response) { - RCLCPP_DEBUG(this->get_logger(), "Recieved service request."); - RCLCPP_DEBUG(this->get_logger(), "Input tensor size: %zu", - request->tensor1.data.size()); - RCLCPP_DEBUG(this->get_logger(), "Expected input size for model %s: %d", - request->model_name.c_str(), input_lengths[request->model_name]); - response->tensor2 = execute(request->model_name, request->tensor1); - - // tensor_publisher_->publish(response->tensor2); -} - -TensorRTEngineNode::TensorRTEngineNode(rclcpp::NodeOptions options) - : Node("TensorRTEngineNode", - options.allow_undeclared_parameters(true) - .automatically_declare_parameters_from_overrides(true)) { - // params - // this->declare_parameter("model_names", ""); //only declared - // parameters - this->declare_parameter("tensor_qos_profile", "default"); - - // set param vars - this->get_parameter("model_names", models_param); // values separated by comma - model_names = string_to_vector(models_param); - - for (std::vector::iterator it = model_names.begin(); - it != model_names.end(); it++) { - std::string m = *it; - this->get_parameter(m + ".model_path", model_paths[m]); - this->get_parameter(m + ".input_dimensions", input_dimensions_strings[m]); - this->get_parameter(m + ".output_dimensions", output_dimensions_strings[m]); - this->get_parameter(m + ".tensor_type", tensor_type_params[m]); - - RCLCPP_DEBUG(this->get_logger(), "Loading parameters."); - RCLCPP_DEBUG(this->get_logger(), "Model Name: %s", m.c_str()); - RCLCPP_DEBUG(this->get_logger(), "Input Dimensions: %s", - input_dimensions_strings[m].c_str()); - RCLCPP_DEBUG(this->get_logger(), "Output Dimensions: %s", - output_dimensions_strings[m].c_str()); - RCLCPP_DEBUG(this->get_logger(), "Tensor Type: %s", - tensor_type_params[m].c_str()); - - // here we set default values - if (!tensor_type_params.count(m)) { - tensor_type_params[m] = "fp32"; - } - - if (!model_paths.count(m) || !input_dimensions_strings.count(m) || - !output_dimensions_strings.count(m)) { - RCLCPP_WARN(this->get_logger(), - "Parameters incomplete. Could not set up model %s", - m.c_str()); - failed_models.push_back(std::distance(model_names.begin(), it)); - - input_dimensions_strings.erase(m); - output_dimensions_strings.erase(m); - - continue; - } - - // expand strings to dims - input_dimensions[m] = string_to_dims(input_dimensions_strings[m]); - output_dimensions[m] = string_to_dims(output_dimensions_strings[m]); - - // set input and output lengths - input_lengths[m] = 1; - for (uint32_t i : input_dimensions[m]) { - input_lengths[m] *= i; - } - - output_lengths[m] = 1; - for (uint32_t i : output_dimensions[m]) { - output_lengths[m] *= i; - } - - // set tensor type - tensor_typelengths[m] = tensor_string_to_typelength(tensor_type_params[m]); - } - - // iterate backwords over failed models backwords to properly erase from - // model_names - for (std::vector::reverse_iterator it = failed_models.rbegin(); - it != failed_models.rend(); it++) { - model_names.erase(model_names.begin() + *it); // erase by index - } - - this->get_parameter("tensor_qos_profile", tensor_qos_param); - auto tensor_qos = - rclcpp::QoS(rclcpp::KeepLast(10), parseQoSString(tensor_qos_param)); - - // // Publisher and subscriber setup - tensor_publisher_ = this->create_publisher( - "tensorrt_output", tensor_qos); - - // tensor_subscription_ = - // this->create_subscription( - // "tensorrt_input", tensor_qos, - // std::bind(&TensorRTEngineNode::tensor_callback, this, - // std::placeholders::_1)); - - service_ = this->create_service( - "tensorrt_request", - std::bind(&TensorRTEngineNode::tensor_request_callback, this, - std::placeholders::_1, std::placeholders::_2)); - - for (std::vector::iterator it = model_names.begin(); - it != model_names.end(); it++) { - std::string m = *it; - // initialize engine - // TODO try/catch engine creation failure - engines[m].reset(new TRTEngine(model_paths[m], input_dimensions[m], - tensor_typelengths[m])); - - if (engines[m].get() == NULL) { - RCLCPP_ERROR(this->get_logger(), "Failed to initialize engine %s.", - m.c_str()); - } - } -} - -neuromesh_interfaces::msg::Tensor TensorRTEngineNode::execute( - const std::string model, - const neuromesh_interfaces::msg::Tensor &tensor_msg) { - RCLCPP_DEBUG(this->get_logger(), "Execute function for model %s", - model.c_str()); - if (input_lengths[model] * tensor_typelengths[model] != - tensor_msg.data.size()) { - RCLCPP_ERROR(this->get_logger(), - "Input tensor size does not match engine input size %i and %i", - input_lengths[model], tensor_msg.data.size()); - - auto output_msg = neuromesh_interfaces::msg::Tensor(); - output_msg.result = 2; - - return output_msg; - } - - // TODO generalize for all types - // process image - const uint8_t *data_ptr = tensor_msg.data.data(); - - std::vector output_data; - output_data.reserve(output_lengths[model] * tensor_typelengths[model]); - - RCLCPP_DEBUG(this->get_logger(), "Input Lengths %i", input_lengths[model]); - - RCLCPP_DEBUG(this->get_logger(), "Output Lengths %i", output_lengths[model]); - - RCLCPP_DEBUG(this->get_logger(), "Tensor Type Lengths %i", - tensor_typelengths[model]); - - engines[model]->runInference( - data_ptr, input_lengths[model] * tensor_typelengths[model], - output_data.data(), output_lengths[model] * tensor_typelengths[model]); - // publish output - RCLCPP_DEBUG(this->get_logger(), "Inference run successfully"); - - // TODO fill more of the output_msg - auto output_msg = neuromesh_interfaces::msg::Tensor(); - - uint8_t *char_ptr = reinterpret_cast(output_data.data()); - output_msg.data = std::vector( - char_ptr, char_ptr + (output_lengths[model] * tensor_typelengths[model])); - - output_msg.name = tensor_msg.name + "_inference"; - output_msg.shape.dims = output_dimensions[model]; - - output_msg.result = 0; - - // float32 - output_msg.data_type = 9; - - // Add header - // output_msg.header = std_msgs::msg::Header(); - // output_msg.header.stamp = this->now(); - // output_msg.header.frame_id = "tensor_frame"; // Set an appropriate frame_id - - return output_msg; -} - -int TensorRTEngineNode::tensor_string_to_typelength(std::string input) { - - if (input == "fp32") - return 4; - else if (input == "uint8") - return 1; - else if (input == "int8") - return 1; - else if (input == "int32") - return 4; - else if (input == "int64") - return 8; - else - return -1; -} - -// Convert string to ROS2 QoS profile -// from -// https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_nvblox/nvblox_ros_common/src/qos.cpp#L26 -rmw_qos_profile_t TensorRTEngineNode::parseQoSString(const std::string &str) { - std::string profile = str; - // Convert to upper case. - std::transform(profile.begin(), profile.end(), profile.begin(), ::toupper); - - if (profile == "SYSTEM_DEFAULT") { - return rmw_qos_profile_system_default; - } - if (profile == "DEFAULT") { - return rmw_qos_profile_default; - } - if (profile == "PARAMETER_EVENTS") { - return rmw_qos_profile_parameter_events; - } - if (profile == "SERVICES_DEFAULT") { - return rmw_qos_profile_services_default; - } - if (profile == "PARAMETERS") { - return rmw_qos_profile_parameters; - } - if (profile == "SENSOR_DATA") { - return rmw_qos_profile_sensor_data; - } - RCLCPP_WARN_STREAM(rclcpp::get_logger("parseQosString"), - "Unknown QoS profile: " << profile - << ". Returning profile: DEFAULT"); - return rmw_qos_profile_default; -} -} // namespace tensorrt_engine_node - -#include "rclcpp_components/register_node_macro.hpp" -RCLCPP_COMPONENTS_REGISTER_NODE(tensorrt_engine_node::TensorRTEngineNode) diff --git a/tensorrt_engine/src/main.cpp b/tensorrt_engine/src/main.cpp deleted file mode 100644 index 7792d20..0000000 --- a/tensorrt_engine/src/main.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include "tensorrt_engine/engine_node.h" - -int main(int argc, char ** argv) -{ - rclcpp::init(argc, argv); - rclcpp::spin(std::make_shared()); - rclcpp::shutdown(); - - printf("shutting down tflite engine node\n"); - return 0; -}; \ No newline at end of file diff --git a/tensorrt_engine/src/trt_engine.cpp b/tensorrt_engine/src/trt_engine.cpp deleted file mode 100644 index 47acca1..0000000 --- a/tensorrt_engine/src/trt_engine.cpp +++ /dev/null @@ -1,121 +0,0 @@ -#include "tensorrt_engine/trt_engine.h" - -#include -#include -#include -#include -#include - -#include "NvInferPlugin.h" - - -nvinfer1::ICudaEngine *createEngine(const std::string &inputFileName, nvinfer1::ILogger &logger) { - - std::string load_file; - std::string engine_path = inputFileName; - //5 is length of ".onnx" and before ".trt" - - if(engine_path.find(".onnx") != std::string::npos){ - engine_path = engine_path.substr(0, engine_path.find(".onnx", engine_path.length() - 5)) + ".trt"; - } - - initLibNvInferPlugins(nullptr, ""); - - //test if conversion already happened and .trt version exists - std::ifstream f(engine_path.c_str()); - if(f.good()){ //".trt" file - load_file = engine_path; - - std::vector buffer; - { - std::ifstream in(engine_path, std::ios::binary | std::ios::ate); - if (!in) - throw std::runtime_error("Cannot open " + engine_path); - std::streamsize ss = in.tellg(); - in.seekg(0, std::ios::beg); - std::cout << "Input file size = " << ss << std::endl; - buffer.resize(ss); - if (0 == ss || !in.read(buffer.data(), ss)) - throw std::runtime_error("Cannot read" + engine_path + ".engine"); - } - - std::unique_ptr runtime(nvinfer1::createInferRuntime(logger)); - assert(runtime != nullptr); - - return runtime->deserializeCudaEngine(buffer.data(), buffer.size(), nullptr); - - }else{ //".onnx" file - load_file = inputFileName; - - std::unique_ptr builder{nvinfer1::createInferBuilder(logger)}; - std::unique_ptr network{ - builder->createNetworkV2(1U << (unsigned) nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)}; - - std::unique_ptr parser{ nvonnxparser::createParser(*network, logger) }; - - if (!parser->parseFromFile(load_file.c_str(), static_cast(nvinfer1::ILogger::Severity::kINFO))) - throw std::runtime_error("ERROR: could not parse model file " + load_file + " !"); - - std::unique_ptr config(builder->createBuilderConfig()); - - nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - - //save to .trt - std::unique_ptr serializedEngine(engine->serialize()); - - std::ofstream out(engine_path, std::ios::binary); - out.write((char *)serializedEngine->data(), serializedEngine->size()); - - - return engine; - } -} - -TRTEngine::TRTEngine(std::string model_filename, std::vector input_dims, int type_length, int batchSize){ - - Logger logger; - logger.log(nvinfer1::ILogger::Severity::kINFO, "Creating engine ..."); - - engine.reset(createEngine(model_filename, logger)); - - if (!engine) - throw std::runtime_error("Engine creation failed !"); - - logger.log(nvinfer1::ILogger::Severity::kINFO, "Engine created."); - - context.reset(engine->createExecutionContext()); - bindings.resize(engine->getNbBindings()); - - // Alloc cuda memory for IO tensors - for (int i = 0; i < engine->getNbBindings(); ++i) { - nvinfer1::Dims dims{engine->getBindingDimensions(i)}; - size_t size = std::accumulate(dims.d, dims.d + dims.nbDims, batchSize, std::multiplies()); - // Create CUDA buffer for Tensor. - cudaMalloc(&bindings[i], size * type_length); - } -} - -TRTEngine::~TRTEngine(){ - - for (int i = 0; i < bindings.size(); i++){ - cudaFree(bindings[i]); - } -} - -void TRTEngine::runInference( const void* inputTensor, int inputSize, void* outputTensor, int outputSize){ - - int inputId = 0; - int outputId = 1; //only supporting 1 input 1 output - - cudaStream_t stream; - cudaStreamCreate(&stream); - - cudaMemcpyAsync(bindings[inputId], inputTensor, inputSize, cudaMemcpyHostToDevice, - stream); - context->enqueueV2(bindings.data(), stream, nullptr); - cudaMemcpyAsync(outputTensor, bindings[outputId], outputSize, - cudaMemcpyDeviceToHost, stream); - - cudaStreamSynchronize(stream); - cudaStreamDestroy(stream); -} diff --git a/tensorrt_engine/src/trt_engine_modified.cpp b/tensorrt_engine/src/trt_engine_modified.cpp index 7af2c90..67908fe 100644 --- a/tensorrt_engine/src/trt_engine_modified.cpp +++ b/tensorrt_engine/src/trt_engine_modified.cpp @@ -1,12 +1,42 @@ #include "tensorrt_engine/trt_engine_modified.h" - #include #include #include #include #include +#include +#include #include "NvInferPlugin.h" +#include + +namespace engine_interface +{ + +class TRTEngine::Impl +{ +public: + Logger logger; + std::unique_ptr engine; + std::unique_ptr runtime; + std::unique_ptr context; + std::vector bindings; + + Impl() : logger(), engine(nullptr), runtime(nullptr), context(nullptr) { + // Initialize the logger + logger.log(nvinfer1::ILogger::Severity::kINFO, "TRTEngine Impl initialized."); + } +}; + +TRTEngine::TRTEngine() + : impl_(std::make_unique()) { +} + +TRTEngine::~TRTEngine() { + for (size_t i = 0; i < impl_->bindings.size(); i++) { + cudaFree(impl_->bindings[i]); + } +} nvinfer1::ICudaEngine *createEngine(const std::string &inputFileName, nvinfer1::ILogger &logger, @@ -82,43 +112,39 @@ nvinfer1::ICudaEngine *createEngine(const std::string &inputFileName, } } -TRTEngine::TRTEngine(std::string model_filename, - std::vector> input_dims, int type_length, - int batchSize) { +bool TRTEngine::loadModel(const std::string& model_path, + const std::vector>& input_dims, + int type_length) +{ + impl_->logger.log(nvinfer1::ILogger::Severity::kINFO, "Creating engine ..."); - logger.log(nvinfer1::ILogger::Severity::kINFO, "Creating engine ..."); - - runtime.reset(nvinfer1::createInferRuntime(logger)); - if (!runtime) + impl_->runtime.reset(nvinfer1::createInferRuntime(impl_->logger)); + if (!impl_->runtime) throw std::runtime_error("Runtime creation failed!"); - engine.reset(createEngine(model_filename, logger, runtime)); + impl_->engine.reset(createEngine(model_path, impl_->logger, impl_->runtime)); - if (!engine) + if (!impl_->engine) throw std::runtime_error("Engine creation failed !"); - logger.log(nvinfer1::ILogger::Severity::kINFO, "Engine created."); + impl_->logger.log(nvinfer1::ILogger::Severity::kINFO, "Engine created."); - context.reset(engine->createExecutionContext()); + impl_->context.reset(impl_->engine->createExecutionContext()); - const int num_io = engine->getNbIOTensors(); - bindings.resize(num_io); + const int num_io = impl_->engine->getNbIOTensors(); + impl_->bindings.resize(num_io); // Alloc cuda memory for IO tensors for (int i = 0; i < num_io; ++i) { - const char *tensor_name = engine->getIOTensorName(i); - nvinfer1::Dims dims = engine->getTensorShape(tensor_name); + const char *tensor_name = impl_->engine->getIOTensorName(i); + nvinfer1::Dims dims = impl_->engine->getTensorShape(tensor_name); size_t size = std::accumulate(dims.d, dims.d + dims.nbDims, 1, std::multiplies()); // Create CUDA buffer for Tensor - cudaMalloc(&bindings[i], size * type_length); + cudaMalloc(&impl_->bindings[i], size * type_length); } -} -TRTEngine::~TRTEngine() { - for (int i = 0; i < bindings.size(); i++) { - cudaFree(bindings[i]); - } + return true; } void TRTEngine::runInference(const std::vector &inputTensors, @@ -126,15 +152,7 @@ void TRTEngine::runInference(const std::vector &inputTensors, std::vector &outputTensors, const std::vector &outputSizes) { - /* - std::cout << "IN ENGINE: my inputTensors.size() looks like " << inputTensors.size() << std::endl; std::cout << "IN ENGINE: my inputSizes.size() looks like " << inputSizes.size() << std::endl; - std::cout << "IN ENGINE: my inputSizes[0] looks like " << inputSizes.at(0) << std::endl; - - std::cout << "IN ENGINE: my outputTensors.size() looks like " << outputTensors.size() << std::endl; std::cout << "IN ENGINE: my outputSizes.size() looks like " << outputSizes.size() << std::endl; - std::cout << "IN ENGINE: my outputSizes[0] looks like " << outputSizes.at(0) << std::endl; - */ - - if (inputTensors.size() + outputTensors.size() != bindings.size()) { + if (inputTensors.size() + outputTensors.size() != impl_->bindings.size()) { throw std::runtime_error( "Number of input and output tensors doesn't match engine bindings"); } @@ -145,26 +163,26 @@ void TRTEngine::runInference(const std::vector &inputTensors, // Copy input data to device for (size_t i = 0; i < inputTensors.size(); ++i) { - cudaMemcpyAsync(bindings[i], inputTensors[i], inputSizes[i], + cudaMemcpyAsync(impl_->bindings[i], inputTensors[i], inputSizes[i], cudaMemcpyHostToDevice, stream); } for (size_t i = 0; i < inputTensors.size(); ++i) { - const char *tensor_name = engine->getIOTensorName(i); - context->setInputTensorAddress(tensor_name, bindings[i]); + const char *tensor_name = impl_->engine->getIOTensorName(i); + impl_->context->setInputTensorAddress(tensor_name, impl_->bindings[i]); } for (size_t i = 0; i < outputTensors.size(); ++i) { - const char *tensor_name = engine->getIOTensorName(inputTensors.size() + i); - context->setOutputTensorAddress(tensor_name, bindings[inputTensors.size() + i]); + const char *tensor_name = impl_->engine->getIOTensorName(inputTensors.size() + i); + impl_->context->setOutputTensorAddress(tensor_name, impl_->bindings[inputTensors.size() + i]); } // Run inference - context->enqueueV3(stream); + impl_->context->enqueueV3(stream); // Copy output data to host for (size_t i = 0; i < outputTensors.size(); ++i) { - cudaMemcpyAsync(outputTensors[i], bindings[inputTensors.size() + i], + cudaMemcpyAsync(outputTensors[i], impl_->bindings[inputTensors.size() + i], outputSizes[i], cudaMemcpyDeviceToHost, stream); } @@ -172,3 +190,7 @@ void TRTEngine::runInference(const std::vector &inputTensors, cudaStreamDestroy(stream); } +} // namespace engine_interface + +#include +PLUGINLIB_EXPORT_CLASS(engine_interface::TRTEngine, engine_interface::BaseEngine) \ No newline at end of file diff --git a/tensorrt_engine/test/test_trt_engine.cpp b/tensorrt_engine/test/test_trt_engine.cpp new file mode 100644 index 0000000..b3a763e --- /dev/null +++ b/tensorrt_engine/test/test_trt_engine.cpp @@ -0,0 +1,47 @@ +#include +#include "tensorrt_engine/trt_engine_modified.h" +#include "tensorrt_engine/logger.h" + +class TRTEngineTest : public ::testing::Test { +protected: + void SetUp() override { + input_dims = {{1, 3, 224, 224}}; + type_length = sizeof(float); + engine_path = "test_engine.trt"; + } + + std::vector> input_dims; + int type_length; + std::string engine_path; +}; + +TEST_F(TRTEngineTest, LoadModelSuccess) { + engine_interface::TRTEngine engine; + EXPECT_TRUE(engine.loadModel(engine_path, input_dims, type_length)); +} + +TEST_F(TRTEngineTest, InferenceProducesOutput) { + engine_interface::TRTEngine engine; + ASSERT_TRUE(engine.loadModel(engine_path, input_dims, type_length)); + + // Prepare dummy input data + size_t num_elements = 1 * 3 * 224 * 224; + std::vector input_data(num_elements, 1.0f); + + std::vector input_tensors = { input_data.data() }; + std::vector input_sizes = { static_cast(num_elements * sizeof(float)) }; + + // Prepare output buffer + std::vector output_data(1000); + std::vector output_tensors = { output_data.data() }; + std::vector output_sizes = { static_cast(output_data.size() * sizeof(float)) }; + + engine.runInference(input_tensors, input_sizes, output_tensors, output_sizes); + + EXPECT_EQ(output_tensors.size(), 1); + EXPECT_NE(output_tensors[0], nullptr); + + float sum = 0.0f; + for (float val : output_data) sum += val; + EXPECT_NE(sum, 0.0f); +}