├── README.md ├── caffe_ros ├── README.md ├── tests │ ├── data │ │ ├── rot_c.jpg │ │ ├── rot_l.jpg │ │ ├── rot_r.jpg │ │ ├── tran_l.jpg │ │ ├── tran_r.jpg │ │ ├── sidewalk.png │ │ ├── yolo_2_obj.png │ │ └── sidewalk-raw-out.png │ ├── tests_basic.launch │ └── tests.cpp ├── src │ ├── caffe_ros_node.cpp │ ├── caffe_ros.cpp │ └── tensor_net.cpp ├── launch │ ├── gscam.launch~ │ ├── gscam.launch │ ├── caffe_gscam_h264.launch │ ├── caffe_gscam_h265.launch │ ├── caffe_gscam.launch │ ├── segnet.launch │ ├── mavros_caffe.launch │ ├── twoDNNs.launch │ └── everything.launch ├── package.xml ├── include │ └── caffe_ros │ │ ├── #segnet.h# │ │ ├── segnet.h │ │ ├── segnet.h~ │ │ ├── caffe_ros.h │ │ ├── tensor_net.h │ │ └── yolo_prediction.h ├── parse_folders.py └── CMakeLists.txt ├── output_view ├── color_mapping.lut ├── launch │ └── segnet-viz.launch ├── package.xml ├── include │ └── output_view │ │ └── segnet.h ├── src │ └── output_view.cpp └── CMakeLists.txt ├── image_pub ├── CMakeLists.txt ├── package.xml └── src │ └── image_pub_node.cpp └── LICENSE.md /README.md: -------------------------------------------------------------------------------- 1 | # caffe_ros 2 | Package containing nodes for deep learning in ROS. 3 | -------------------------------------------------------------------------------- /caffe_ros/README.md: -------------------------------------------------------------------------------- 1 | # custom-caffe_ros 2 | A clean commit of the caffe_ros node created by Alexey. 3 | -------------------------------------------------------------------------------- /caffe_ros/tests/data/rot_c.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/rot_c.jpg -------------------------------------------------------------------------------- /caffe_ros/tests/data/rot_l.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/rot_l.jpg -------------------------------------------------------------------------------- /caffe_ros/tests/data/rot_r.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/rot_r.jpg -------------------------------------------------------------------------------- /caffe_ros/tests/data/tran_l.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/tran_l.jpg -------------------------------------------------------------------------------- /caffe_ros/tests/data/tran_r.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/tran_r.jpg -------------------------------------------------------------------------------- /caffe_ros/tests/data/sidewalk.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/sidewalk.png -------------------------------------------------------------------------------- /caffe_ros/tests/data/yolo_2_obj.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/yolo_2_obj.png -------------------------------------------------------------------------------- /caffe_ros/tests/data/sidewalk-raw-out.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVIDIA-AI-IOT/caffe_ros/HEAD/caffe_ros/tests/data/sidewalk-raw-out.png -------------------------------------------------------------------------------- /output_view/color_mapping.lut: -------------------------------------------------------------------------------- 1 | -- Look up table giving colors for each channel, 2 | -- input output, 3 | 0, 10 4 | 1, 20 5 | 3, 30 6 | 4, 40 7 | 5, 50 8 | 6, 60 9 | 7, 70 10 | 8, 80 11 | 9, 90 12 | 10, 100 13 | 11, 110 14 | 12, 120 15 | 13, 130 16 | 14, 140 17 | 15, 150 18 | 16, 160 19 | 17, 170 20 | 18, 180 21 | 19, 190 22 | 20, 200 23 | 21, 210 -------------------------------------------------------------------------------- /caffe_ros/src/caffe_ros_node.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #include "caffe_ros/caffe_ros.h" 5 | 6 | int main(int argc, char **argv) 7 | { 8 | ros::init(argc, argv, "caffe_ros"); 9 | 10 | caffe_ros::CaffeRos caffe_r; 11 | caffe_r.spin(); 12 | return 0; 13 | } 14 | -------------------------------------------------------------------------------- /output_view/launch/segnet-viz.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /caffe_ros/launch/gscam.launch~: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /caffe_ros/launch/gscam.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | -------------------------------------------------------------------------------- /image_pub/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | project(image_pub) 3 | 4 | find_package(catkin REQUIRED COMPONENTS 5 | roscpp 6 | std_msgs 7 | ) 8 | 9 | ########### 10 | ## Build ## 11 | ########### 12 | 13 | set (CMAKE_CXX_STANDARD 14) 14 | set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall") 15 | 16 | include_directories( 17 | include 18 | ${catkin_INCLUDE_DIRS} 19 | ) 20 | 21 | file(GLOB image_pub_sources src/*.cpp) 22 | 23 | add_executable(image_pub_node ${image_pub_sources}) 24 | 25 | target_link_libraries(image_pub_node 26 | ${catkin_LIBRARIES} 27 | opencv_core 28 | opencv_imgproc 29 | opencv_highgui 30 | ) 31 | -------------------------------------------------------------------------------- /image_pub/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | image_pub 4 | 0.1.0 5 | Image publisher node to be used to test caffe_ros node. 6 | Alexey Kamenev 7 | Alexey Kamenev 8 | BSD 9 | 10 | catkin 11 | 12 | roscpp 13 | std_msgs 14 | 15 | roscpp 16 | std_msgs 17 | 18 | 19 | 20 | -------------------------------------------------------------------------------- /caffe_ros/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | caffe_ros 4 | 0.1.0 5 | ROS wrapper for Caffe inference API with TensorRT. 6 | Alexey Kamenev 7 | Alexey Kamenev 8 | BSD 9 | 10 | catkin 11 | 12 | roscpp 13 | std_msgs 14 | cuda-toolkit-8-0 15 | 16 | 17 | roscpp 18 | std_msgs 19 | 20 | gtest 21 | 22 | 23 | 24 | -------------------------------------------------------------------------------- /caffe_ros/include/caffe_ros/#segnet.h#: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | namespace caffe_ros 4 | { 5 | void getArgMax(const float *scores, size_t num_channels, cv::Mat &classMap) 6 | { 7 | uint32_t height = classMap.rows; 8 | uint32_t width = classMap.cols; 9 | for( uint32_t y=0; y < height; y++ ) 10 | { 11 | for( uint32_t x=0; x < width; x++ ) 12 | { 13 | 14 | // find max probability across the channels 15 | float p_max[3] = {-100000.0f, -100000.0f, -100000.0f }; 16 | int c_max[3] = { -1, -1, -1 }; 17 | 18 | for( uint32_t c=0; c < num_channels; c++ )// classes 19 | { 20 | // get current score 21 | const float p = scores[c * width * height + y * width + x]; 22 | 23 | // choose top three classes 24 | if( c_max[0] < 0 || p > p_max[0] ) 25 | { 26 | p_max[0] = p; 27 | c_max[0] = c; 28 | } 29 | else if( c_max[1] < 0 || p > p_max[1] ) 30 | { 31 | p_max[1] = p; 32 | c_max[1] = c; 33 | } 34 | else if( c_max[2] < 0 || p > p_max[2] ) 35 | { 36 | p_max[2] = p; 37 | c_max[2] = c; 38 | } 39 | } 40 | 41 | classMap.at(y, x) = (uint8_t) c_max[0]; 42 | } 43 | } 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /caffe_ros/include/caffe_ros/segnet.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #ifndef CAFFE_ROS_SEGMENTATION_H 5 | #define CAFFE_ROS_SEGMENTATION_H 6 | 7 | namespace caffe_ros 8 | { 9 | void getArgMax(const float *scores, size_t num_channels, cv::Mat &class_map) 10 | { 11 | size_t height = class_map.rows; 12 | size_t width = class_map.cols; 13 | 14 | for (size_t y = 0; y < height; y++) 15 | { 16 | for (size_t x = 0; x < width; x++) 17 | { 18 | // find max probability across the channels 19 | float p_max = -1.0f; 20 | int c_max = -1; 21 | 22 | // loop through channels (each representing a class) 23 | for (size_t c = 0; c < num_channels; c++) 24 | { 25 | // get current score 26 | const float p = scores[c * width * height + y * width + x]; 27 | 28 | if (c_max < 0 || p > p_max) 29 | { 30 | p_max = p; 31 | c_max = c; 32 | } 33 | } 34 | 35 | // set most likely class at proper pixel location 36 | class_map.at(y, x) = (uint8_t) c_max; 37 | } 38 | } 39 | } 40 | } 41 | 42 | #endif 43 | -------------------------------------------------------------------------------- /output_view/package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | output_view 4 | 0.0.0 5 | The output_view package 6 | 7 | Rachel Gardner 8 | BSD 9 | 10 | 11 | 12 | 13 | 14 | 15 | Rachel Gardner --> 16 | 17 | catkin 18 | cv_bridge 19 | image_transport 20 | roscpp 21 | sensor_msgs 22 | std_msgs 23 | message_filters 24 | cv_bridge 25 | image_transport 26 | roscpp 27 | sensor_msgs 28 | std_msgs 29 | message_filters 30 | 31 | 32 | -------------------------------------------------------------------------------- /caffe_ros/include/caffe_ros/segnet.h~: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | namespace caffe_ros 4 | { 5 | uint8_t* getArgMax(const float *scores, size_t width, size_t height, size_t num_channels) 6 | { 7 | uint8_t *classMap[width*height]; 8 | 9 | for( uint32_t y=0; y < height; y++ ) 10 | { 11 | for( uint32_t x=0; x < width; x++ ) 12 | { 13 | 14 | // find max probability across the channels 15 | float p_max[3] = {-100000.0f, -100000.0f, -100000.0f }; 16 | int c_max[3] = { -1, -1, -1 }; 17 | 18 | for( uint32_t c=0; c < num_channels; c++ )// classes 19 | { 20 | // get current score 21 | const float p = scores[c * width * height + y * width + x]; 22 | 23 | // choose top three classes 24 | if( c_max[0] < 0 || p > p_max[0] ) 25 | { 26 | p_max[0] = p; 27 | c_max[0] = c; 28 | } 29 | else if( c_max[1] < 0 || p > p_max[1] ) 30 | { 31 | p_max[1] = p; 32 | c_max[1] = c; 33 | } 34 | else if( c_max[2] < 0 || p > p_max[2] ) 35 | { 36 | p_max[2] = p; 37 | c_max[2] = c; 38 | } 39 | } 40 | 41 | // const int argmax = (c_max[0] == ignoreID) ? c_max[1] : c_max[0]; 42 | 43 | // classMap[y * s_w + x] = argmax; 44 | classMap[y * width + x] = c_max[0]; 45 | } 46 | } 47 | return classMap; 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /output_view/include/output_view/segnet.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #ifndef SEGNET 5 | #define SEGNET 6 | 7 | bool loadClassColors(const char* filename, cv::Mat &lookUpTable) 8 | { 9 | 10 | if(!filename) 11 | return false; 12 | 13 | FILE* f = fopen(filename, "r"); 14 | 15 | if(!f) 16 | { 17 | printf("failed to open %s\n", filename); 18 | return false; 19 | } 20 | 21 | char str[512]; 22 | int idx = 0; 23 | 24 | while( fgets(str, 512, f) != NULL ) 25 | { 26 | const int len = strlen(str); 27 | 28 | if( len > 0 ) 29 | { 30 | if( str[len-1] == '\n' ) 31 | str[len-1] = 0; 32 | 33 | int r = 255; 34 | int g = 255; 35 | int b = 255; 36 | 37 | sscanf(str, "%i %i %i ", &r, &g, &b); 38 | lookUpTable.at(idx) = cv::Vec3b(r, g, b); 39 | idx++; 40 | } 41 | } 42 | 43 | fclose(f); 44 | 45 | if( idx == 0 ) 46 | return false; 47 | 48 | return true; 49 | } 50 | 51 | void colorImage(const cv::Mat &lookUpTable, const cv::Mat &values, cv::Mat &finalImage) 52 | { 53 | for(int y=0; y(cv::Point(x,y)); 56 | cv::Vec3b color = lookUpTable.at(class_id); 57 | finalImage.at(cv::Point(x, y)) = color; 58 | } 59 | } 60 | } 61 | #endif 62 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions 5 | are met: 6 | * Redistributions of source code must retain the above copyright 7 | notice, this list of conditions and the following disclaimer. 8 | * Redistributions in binary form must reproduce the above copyright 9 | notice, this list of conditions and the following disclaimer in the 10 | documentation and/or other materials provided with the distribution. 11 | * Neither the name of NVIDIA CORPORATION nor the names of its 12 | contributors may be used to endorse or promote products derived 13 | from this software without specific prior written permission. 14 | 15 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY 16 | EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 17 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR 18 | PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR 19 | CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 20 | EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 21 | PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 22 | PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY 23 | OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /caffe_ros/launch/caffe_gscam_h264.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | -------------------------------------------------------------------------------- /caffe_ros/launch/caffe_gscam_h265.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | -------------------------------------------------------------------------------- /caffe_ros/launch/caffe_gscam.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | -------------------------------------------------------------------------------- /caffe_ros/launch/segnet.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /caffe_ros/launch/mavros_caffe.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | -------------------------------------------------------------------------------- /caffe_ros/include/caffe_ros/caffe_ros.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #ifndef CAFFE_ROS_CAFFE_ROS_H 5 | #define CAFFE_ROS_CAFFE_ROS_H 6 | 7 | #include 8 | #include 9 | #include "caffe_ros/tensor_net.h" 10 | 11 | namespace caffe_ros 12 | { 13 | // Implements caffe_ros node. 14 | class CaffeRos 15 | { 16 | public: 17 | CaffeRos(); 18 | ~CaffeRos() = default; 19 | 20 | void spin(); 21 | 22 | private: 23 | // Specifies whether to apply any post processing to the output of the DNN. 24 | enum class PostProc 25 | { 26 | None = 0, 27 | YOLO, // Compute object boxes from the output of YOLO DNN. 28 | Segmentation // Collapse multidimensional softmax into a matrix of class IDs. 29 | }; 30 | 31 | // Default camera queue size. Recommended value is one as to make 32 | // sure we process most recent image from the camera. 33 | const int DEFAULT_CAMERA_QUEUE_SIZE = 1; 34 | // DNN output (publisher) queue. Value of 1 makes sure only most recent 35 | // output gets published. 36 | const int DEFAULT_DNN_QUEUE_SIZE = 1; 37 | 38 | // Current image being worked on. 39 | sensor_msgs::Image::ConstPtr cur_img_; 40 | 41 | // Publisher for the DNN output. 42 | ros::Publisher output_pub_; 43 | // Subscriber to camera capture topic (gscam). 44 | ros::Subscriber image_sub_; 45 | // DNN predictor. 46 | TensorNet net_; 47 | 48 | bool debug_mode_; 49 | std::string debug_dir_; 50 | 51 | PostProc post_proc_; 52 | 53 | // Probability and IOU thresholds used in object detection net (YOLO). 54 | float obj_det_threshold_; 55 | float iou_threshold_; 56 | 57 | // Max rate to run the node at (in Hz). 58 | float max_rate_hz_; 59 | 60 | private: 61 | sensor_msgs::Image::ConstPtr computeOutputs(); 62 | 63 | void imageCallback(const sensor_msgs::Image::ConstPtr& msg); 64 | 65 | void setPostProcessing(const std::string& postProc) 66 | { 67 | if (postProc.size() == 0) 68 | post_proc_ = PostProc::None; 69 | else if (postProc == "YOLO") 70 | post_proc_ = PostProc::YOLO; 71 | else if (postProc == "Segmentation") 72 | post_proc_ = PostProc::Segmentation; 73 | else 74 | { 75 | ROS_FATAL("Post processing %s is not supported. Supported: YOLO, Segmentation", postProc.c_str()); 76 | ros::shutdown(); 77 | } 78 | } 79 | 80 | }; 81 | } 82 | 83 | #endif 84 | -------------------------------------------------------------------------------- /image_pub/src/image_pub_node.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | int main(int argc, char **argv) 9 | { 10 | ROS_INFO("Starting image_pub ROS node...\n"); 11 | 12 | ros::init(argc, argv, "image_pub"); 13 | ros::NodeHandle nh("~"); 14 | 15 | std::string camera_topic; 16 | std::string img_path; 17 | float pub_rate; 18 | int start_sec; 19 | bool repeat; 20 | nh.param("camera_topic", camera_topic, "/camera/image_raw"); 21 | nh.param("img_path", img_path, ""); 22 | nh.param("pub_rate", pub_rate, 30.0f); 23 | nh.param("start_sec", start_sec, 0); 24 | nh.param("repeat", repeat, false); 25 | 26 | ROS_INFO("Topic : %s", camera_topic.c_str()); 27 | ROS_INFO("Source: %s", img_path.c_str()); 28 | ROS_INFO("Rate : %.1f", pub_rate); 29 | ROS_INFO("Start : %d", start_sec); 30 | ROS_INFO("Repeat: %s", repeat ? "yes" : "no"); 31 | 32 | ros::Publisher img_pub = nh.advertise(camera_topic, 100); 33 | 34 | cv::VideoCapture vid_cap(img_path.c_str()); 35 | if (start_sec > 0) 36 | vid_cap.set(CV_CAP_PROP_POS_MSEC, 1000.0 * start_sec); 37 | 38 | ros::Rate rate(pub_rate); 39 | while (ros::ok()) 40 | { 41 | cv::Mat img; 42 | if (!vid_cap.read(img)) 43 | { 44 | if (repeat) 45 | { 46 | vid_cap.open(img_path.c_str()); 47 | if (start_sec > 0) 48 | vid_cap.set(CV_CAP_PROP_POS_MSEC, 1000.0 * start_sec); 49 | continue; 50 | } 51 | ROS_ERROR("Failed to capture frame."); 52 | ros::shutdown(); 53 | } 54 | else 55 | { 56 | //ROS_DEBUG("Image: %dx%dx%d, %zu, %d", img.rows, img.cols, img.channels(), img.elemSize(), img.type() == CV_8UC3); 57 | if (img.type() != CV_8UC3) 58 | img.convertTo(img, CV_8UC3); 59 | // Convert image from BGR format used by OpenCV to RGB. 60 | cv::cvtColor(img, img, CV_BGR2RGB); 61 | 62 | auto img_msg = boost::make_shared(); 63 | img_msg->header.stamp = ros::Time::now(); 64 | img_msg->encoding = "rgb8"; 65 | img_msg->width = img.cols; 66 | img_msg->height = img.rows; 67 | img_msg->step = img_msg->width * img.channels(); 68 | auto ptr = img.ptr(0); 69 | img_msg->data = std::vector(ptr, ptr + img_msg->step * img_msg->height); 70 | img_pub.publish(img_msg); 71 | } 72 | ros::spinOnce(); 73 | rate.sleep(); 74 | } 75 | 76 | return 0; 77 | } 78 | -------------------------------------------------------------------------------- /caffe_ros/parse_folders.py: -------------------------------------------------------------------------------- 1 | ''' 2 | parse_folders.py 3 | 4 | This script is used to convert a directory structure into the appropriate map files for use 5 | in DIGITS. The script randomizes and partitions the data, then outputs the appropriate label 6 | indices (using the name of the folder the image was found in as the image's label). 7 | ''' 8 | 9 | import argparse 10 | import os 11 | import sys 12 | import numpy as np 13 | import random 14 | 15 | def main(root_dir, path_to_output_files, labels, percent_train): 16 | 17 | train_map_path = os.path.join(path_to_output_files, 'train_map.txt') 18 | validate_map_path = os.path.join(path_to_output_files, 'val_map.txt') 19 | labels_map_path = os.path.join(path_to_output_files, 'labels_map.txt') 20 | 21 | with open(train_map_path, 'w') as train_map, \ 22 | open(validate_map_path, 'w') as validate_map, \ 23 | open(labels_map_path, 'w') as labels_map: 24 | 25 | # generate list of all files 26 | files = [os.path.join(dp, f) for dp, dn, file_names in os.walk(root_dir) for f in file_names] 27 | 28 | # enumerate all possible labels 29 | if not labels: 30 | subdirs = [name for name in os.listdir(root_dir)] 31 | labels = {subdirs[i] : i for i in range(len(subdirs))} 32 | else: 33 | labels = {labels[i] : i for i in range(len(labels))} 34 | 35 | # randomly split the data 36 | random.shuffle(files) 37 | num_in_train = int((percent_train/float(100))*len(files)) 38 | train, validate = np.split(files, [num_in_train]) 39 | 40 | for file in train: 41 | base_path, subdir = os.path.split(os.path.dirname(file)) 42 | fileName = os.path.join(subdir, os.path.basename(file)) 43 | train_map.write('{} {}\n'.format(fileName, labels[subdir])) 44 | 45 | for file in validate: 46 | base_path, subdir = os.path.split(os.path.dirname(file)) 47 | fileName = os.path.join(subdir, os.path.basename(file)) 48 | validate_map.write('{} {}\n'.format(fileName, labels[subdir])) 49 | 50 | # sort labels for printing them out 51 | labels = sorted(labels, key=labels.get) 52 | for label in labels: 53 | labels_map.write('{}\n'.format(label)) 54 | 55 | if __name__ == '__main__': 56 | parser = argparse.ArgumentParser(description='Create map files from custom directory structure.') 57 | parser.add_argument('root_dir') 58 | parser.add_argument('path_to_output_files') 59 | parser.add_argument('--labels') 60 | parser.add_argument('percent_train', type=int) # i.e. 80 to mean 80% 61 | args = parser.parse_args() 62 | labels = [] 63 | if args.labels: 64 | labels = args.labels.split(",") 65 | print(labels) 66 | main(args.root_dir, args.path_to_output_files, labels, args.percent_train) 67 | -------------------------------------------------------------------------------- /caffe_ros/launch/twoDNNs.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | -------------------------------------------------------------------------------- /caffe_ros/include/caffe_ros/tensor_net.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #ifndef CAFFE_ROS_TENSOR_NET_H 5 | #define CAFFE_ROS_TENSOR_NET_H 6 | 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | namespace caffe_ros 13 | { 14 | 15 | class TensorNet 16 | { 17 | public: 18 | using ConstStr = const std::string; 19 | 20 | TensorNet(); 21 | virtual ~TensorNet(); 22 | 23 | void loadNetwork(ConstStr& prototxtPath, ConstStr& modelPath, 24 | ConstStr& inputBlob = "data", ConstStr& outputBlob = "prob", 25 | bool useFP16 = true, bool use_cached_model = true); 26 | 27 | void forward(const unsigned char* input, size_t w, size_t h, size_t c); 28 | 29 | int getInWidth() const { return in_dims_.w(); } 30 | int getInHeight() const { return in_dims_.h(); } 31 | int getInChannels() const { return in_dims_.c(); } 32 | 33 | int getOutWidth() const { return out_dims_.w(); } 34 | int getOutHeight() const { return out_dims_.h(); } 35 | int getOutChannels() const { return out_dims_.c(); } 36 | 37 | const float* getOutput() const { return out_h_; } 38 | 39 | void setInputFormat(ConstStr& inputFormat) 40 | { 41 | if (inputFormat == "BGR") 42 | inp_fmt_ = InputFormat::BGR; 43 | else if (inputFormat == "RGB") 44 | inp_fmt_ = InputFormat::RGB; 45 | else 46 | { 47 | ROS_FATAL("Input format %s is not supported. Supported formats: BGR and RGB", inputFormat.c_str()); 48 | ros::shutdown(); 49 | } 50 | } 51 | 52 | void setShift(float shift) 53 | { 54 | assert(std::isfinite(shift)); 55 | inp_shift_ = shift; 56 | } 57 | 58 | void setScale(float scale) 59 | { 60 | assert(std::isfinite(scale)); 61 | inp_scale_ = scale; 62 | } 63 | 64 | void showProfile(bool on) 65 | { 66 | assert(context_ != nullptr); 67 | debug_mode_ = on; 68 | context_->setProfiler(on ? &s_profiler : nullptr); 69 | } 70 | 71 | protected: 72 | 73 | // Formats of the input layer. BGR is usually used by most of the frameworks that use OpenCV. 74 | enum class InputFormat 75 | { 76 | BGR = 0, 77 | RGB 78 | }; 79 | 80 | void profileModel(ConstStr& prototxtPath, ConstStr& modelPath, bool useFP16, ConstStr& outputBlob, std::ostream& model); 81 | 82 | class Logger : public nvinfer1::ILogger 83 | { 84 | void log(Severity severity, const char *msg) override; 85 | }; 86 | static Logger s_log; 87 | 88 | class Profiler : public nvinfer1::IProfiler 89 | { 90 | public: 91 | void printLayerTimes(); 92 | protected: 93 | void reportLayerTime(const char *layerName, float ms) override; 94 | private: 95 | using Record = std::pair; 96 | std::vector profile_; 97 | }; 98 | static Profiler s_profiler; 99 | 100 | nvinfer1::IRuntime* infer_; 101 | nvinfer1::ICudaEngine* engine_; 102 | nvinfer1::IExecutionContext* context_; 103 | 104 | nvinfer1::DimsCHW in_dims_; 105 | nvinfer1::DimsCHW out_dims_; 106 | 107 | InputFormat inp_fmt_ = InputFormat::BGR; 108 | 109 | cv::Mat in_h_; 110 | cv::Mat in_final_h_; 111 | // cv::gpu::GpuMat m_inOrigD; 112 | // cv::gpu::GpuMat m_inD; 113 | cv::gpu::GpuMat in_d_; 114 | float* out_h_ = nullptr; 115 | float* out_d_ = nullptr; 116 | 117 | float inp_shift_ = 0; 118 | float inp_scale_ = 1; 119 | 120 | // This is a separate flag from ROS_DEBUG to enable only specific profiling 121 | // of data preparation and DNN feed forward. 122 | bool debug_mode_ = false; 123 | }; 124 | 125 | } 126 | 127 | #endif 128 | -------------------------------------------------------------------------------- /caffe_ros/launch/everything.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | -------------------------------------------------------------------------------- /output_view/src/output_view.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include "output_view/segnet.h" 14 | 15 | static const std::string OPENCV_WINDOW = "Image window"; 16 | 17 | using namespace sensor_msgs; 18 | using namespace message_filters; 19 | 20 | class ImageConverter 21 | { 22 | ros::Publisher image_pub; 23 | double alpha; 24 | cv::Mat lookUpTable; 25 | 26 | public: 27 | 28 | ImageConverter(const std::string &colors_file, double a, const std::string &image_output) 29 | { 30 | /* 31 | ROS_INFO("Now publishing to %s topic", image_output.c_str()); 32 | image_pub = nh.advertise(image_output, 1); 33 | */ 34 | alpha = a; 35 | 36 | // TODO - undefined behavior when getting in more colors than it expects 37 | lookUpTable = cv::Mat(1, 256, CV_8UC3); 38 | 39 | bool result = loadClassColors(colors_file.c_str(), lookUpTable); 40 | ROS_ASSERT(result == true); // NOTE: expressions inside ROS_ASSERT will not execute 41 | ROS_INFO("Loaded colors from %s", colors_file.c_str()); 42 | } 43 | 44 | ~ImageConverter() 45 | { 46 | cv::destroyWindow(OPENCV_WINDOW); 47 | } 48 | 49 | void segmentationCb(const ImageConstPtr& raw_image, const ImageConstPtr& seg_info) 50 | { 51 | cv_bridge::CvImageConstPtr raw_image_ptr; 52 | cv_bridge::CvImageConstPtr seg_info_ptr; 53 | try 54 | { 55 | raw_image_ptr = cv_bridge::toCvShare(raw_image, sensor_msgs::image_encodings::BGR8); 56 | seg_info_ptr = cv_bridge::toCvShare(seg_info, sensor_msgs::image_encodings::MONO8); 57 | } 58 | catch (cv_bridge::Exception& e) 59 | { 60 | ROS_ERROR("cv_bridge exception: %s", e.what()); 61 | return; 62 | } 63 | 64 | cv::Mat colored_image(seg_info_ptr->image.size(), CV_8UC3); 65 | colorImage(lookUpTable, seg_info_ptr->image, colored_image); 66 | 67 | cv::Mat scaled_image(raw_image_ptr->image.size(), CV_8UC3); 68 | // TODO - add default scaling behavior that preserves ratio 69 | cv::resize(colored_image, scaled_image, scaled_image.size()); 70 | 71 | cv::Mat final_image; 72 | cv::addWeighted(scaled_image, alpha, raw_image_ptr->image, 1 - alpha, 0.0, final_image); 73 | 74 | // Update GUI Window 75 | cv::namedWindow(OPENCV_WINDOW, CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO | CV_GUI_NORMAL); 76 | cv::imshow(OPENCV_WINDOW, final_image); 77 | cv::waitKey(3); 78 | 79 | // ImagePtr output = cv_bridge::CvImage(std_msgs::Header(), "bgr8", scaled_image).toImageMsg(); 80 | 81 | // Output modified video stream 82 | // image_pub.publish(output); 83 | } 84 | }; 85 | 86 | int main(int argc, char** argv) 87 | { 88 | ros::init(argc, argv, "image_converter"); 89 | 90 | ros::NodeHandle nh("~"); 91 | 92 | std::string image_input, image_output; 93 | std::string annotation_input, annotation_type; 94 | std::string colors_file; 95 | double alpha; 96 | 97 | nh.param("image_input", image_input, "/camera/image_raw"); 98 | nh.param("annotation_input", annotation_input, "/segnet/network/output"); 99 | nh.param("annotation_type", annotation_type, "segmentation"); 100 | nh.param("image_output", image_output, "/dnn_data/out"); 101 | nh.param("colors_file", colors_file, "/path/to/colors_file.lut"); 102 | nh.param("alpha", alpha, 0.8); 103 | 104 | ImageConverter img(colors_file, alpha, image_output); 105 | 106 | // TODO - time sync code cannot be put into the constructor of ImageView...why? 107 | // TODO - this needs to be a class only because of cv::destroyWindow 108 | ROS_INFO("Now subscribing to raw image topic: %s", image_input.c_str()); 109 | message_filters::Subscriber image_sub(nh, image_input, 1); 110 | 111 | ROS_INFO("Now subscribing to data topic: %s (%s)", annotation_input.c_str(), annotation_type.c_str()); 112 | message_filters::Subscriber info_sub(nh, annotation_input, 1); 113 | 114 | TimeSynchronizer sync(image_sub, info_sub, 10); 115 | 116 | if (annotation_type == "segmentation") { 117 | sync.registerCallback(boost::bind(&ImageConverter::segmentationCb, &img, _1, _2)); 118 | } else { 119 | ROS_ERROR("Annotation type not supported"); 120 | } 121 | 122 | ros::spin(); 123 | return 0; 124 | } 125 | -------------------------------------------------------------------------------- /caffe_ros/tests/tests_basic.launch: -------------------------------------------------------------------------------- 1 | 2 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 91 | 92 | 93 | 94 | 95 | -------------------------------------------------------------------------------- /caffe_ros/include/caffe_ros/yolo_prediction.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #ifndef CAFFE_ROS_YOLO_PREDICTION_H 5 | #define CAFFE_ROS_YOLO_PREDICTION_H 6 | 7 | namespace caffe_ros 8 | { 9 | struct ObjectPrediction 10 | { 11 | int label; 12 | float prob; 13 | int x; 14 | int y; 15 | int w; 16 | int h; 17 | }; 18 | 19 | std::vector getYoloPredictions(const float* predictions, size_t num_preds, int w_in, int h_in, 20 | float prob_threshold = 0.1) 21 | { 22 | assert(predictions != nullptr); 23 | assert(num_preds > 0); 24 | assert(w_in > 0); 25 | assert(h_in > 0); 26 | assert(0 < prob_threshold && prob_threshold <= 1); 27 | 28 | const int grid_size = 7; 29 | const int num_lab = 20; 30 | const int num_box = 2; 31 | const int num_box_coord = 4; 32 | 33 | assert(num_preds == grid_size * grid_size * (num_box * (num_box_coord + 1) + num_lab)); 34 | 35 | std::vector res; 36 | size_t icell = 0; 37 | for (int row = 0; row < grid_size; row++) 38 | { 39 | for (int col = 0; col < grid_size; col++, icell++) 40 | { 41 | // Find max conditional class probability for the current cell. 42 | auto cell_preds = predictions + icell * num_lab; 43 | auto it_max = std::max_element(cell_preds, cell_preds + num_lab); 44 | int imax_p = it_max - cell_preds; 45 | assert(0 <= imax_p && imax_p < num_lab); 46 | float max_p = *(it_max); 47 | // Find a box with a max condidence prediction. 48 | auto cell_box_scores = predictions + grid_size * grid_size * num_lab + icell * num_box; 49 | auto it_box_max = std::max_element(cell_box_scores, cell_box_scores + num_box); 50 | int imax_box = it_box_max - cell_box_scores; 51 | assert(0 <= imax_box && imax_box < num_box); 52 | float box_score = *(it_box_max); 53 | // Skip entries with conditional class probability below the threshold. 54 | if (box_score * max_p < prob_threshold) 55 | continue; 56 | // Save box for the current cell. 57 | auto cell_box_coords = predictions + grid_size * grid_size * (num_lab + num_box) + (icell * num_box + imax_box) * num_box_coord; 58 | float x = (cell_box_coords[0] + col) / grid_size * w_in; 59 | float y = (cell_box_coords[1] + row) / grid_size * h_in; 60 | float w = std::max(cell_box_coords[2], 0.0f); 61 | float h = std::max(cell_box_coords[3], 0.0f); 62 | // Square the w/h as it was trained like that in YOLO. 63 | w *= w * w_in; 64 | h *= h * h_in; 65 | // x,y is the center of the box, find top left corner. 66 | x -= w / 2; 67 | y -= h / 2; 68 | // Make sure box coordinates are in the valid range. 69 | x = std::min(std::max(x, 0.0f), (float)w_in - 1); 70 | y = std::min(std::max(y, 0.0f), (float)h_in - 1); 71 | w = std::min(w, w_in - x); 72 | h = std::min(h, h_in - y); 73 | ObjectPrediction cur; 74 | cur.label = imax_p; 75 | cur.prob = box_score * max_p; 76 | cur.x = (int)x; 77 | cur.y = (int)y; 78 | cur.w = (int)w; 79 | cur.h = (int)h; 80 | assert(0 <= cur.x && cur.x < w_in); 81 | assert(0 <= cur.y && cur.y < h_in); 82 | assert(0 < cur.x + cur.w && cur.x + cur.w <= w_in); 83 | assert(0 < cur.y + cur.h && cur.y + cur.h <= h_in); 84 | res.push_back(cur); 85 | } 86 | } 87 | assert(icell == grid_size * grid_size); 88 | return res; 89 | } 90 | 91 | // Filter objects according to IOU threshold. 92 | std::vector filterByIOU(std::vector src, float iou_threshold = 0.5) 93 | { 94 | assert(0 < iou_threshold && iou_threshold <= 1); 95 | 96 | // Filter out overlapping boxes. 97 | size_t i1 = 0; 98 | while (i1 < src.size()) 99 | { 100 | auto b1 = src[i1]; 101 | size_t i2 = i1 + 1; 102 | while (i2 < src.size()) 103 | { 104 | auto b2 = src[i2]; 105 | float b_union = b1.w * b1.h + b2.w * b2.h; 106 | assert(b_union > 0); 107 | float wi = std::max(std::min(b1.x + b1.w - b2.x, b2.x + b2.w - b1.x), 0); 108 | float hi = std::max(std::min(b1.y + b1.h - b2.y, b2.y + b2.h - b1.y), 0); 109 | float b_intersect = wi * hi; 110 | assert(0 <= b_intersect && b_intersect <= b_union); 111 | float iou = b_intersect / (b_union - b_intersect); 112 | // Remove box with IOU above threshold (e.g. "ambiguous" or "duplicate" box). 113 | if (iou > iou_threshold) 114 | src.erase(src.begin() + i2); 115 | else 116 | i2++; 117 | } 118 | i1++; 119 | } 120 | 121 | return src; 122 | } 123 | } 124 | 125 | #endif 126 | -------------------------------------------------------------------------------- /caffe_ros/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.1) 2 | project(caffe_ros) 3 | 4 | ## Find catkin macros and libraries 5 | ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz) 6 | ## is used, also find other catkin packages 7 | find_package(catkin REQUIRED COMPONENTS 8 | roscpp 9 | std_msgs 10 | ) 11 | 12 | ## System dependencies are found with CMake's conventions 13 | #find_package(Boost REQUIRED COMPONENTS system) 14 | find_package(CUDA) 15 | # OpenCV 2.4.13 is default on Jetson. Might be required to install in simulation environment. 16 | find_package(OpenCV 2.4.13 REQUIRED) 17 | 18 | ## Uncomment this if the package has a setup.py. This macro ensures 19 | ## modules and global scripts declared therein get installed 20 | ## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html 21 | # catkin_python_setup() 22 | 23 | ################################################ 24 | ## Declare ROS messages, services and actions ## 25 | ################################################ 26 | 27 | ## To declare and build messages, services or actions from within this 28 | ## package, follow these steps: 29 | ## * Let MSG_DEP_SET be the set of packages whose message types you use in 30 | ## your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...). 31 | ## * In the file package.xml: 32 | ## * add a build_depend tag for "message_generation" 33 | ## * add a build_depend and a run_depend tag for each package in MSG_DEP_SET 34 | ## * If MSG_DEP_SET isn't empty the following dependency has been pulled in 35 | ## but can be declared for certainty nonetheless: 36 | ## * add a run_depend tag for "message_runtime" 37 | ## * In this file (CMakeLists.txt): 38 | ## * add "message_generation" and every package in MSG_DEP_SET to 39 | ## find_package(catkin REQUIRED COMPONENTS ...) 40 | ## * add "message_runtime" and every package in MSG_DEP_SET to 41 | ## catkin_package(CATKIN_DEPENDS ...) 42 | ## * uncomment the add_*_files sections below as needed 43 | ## and list every .msg/.srv/.action file to be processed 44 | ## * uncomment the generate_messages entry below 45 | ## * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...) 46 | 47 | ## Generate messages in the 'msg' folder 48 | # add_message_files( 49 | # FILES 50 | # Message1.msg 51 | # Message2.msg 52 | # ) 53 | 54 | ## Generate services in the 'srv' folder 55 | # add_service_files( 56 | # FILES 57 | # Service1.srv 58 | # Service2.srv 59 | # ) 60 | 61 | ## Generate actions in the 'action' folder 62 | # add_action_files( 63 | # FILES 64 | # Action1.action 65 | # Action2.action 66 | # ) 67 | 68 | ## Generate added messages and services with any dependencies listed here 69 | # generate_messages( 70 | # DEPENDENCIES 71 | # std_msgs 72 | # ) 73 | 74 | ################################### 75 | ## catkin specific configuration ## 76 | ################################### 77 | ## The catkin_package macro generates cmake config files for your package 78 | ## Declare things to be passed to dependent projects 79 | ## INCLUDE_DIRS: uncomment this if you package contains header files 80 | ## LIBRARIES: libraries you create in this project that dependent projects also need 81 | ## CATKIN_DEPENDS: catkin_packages dependent projects also need 82 | ## DEPENDS: system dependencies of this project that dependent projects also need 83 | catkin_package( 84 | INCLUDE_DIRS include 85 | # LIBRARIES caffe_ros 86 | # CATKIN_DEPENDS roscpp std_msgs 87 | # DEPENDS system_lib 88 | ) 89 | 90 | ########### 91 | ## Build ## 92 | ########### 93 | 94 | set (CMAKE_CXX_STANDARD 14) 95 | set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall") 96 | 97 | set( 98 | CUDA_NVCC_FLAGS 99 | ${CUDA_NVCC_FLAGS}; 100 | -O3 -gencode arch=compute_53,code=sm_53 101 | ) 102 | 103 | ## Specify additional locations of header files 104 | ## Your package locations should be listed before other locations 105 | # include_directories(include) 106 | include_directories( 107 | include 108 | ${catkin_INCLUDE_DIRS} 109 | ) 110 | 111 | ## Declare a C++ library 112 | # add_library(caffe_ros 113 | # src/${PROJECT_NAME}/caffe_ros.cpp 114 | # ) 115 | 116 | ## Add cmake target dependencies of the library 117 | ## as an example, code may need to be generated before libraries 118 | ## either from message generation or dynamic reconfigure 119 | # add_dependencies(caffe_ros ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) 120 | 121 | ## Declare a C++ executable 122 | file(GLOB caffe_ros_sources src/*.cpp) 123 | 124 | cuda_add_executable(caffe_ros_node ${caffe_ros_sources}) 125 | 126 | ## Add cmake target dependencies of the executable 127 | ## same as for the library above 128 | # add_dependencies(caffe_ros_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) 129 | 130 | ## Specify libraries to link a library or executable target against 131 | target_link_libraries(caffe_ros_node 132 | ${catkin_LIBRARIES} 133 | nvcaffe_parser 134 | nvinfer 135 | opencv_core 136 | opencv_imgproc 137 | opencv_gpu 138 | opencv_highgui 139 | ) 140 | 141 | ############# 142 | ## Install ## 143 | ############# 144 | 145 | # all install targets should use catkin DESTINATION variables 146 | # See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html 147 | 148 | ## Mark executable scripts (Python etc.) for installation 149 | ## in contrast to setup.py, you can choose the destination 150 | # install(PROGRAMS 151 | # scripts/my_python_script 152 | # DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 153 | # ) 154 | 155 | ## Mark executables and/or libraries for installation 156 | # install(TARGETS caffe_ros caffe_ros_node 157 | # ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} 158 | # LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} 159 | # RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 160 | # ) 161 | 162 | ## Mark cpp header files for installation 163 | # install(DIRECTORY include/${PROJECT_NAME}/ 164 | # DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} 165 | # FILES_MATCHING PATTERN "*.h" 166 | # PATTERN ".svn" EXCLUDE 167 | # ) 168 | 169 | ## Mark other files for installation (e.g. launch and bag files, etc.) 170 | # install(FILES 171 | # # myfile1 172 | # # myfile2 173 | # DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION} 174 | # ) 175 | 176 | ############# 177 | ## Testing ## 178 | ############# 179 | 180 | if(CATKIN_ENABLE_TESTING) 181 | find_package(rostest REQUIRED) 182 | add_rostest_gtest(${PROJECT_NAME}_tests tests/tests_basic.launch tests/tests.cpp) 183 | target_link_libraries(${PROJECT_NAME}_tests 184 | ${catkin_LIBRARIES} 185 | opencv_core 186 | opencv_imgproc 187 | opencv_highgui 188 | ) 189 | 190 | ## Add gtest based cpp test target and link libraries 191 | #catkin_add_gtest(${PROJECT_NAME}_test test/test_caffe_ros.cpp) 192 | #if(TARGET ${PROJECT_NAME}-test) 193 | # target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME}) 194 | #endif() 195 | endif() 196 | 197 | ## Add folders to be run by python nosetests 198 | # catkin_add_nosetests(test) 199 | -------------------------------------------------------------------------------- /output_view/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8.3) 2 | project(output_view) 3 | 4 | ## Compile as C++11, supported in ROS Kinetic and newer 5 | # add_compile_options(-std=c++11) 6 | 7 | ## Find catkin macros and libraries 8 | ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz) 9 | ## is used, also find other catkin packages 10 | find_package(catkin REQUIRED COMPONENTS 11 | cv_bridge 12 | image_transport 13 | roscpp 14 | sensor_msgs 15 | message_filters 16 | std_msgs 17 | ) 18 | 19 | ## System dependencies are found with CMake's conventions 20 | # find_package(Boost REQUIRED COMPONENTS system) 21 | 22 | 23 | ## Uncomment this if the package has a setup.py. This macro ensures 24 | ## modules and global scripts declared therein get installed 25 | ## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html 26 | # catkin_python_setup() 27 | 28 | ################################################ 29 | ## Declare ROS messages, services and actions ## 30 | ################################################ 31 | 32 | ## To declare and build messages, services or actions from within this 33 | ## package, follow these steps: 34 | ## * Let MSG_DEP_SET be the set of packages whose message types you use in 35 | ## your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...). 36 | ## * In the file package.xml: 37 | ## * add a build_depend tag for "message_generation" 38 | ## * add a build_depend and a run_depend tag for each package in MSG_DEP_SET 39 | ## * If MSG_DEP_SET isn't empty the following dependency has been pulled in 40 | ## but can be declared for certainty nonetheless: 41 | ## * add a run_depend tag for "message_runtime" 42 | ## * In this file (CMakeLists.txt): 43 | ## * add "message_generation" and every package in MSG_DEP_SET to 44 | ## find_package(catkin REQUIRED COMPONENTS ...) 45 | ## * add "message_runtime" and every package in MSG_DEP_SET to 46 | ## catkin_package(CATKIN_DEPENDS ...) 47 | ## * uncomment the add_*_files sections below as needed 48 | ## and list every .msg/.srv/.action file to be processed 49 | ## * uncomment the generate_messages entry below 50 | ## * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...) 51 | 52 | ## Generate messages in the 'msg' folder 53 | # add_message_files( 54 | # FILES 55 | # Message1.msg 56 | # Message2.msg 57 | # ) 58 | 59 | ## Generate services in the 'srv' folder 60 | # add_service_files( 61 | # FILES 62 | # Service1.srv 63 | # Service2.srv 64 | # ) 65 | 66 | ## Generate actions in the 'action' folder 67 | # add_action_files( 68 | # FILES 69 | # Action1.action 70 | # Action2.action 71 | # ) 72 | 73 | ## Generate added messages and services with any dependencies listed here 74 | # generate_messages( 75 | # DEPENDENCIES 76 | # sensor_msgs# std_msgs 77 | # ) 78 | 79 | ################################################ 80 | ## Declare ROS dynamic reconfigure parameters ## 81 | ################################################ 82 | 83 | ## To declare and build dynamic reconfigure parameters within this 84 | ## package, follow these steps: 85 | ## * In the file package.xml: 86 | ## * add a build_depend and a run_depend tag for "dynamic_reconfigure" 87 | ## * In this file (CMakeLists.txt): 88 | ## * add "dynamic_reconfigure" to 89 | ## find_package(catkin REQUIRED COMPONENTS ...) 90 | ## * uncomment the "generate_dynamic_reconfigure_options" section below 91 | ## and list every .cfg file to be processed 92 | 93 | ## Generate dynamic reconfigure parameters in the 'cfg' folder 94 | # generate_dynamic_reconfigure_options( 95 | # cfg/DynReconf1.cfg 96 | # cfg/DynReconf2.cfg 97 | # ) 98 | 99 | ################################### 100 | ## catkin specific configuration ## 101 | ################################### 102 | ## The catkin_package macro generates cmake config files for your package 103 | ## Declare things to be passed to dependent projects 104 | ## INCLUDE_DIRS: uncomment this if you package contains header files 105 | ## LIBRARIES: libraries you create in this project that dependent projects also need 106 | ## CATKIN_DEPENDS: catkin_packages dependent projects also need 107 | ## DEPENDS: system dependencies of this project that dependent projects also need 108 | catkin_package( 109 | INCLUDE_DIRS include 110 | # LIBRARIES output_view 111 | CATKIN_DEPENDS cv_bridge image_transport roscpp sensor_msgs std_msgs message_filters 112 | DEPENDS system_lib 113 | ) 114 | 115 | ########### 116 | ## Build ## 117 | ########### 118 | 119 | ## Specify additional locations of header files 120 | ## Your package locations should be listed before other locations 121 | include_directories( 122 | include 123 | ${catkin_INCLUDE_DIRS} 124 | ) 125 | 126 | ## Declare a C++ library 127 | # add_library(${PROJECT_NAME} 128 | # src/${PROJECT_NAME}/output_view.cpp 129 | # ) 130 | 131 | ## Add cmake target dependencies of the library 132 | ## as an example, code may need to be generated before libraries 133 | ## either from message generation or dynamic reconfigure 134 | # add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) 135 | 136 | ## Declare a C++ executable 137 | ## With catkin_make all packages are built within a single CMake context 138 | ## The recommended prefix ensures that target names across packages don't collide 139 | add_executable(${PROJECT_NAME} src/output_view.cpp) 140 | 141 | ## Rename C++ executable without prefix 142 | ## The above recommended prefix causes long target names, the following renames the 143 | ## target back to the shorter version for ease of user use 144 | ## e.g. "rosrun someones_pkg node" instead of "rosrun someones_pkg someones_pkg_node" 145 | # set_target_properties(${PROJECT_NAME}_node PROPERTIES OUTPUT_NAME node PREFIX "") 146 | 147 | ## Add cmake target dependencies of the executable 148 | ## same as for the library above 149 | # add_dependencies(${PROJECT_NAME}_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) 150 | 151 | ## Specify libraries to link a library or executable target against 152 | target_link_libraries(${PROJECT_NAME} 153 | ${catkin_LIBRARIES} 154 | ) 155 | 156 | ############# 157 | ## Install ## 158 | ############# 159 | 160 | # all install targets should use catkin DESTINATION variables 161 | # See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html 162 | 163 | ## Mark executable scripts (Python etc.) for installation 164 | ## in contrast to setup.py, you can choose the destination 165 | # install(PROGRAMS 166 | # scripts/my_python_script 167 | # DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 168 | # ) 169 | 170 | ## Mark executables and/or libraries for installation 171 | # install(TARGETS ${PROJECT_NAME} ${PROJECT_NAME}_node 172 | # ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} 173 | # LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} 174 | # RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 175 | # ) 176 | 177 | ## Mark cpp header files for installation 178 | # install(DIRECTORY include/${PROJECT_NAME}/ 179 | # DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} 180 | # FILES_MATCHING PATTERN "*.h" 181 | # PATTERN ".svn" EXCLUDE 182 | # ) 183 | 184 | ## Mark other files for installation (e.g. launch and bag files, etc.) 185 | # install(FILES 186 | # # myfile1 187 | # # myfile2 188 | # DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION} 189 | # ) 190 | 191 | ############# 192 | ## Testing ## 193 | ############# 194 | 195 | ## Add gtest based cpp test target and link libraries 196 | # catkin_add_gtest(${PROJECT_NAME}-test test/test_output_view.cpp) 197 | # if(TARGET ${PROJECT_NAME}-test) 198 | # target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME}) 199 | # endif() 200 | 201 | ## Add folders to be run by python nosetests 202 | # catkin_add_nosetests(test) 203 | -------------------------------------------------------------------------------- /caffe_ros/src/caffe_ros.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #include "caffe_ros/caffe_ros.h" 5 | #include "caffe_ros/tensor_net.h" 6 | 7 | #include 8 | #include 9 | 10 | #include "caffe_ros/yolo_prediction.h" 11 | #include "caffe_ros/segnet.h" 12 | 13 | namespace caffe_ros 14 | { 15 | CaffeRos::CaffeRos() 16 | { 17 | ROS_INFO("Starting Caffe ROS node..."); 18 | ros::NodeHandle nh("~"); 19 | 20 | std::string camera_topic; 21 | std::string prototxt_path; 22 | std::string model_path; 23 | std::string input_layer; 24 | std::string output_layer; 25 | std::string inp_fmt; 26 | std::string post_proc; 27 | bool use_FP16; 28 | float inp_scale; 29 | float inp_shift; 30 | int camera_queue_size; 31 | int dnn_queue_size; 32 | bool use_cached_model; 33 | 34 | nh.param("camera_topic", camera_topic, "/camera/image_raw"); 35 | nh.param("prototxt_path", prototxt_path, ""); 36 | nh.param("model_path", model_path, ""); 37 | nh.param("input_layer", input_layer, "data"); 38 | nh.param("output_layer", output_layer, "prob"); 39 | nh.param("inp_fmt", inp_fmt, "BGR"); 40 | nh.param("post_proc", post_proc, ""); 41 | nh.param("use_fp16", use_FP16, true); 42 | nh.param("inp_scale", inp_scale, 1.0f); 43 | nh.param("inp_shift", inp_shift, 0.0f); 44 | nh.param("camera_queue_size", camera_queue_size, DEFAULT_CAMERA_QUEUE_SIZE); 45 | nh.param("dnn_queue_size", dnn_queue_size, DEFAULT_DNN_QUEUE_SIZE); 46 | nh.param("obj_det_threshold", obj_det_threshold_, 0.15f); 47 | nh.param("iou_threshold", iou_threshold_, 0.2f); 48 | nh.param("max_rate_hz", max_rate_hz_, 30.0f); 49 | nh.param("debug_mode", debug_mode_, false); 50 | nh.param("use_cached_model", use_cached_model, true); 51 | 52 | ROS_INFO("Camera: %s", camera_topic.c_str()); 53 | ROS_INFO("Proto : %s", prototxt_path.c_str()); 54 | ROS_INFO("Model : %s", model_path.c_str()); 55 | ROS_INFO("Input : %s", input_layer.c_str()); 56 | ROS_INFO("Output: %s", output_layer.c_str()); 57 | ROS_INFO("In Fmt: %s", inp_fmt.c_str()); 58 | ROS_INFO("FP16 : %s", use_FP16 ? "yes" : "no"); 59 | ROS_INFO("Scale : %.4f", inp_scale); 60 | ROS_INFO("Shift : %.2f", inp_shift); 61 | ROS_INFO("Cam Q : %d", camera_queue_size); 62 | ROS_INFO("DNN Q : %d", dnn_queue_size); 63 | ROS_INFO("Post P: %s", post_proc.c_str()); 64 | ROS_INFO("Obj T : %.2f", obj_det_threshold_); 65 | ROS_INFO("IOU T : %.2f", iou_threshold_); 66 | ROS_INFO("Rate : %.1f", max_rate_hz_); 67 | ROS_INFO("Debug : %s", debug_mode_ ? "yes" : "no"); 68 | 69 | setPostProcessing(post_proc); 70 | 71 | net_.loadNetwork(prototxt_path, model_path, input_layer, output_layer, use_FP16, use_cached_model); 72 | net_.setInputFormat(inp_fmt); 73 | net_.setScale(inp_scale); 74 | net_.setShift(inp_shift); 75 | if (debug_mode_) 76 | net_.showProfile(true); 77 | 78 | image_sub_ = nh.subscribe(camera_topic, camera_queue_size, &CaffeRos::imageCallback, this); 79 | output_pub_ = nh.advertise("network/output", dnn_queue_size); 80 | } 81 | 82 | void CaffeRos::spin() 83 | { 84 | // We currently use single-threaded version of spinner as we have only one topic 85 | // that we subscribe to: image feed from camera. 86 | // DNN needs only the most recent image and GPU can process only one batch at a time 87 | // so there is no need for more than one thread or async processing. 88 | // The async code is provided anyway but commented out. 89 | 90 | // The code is based on tutorial: http://wiki.ros.org/roscpp/Overview/Callbacks%20and%20Spinning 91 | // ros::AsyncSpinner spinner(4); // Use 4 threads 92 | // spinner.start(); 93 | // ros::waitForShutdown(); 94 | // ROS_INFO("Caffe ROS node is stopped."); 95 | 96 | ros::Rate rate(max_rate_hz_); 97 | ros::spinOnce(); 98 | while (ros::ok()) 99 | { 100 | auto out_msg = computeOutputs(); 101 | if (out_msg != nullptr) 102 | output_pub_.publish(out_msg); 103 | ros::spinOnce(); 104 | rate.sleep(); 105 | } 106 | } 107 | 108 | sensor_msgs::Image::ConstPtr CaffeRos::computeOutputs() 109 | { 110 | if (cur_img_ == nullptr) 111 | return nullptr; 112 | 113 | auto img = *cur_img_; 114 | net_.forward(img.data.data(), img.width, img.height, 3); 115 | auto out_msg = boost::make_shared(); 116 | // Set stamp to the same value as source image so we can synchronize with other nodes if needed. 117 | out_msg->header.stamp.sec = img.header.stamp.sec; 118 | out_msg->header.stamp.nsec = img.header.stamp.nsec; 119 | 120 | // Use single precision multidimensional array to represent outputs. 121 | // This can be useful in case DNN output is multidimensional such as in segmentation networks. 122 | // Note that encoding may not be compatible with other ROS code that uses Image in case number of channels > 4. 123 | // For classification nets, TensorRT uses 'c' dimension to reperesent number of classes. 124 | if (post_proc_ == PostProc::None) 125 | { 126 | out_msg->encoding = "32FC" + std::to_string(net_.getOutChannels()); 127 | out_msg->width = net_.getOutWidth(); 128 | out_msg->height = net_.getOutHeight(); 129 | out_msg->step = out_msg->width * net_.getOutChannels() * sizeof(float); 130 | size_t count = out_msg->step * out_msg->height; 131 | auto ptr = reinterpret_cast(net_.getOutput()); 132 | out_msg->data = std::vector(ptr, ptr + count); 133 | } 134 | else if (post_proc_ == PostProc::Segmentation) 135 | { 136 | cv::Mat class_map(net_.getOutHeight(), net_.getOutWidth(), CV_8UC1); 137 | getArgMax(net_.getOutput(), net_.getOutChannels(), class_map); 138 | 139 | out_msg->encoding = "mono8"; 140 | out_msg->width = class_map.cols; 141 | out_msg->height = class_map.rows; 142 | out_msg->step = class_map.cols; 143 | size_t count = out_msg->step * out_msg->height; 144 | const unsigned char *ptr = class_map.data; 145 | out_msg->data = std::vector(ptr, ptr + count); 146 | } 147 | else if (post_proc_ == PostProc::YOLO) 148 | { 149 | // Width and height are 1 in case of YOLO. 150 | ROS_ASSERT(net_.getOutWidth() == 1); 151 | ROS_ASSERT(net_.getOutHeight() == 1); 152 | // Get bounding boxes and apply IOU filter. 153 | // REVIEW alexeyk: move magic constants to node arguments. 154 | auto preds = getYoloPredictions(net_.getOutput(), net_.getOutChannels(), img.width, img.height, obj_det_threshold_); 155 | preds = filterByIOU(preds, iou_threshold_); 156 | // Copy results to float array. Label and coords will be converted to float. 157 | const int num_col = 6; 158 | std::vector msg_data(preds.size() * num_col); 159 | size_t i = 0; 160 | for (size_t row = 0; row < preds.size(); row++, i += num_col) 161 | { 162 | msg_data[i] = (float)preds[row].label; 163 | msg_data[i + 1] = preds[row].prob; 164 | msg_data[i + 2] = (float)preds[row].x; 165 | msg_data[i + 3] = (float)preds[row].y; 166 | msg_data[i + 4] = (float)preds[row].w; 167 | msg_data[i + 5] = (float)preds[row].h; 168 | } 169 | ROS_ASSERT(i == msg_data.size()); 170 | // Create message. 171 | // YOLO output is represented as a matrix where each row is 172 | // a predicted object vector of size 6: label, prob and 4 bounding box coordinates. 173 | out_msg->encoding = "32FC1"; 174 | out_msg->width = num_col; 175 | out_msg->height = preds.size(); 176 | out_msg->step = out_msg->width * sizeof(float); 177 | size_t count = out_msg->step * out_msg->height; 178 | auto ptr = reinterpret_cast(msg_data.data()); 179 | out_msg->data = std::vector(ptr, ptr + count); 180 | ROS_ASSERT(msg_data.size() * sizeof(float) == count); 181 | } 182 | else 183 | { 184 | // Should not happen, yeah... 185 | ROS_FATAL("Invalid post processing."); 186 | ros::shutdown(); 187 | } 188 | 189 | // Set to null to mark as completed. 190 | cur_img_ = nullptr; 191 | 192 | return out_msg; 193 | } 194 | 195 | void CaffeRos::imageCallback(const sensor_msgs::Image::ConstPtr& msg) 196 | { 197 | auto img = *msg; 198 | //ROS_DEBUG("imageCallback: %u, %u, %s", img.width, img.height, img.encoding.c_str()); 199 | // Only RGB8 is currently supported. 200 | if (img.encoding != "rgb8") 201 | { 202 | ROS_FATAL("Image encoding %s is not yet supported. Supported encodings: rgb8", img.encoding.c_str()); 203 | ros::shutdown(); 204 | } 205 | cur_img_ = msg; 206 | } 207 | 208 | } 209 | -------------------------------------------------------------------------------- /caffe_ros/src/tensor_net.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include "caffe_ros/tensor_net.h" 12 | 13 | namespace caffe_ros 14 | { 15 | 16 | TensorNet::Logger TensorNet::s_log; 17 | TensorNet::Profiler TensorNet::s_profiler; 18 | 19 | static nvinfer1::DimsCHW DimsToCHW(nvinfer1::Dims dims) 20 | { 21 | ROS_ASSERT(dims.nbDims == 3); 22 | ROS_ASSERT(dims.type[0] == nvinfer1::DimensionType::kCHANNEL); 23 | ROS_ASSERT(dims.type[1] == nvinfer1::DimensionType::kSPATIAL); 24 | ROS_ASSERT(dims.type[2] == nvinfer1::DimensionType::kSPATIAL); 25 | return nvinfer1::DimsCHW(dims.d[0], dims.d[1], dims.d[2]); 26 | } 27 | 28 | void TensorNet::Logger::log(Severity severity, const char *msg) 29 | { 30 | if (severity != Severity::kINFO) 31 | ROS_INFO("[TensorRT] %s", msg); 32 | } 33 | 34 | void TensorNet::Profiler::reportLayerTime(const char *layerName, float ms) 35 | { 36 | auto record = std::find_if(profile_.begin(), profile_.end(), [&](const Record &r) { return r.first == layerName; }); 37 | if (record == profile_.end()) 38 | profile_.push_back(std::make_pair(layerName, ms)); 39 | else 40 | record->second = ms; 41 | } 42 | 43 | void TensorNet::Profiler::printLayerTimes() 44 | { 45 | float total_time = 0; 46 | for (size_t i = 0; i < profile_.size(); i++) 47 | { 48 | //ROS_INFO("%-40.40s %4.3fms", profile_[i].first.c_str(), profile_[i].second); 49 | total_time += profile_[i].second; 50 | } 51 | ROS_INFO("All layers : %4.3f", total_time); 52 | } 53 | 54 | TensorNet::TensorNet() 55 | { 56 | } 57 | 58 | TensorNet::~TensorNet() 59 | { 60 | if (out_h_ != nullptr) 61 | { 62 | cudaError_t err = cudaFreeHost(out_h_); 63 | if (err != cudaSuccess) 64 | ROS_WARN("cudaFreeHost returned %d", (int)err); 65 | out_h_ = nullptr; 66 | out_d_ = nullptr; 67 | } 68 | } 69 | 70 | void TensorNet::profileModel(ConstStr& prototxt_path, ConstStr& model_path, bool use_FP16, 71 | ConstStr& output_blob, std::ostream& model) 72 | { 73 | auto builder = nvinfer1::createInferBuilder(s_log); 74 | auto network = builder->createNetwork(); 75 | 76 | builder->setMinFindIterations(3); // allow time for TX1 GPU to spin up 77 | builder->setAverageFindIterations(2); 78 | 79 | auto parser = nvcaffeparser1::createCaffeParser(); 80 | bool has_fast_FP16 = builder->platformHasFastFp16(); 81 | ROS_INFO("Hardware support of fast FP16: %s.", has_fast_FP16 ? "yes" : "no"); 82 | if (has_fast_FP16 && !use_FP16) 83 | ROS_INFO("... however, the model will be loaded as FP32."); 84 | 85 | nvinfer1::DataType model_data_type = (has_fast_FP16 && use_FP16) ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT; 86 | auto blob_finder = parser->parse(prototxt_path.c_str(), model_path.c_str(), *network, model_data_type); 87 | if (blob_finder == nullptr) 88 | { 89 | ROS_FATAL("Failed to parse network: %s, %s", prototxt_path.c_str(), model_path.c_str()); 90 | ros::shutdown(); 91 | } 92 | ROS_INFO("Loaded model from: %s, %s", prototxt_path.c_str(), model_path.c_str()); 93 | 94 | auto out_b = blob_finder->find(output_blob.c_str()); 95 | if (out_b == nullptr) 96 | { 97 | ROS_FATAL("Could not find output blob: %s", output_blob.c_str()); 98 | ros::shutdown(); 99 | } 100 | network->markOutput(*out_b); 101 | 102 | // Build model. 103 | // REVIEW alexeyk: make configurable? 104 | // Note: FP16 requires batch size to be even, TensorRT will switch automatically when building an engine. 105 | builder->setMaxBatchSize(1); 106 | builder->setMaxWorkspaceSize(16 * 1024 * 1024); 107 | 108 | builder->setHalf2Mode(has_fast_FP16 && use_FP16); 109 | 110 | ROS_INFO("Building CUDA engine..."); 111 | auto engine = builder->buildCudaEngine(*network); 112 | if (engine == nullptr) 113 | { 114 | ROS_FATAL("Failed to build CUDA engine."); 115 | ros::shutdown(); 116 | } 117 | 118 | // Save model. 119 | nvinfer1::IHostMemory* model_ptr = engine->serialize(); 120 | ROS_ASSERT(model_ptr != nullptr); 121 | model.write(reinterpret_cast(model_ptr->data()), model_ptr->size()); 122 | model_ptr->destroy(); 123 | 124 | ROS_INFO("Done building."); 125 | 126 | // Cleanup. 127 | network->destroy(); 128 | parser->destroy(); 129 | engine->destroy(); 130 | builder->destroy(); 131 | } 132 | 133 | void TensorNet::loadNetwork(ConstStr& prototxt_path, ConstStr& model_path, 134 | ConstStr& input_blob, ConstStr& output_blob, 135 | bool use_FP16, bool use_cached_model) 136 | { 137 | infer_ = nvinfer1::createInferRuntime(s_log); 138 | if (infer_ == nullptr) 139 | { 140 | ROS_FATAL("Failed to create inference runtime."); 141 | ros::shutdown(); 142 | } 143 | 144 | std::stringstream model; 145 | if (!use_cached_model) 146 | profileModel(prototxt_path, model_path, use_FP16, output_blob, model); 147 | else 148 | { 149 | std::string cached_model_path(model_path + ".cache"); 150 | std::ifstream cached_model(cached_model_path, std::ios::binary); 151 | 152 | if (cached_model.good()) 153 | { 154 | ROS_INFO("Loading cached model from: %s", cached_model_path.c_str()); 155 | model << cached_model.rdbuf(); 156 | } 157 | else 158 | { 159 | profileModel(prototxt_path, model_path, use_FP16, output_blob, model); 160 | ROS_INFO("Saving cached model to: %s", cached_model_path.c_str()); 161 | std::ofstream cacheFile(cached_model_path, std::ios::binary); 162 | cacheFile << model.rdbuf(); 163 | } 164 | } 165 | 166 | model.seekg(0, model.beg); 167 | const auto& model_final = model.str(); 168 | 169 | engine_ = infer_->deserializeCudaEngine(model_final.c_str(), model_final.size(), nullptr); 170 | if (engine_ == nullptr) 171 | { 172 | ROS_FATAL("Failed to deserialize engine."); 173 | ros::shutdown(); 174 | } 175 | 176 | context_ = engine_->createExecutionContext(); 177 | if (context_ == nullptr) 178 | { 179 | ROS_FATAL("Failed to create execution context."); 180 | ros::shutdown(); 181 | } 182 | ROS_INFO("Created CUDA engine and context."); 183 | 184 | int iinp = engine_->getBindingIndex(input_blob.c_str()); 185 | in_dims_ = DimsToCHW(engine_->getBindingDimensions(iinp)); 186 | ROS_INFO("Input : (W:%4u, H:%4u, C:%4u).", in_dims_.w(), in_dims_.h(), in_dims_.c()); 187 | //cv::gpu::ensureSizeIsEnough(in_dims_.h(), in_dims_.w(), CV_8UC3, in_d_); 188 | in_d_ = cv::gpu::createContinuous(in_dims_.c(), in_dims_.w() * in_dims_.h(), CV_32FC1); 189 | assert(in_d_.isContinuous()); 190 | 191 | int iout = engine_->getBindingIndex(output_blob.c_str()); 192 | out_dims_ = DimsToCHW(engine_->getBindingDimensions(iout)); 193 | ROS_INFO("Output: (W:%4u, H:%4u, C:%4u).", out_dims_.w(), out_dims_.h(), out_dims_.c()); 194 | 195 | // Allocate mapped memory for the outputs. 196 | size_t outSizeBytes = out_dims_.w() * out_dims_.h() * out_dims_.c() * sizeof(float); 197 | if (cudaHostAlloc(&out_h_, outSizeBytes, cudaHostAllocMapped) != cudaSuccess) 198 | { 199 | ROS_FATAL("Could not allocate %zu bytes for the output, error: %u.", outSizeBytes, cudaGetLastError()); 200 | ros::shutdown(); 201 | } 202 | if (cudaHostGetDevicePointer(&out_d_, out_h_, 0) != cudaSuccess) 203 | { 204 | ROS_FATAL("Could not get device pointer for the output, error: %u.", cudaGetLastError()); 205 | ros::shutdown(); 206 | } 207 | } 208 | 209 | void TensorNet::forward(const unsigned char* input, size_t w, size_t h, size_t c) 210 | { 211 | ROS_ASSERT(c == (size_t)in_dims_.c()); 212 | //ROS_DEBUG("Forward: input image is (%zu, %zu, %zu), network input is (%u, %u, %u)", w, h, c, in_dims_.w(), in_dims_.h(), in_dims_.c()); 213 | 214 | // REVIEW alexeyk: extract to a separate methog/transformer class. 215 | // Perform image pre-processing (scaling, conversion etc). 216 | 217 | ros::Time start = ros::Time::now(); 218 | 219 | in_h_ = cv::Mat((int)h, (int)w, CV_8UC3, (void*)input); 220 | // Convert image from RGB to BGR format used by OpenCV if needed. 221 | if (inp_fmt_ == InputFormat::BGR) 222 | cv::cvtColor(in_h_, in_h_, CV_RGB2BGR); 223 | //ROS_INFO("Dims: (%zu, %zu) -> (%zu, %zu)", w, h, (size_t)in_dims_.w(), (size_t)in_dims_.h()); 224 | // Convert to floating point type. 225 | in_h_.convertTo(in_h_, CV_32F); 226 | // Resize (anisotropically) to input layer size. 227 | cv::resize(in_h_, in_h_, cv::Size(in_dims_.w(), in_dims_.h()), 0, 0, cv::INTER_CUBIC); 228 | // Scale if needed. 229 | if (inp_scale_ != 1) 230 | in_h_ *= inp_scale_; 231 | // Shift if needed. 232 | if (inp_shift_ != 0) 233 | in_h_ = in_h_ + inp_shift_; 234 | // Transpose to get CHW format. 235 | in_final_h_ = in_h_.reshape(1, in_dims_.w() * in_dims_.h()).t(); 236 | // Copy to the device. 237 | ROS_ASSERT(in_final_h_.isContinuous()); 238 | ROS_ASSERT(in_d_.isContinuous()); 239 | ROS_ASSERT(in_d_.size().area() * in_d_.channels() == in_final_h_.size().area() * in_final_h_.channels()); 240 | if (cudaMemcpy(in_d_.ptr(0), in_final_h_.ptr(0), 241 | in_d_.size().area() * in_d_.channels() * sizeof(float), cudaMemcpyHostToDevice) != cudaSuccess) 242 | { 243 | ROS_FATAL("Could not copy data to device, error: %u.", cudaGetLastError()); 244 | ros::shutdown(); 245 | } 246 | 247 | if (debug_mode_) 248 | ROS_INFO("Preproc time: %.3f", (ros::Time::now() - start).toSec() * 1000); 249 | 250 | // GPU version: 251 | // in_orig_d_.upload(cv::Mat((int)h, (int)w, CV_8UC3, (void*)input)); 252 | // // Resize to input layer size. 253 | // cv::gpu::resize(in_orig_d_, in_d_, cv::Size(in_dims_.w(), in_dims_.h()), 0, 0, cv::INTER_CUBIC); 254 | // // Convert to floating point type. 255 | // in_d_.convertTo(in_f_d_, CV_32FC3); 256 | // // Subtract shift. 257 | // // REVIEW alexeyk: should be configurable as some models already have this in prototxt. 258 | // cv::gpu::subtract(in_f_d_, 128.0f, in_f_d_); 259 | // // Transpose to get CHW format. 260 | // ROS_DEBUG("in_f_d_: %zu", in_f_d_.elemSize()); 261 | // cv::gpu::transpose(in_f_d_, in_f_d_); 262 | 263 | // cv::Mat cpuM; 264 | // in_d_.download(cpuM); 265 | // std::ofstream file("/home/ubuntu/tmp.raw", std::ios::binary); 266 | // file.write(cpuM.ptr(0), std::streamsize(in_dims_.w() * in_dims_.h() * c)); 267 | // file.close(); 268 | 269 | void* bufs[] = {in_d_.ptr(0), out_d_}; 270 | context_->execute(1, bufs); 271 | if (debug_mode_) 272 | s_profiler.printLayerTimes(); 273 | ROS_DEBUG("Forward out (first 3 values): [%.4f, %.4f, %.4f]", out_h_[0], out_h_[1], out_h_[2]); 274 | } 275 | 276 | } 277 | 278 | -------------------------------------------------------------------------------- /caffe_ros/tests/tests.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2017, NVIDIA CORPORATION. All rights reserved. 2 | // Full license terms provided in LICENSE.md file. 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | namespace fs = boost::filesystem; 12 | 13 | // REVIEW alexeyk: refactor to functor or something prettier? 14 | class CaffeRosTestsCallback 15 | { 16 | public: 17 | CaffeRosTestsCallback(): dnn_out_(nullptr) 18 | { 19 | } 20 | 21 | void dnnCallback(const sensor_msgs::Image::ConstPtr &msg) 22 | { 23 | dnn_out_ = msg; 24 | } 25 | 26 | sensor_msgs::Image::ConstPtr dnn_out_; 27 | }; 28 | 29 | static boost::shared_ptr readImage(const std::string& filename) 30 | { 31 | auto img = cv::imread(filename); 32 | SCOPED_TRACE(filename); 33 | EXPECT_TRUE(img.cols > 0 && img.rows > 0); 34 | // Convert image from BGR format used by OpenCV to RGB. 35 | cv::cvtColor(img, img, CV_BGR2RGB); 36 | auto img_msg = boost::make_shared(); 37 | img_msg->encoding = "rgb8"; 38 | img_msg->width = img.cols; 39 | img_msg->height = img.rows; 40 | img_msg->step = img_msg->width * img.channels(); 41 | auto ptr = img.ptr(0); 42 | img_msg->data = std::vector(ptr, ptr + img_msg->step * img_msg->height); 43 | return img_msg; 44 | } 45 | 46 | TEST(CaffeRosTests, TrailNetPredictions) 47 | { 48 | ros::NodeHandle nh("~"); 49 | std::string test_data_dir; 50 | nh.param("test_data_dir", test_data_dir, ""); 51 | ASSERT_TRUE(fs::exists(test_data_dir)); 52 | 53 | CaffeRosTestsCallback t; 54 | auto dnn_sub = nh.subscribe("/trails_dnn/network/output", 1, 55 | &CaffeRosTestsCallback::dnnCallback, &t); 56 | const char* camera_topic = "/camera_trails/image_raw"; 57 | auto img_pub = nh.advertise(camera_topic, 1); 58 | 59 | // Test images and expected predictions. 60 | auto images = std::vector{"rot_l.jpg", "rot_c.jpg", "rot_r.jpg", "tran_l.jpg", "tran_r.jpg"}; 61 | float predictions[][6] = {{0.932, 0.060, 0.006, 0.080, 0.848, 0.071}, 62 | {0.040, 0.958, 0.001, 0.488, 0.375, 0.135}, 63 | {0.000, 0.027, 0.971, 0.036, 0.407, 0.555}, 64 | {0.011, 0.988, 0.000, 0.981, 0.008, 0.009}, 65 | {0.000, 0.855, 0.144, 0.013, 0.031, 0.954}}; 66 | 67 | // When running using rostest, current directory is $HOME/.ros 68 | fs::path data_dir{test_data_dir}; 69 | 70 | for (size_t i = 0; i < images.size(); i++) 71 | { 72 | auto img_msg = readImage((data_dir / images[i]).string()); 73 | // Use image index as a unique timestamp. 74 | img_msg->header.stamp.sec = 0; 75 | img_msg->header.stamp.nsec = (int)i; 76 | 77 | ros::Rate rate(1000); 78 | // Wait until DNN processes the current messages. There might be multiple messages 79 | // in the queue so make sure to select the right one based on current index. 80 | while (ros::ok() && (t.dnn_out_ == nullptr || t.dnn_out_->header.stamp.nsec != i)) 81 | { 82 | img_pub.publish(img_msg); 83 | ros::spinOnce(); 84 | rate.sleep(); 85 | } 86 | 87 | EXPECT_TRUE(t.dnn_out_ != nullptr); 88 | auto dnn_out = *t.dnn_out_; 89 | // The output should be 1x1x6 (HxWxC). 90 | EXPECT_EQ(dnn_out.width, 1); 91 | EXPECT_EQ(dnn_out.height, 1); 92 | // float32, channels == 6. 93 | EXPECT_EQ(dnn_out.encoding, "32FC6"); 94 | 95 | auto data = reinterpret_cast(dnn_out.data.data()); 96 | for (int col = 0; col < 6; col++) 97 | { 98 | // Must use proper floating point comparison. 99 | EXPECT_NEAR(data[col], predictions[i][col], 0.001f) << "Values are not equal at (" << i << ", " << col <<")"; 100 | } 101 | } 102 | } 103 | 104 | TEST(CaffeRosTests, TrailNetPredictionsFP16) 105 | { 106 | ros::NodeHandle nh("~"); 107 | std::string test_data_dir; 108 | nh.param("test_data_dir", test_data_dir, ""); 109 | ASSERT_TRUE(fs::exists(test_data_dir)); 110 | 111 | CaffeRosTestsCallback t; 112 | auto dnn_sub = nh.subscribe("/trails_dnn_fp16/network/output", 1, 113 | &CaffeRosTestsCallback::dnnCallback, &t); 114 | const char* camera_topic = "/camera_trails_fp16/image_raw"; 115 | auto img_pub = nh.advertise(camera_topic, 1); 116 | 117 | // Test images and expected predictions. 118 | auto images = std::vector{"rot_l.jpg", "rot_c.jpg", "rot_r.jpg", "tran_l.jpg", "tran_r.jpg"}; 119 | float predictions[][6] = {{0.932, 0.060, 0.006, 0.080, 0.848, 0.071}, 120 | {0.040, 0.958, 0.001, 0.488, 0.375, 0.135}, 121 | {0.000, 0.027, 0.971, 0.036, 0.407, 0.555}, 122 | {0.011, 0.988, 0.000, 0.981, 0.008, 0.009}, 123 | {0.000, 0.855, 0.144, 0.013, 0.031, 0.954}}; 124 | 125 | // When running using rostest, current directory is $HOME/.ros 126 | fs::path data_dir{test_data_dir}; 127 | 128 | for (size_t i = 0; i < images.size(); i++) 129 | { 130 | auto img_msg = readImage((data_dir / images[i]).string()); 131 | // Use image index as a unique timestamp. 132 | img_msg->header.stamp.sec = 0; 133 | img_msg->header.stamp.nsec = (int)i; 134 | 135 | ros::Rate rate(1000); 136 | // Wait until DNN processes the current messages. There might be multiple messages 137 | // in the queue so make sure to select the right one based on current index. 138 | while (ros::ok() && (t.dnn_out_ == nullptr || t.dnn_out_->header.stamp.nsec != i)) 139 | { 140 | img_pub.publish(img_msg); 141 | ros::spinOnce(); 142 | rate.sleep(); 143 | } 144 | 145 | EXPECT_TRUE(t.dnn_out_ != nullptr); 146 | auto dnn_out = *t.dnn_out_; 147 | // The output should be 1x1x6 (HxWxC). 148 | EXPECT_EQ(dnn_out.width, 1); 149 | EXPECT_EQ(dnn_out.height, 1); 150 | // float32, channels == 6. 151 | EXPECT_EQ(dnn_out.encoding, "32FC6"); 152 | 153 | auto data = reinterpret_cast(dnn_out.data.data()); 154 | for (int col = 0; col < 6; col++) 155 | { 156 | // Must use proper floating point comparison. 157 | // For FP16 the tolerance is higher than for FP32. 158 | EXPECT_NEAR(data[col], predictions[i][col], 0.02f) << "Values are not equal at (" << i << ", " << col <<")"; 159 | } 160 | } 161 | } 162 | 163 | TEST(CaffeRosTests, YoloNetPredictions) 164 | { 165 | ros::NodeHandle nh("~"); 166 | std::string test_data_dir; 167 | nh.param("test_data_dir", test_data_dir, ""); 168 | ASSERT_TRUE(fs::exists(test_data_dir)); 169 | 170 | CaffeRosTestsCallback t; 171 | auto dnn_sub = nh.subscribe("/object_dnn/network/output", 1, 172 | &CaffeRosTestsCallback::dnnCallback, &t); 173 | const char* camera_topic = "/camera_object/image_raw"; 174 | auto img_pub = nh.advertise(camera_topic, 1); 175 | 176 | // Test images and expected predictions. 177 | std::string image{"yolo_2_obj.png"}; 178 | float predictions[][6] = {{14, 0.290, 184, 128, 72, 158}, 179 | {14, 0.660, 529, 84, 105, 239}}; 180 | 181 | // When running using rostest, current directory is $HOME/.ros 182 | fs::path data_dir{test_data_dir}; 183 | 184 | auto img_msg = readImage((data_dir / image).string()); 185 | 186 | ros::Rate rate(1000); 187 | // Wait until DNN processes the current messages. There might be multiple messages 188 | // in the queue so take the first one. 189 | while (ros::ok() && t.dnn_out_ == nullptr) 190 | { 191 | img_pub.publish(img_msg); 192 | ros::spinOnce(); 193 | rate.sleep(); 194 | } 195 | 196 | EXPECT_TRUE(t.dnn_out_ != nullptr); 197 | auto dnn_out = *t.dnn_out_; 198 | // YOLO output is matrix of n x 6. 199 | EXPECT_EQ(dnn_out.width, 6); 200 | // There should be 2 objects detected in the test image. 201 | EXPECT_EQ(dnn_out.height, 2); 202 | // Single channel, float32. 203 | EXPECT_EQ(dnn_out.encoding, "32FC1"); 204 | 205 | auto data = reinterpret_cast(dnn_out.data.data()); 206 | for (size_t row = 0; row < dnn_out.height; row++) 207 | { 208 | for (size_t col = 0; col < 6; col++) 209 | { 210 | // Must use proper floating point comparison. 211 | EXPECT_NEAR(data[row * 6 + col], predictions[row][col], 0.001f) << "Values are not equal at (" << row << ", " << col <<")"; 212 | } 213 | } 214 | } 215 | 216 | TEST(CaffeRosTests, YoloNetPredictionsFP16) 217 | { 218 | ros::NodeHandle nh("~"); 219 | std::string test_data_dir; 220 | nh.param("test_data_dir", test_data_dir, ""); 221 | ASSERT_TRUE(fs::exists(test_data_dir)); 222 | 223 | CaffeRosTestsCallback t; 224 | auto dnn_sub = nh.subscribe("/object_dnn_fp16/network/output", 1, 225 | &CaffeRosTestsCallback::dnnCallback, &t); 226 | const char* camera_topic = "/camera_object_fp16/image_raw"; 227 | auto img_pub = nh.advertise(camera_topic, 1); 228 | 229 | // Test images and expected predictions. 230 | std::string image{"yolo_2_obj.png"}; 231 | float predictions[][6] = {{14, 0.290, 184, 128, 72, 158}, 232 | {14, 0.660, 529, 84, 105, 239}}; 233 | 234 | // When running using rostest, current directory is $HOME/.ros 235 | fs::path data_dir{test_data_dir}; 236 | 237 | auto img_msg = readImage((data_dir / image).string()); 238 | 239 | ros::Rate rate(1000); 240 | // Wait until DNN processes the current messages. There might be multiple messages 241 | // in the queue so take the first one. 242 | while (ros::ok() && t.dnn_out_ == nullptr) 243 | { 244 | img_pub.publish(img_msg); 245 | ros::spinOnce(); 246 | rate.sleep(); 247 | } 248 | 249 | EXPECT_TRUE(t.dnn_out_ != nullptr); 250 | auto dnn_out = *t.dnn_out_; 251 | // YOLO output is matrix of n x 6. 252 | EXPECT_EQ(dnn_out.width, 6); 253 | // There should be 2 objects detected in the test image. 254 | EXPECT_EQ(dnn_out.height, 2); 255 | // Single channel, float32. 256 | EXPECT_EQ(dnn_out.encoding, "32FC1"); 257 | 258 | auto data = reinterpret_cast(dnn_out.data.data()); 259 | for (size_t row = 0; row < dnn_out.height; row++) 260 | { 261 | for (size_t col = 0; col < 6; col++) 262 | { 263 | // Must use proper floating point comparison. 264 | // For FP16 the tolerance is higher than for FP32 and will be different 265 | // for label and probability and coordinates. 266 | float tolerance = col <= 1 ? 0.02f : 1.0f; 267 | EXPECT_NEAR(data[row * 6 + col], predictions[row][col], tolerance) << "Values are not equal at (" << row << ", " << col <<")"; 268 | } 269 | } 270 | } 271 | /* 272 | TEST(CaffeRosTests, FCNAlexNetPredictions) 273 | { 274 | ros::NodeHandle nh("~"); 275 | std::string test_data_dir; 276 | nh.param("test_data_dir", test_data_dir, ""); 277 | ASSERT_TRUE(fs::exists(test_data_dir)); 278 | 279 | CaffeRosTestsCallback t; 280 | auto dnn_sub = nh.subscribe("/fcn_alexnet_dnn/network/output", 1, 281 | &CaffeRosTestsCallback::dnnCallback, &t); 282 | const char* camera_topic = "/camera_sidewalk/image_raw"; 283 | auto img_pub = nh.advertise(camera_topic, 1); 284 | 285 | // Test images and expected predictions. 286 | std::string image{"sidewalk.png"}; 287 | cv::Mat predictions = cv::imread("sidewalk-raw-out.png"); 288 | 289 | // When running using rostest, current directory is $HOME/.ros 290 | fs::path data_dir{test_data_dir}; 291 | 292 | auto img_msg = readImage((data_dir / image).string()); 293 | 294 | ros::Rate rate(1000); 295 | // Wait until DNN processes the current messages. There might be multiple messages 296 | // in the queue so take the first one. 297 | while (ros::ok() && t.dnn_out_ == nullptr) 298 | { 299 | img_pub.publish(img_msg); 300 | ros::spinOnce(); 301 | rate.sleep(); 302 | } 303 | 304 | EXPECT_TRUE(t.dnn_out_ != nullptr); 305 | auto dnn_out = *t.dnn_out_; 306 | // For this size input, the output will be 54 x 11. 307 | EXPECT_EQ(dnn_out.width, 54); 308 | EXPECT_EQ(dnn_out.height, 11); 309 | // Single channel, unsigned char. 310 | EXPECT_EQ(dnn_out.encoding, "mono8"); 311 | 312 | auto data = reinterpret_cast(dnn_out.data.data()); 313 | for (size_t y = 0; y < dnn_out.height; y++) 314 | { 315 | for (size_t x = 0; x < dnn_out.width; x++) 316 | { 317 | unsigned char expected = predictions.at(y, x); 318 | unsigned char actual = data[y * dnn_out.height + x]; 319 | 320 | EXPECT_EQ(expected, actual) << "Values are not equal at (" << x << ", " << y <<")"; 321 | } 322 | } 323 | } 324 | */ 325 | int main(int argc, char **argv) 326 | { 327 | ros::init(argc, argv, "CaffeRosTests"); 328 | testing::InitGoogleTest(&argc, argv); 329 | return RUN_ALL_TESTS(); 330 | } 331 | --------------------------------------------------------------------------------