diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index ba7d683e..e3f2eaf4 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -2,6 +2,7 @@ FROM ros:noetic-ros-core AS base ARG TZ=Europe/Zurich +ARG VGLUSERS_GID=1004 ENV DEBIAN_FRONTEND=noninteractive ENV PATH="${PATH}:/opt/hpcx/ompi/bin" @@ -15,25 +16,36 @@ ENV LANG="en_US.UTF-8" \ # Add user "asl" with sudo rights RUN groupadd -r asl && \ - useradd --create-home --gid asl --groups dialout,plugdev --shell /bin/bash asl && \ + groupadd -g ${VGLUSERS_GID} vglusers && \ + useradd --create-home --gid asl --groups dialout,plugdev,vglusers --shell /bin/bash asl && \ mkdir -p /etc/sudoers.d && \ echo 'asl ALL=NOPASSWD: ALL' > /etc/sudoers.d/asl -# CUDA: Install +# CUDA: Install (NOTE: libcublas 12.6 is needed for trtexec) RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb && \ sudo dpkg -i cuda-keyring_1.1-1_all.deb && \ sudo apt update && \ - sudo apt -y install cuda-toolkit-12-6 cudnn9-cuda-12 + sudo apt -y install cuda-toolkit-11-8 cudnn9-cuda-11 libcublas-12-6 # CUDA: Add PATH and LD_LIBRARY_PATH to .bash_aliases -RUN echo 'export PATH=$PATH:/usr/local/cuda/bin' >> /home/asl/.bash_aliases && \ - echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64' >> /home/asl/.bash_aliases +RUN echo 'export PATH=$PATH:/usr/local/cuda-11.8/bin' >> /home/asl/.bash_aliases && \ + echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.8/lib64' >> /home/asl/.bash_aliases -# TensorRT: Install from .deb file -COPY resources/nv-tensorrt-repo-ubuntu2004-cuda11.6-trt8.4.3.1-ga-20220813_1-1_amd64.deb /tmp/ -RUN sudo dpkg -i /tmp/resources/nv-tensorrt-repo-ubuntu2004-cuda11.6-trt8.4.3.1-ga-20220813_1-1_amd64.deb || sudo apt-get install -f -y -# RUN sudo cp /var/nv-tensorrt-repo-ubuntu2004-cuda11.6-trt8.4.3.1-ga-20220813/*-keyring.gpg /usr/share/keyrings/ && sudo apt update && sudo apt install -y tensorrt -RUN sudo apt update && sudo apt install -y tensorrt +# TensorRT: Install from .deb file: Seems we run 8.5.2.2 (which is bundled with 8.5.3) +COPY .devcontainer/nv-tensorrt-local-repo-ubuntu2004-8.5.3-cuda-11.8_1.0-1_amd64.deb /tmp/tensorrt.deb +RUN sudo dpkg -i /tmp/tensorrt.deb +RUN sudo cp /var/nv-tensorrt-local-repo-ubuntu2004-8.5.3-cuda-11.8/nv-tensorrt-local-3EFA7C6A-keyring.gpg /usr/share/keyrings/ +RUN sudo apt update && sudo apt install -y tensorrt=8.5.2.2-1+cuda11.8 \ + libnvinfer8=8.5.2-1+cuda11.8 \ + libnvinfer-plugin8=8.5.2-1+cuda11.8 \ + libnvparsers8=8.5.2-1+cuda11.8 \ + libnvonnxparsers8=8.5.2-1+cuda11.8 \ + libnvinfer-bin=8.5.2-1+cuda11.8 \ + libnvinfer-dev=8.5.2-1+cuda11.8 \ + libnvinfer-plugin-dev=8.5.2-1+cuda11.8 \ + libnvparsers-dev=8.5.2-1+cuda11.8 \ + libnvonnxparsers-dev=8.5.2-1+cuda11.8 \ + libnvinfer-samples=8.5.2-1+cuda11.8 # OpenCV Build Stage FROM base AS opencv_build diff --git a/.devcontainer/build_opencv.sh b/.devcontainer/build_opencv.sh index 208995b9..d5b50087 100755 --- a/.devcontainer/build_opencv.sh +++ b/.devcontainer/build_opencv.sh @@ -104,8 +104,8 @@ install_dependencies () { # Automatically detect installed CUDA version detect_cuda_version() { - if command -v /usr/local/cuda/bin/nvcc &> /dev/null; then - CUDA_VERSION=$(/usr/local/cuda/bin/nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+') + if command -v /usr/local/cuda-11.8/bin/nvcc &> /dev/null; then + CUDA_VERSION=$(/usr/local/cuda-11.8/bin/nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+') else echo "CUDA not found. Please install CUDA before running this script." exit 1 @@ -153,6 +153,7 @@ configure() { -D BUILD_opencv_python3=OFF -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=${PREFIX} + -D CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda-11.8 -D CUDA_ARCH_BIN=${CUDA_ARCH_BIN} -D CUDA_FAST_MATH=ON -D CUDNN_VERSION='${CUDNN_VERSION}' @@ -170,8 +171,7 @@ configure() { -D WITH_OPENGL=ON" if [[ "$1" != "test" ]] ; then - CMAKEFLAGS=" - ${CMAKEFLAGS} + CMAKEFLAGS="${CMAKEFLAGS} -D BUILD_PERF_TESTS=OFF -D BUILD_TESTS=OFF" fi diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 61a37ae1..74169662 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -3,7 +3,7 @@ "customizations": { "vscode": { "extensions": [ - "ms-python.python" + "ms-vscode.cpptools" ], "settings": { "files.hotExit": "off", diff --git a/.devcontainer/devcontainer_all_packages.sh b/.devcontainer/devcontainer_all_packages.sh index cdf99f60..dc37609f 100755 --- a/.devcontainer/devcontainer_all_packages.sh +++ b/.devcontainer/devcontainer_all_packages.sh @@ -11,6 +11,7 @@ main() { ccache cmake curl + gdb gfortran git gnupg @@ -52,6 +53,7 @@ main() { python3-dev python3-matplotlib python3-numpy + python3-pip python3-rosdep qv4l2 software-properties-common diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6a4b6bf3..f9e6ab43 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,7 +9,7 @@ jobs: strategy: matrix: - ci_script: [pr_compile] + ci_script: [pr_compile, pr_run_tests] steps: - name: Checkout diff --git a/.gitignore b/.gitignore index b9fc4b5c..e167d158 100644 --- a/.gitignore +++ b/.gitignore @@ -8,5 +8,9 @@ docs/xml site/ -resources/* +.devcontainer/*.deb +*.out +test/resources/*.engine +test/resources/*.onnx build/* +test/resources/*.plan diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..c4cbee9f --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,26 @@ +{ + // Use IntelliSense to learn about possible attributes. + // Hover to view descriptions of existing attributes. + // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 + "version": "0.2.0", + "configurations": [ + { + "name": "Debug test_depth_anything_v2", + "type": "cppdbg", + "request": "launch", + "program": "${workspaceFolder}/build/devel/lib/usb_cam/test_depth_anything_v2", + "cwd": "${workspaceFolder}", + "environment": [], + "externalConsole": false, + "MIMode": "gdb", // Use gdb for debugging + "setupCommands": [ + { + "description": "Enable pretty-printing for gdb", + "text": "-enable-pretty-printing", + "ignoreFailures": true + } + ], + "miDebuggerPath": "/usr/bin/gdb" + } + ] +} \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 733b3bd4..2cd1cf54 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -46,10 +46,15 @@ endif() set(TENSORRT_LIB_DIR /usr/lib/x86_64-linux-gnu CACHE PATH "Path to TensorRT libraries") find_library(NVINFER nvinfer PATHS ${TENSORRT_LIB_DIR}) find_library(NVINFER_PLUGIN nvinfer_plugin PATHS ${TENSORRT_LIB_DIR}) +find_library(NVONNXPARSER nvonnxparser PATHS ${TENSORRT_LIB_DIR}) + +message(STATUS "TensorRT NVINFER library found at: ${NVINFER}") +message(STATUS "TensorRT NVINFER_PLUGIN library found at: ${NVINFER_PLUGIN}") +message(STATUS "TensorRT NVONNXPARSER library found at: ${NVONNXPARSER}") # Check if TensorRT libraries are found -if(NOT NVINFER OR NOT NVINFER_PLUGIN) - message(FATAL_ERROR "TensorRT libraries not found. Set TENSORRT_LIB_DIR correctly.") +if(NOT NVINFER OR NOT NVINFER_PLUGIN OR NOT NVONNXPARSER) + message(FATAL_ERROR "TensorRT libraries not found. Ensure TENSORRT_LIB_DIR is set correctly.") endif() ## Build the USB camera library @@ -79,6 +84,7 @@ target_link_libraries(${PROJECT_NAME} ${swscale_LIBRARIES} ${NVINFER} ${NVINFER_PLUGIN} + ${NVONNXPARSER} ) # Define catkin package @@ -90,12 +96,16 @@ catkin_package( # Build the USB camera node executable add_executable(${PROJECT_NAME}_node src/usb_cam_node.cpp) target_link_libraries(${PROJECT_NAME}_node - ${PROJECT_NAME} - ${catkin_LIBRARIES} - ${avcodec_LIBRARIES} - ${avutil_LIBRARIES} - ${swscale_LIBRARIES} + ${PROJECT_NAME} + ${catkin_LIBRARIES} + ${avcodec_LIBRARIES} + ${avutil_LIBRARIES} + ${swscale_LIBRARIES} + ${NVINFER} + ${NVINFER_PLUGIN} + ${NVONNXPARSER} ) +set_target_properties(${PROJECT_NAME}_node PROPERTIES LINK_FLAGS "-Wl,--no-as-needed") # Ensure include directories are set target_include_directories(${PROJECT_NAME}_node PUBLIC @@ -103,22 +113,22 @@ target_include_directories(${PROJECT_NAME}_node PUBLIC # Testing if(BUILD_TESTING) - if($ENV{ROS_VERSION} EQUAL 2) - find_package(ament_lint_auto REQUIRED) - ament_lint_auto_find_test_dependencies() - find_package(ament_cmake_gtest) - - # Unit tests - ament_add_gtest(test_usb_cam_utils test/test_usb_cam_utils.cpp) - target_link_libraries(test_usb_cam_utils ${PROJECT_NAME}) - - ament_add_gtest(test_pixel_formats test/test_pixel_formats.cpp) - target_link_libraries(test_pixel_formats ${PROJECT_NAME}) - - # Integration tests (commented out as per TODO) - # ament_add_gtest(test_usb_cam_lib test/test_usb_cam_lib.cpp) - # target_link_libraries(test_usb_cam_lib ${PROJECT_NAME}) - endif() + # Find GTest package + set(CMAKE_BUILD_TYPE Debug) + find_package(GTest REQUIRED) + include_directories(${GTEST_INCLUDE_DIRS}) + + catkin_add_gtest(test_depth_anything_v2 test/test_depth_anything_v2.cpp) + target_link_libraries(test_depth_anything_v2 + ${PROJECT_NAME} + ${catkin_LIBRARIES} + GTest::gtest_main + ${NVINFER} + ${NVINFER_PLUGIN} + ${NVONNXPARSER} + ${CUDA_LIBRARIES} + ) + set_target_properties(test_depth_anything_v2 PROPERTIES LINK_FLAGS "-Wl,--no-as-needed") endif() # Installation rules diff --git a/ci/pr_run_tests.sh b/ci/pr_run_tests.sh new file mode 100755 index 00000000..d95acff8 --- /dev/null +++ b/ci/pr_run_tests.sh @@ -0,0 +1,56 @@ +#!/bin/bash +source /opt/ros/noetic/setup.bash + +# Set paths for the model and plan files +MODEL_PATH="test/resources/depth_anything_v2_vits_16.onnx" + +# Check if the ONNX model file exists +if [ ! -f "$MODEL_PATH" ]; then + # If the model file doesn't exist, check if the environment variable is set + if [ -z "$DEPTH_ANYTHING_V2_VITS_16_LINK" ]; then + echo "The model file does not exist, and the environment variable DEPTH_ANYTHING_V2_VITS_16_LINK is not set." + exit 1 + else + # If the environment variable is set, download the model + echo "ONNX model file not found. Attempting to download..." + + # Create the directory if it doesn't exist + mkdir -p "$(dirname "$MODEL_PATH")" + + # Download the file + if wget -O "$MODEL_PATH" "$DEPTH_ANYTHING_V2_VITS_16_LINK"; then + echo "Download successful." + else + echo "Download failed." + exit 1 + fi + fi +else + echo "ONNX model file already exists." +fi + +# Build the project and run tests +rm -rf build +mkdir -p build +cd build + +if cmake .. -DBUILD_TESTING=ON; then + echo "CMake successful." + if make test_depth_anything_v2; then + echo "Make successful." + else + echo "Make failed." + exit 1 + fi +else + echo "CMake failed." + exit 1 +fi + +# Run the test executable +if ./devel/lib/usb_cam/test_depth_anything_v2; then + echo "Tests successful." +else + echo "Tests failed." + exit 1 +fi diff --git a/include/usb_cam/formats/bayer.hpp b/include/usb_cam/formats/bayer.hpp index 3657f8ea..1264ee6c 100644 --- a/include/usb_cam/formats/bayer.hpp +++ b/include/usb_cam/formats/bayer.hpp @@ -113,7 +113,7 @@ class BAYER_GRBG10 : public pixel_format_base } for (uint8_t i = 0; i < 3; ++i) { - channels[i].convertTo(channels[i], CV_8U, avg_gains[i] / 256.0); + channels[i].convertTo(channels[i], CV_8U, avg_gains[i] / 255.0); } ROS_DEBUG("BGR gains: %.3f %.3f %.3f", avg_gains[0], avg_gains[1], avg_gains[2]); cv::merge(channels, _rgb_image_out); @@ -123,7 +123,7 @@ class BAYER_GRBG10 : public pixel_format_base cv::cuda::split(_gpu_rgb_image, _gpu_rgb_channels); for (int i = 0; i < 3; ++i) { - _gpu_rgb_channels[i].convertTo(_gpu_rgb_channels[i], CV_8U, _wb_gains[i] / 256.0); + _gpu_rgb_channels[i].convertTo(_gpu_rgb_channels[i], CV_8U, _wb_gains[i] / 255.0); } cv::cuda::merge(_gpu_rgb_channels, _gpu_rgb_image_8bit); diff --git a/include/usb_cam/learning/depth_anything_v2.hpp b/include/usb_cam/learning/depth_anything_v2.hpp new file mode 100644 index 00000000..3e348fa4 --- /dev/null +++ b/include/usb_cam/learning/depth_anything_v2.hpp @@ -0,0 +1,60 @@ +#ifndef DEPTH_ANYTHING_HPP_ +#define DEPTH_ANYTHING_HPP_ + +#include "interface.hpp" +#include "ros/ros.h" +#include +#include +#include + +class DepthAnythingV2 : public LearningInterface { +public: + DepthAnythingV2(ros::NodeHandle* nh, std::string model_path) { + _INPUT_SIZE = cv::Size(_HEIGHT, _WIDTH); + _model_path = model_path; + + if (nh != nullptr) { + _depth_publication = nh->advertise("depth_anything_v2", 1); + } + } + + void set_input(sensor_msgs::Image& msg) override { + cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::RGB8); + cv::Mat image = cv_ptr->image; + + // Change size to 518x518 (still uint8) + cv::Mat resized_image; + cv::resize(image, resized_image, _INPUT_SIZE); + + // Change to float32 between 0 and 1 + std::vector channels; + cv::split(resized_image, channels); + for (uint8_t i = 0; i < 3; ++i) { + channels[i].convertTo(channels[i], CV_32F, 1.0f / 255.0f); + } + cv::Mat float_image; + cv::merge(channels, float_image); + _input_data = float_image.reshape(1, 1).ptr(0); + } + + void publish() override { + if (_depth_publication.getTopic() != "") { + cv::Mat depth_prediction = cv::Mat(_HEIGHT, _WIDTH, CV_32FC1, _output_data); + + cv_bridge::CvImage depth_image; + depth_image.header.stamp = ros::Time::now(); + depth_image.header.frame_id = "depth_frame"; + depth_image.encoding = sensor_msgs::image_encodings::TYPE_32FC1; + depth_image.image = depth_prediction; + _depth_publication.publish(depth_image.toImageMsg()); + } + } + +private: + const size_t _HEIGHT = 518; + const size_t _WIDTH = 518; + cv::Size _INPUT_SIZE; + ros::Publisher _depth_publication; +}; + +#endif // DEPTH_ANYTHING_HPP_ diff --git a/include/usb_cam/learning/interface.hpp b/include/usb_cam/learning/interface.hpp index 30e75dcc..8d1d1e0d 100644 --- a/include/usb_cam/learning/interface.hpp +++ b/include/usb_cam/learning/interface.hpp @@ -1,55 +1,57 @@ #ifndef LEARNING_INTERFACE_HPP_ #define LEARNING_INTERFACE_HPP_ -#include -#include +#include #include #include #include -#include -#include -#include -#include +#include +#include +#include +#include #include class LearningInterface { public: LearningInterface() : _model_path("") {} - virtual void set_input(const uint8_t* input_buffer, size_t height, size_t width) = 0; - virtual void get_output(uint8_t* output_buffer) = 0; + virtual void set_input(sensor_msgs::Image& image) = 0; virtual void publish() = 0; void load_model(); - bool run_inference(size_t batch_size); - + void predict(); - virtual ~LearningInterface() { - // if (_context) _context->destroy(); - // if (_engine) _engine->destroy(); - // if (_runtime) _runtime->destroy(); + nvinfer1::ICudaEngine* get_engine() { return _engine; } + nvinfer1::IExecutionContext* get_context() { return _context; } + nvinfer1::IRuntime* get_runtime() { return _runtime; } - if (_buffers[0]) cudaFree(_buffers[0]); - if (_buffers[1]) cudaFree(_buffers[1]); - - delete[] _input_buffer; - delete[] _output_buffer; - } + ~LearningInterface(); protected: - float* _input_buffer = nullptr; - float* _output_buffer = nullptr; - size_t input_height; - size_t input_width; - size_t output_height; - size_t output_width; + cudaStream_t _stream; + float* _input_data = nullptr; + float* _output_data = nullptr; + nvinfer1::ICudaEngine* _engine; + nvinfer1::IExecutionContext* _context; + nvinfer1::INetworkDefinition* _network; + nvinfer1::IRuntime* _runtime; std::string _model_path; private: - nvinfer1::ICudaEngine* _engine = nullptr; - nvinfer1::IExecutionContext* _context = nullptr; - nvinfer1::IRuntime* _runtime = nullptr; void* _buffers[2] = { nullptr, nullptr }; + + // TODO: static? + class Logger : public nvinfer1::ILogger { + void log(Severity severity, const char* msg) noexcept override { + // Only output logs with severity greater than warning + if (severity <= Severity::kWARNING) { + std::cout << msg << std::endl; + } + } + } _logger; + + bool _save_engine(const std::string& onnx_path); + void _build(std::string onnx_path); }; #endif // LEARNING_INTERFACE_HPP_ diff --git a/include/usb_cam/learning/raft.hpp b/include/usb_cam/learning/raft.hpp deleted file mode 100644 index 1cf9d439..00000000 --- a/include/usb_cam/learning/raft.hpp +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef RAFT_HPP_ -#define RAFT_HPP_ - -#include "interface.hpp" -#include -#include - -class Raft : public LearningInterface { -public: - Raft(std::string model_path, size_t network_height, size_t network_width) : _network_height(network_height), _network_width(network_width) { - _model_path = model_path; - _network_size = cv::Size(_network_width, _network_height); - } - - void set_input(const uint8_t* input_buffer, size_t height, size_t width) override { - // Resize frame to network size - cv::Mat input_frame(height, width, CV_8UC1, (void*)input_buffer); - cv::Mat resized_frame; - cv::resize(input_frame, resized_frame, _network_size); - - cv::Mat float_frame; - resized_frame.convertTo(float_frame, CV_32FC1, _uint8_to_float); - - cudaMemcpy(_input_buffer, float_frame.ptr(), _network_width * _network_height * sizeof(float), cudaMemcpyHostToDevice); - } - - void get_output(uint8_t* output_buffer) override { - // TODO - } - - void publish() override { - // TODO - } - - -private: - const size_t _network_height; - const size_t _network_width; - cv::Size _network_size; - static constexpr float _uint8_to_float = 1.0f / 255.0f; -}; - -#endif // RAFT_HPP_ \ No newline at end of file diff --git a/scripts/setup_ros_client.sh b/scripts/setup_ros_client.sh new file mode 100644 index 00000000..845e2fe3 --- /dev/null +++ b/scripts/setup_ros_client.sh @@ -0,0 +1,2 @@ +export ROS_MASTER_URI=http://169.254.38.82:11311 +export ROS_HOSTNAME=169.254.229.54 diff --git a/scripts/setup_ros_master.sh b/scripts/setup_ros_master.sh new file mode 100644 index 00000000..06c0645e --- /dev/null +++ b/scripts/setup_ros_master.sh @@ -0,0 +1,2 @@ +export ROS_MASTER_URI=http://169.254.38.82:11311 +export ROS_HOSTNAME=169.254.38.82 diff --git a/scripts/show_image.py b/scripts/show_image.py deleted file mode 100644 index 8e29a7aa..00000000 --- a/scripts/show_image.py +++ /dev/null @@ -1,97 +0,0 @@ -#!/usr/bin/env python3 - -# Copyright 2021 Evan Flynn, Lucas Walter -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# * Redistributions of source code must retain the above copyright -# notice, this list of conditions and the following disclaimer. -# -# * Redistributions in binary form must reproduce the above copyright -# notice, this list of conditions and the following disclaimer in the -# documentation and/or other materials provided with the distribution. -# -# * Neither the name of the Evan Flynn nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -# POSSIBILITY OF SUCH DAMAGE. - - -import cv2 -import numpy as np - -import rclpy -from rclpy.node import Node -from sensor_msgs.msg import Image - - -class ExamineImage(Node): - - def __init__(self): - super().__init__('examine_image') - - self.mat = None - self.sub = self.create_subscription( - Image, - 'image_raw', - self.image_callback, - 100) - - def image_callback(self, msg): - sz = (msg.height, msg.width) - # print(msg.header.stamp) - if False: - print('{encoding} {width} {height} {step} {data_size}'.format( - encoding=msg.encoding, width=msg.width, height=msg.height, - step=msg.step, data_size=len(msg.data))) - if msg.step * msg.height != len(msg.data): - print('bad step/height/data size') - return - - if msg.encoding == 'rgb8': - dirty = (self.mat is None or msg.width != self.mat.shape[1] or - msg.height != self.mat.shape[0] or len(self.mat.shape) < 2 or - self.mat.shape[2] != 3) - if dirty: - self.mat = np.zeros([msg.height, msg.width, 3], dtype=np.uint8) - self.mat[:, :, 2] = np.array(msg.data[0::3]).reshape(sz) - self.mat[:, :, 1] = np.array(msg.data[1::3]).reshape(sz) - self.mat[:, :, 0] = np.array(msg.data[2::3]).reshape(sz) - elif msg.encoding == 'mono8': - self.mat = np.array(msg.data).reshape(sz) - else: - print('unsupported encoding {}'.format(msg.encoding)) - return - if self.mat is not None: - cv2.imshow('image', self.mat) - cv2.waitKey(5) - - -def main(args=None): - rclpy.init(args=args) - - examine_image = ExamineImage() - - try: - rclpy.spin(examine_image) - except KeyboardInterrupt: - pass - - examine_image.destroy_node() - rclpy.shutdown() - - -if __name__ == '__main__': - main() diff --git a/src/interface.cpp b/src/interface.cpp index 1a402f48..2e64109f 100644 --- a/src/interface.cpp +++ b/src/interface.cpp @@ -1,80 +1,101 @@ #include "usb_cam/learning/interface.hpp" +#include + +using namespace nvinfer1; void LearningInterface::load_model() { - // Open and try to read the file - std::ifstream file(_model_path, std::ios::binary); - if (file.good()) { - file.seekg(0, std::ios::end); - const size_t model_size = file.tellg(); - file.seekg(0, std::ios::beg); - - // Read the model data - std::vector model_data(model_size); - file.read(model_data.data(), model_size); - file.close(); + if (_model_path.find(".onnx") == std::string::npos) { + std::ifstream engine_stream(_model_path, std::ios::binary); + engine_stream.seekg(0, std::ios::end); + + const size_t model_size = engine_stream.tellg(); + engine_stream.seekg(0, std::ios::beg); + + std::unique_ptr engine_data(new char[model_size]); + engine_stream.read(engine_data.get(), model_size); + engine_stream.close(); + + // Create tensorrt model + _runtime = nvinfer1::createInferRuntime(_logger); + _engine = _runtime->deserializeCudaEngine(engine_data.get(), model_size); + _context = _engine->createExecutionContext(); - // Create logger instance - class Logger : public nvinfer1::ILogger { - public: - void log(Severity severity, const char* msg) noexcept override { - std::cout << msg << std::endl; // Log the message - } - } logger; // Create a logger instance - - _runtime = nvinfer1::createInferRuntime(logger); - if (_runtime != nullptr) { - _engine = _runtime->deserializeCudaEngine(model_data.data(), model_size); - if (_engine != nullptr) { - _context = _engine->createExecutionContext(); - if (_context != nullptr) { - // Allocate buffers for input and output - size_t input_size; - size_t output_size; - // for (int io = 0; io < _engine->getNbIOTensors(); io++) { - // const char* name = _engine->getIOTensorName(io); - // std::cout << io << ": " << name; - // const nvinfer1::Dims dims = _engine->getTensorShape(name); - - // size_t total_dims = 1; - // for (int d = 0; d < dims.nbDims; d++) { - // total_dims *= dims.d[d]; - // } - - // std::cout << " size: " << total_dims << std::endl; - // if (io == 0) { - // input_size = total_dims * sizeof(float); - // } else if (io == 1) { - // output_size = total_dims * sizeof(float); - // } - // } - - // Allocate device buffers - cudaMalloc(&_buffers[0], input_size); - cudaMalloc(&_buffers[1], output_size); - - // Allocate CPU buffers - _input_buffer = new float[input_size / sizeof(float)]; - _output_buffer = new float[output_size / sizeof(float)]; - - std::cout << "TensorRT model loaded successfully from: " << _model_path << std::endl; - } else { - std::cout << "Failed to create execution context." << std::endl; - } - } else { - std::cout << "Failed to create TensorRT engine." << std::endl; - } - } else { - std::cout << "Failed to create TensorRT runtime." << std::endl; - } } else { - std::cout << "Failed to open model file." << std::endl; + // Build an engine from an onnx model + _build(_model_path); + _save_engine(_model_path); } + + // Define input dimensions + const auto input_dims = _engine->getTensorShape(_engine->getIOTensorName(0)); + const int input_h = input_dims.d[2]; + const int input_w = input_dims.d[3]; + + // Create CUDA stream + cudaStreamCreate(&_stream); + + cudaMalloc(&_buffers[0], 3 * input_h * input_w * sizeof(float)); + cudaMalloc(&_buffers[1], input_h * input_w * sizeof(float)); + + _output_data = new float[input_h * input_w]; +} + +void LearningInterface::_build(std::string onnx_path) { + auto builder = createInferBuilder(_logger); + const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); + INetworkDefinition* network = builder->createNetworkV2(explicitBatch); + IBuilderConfig* config = builder->createBuilderConfig(); + config->setFlag(BuilderFlag::kFP16); + nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, _logger); + bool parsed = parser->parseFromFile(onnx_path.c_str(), static_cast(nvinfer1::ILogger::Severity::kINFO)); + IHostMemory* plan{ builder->buildSerializedNetwork(*network, *config) }; + + _runtime = createInferRuntime(_logger); + _engine = _runtime->deserializeCudaEngine(plan->data(), plan->size()); + _context = _engine->createExecutionContext(); + + delete network; + delete config; + delete parser; + delete plan; } -bool LearningInterface::run_inference(size_t batch_size) { - if (!_context->executeV2(_buffers)) { - std::cerr << "Failed to execute inference." << std::endl; +bool LearningInterface::_save_engine(const std::string& onnx_path) { + std::string engine_path; + size_t dot_index = onnx_path.find_last_of("."); + if (dot_index != std::string::npos) { + engine_path = onnx_path.substr(0, dot_index) + ".engine"; + + } else { return false; } + + if (_engine) { + nvinfer1::IHostMemory* data = _engine->serialize(); + std::ofstream file; + file.open(engine_path, std::ios::binary | std::ios::out); + if (!file.is_open()) { + std::cout << "Create engine file" << engine_path << " failed" << std::endl; + return 0; + } + + file.write((const char*)data->data(), data->size()); + file.close(); + + delete data; + } return true; } + +void LearningInterface::predict() { + cudaMemcpyAsync(_buffers[0], _input_data, sizeof(_input_data) * sizeof(float), cudaMemcpyHostToDevice, _stream); + _context->executeV2(_buffers); + cudaStreamSynchronize(_stream); + cudaMemcpyAsync(_output_data, _buffers[1], sizeof(_input_data) * sizeof(float), cudaMemcpyDeviceToHost); +} + +LearningInterface::~LearningInterface() { + cudaFree(_stream); + cudaFree(_buffers[0]); + cudaFree(_buffers[1]); +} diff --git a/src/usb_cam_node.cpp b/src/usb_cam_node.cpp index 388fa13b..517a43e7 100644 --- a/src/usb_cam_node.cpp +++ b/src/usb_cam_node.cpp @@ -38,7 +38,7 @@ #include "usb_cam/utils.hpp" #include "usb_cam/learning/interface.hpp" -#include "usb_cam/learning/raft.hpp" +#include "usb_cam/learning/depth_anything_v2.hpp" namespace usb_cam { class UsbCamNode { @@ -72,7 +72,8 @@ class UsbCamNode { UsbCamNode() : m_node("~") { // Setup the network that outputs derivates of the image captured - // networks.push_back(std::make_unique("resources/raft.onnx", 240, 320)); + // TODO: Actual network + networks.push_back(std::make_unique(&m_node, "depth_anything_v2_vitb.onnx")); // Advertise the main image topic image_transport::ImageTransport it(m_node); @@ -177,12 +178,11 @@ class UsbCamNode { m_image_pub.publish(m_image, *ci); // Run all the networks - // for (const auto& net : networks) { - // net->set_input(m_image.data.data(), 1920, 1200); - // if (net->run_inference(1)) { - // net->publish(); - // } - // } + for (const auto& net : networks) { + net->set_input(m_image); + net->predict(); + net->publish(); + } return true; } diff --git a/test/resources/maschinenhalle_example_frame.jpg b/test/resources/maschinenhalle_example_frame.jpg new file mode 100644 index 00000000..ce7f7b76 Binary files /dev/null and b/test/resources/maschinenhalle_example_frame.jpg differ diff --git a/test/test_depth_anything_v2.cpp b/test/test_depth_anything_v2.cpp new file mode 100644 index 00000000..a665cdae --- /dev/null +++ b/test/test_depth_anything_v2.cpp @@ -0,0 +1,63 @@ +#include "usb_cam/learning/depth_anything_v2.hpp" +#include +#include +#include +#include +#include +#include + +// This class provides access to protected members that we normally don't want to expose +class DepthAnythingV2Test : public DepthAnythingV2 { + public: + DepthAnythingV2Test(const std::string& model_path) : DepthAnythingV2(nullptr, model_path) {} + float* get_input_data() { return _input_data; } +}; + +class TestDepthAnythingV2 : public ::testing::Test { +protected: + const std::string model_path = "/workspaces/v4l2_camera/test/resources/depth_anything_v2_vits_16.onnx"; + const std::string test_image_path = "/workspaces/v4l2_camera/test/resources/maschinenhalle_example_frame.jpg"; + + cv_bridge::CvImage cv_image; + cv::Mat img; + DepthAnythingV2Test* depth_anything_v2; + sensor_msgs::Image msg; + + void SetUp() override { + img = cv::imread(test_image_path, cv::IMREAD_COLOR); + cv_image.image = img; + cv_image.encoding = sensor_msgs::image_encodings::RGB8; + cv_image.toImageMsg(msg); + + depth_anything_v2 = new DepthAnythingV2Test(model_path); + depth_anything_v2->load_model(); + } + + void TearDown() override { + delete depth_anything_v2; + } +}; + +TEST_F(TestDepthAnythingV2, TestModelLoad) { + ASSERT_NE(depth_anything_v2->get_engine(), nullptr); + ASSERT_NE(depth_anything_v2->get_context(), nullptr); + ASSERT_NE(depth_anything_v2->get_runtime(), nullptr); +} + +TEST_F(TestDepthAnythingV2, TestSetInput) { + depth_anything_v2->set_input(msg); + ASSERT_NE(depth_anything_v2->get_input_data(), nullptr); + + const size_t expected_size = 518 * 518 * 3; + float* input_data = depth_anything_v2->get_input_data(); + ASSERT_NE(input_data, nullptr); + ASSERT_FLOAT_EQ(input_data[0], img.at(0, 0)[0] / 255.0f); +} + +TEST_F(TestDepthAnythingV2, TestPredictPublish) { + for (size_t i = 0; i < 10; i++) { + depth_anything_v2->set_input(msg); + depth_anything_v2->predict(); + depth_anything_v2->publish(); + } +} diff --git a/test/test_pixel_formats.cpp b/test/test_pixel_formats.cpp deleted file mode 100644 index 44a9a566..00000000 --- a/test/test_pixel_formats.cpp +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright 2023 Evan Flynn -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// * Redistributions of source code must retain the above copyright -// notice, this list of conditions and the following disclaimer. -// -// * Redistributions in binary form must reproduce the above copyright -// notice, this list of conditions and the following disclaimer in the -// documentation and/or other materials provided with the distribution. -// -// * Neither the name of the Evan Flynn nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -// POSSIBILITY OF SUCH DAMAGE. - - -#include - -#include - -#include "usb_cam/formats/pixel_format_base.hpp" - -TEST(test_pixel_formats, pixel_format_base_class) { - auto test_pix_fmt = usb_cam::formats::default_pixel_format(); - - EXPECT_EQ(test_pix_fmt.name(), "yuyv"); - EXPECT_EQ(test_pix_fmt.v4l2(), V4L2_PIX_FMT_YUYV); - EXPECT_EQ(test_pix_fmt.channels(), 2); - EXPECT_EQ(test_pix_fmt.bit_depth(), 8); - EXPECT_EQ(test_pix_fmt.requires_conversion(), false); - - EXPECT_EQ(test_pix_fmt.is_bayer(), false); - // TOOD(flynneva): should this be true for `yuyv`? - EXPECT_EQ(test_pix_fmt.is_color(), false); - EXPECT_EQ(test_pix_fmt.is_mono(), false); -} diff --git a/test/test_usb_cam_lib.cpp b/test/test_usb_cam_lib.cpp deleted file mode 100644 index f0801586..00000000 --- a/test/test_usb_cam_lib.cpp +++ /dev/null @@ -1,59 +0,0 @@ -// Copyright 2023 Evan Flynn -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// * Redistributions of source code must retain the above copyright -// notice, this list of conditions and the following disclaimer. -// -// * Redistributions in binary form must reproduce the above copyright -// notice, this list of conditions and the following disclaimer in the -// documentation and/or other materials provided with the distribution. -// -// * Neither the name of the Evan Flynn nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -// POSSIBILITY OF SUCH DAMAGE. - -#include - -#include -#include -#include - -#include "usb_cam/usb_cam.hpp" -#include "usb_cam/utils.hpp" - -TEST(test_usb_cam_lib, test_usb_cam_class) { - usb_cam::UsbCam test_usb_cam; - - auto supported_fmts = test_usb_cam.get_supported_formats(); - - // TODO(flynneva): iterate over availble formats with test_usb_cam obj - for (auto fmt : supported_fmts) { - std::cerr << "format: " << fmt.format.type << std::endl; - } - - // TODO(flynneva): rework these tests in another MR - { - // test_usb_cam.configure( - // "/dev/video0", - // usb_cam::utils::IO_METHOD_MMAP, - // "yuyv2rgb", 640, 480, 30); - // test_usb_cam.start(); - // TODO(flynneva): uncomment once /dev/video0 can be simulated in CI - // EXPECT_TRUE(test_usb_cam.is_capturing()); - // test_usb_cam.shutdown(); - } -} diff --git a/test/test_usb_cam_utils.cpp b/test/test_usb_cam_utils.cpp deleted file mode 100644 index a0ac6fcb..00000000 --- a/test/test_usb_cam_utils.cpp +++ /dev/null @@ -1,83 +0,0 @@ -// Copyright 2023 Evan Flynn -// -// Redistribution and use in source and binary forms, with or without -// modification, are permitted provided that the following conditions are met: -// -// * Redistributions of source code must retain the above copyright -// notice, this list of conditions and the following disclaimer. -// -// * Redistributions in binary form must reproduce the above copyright -// notice, this list of conditions and the following disclaimer in the -// documentation and/or other materials provided with the distribution. -// -// * Neither the name of the Evan Flynn nor the names of its -// contributors may be used to endorse or promote products derived from -// this software without specific prior written permission. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE -// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -// POSSIBILITY OF SUCH DAMAGE. - -#include -#include - -#include - -#include "usb_cam/utils.hpp" -#include "usb_cam/formats/utils.hpp" - - -using usb_cam::utils::io_method_from_string; - - -TEST(test_usb_cam_utils, test_io_method_from_string) { - usb_cam::utils::io_method_t test_io; - - test_io = io_method_from_string("mmap"); - EXPECT_EQ(test_io, usb_cam::utils::IO_METHOD_MMAP); - - test_io = io_method_from_string("read"); - EXPECT_EQ(test_io, usb_cam::utils::IO_METHOD_READ); - - test_io = io_method_from_string("userptr"); - EXPECT_EQ(test_io, usb_cam::utils::IO_METHOD_USERPTR); - - test_io = io_method_from_string("bananas"); - EXPECT_EQ(test_io, usb_cam::utils::IO_METHOD_UNKNOWN); -} - -TEST(test_usb_cam_utils, test_clip_value) { - // Clip values to 0 if -128<=val<0 - for (int i = -128; i < 0; i++) { - EXPECT_EQ(0, usb_cam::formats::CLIPVALUE(i)); - } - // Do not clip values between 0 383 - // these will use the old method (non-array method) - EXPECT_EQ(0, usb_cam::formats::CLIPVALUE(-129)); - EXPECT_EQ(255, usb_cam::formats::CLIPVALUE(400)); -} - -TEST(test_usb_cam_utils, test_monotonic_to_real_time) { - // Get timeval to use for test - - const time_t test_time_t = usb_cam::utils::get_epoch_time_shift_us(); - - EXPECT_NE(test_time_t, 0); -}