Merge pull request NVIDIA-AI-IOT#95 from NVIDIA-Jetson/alexeyk/bgra

Added support for BGRA image topics (ZED). Fixes NVIDIA-AI-IOT#94.

ros/packages/caffe_ros/include/caffe_ros/tensor_net.h

@@ -25,7 +25,7 @@ class TensorNet
                      ConstStr& input_blob, ConstStr& output_blob,
                      nvinfer1::DataType data_type, bool use_cached_model);
 
-    void forward(const unsigned char* input, size_t w, size_t h, size_t c, const std::string& encoding);
+    void forward(const unsigned char* input, size_t w, size_t h, const std::string& encoding);
 
     int getInWidth() const    { return in_dims_.w(); }
     int getInHeight() const   { return in_dims_.h(); }

ros/packages/caffe_ros/src/caffe_ros.cpp

@@ -131,7 +131,7 @@ sensor_msgs::Image::ConstPtr CaffeRos::computeOutputs()
         return nullptr;
 
     auto img = *cur_img_;
-    net_.forward(img.data.data(), img.width, img.height, 3, img.encoding);
+    net_.forward(img.data.data(), img.width, img.height, img.encoding);
     auto out_msg = boost::make_shared<sensor_msgs::Image>();
     // Set stamp and frame id to the same value as source image so we can synchronize with other nodes if needed.
     out_msg->header.stamp.sec  = img.header.stamp.sec;
@@ -205,9 +205,9 @@ void CaffeRos::imageCallback(const sensor_msgs::Image::ConstPtr& msg)
     auto img = *msg;
     //ROS_DEBUG("imageCallback: %u, %u, %s", img.width, img.height, img.encoding.c_str());
     // Only RGB8 is currently supported.
-    if (img.encoding != "rgb8" && img.encoding != "bgr8")
+    if (img.encoding != "rgb8" && img.encoding != "bgr8" && img.encoding != "bgra8")
     {
-        ROS_FATAL("Image encoding %s is not yet supported. Supported encodings: rgb8, bgr8", img.encoding.c_str());
+        ROS_FATAL("Image encoding %s is not yet supported. Supported encodings: rgb8, bgr8, bgra8", img.encoding.c_str());
         ros::shutdown();
     }
     cur_img_ = msg;

ros/packages/caffe_ros/src/tensor_net.cpp

@@ -259,18 +259,17 @@ void TensorNet::loadNetwork(ConstStr& prototxt_path, ConstStr& model_path,
     }
 }
 
