[update] Add probability module

Author: ShigemichiMatsuzaki

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper.h

@@ -25,7 +25,8 @@ private :
   bool import_module(const std::string filename);
   void img2tensor(cv::Mat & img, at::Tensor & tensor, const bool use_gpu = true);
   void tensor2img(at::Tensor tensor, cv::Mat & img);
-  at::Tensor get_output(at::Tensor input_tensor);
+//  at::Tensor get_output(at::Tensor input_tensor);
+  std::tuple<at::Tensor, at::Tensor> get_output(at::Tensor input_tensor);
   at::Tensor get_argmax(at::Tensor input_tensor);
 };
 //}

launch/pytorch_enet_ros.launch

@@ -9,7 +9,7 @@
     <param name="colormap" value="$(find pytorch_enet_ros)/images/camvid12.png" />
     -->
     <remap from="~image" to="$(arg image)" />
-    <param name="model_file" value="$(find pytorch_enet_ros)/models/espdnet_ue_uest_trav.pt" />
+    <param name="model_file" value="$(find pytorch_enet_ros)/models/espdnet_ue_trav_20210115-151110.pt" />
     <!--
     <param name="model_file" value="$(find pytorch_enet_ros)/models/ENet_greenhouse.pt" />
     -->

src/pytorch_cpp_wrapper.cpp

@@ -5,6 +5,7 @@
  */
 
 
+#include <torch/torch.h>
 #include "pytorch_cpp_wrapper/pytorch_cpp_wrapper.h"
 #include <torch/script.h> // One-stop header.
 #include <torch/data/transforms/tensor.h> // One-stop header.
@@ -52,7 +53,7 @@ PyTorchCppWrapper::import_module(const std::string filename)
     return true;
   }
   catch (const c10::Error& e) {
-    std::cerr << "error loading the model\n";
+    std::cerr << e.what();
     return false;
   }
 }
@@ -78,16 +79,18 @@ PyTorchCppWrapper::img2tensor(cv::Mat & img, at::Tensor & tensor, const bool use
 void
 PyTorchCppWrapper::tensor2img(at::Tensor tensor, cv::Mat & img)
 {
-  std::cout << tensor.sizes() << std::endl;
   // Get the size of the input image
   int height = tensor.sizes()[0];
   int width  = tensor.sizes()[1];
 
+  tensor = tensor.to(torch::kCPU);
+
   // Convert to OpenCV
   img = cv::Mat(height, width, CV_8U, tensor. template data<uint8_t>());
 }
 
-at::Tensor
+//at::Tensor
+std::tuple<at::Tensor, at::Tensor>
 PyTorchCppWrapper::get_output(at::Tensor input_tensor)
 {
   // Execute the model and turn its output into a tensor.
@@ -97,8 +100,17 @@ PyTorchCppWrapper::get_output(at::Tensor input_tensor)
 
   at::Tensor output1 = outputs->elements()[0].toTensor();
   at::Tensor output2 = outputs->elements()[1].toTensor();
+  at::Tensor prob = outputs->elements()[2].toTensor();
+
+  // Divide probability by c
+  prob = torch::sigmoid(prob) / 0.3;
+  // Limit the values in range [0, 1]
+  prob = at::clamp(prob, 0.0, 1.0);
+
+//  return output1 + 0.5 * output2;
+  at::Tensor segmentation = output1 + 0.5 * output2;
 
-  return output1 + 0.5 * output2;
+  return std::forward_as_tuple(segmentation, prob);
 }
 
 at::Tensor 

src/pytorch_enet_ros.cpp

@@ -18,7 +18,7 @@ PyTorchENetROS::PyTorchENetROS(ros::NodeHandle & nh)
   std::string filename;
   nh_.param<std::string>("model_file", filename, "");
   if(!pt_wrapper_.import_module(filename)) {
-    ROS_ERROR("Failed to import the  model file [%s]", filename.c_str());
+    ROS_ERROR("Failed to import the model file [%s]", filename.c_str());
     ros::shutdown();
   }
 
@@ -44,6 +44,7 @@ PyTorchENetROS::image_callback(const sensor_msgs::ImageConstPtr& msg)
   // Run inference
   sensor_msgs::ImagePtr label_msg;
   sensor_msgs::ImagePtr color_label_msg;
+  sensor_msgs::ImagePtr prob_msg;
   std::tie(label_msg, color_label_msg) = inference(cv_ptr->image);
 
   // Set header
@@ -69,6 +70,7 @@ PyTorchENetROS::image_inference_srv_callback(semantic_segmentation_srvs::GetLabe
   // Run inference
   sensor_msgs::ImagePtr label_msg;
   sensor_msgs::ImagePtr color_label_msg;
+  sensor_msgs::ImagePtr prob_msg;
   std::tie(label_msg, color_label_msg) = inference(cv_ptr->image);
 
   res.label_img = *label_msg;
@@ -103,27 +105,35 @@ PyTorchENetROS::inference(cv::Mat & input_img)
   for(int i = 0; i < mean_vec.size(); i++) {
     input_tensor[0][i] = (input_tensor[0][i] - mean_vec[i]) / std_vec[i];
   }
-  std::cout << input_tensor.sizes() << std::endl;
+//  std::cout << input_tensor.sizes() << std::endl;
 
   // Execute the model and turn its output into a tensor.
-  at::Tensor output = pt_wrapper_.get_output(input_tensor);
+  at::Tensor segmentation;
+  at::Tensor prob;
+  std::tie(segmentation, prob) = pt_wrapper_.get_output(input_tensor);
+//  at::Tensor output = pt_wrapper_.get_output(input_tensor);
   // Calculate argmax to get a label on each pixel
-  at::Tensor output_args = pt_wrapper_.get_argmax(output);
+//  at::Tensor output_args = pt_wrapper_.get_argmax(output);
+
+  at::Tensor output_args = pt_wrapper_.get_argmax(segmentation);
 
   // Convert to OpenCV
   cv::Mat label;
+  cv::Mat prob_cv;
   pt_wrapper_.tensor2img(output_args[0], label);
+  pt_wrapper_.tensor2img((prob[0][0]*255).to(torch::kByte), prob_cv);
 
   // Set the size
   cv::Size s_orig(width_orig, height_orig);
   // Resize the input image back to the original size
   cv::resize(label, label, s_orig, cv::INTER_NEAREST);
-
+  cv::resize(prob_cv, prob_cv, s_orig, cv::INTER_LINEAR);
   cv::Mat color_label;
   label_to_color(label, color_label);
 
   // Generate an image message
-  sensor_msgs::ImagePtr label_msg = cv_bridge::CvImage(std_msgs::Header(), "mono8", label).toImageMsg();
+//  sensor_msgs::ImagePtr label_msg = cv_bridge::CvImage(std_msgs::Header(), "mono8", label).toImageMsg();
+  sensor_msgs::ImagePtr label_msg = cv_bridge::CvImage(std_msgs::Header(), "mono8", prob_cv).toImageMsg();
   sensor_msgs::ImagePtr color_label_msg = cv_bridge::CvImage(std_msgs::Header(), "rgb8", color_label).toImageMsg();
 
   return std::forward_as_tuple(label_msg, color_label_msg);