-void TensorNet::forward(const unsigned char* input, size_t w, size_t h, size_t c, const std::string& encoding)
+void TensorNet::forward(const unsigned char* input, size_t w, size_t h, const std::string& encoding)
 {
-    ROS_ASSERT(encoding == "rgb8" || encoding == "bgr8");
-    ROS_ASSERT(c == (size_t)in_dims_.c());
+    ROS_ASSERT(encoding == "rgb8" || encoding == "bgr8" || encoding == "bgra8");
     //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());
 
     // REVIEW alexeyk: extract to a separate methog/transformer class.
     // Perform image pre-processing (scaling, conversion etc).
 
     ros::Time start = ros::Time::now();
 
-    in_h_ = cv::Mat((int)h, (int)w, CV_8UC3, (void*)input);
+    in_h_ = cv::Mat((int)h, (int)w, encoding == "bgra8" ? CV_8UC4 : CV_8UC3, (void*)input);
     in_final_h_ = preprocessImage(in_h_, in_dims_.w(), in_dims_.h(), inp_fmt_, encoding, inp_scale_, inp_shift_);
     // Copy to the device.
     ROS_ASSERT(in_final_h_.isContinuous());
@@ -307,15 +306,19 @@ cv::Mat preprocessImage(cv::Mat img, int dst_img_w, int dst_img_h, InputFormat i
     // Handle encodings.
     if (inp_fmt == InputFormat::BGR)
     {
-        // Convert image from RGB to BGR format used by OpenCV if needed.
+        // Convert image from RGB/BGRA to BGR format.
         if (encoding == "rgb8")
             cv::cvtColor(img, img, CV_RGB2BGR);
+        else if (encoding == "bgra8")
+            cv::cvtColor(img, img, CV_BGRA2BGR);
     }
     else if (inp_fmt == InputFormat::RGB)
     {
         // Input image in OpenCV BGR, convert to RGB.
         if (encoding == "bgr8")
             cv::cvtColor(img, img, CV_BGR2RGB);
+        else if (encoding == "bgra8")
+            cv::cvtColor(img, img, CV_BGRA2RGB);
     }
     //ROS_INFO("Dims: (%zu, %zu) -> (%zu, %zu)", w, h, (size_t)dst_img_w, (size_t)dst_img_h);
     // Convert to floating point type.

ros/packages/caffe_ros/tests/data/trailnet_int8_calib.cache

@@ -1,54 +1,56 @@
 1
-(Unnamed ITensor* 12): 3c66f92e
-conv1: 3c04519f
-pool1: 3c016718
-res1_1_1: 3c010a14
-(Unnamed ITensor* 21): 3ce8219b
-(Unnamed ITensor* 23): 3d1474bf
-(Unnamed ITensor* 31): 3d8316b5
-res2_1_1: 3c01cfc8
-softmax_t: 3be25d0c
-(Unnamed ITensor* 56): 3e0c8ef5
-(Unnamed ITensor* 37): 3d76540e
-(Unnamed ITensor* 14): 3cb3e96f
-(Unnamed ITensor* 18): 3ca98b20
-res4_2_1: 3c023151
-res1_2: 3cbbbdf3
-(Unnamed ITensor* 80): 3dcdbb3f
-fc3_t: 3d8e6f0d
-(Unnamed ITensor* 33): 3d9aba8f
-res2_1: 3c01a251
-(Unnamed ITensor* 42): 3dc7db29
-(Unnamed ITensor* 50): 3e30779d
-(Unnamed ITensor* 59): 3e49a689
-sub_mean: 3b9e3ebb
-data: 400fabe8
-(Unnamed ITensor* 46): 3daa87d1
-res2_1_2: 3d2c4b58
-(Unnamed ITensor* 40): 3db93d03
-res3_1: 3d5a4a5b
-res2_2_1: 3c01ca05
-(Unnamed ITensor* 78): 3e4fb8a4
-(Unnamed ITensor* 61): 3e51b72a
-(Unnamed ITensor* 52): 3e368ce7
-res3_1_2: 3db36048
-res1_2_1: 3c497c69
-fc3: 3e91f7f1
-(Unnamed ITensor* 69): 3e1a721f
-res2_2: 3c016a8c
-res4_1_2: 3de4c1d9
-res1_1: 3c66f92f
-res3_2_1: 3d455a9e
-softmax: 3c00e5d4
-res4_1: 3d1a681c
-(Unnamed ITensor* 27): 3d3b70b6
-res3_2: 3d3b9f83
-res4_2: 3c023e26
-(Unnamed ITensor* 9): 3c7091cc
-(Unnamed ITensor* 75): 3e09c85e
-res3_1_1: 3c01b81e
-pool_avg: 3c159346
-(Unnamed ITensor* 65): 3e27c4e6
-(Unnamed ITensor* 4): 3c81c69e
-res4_1_1: 3d669c33
-(Unnamed ITensor* 71): 3e18ef86
+(Unnamed ITensor* 12): 3c99e2ac
+(Unnamed ITensor* 40): 3dca9489
+pool1: 3c296417
+res4_1_proj: 3df4e5c1
+res1_1_1: 3c015a58
+(Unnamed ITensor* 21): 3d032105
+(Unnamed ITensor* 23): 3d1f25de
+(Unnamed ITensor* 31): 3d892713
+res2_1_1: 3c019f21
+softmax_t: 3c02c000
+(Unnamed ITensor* 56): 3e100a0c
+(Unnamed ITensor* 37): 3d75caa3
+(Unnamed ITensor* 14): 3cd54ba3
+res1_2: 3cb3f14b
+(Unnamed ITensor* 80): 3e225f75
+fc3_t: 3cb15991
+(Unnamed ITensor* 4): 3c8ac3bf
+(Unnamed ITensor* 33): 3d98cd1b
+res2_1: 3c025baf
+(Unnamed ITensor* 42): 3dd8d7cc
+(Unnamed ITensor* 50): 3e38d6fe
+(Unnamed ITensor* 59): 3e57e962
+sub_mean: 3baad922
+data: 401ce88d
+(Unnamed ITensor* 46): 3db26938
+conv1: 3c01440f
+(Unnamed ITensor* 69): 3e2d0fe0
+res2_2: 3d2ce289
+res3_1: 3d21580a
+res2_2_1: 3c01bdb2
+(Unnamed ITensor* 61): 3e5edd05
+(Unnamed ITensor* 78): 3e818215
+(Unnamed ITensor* 52): 3e44452b
+res1_2_1: 3c6fbcaf
+fc3: 3e350887
+res3_1_proj: 3db4455f
+(Unnamed ITensor* 18): 3cd5a116
+res2_1_proj: 3cf6aa94
+res4_2_1: 3d668360
+res1_1: 3c963591
+(Unnamed ITensor* 9): 3c880658
+res4_2: 3d0ef404
+res3_2_1: 3d3ccadc
+softmax copy: 3c01eeb9
+softmax: 3c01eeb9
+res4_1: 3d1f8170
+(Unnamed ITensor* 27): 3d315523
+res3_2: 3d3da341
+res4_1_1: 3d658b6e
+softmax_t copy: 3c02c000
+(Unnamed ITensor* 71): 3e327e8e
+(Unnamed ITensor* 75): 3e0f1a0c
+(Unnamed ITensor* 65): 3e302791
+res3_1_1: 3d1f303c
+pool_avg: 3c01d9c7

ros/packages/caffe_ros/tests/tests.cpp

@@ -28,13 +28,15 @@ class CaffeRosTestsCallback
 
 static boost::shared_ptr<sensor_msgs::Image> readImage(const std::string& filename, const std::string& encoding = "rgb8")
 {
-    EXPECT_TRUE(encoding == "rgb8" || encoding == "bgr8");
+    EXPECT_TRUE(encoding == "rgb8" || encoding == "bgr8" || encoding == "bgra8");
     auto img = cv::imread(filename);
     SCOPED_TRACE(filename);
     EXPECT_TRUE(img.cols > 0 && img.rows > 0);
-    // Convert image from BGR format used by OpenCV to RGB.
+    // Convert image from BGR format used by OpenCV to RGB or BGRA.
     if (encoding == "rgb8")
         cv::cvtColor(img, img, CV_BGR2RGB);
+    else if (encoding == "bgra8")
+        cv::cvtColor(img, img, CV_BGR2BGRA);
     auto img_msg = boost::make_shared<sensor_msgs::Image>();
     img_msg->encoding = encoding;
     img_msg->width = img.cols;
@@ -161,6 +163,64 @@ TEST(CaffeRosTests, TrailNetPredictionsBGR8)
     }
 }
 
+TEST(CaffeRosTests, TrailNetPredictionsBGRA8)
+{
+    ros::NodeHandle nh("~");
+    std::string test_data_dir;
+    nh.param<std::string>("test_data_dir", test_data_dir, "");
+    ASSERT_TRUE(fs::exists(test_data_dir));
+
+    CaffeRosTestsCallback t;
+    auto dnn_sub = nh.subscribe<sensor_msgs::Image>("/trailnet/dnn/network/output", 1,
+                                                    &CaffeRosTestsCallback::dnnCallback, &t);
+    const char* camera_topic = "/trailnet/camera/image_raw";
+    auto img_pub = nh.advertise<sensor_msgs::Image>(camera_topic, 1);
+
+    // Test images and expected predictions.
+    auto images            = std::vector<std::string>{"rot_l.jpg", "rot_c.jpg", "rot_r.jpg", "tran_l.jpg", "tran_r.jpg"};
+    float predictions[][6] = {{0.932, 0.060, 0.006, 0.080, 0.848, 0.071},
+                              {0.040, 0.958, 0.001, 0.488, 0.375, 0.135},
+                              {0.000, 0.027, 0.971, 0.036, 0.407, 0.555},
+                              {0.011, 0.988, 0.000, 0.981, 0.008, 0.009},
+                              {0.000, 0.855, 0.144, 0.013, 0.031, 0.954}};
+
+    // When running using rostest, current directory is $HOME/.ros
+    fs::path data_dir{test_data_dir};
+
+    for (size_t i = 0; i < images.size(); i++)
+    {
+        auto img_msg = readImage((data_dir / images[i]).string(), "bgra8");
+        // Use image index as a unique timestamp.
+        img_msg->header.stamp.sec  = 0;
+        img_msg->header.stamp.nsec = (int)i;
+
+        ros::Rate rate(1000);
+        // Wait until DNN processes the current messages. There might be multiple messages
+        // in the queue so make sure to select the right one based on current index.
+        while (ros::ok() && (t.dnn_out_ == nullptr || t.dnn_out_->header.stamp.nsec != i))
+        {
+            img_pub.publish(img_msg);
+            ros::spinOnce();
+            rate.sleep();
+        }
+
+        EXPECT_TRUE(t.dnn_out_ != nullptr);
+        auto dnn_out = *t.dnn_out_;
+        // The output should be 1x1x6 (HxWxC).
+        EXPECT_EQ(dnn_out.width,  1);
+        EXPECT_EQ(dnn_out.height, 1);
+        // float32, channels == 6.
+        EXPECT_EQ(dnn_out.encoding, "32FC6");
+
+        auto data  = reinterpret_cast<const float*>(dnn_out.data.data());
+        for (int col = 0; col < 6; col++)
+        {
+            // Must use proper floating point comparison.
+            EXPECT_NEAR(data[col], predictions[i][col], 0.001f) << "Values are not equal at (" << i << ", " << col <<")";
+        }
+    }
+}
+
 TEST(CaffeRosTests, TrailNetPredictionsFP16)
 {
     ros::NodeHandle nh("~");