[update] Implement uncertainty calculation based on the entropy

Author: ShigemichiMatsuzaki

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h

@@ -15,11 +15,13 @@
 class PyTorchCppWrapperBase {
 protected :
   torch::jit::script::Module module_;
+  int class_num_;
+  float max_entropy_;
 
 public:
   PyTorchCppWrapperBase();
-  PyTorchCppWrapperBase(const std::string & filename);
-  PyTorchCppWrapperBase(const char* filename);
+  PyTorchCppWrapperBase(const std::string & filename, const int class_num);
+  PyTorchCppWrapperBase(const char* filename, const int class_num);
 
   /**
    * @brief import a network 
@@ -55,7 +57,7 @@ protected :
    * @param[in]  tensor
    * @param[out] tensor that has index of max value in each element
    */
-  at::Tensor get_entropy(at::Tensor input_tensor);
+  at::Tensor get_entropy(at::Tensor input_tensor, const bool normalize);
 
 };
 //}

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper_seg_trav_path.h

@@ -18,6 +18,9 @@ private :
   float c_{0.3};
 
 public:
+  PyTorchCppWrapperSegTravPath(const std::string & filename, const int class_num);
+  PyTorchCppWrapperSegTravPath(const char* filename, const int class_num);
+
   /**
    * @brief Get outputs from the model
    * @param[in] input_tensor Input tensor

include/pytorch_ros/pytorch_seg_trav_path_ros.h

@@ -36,11 +36,12 @@ class PyTorchSegTravPathROS {
   image_transport::Publisher  pub_label_image_;
   image_transport::Publisher  pub_color_image_;
   image_transport::Publisher  pub_prob_image_;
+  image_transport::Publisher  pub_uncertainty_image_;
   ros::Publisher pub_start_point_;
   ros::Publisher pub_end_point_;
   ros::Time stamp_of_current_image_;
 
-  PyTorchCppWrapperSegTravPath pt_wrapper_;
+  std::shared_ptr<PyTorchCppWrapperSegTravPath> pt_wrapper_ptr_;
 
   // Used to convert a label image to a color image
   cv::Mat colormap_;
@@ -59,7 +60,8 @@ class PyTorchSegTravPathROS {
    * @param[in] input_image OpenCV image 
    * @return    A tuple of messages of the inference results
    */
-  std::tuple<sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, geometry_msgs::PointStampedPtr, geometry_msgs::PointStampedPtr> inference(cv::Mat & input_image);
+  std::tuple<sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, geometry_msgs::PointStampedPtr, geometry_msgs::PointStampedPtr>
+    inference(cv::Mat & input_image);
 
   /** 
    * @brief Service callback
@@ -99,6 +101,13 @@ class PyTorchSegTravPathROS {
    * @return                  A tuple of start and end points as geometry_msgs::PointStampedPtr
    */
   std::tuple<geometry_msgs::PointStampedPtr, geometry_msgs::PointStampedPtr> tensor_to_points(const at::Tensor point_tensor, const int & width, const int & height);
+
+  /** 
+   * @brief Normalize a tensor to feed in a model
+   * @param[in]  input        Tensor
+   */
+  void normalize_tensor(at::Tensor & input_tensor);
+
 };
 
 #endif

src/pytorch_cpp_wrapper_base.cpp

@@ -4,7 +4,6 @@
  *
  */
 
-
 #include <torch/torch.h>
 #include "pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h"
 #include <torch/script.h> // One-stop header.
@@ -13,17 +12,38 @@
 #include <opencv2/opencv.hpp>
 #include "opencv2/highgui/highgui.hpp"
 #include <typeinfo>
+#include <cmath>
 
 PyTorchCppWrapperBase::PyTorchCppWrapperBase() {}
 
-PyTorchCppWrapperBase::PyTorchCppWrapperBase(const std::string & filename) {
+PyTorchCppWrapperBase::PyTorchCppWrapperBase(const std::string & filename, const int class_num)
+  : class_num_(class_num)
+{
   // Import model
   import_module(filename);
+
+  // Calculate the maximum possible entropy 
+  //  to normalize the entropy value in [0, 1].
+  max_entropy_ = 0;
+  const float prob = (float) 1.0 / class_num_;
+  for(int i = 0; i < class_num_; ++i) {
+    max_entropy_ += -prob * std::log(prob);
+  }
 }
 
-PyTorchCppWrapperBase::PyTorchCppWrapperBase(const char* filename) {
+PyTorchCppWrapperBase::PyTorchCppWrapperBase(const char* filename, const int class_num) 
+  : class_num_(class_num)
+{
   // Import model
   import_module(std::string(filename));
+
+  // Calculate the maximum possible entropy 
+  //  to normalize the entropy value in [0, 1].
+  max_entropy_ = 0;
+  const float prob = (float) 1.0 / class_num_;
+  for(int i = 0; i < class_num_; ++i) {
+    max_entropy_ += -prob * std::log(prob);
+  }
 }
 
 /**
@@ -101,7 +121,7 @@ at::Tensor
 PyTorchCppWrapperBase::get_argmax(at::Tensor input_tensor)
 {
   // Calculate argmax to get a label on each pixel
-  at::Tensor output = at::argmax(input_tensor, 1).to(torch::kCPU).to(at::kByte);
+  at::Tensor output = at::argmax(input_tensor, /*dim=*/1).to(torch::kCPU).to(at::kByte);
 
   return output;
 }
@@ -112,7 +132,7 @@ PyTorchCppWrapperBase::get_argmax(at::Tensor input_tensor)
    * @param[out] tensor that has index of max value in each element
    */
 at::Tensor
-PyTorchCppWrapperBase::get_entropy(at::Tensor input_tensor)
+PyTorchCppWrapperBase::get_entropy(at::Tensor input_tensor, const bool normalize = true)
 {
   input_tensor.to(torch::kCUDA);
   // Calculate the entropy at each pixel
@@ -121,5 +141,11 @@ PyTorchCppWrapperBase::get_entropy(at::Tensor input_tensor)
 
   at::Tensor entropy = -torch::sum(p * log_p, /*dim=*/1);
 
+  if(normalize)
+    entropy = entropy / max_entropy_;
+
   return entropy;
 }
+
+
+

src/pytorch_cpp_wrapper_seg_trav_path.cpp

@@ -14,6 +14,14 @@
 #include "opencv2/highgui/highgui.hpp"
 #include <typeinfo>
 
+PyTorchCppWrapperSegTravPath::PyTorchCppWrapperSegTravPath(const std::string & filename, const int class_num)
+  : PyTorchCppWrapperBase(filename, class_num)
+{ }
+
+PyTorchCppWrapperSegTravPath::PyTorchCppWrapperSegTravPath(const char* filename, const int class_num)
+  : PyTorchCppWrapperBase(filename, class_num)
+{ }
+
 /**
  * @brief Get outputs from the model
  * @param[in] input_tensor Input tensor

src/pytorch_seg_trav_path_ros.cpp

@@ -13,14 +13,16 @@ PyTorchSegTravPathROS::PyTorchSegTravPathROS(ros::NodeHandle & nh)
   pub_label_image_ = it_.advertise("label", 1);
   pub_color_image_ = it_.advertise("color_label", 1);
   pub_prob_image_ = it_.advertise("prob", 1);
+  pub_uncertainty_image_ = it_.advertise("uncertainty", 1);
   pub_start_point_ = nh_.advertise<geometry_msgs::PointStamped>("start_point", 1);
   pub_end_point_ = nh_.advertise<geometry_msgs::PointStamped>("end_point", 1);
   get_label_image_server_ = nh_.advertiseService("get_label_image", &PyTorchSegTravPathROS::image_inference_srv_callback, this);
 
   // Import the model
   std::string filename;
   nh_.param<std::string>("model_file", filename, "");
-  if(!pt_wrapper_.import_module(filename)) {
+  pt_wrapper_ptr_.reset(new PyTorchCppWrapperSegTravPath(filename, 4));
+  if(!pt_wrapper_ptr_->import_module(filename)) {
     ROS_ERROR("Failed to import the model file [%s]", filename.c_str());
     ros::shutdown();
   }
@@ -53,19 +55,22 @@ PyTorchSegTravPathROS::image_callback(const sensor_msgs::ImageConstPtr& msg)
   sensor_msgs::ImagePtr label_msg;
   sensor_msgs::ImagePtr color_label_msg;
   sensor_msgs::ImagePtr prob_msg;
+  sensor_msgs::ImagePtr uncertainty_msg;
   geometry_msgs::PointStampedPtr start_point_msg;
   geometry_msgs::PointStampedPtr end_point_msg;
-  std::tie(label_msg, color_label_msg, prob_msg, start_point_msg, end_point_msg) = inference(cv_ptr->image);
+  std::tie(label_msg, color_label_msg, prob_msg, uncertainty_msg, start_point_msg, end_point_msg) = inference(cv_ptr->image);
 
   // Set header
   label_msg->header = msg->header;
   color_label_msg->header = msg->header;
   prob_msg->header = msg->header;
+  uncertainty_msg->header = msg->header;
 
   // Publish the messages
   pub_label_image_.publish(label_msg);
   pub_color_image_.publish(color_label_msg);
   pub_prob_image_.publish(prob_msg);
+  pub_uncertainty_image_.publish(uncertainty_msg);
   pub_start_point_.publish(start_point_msg);
   pub_end_point_.publish(end_point_msg);
 }
@@ -88,9 +93,10 @@ PyTorchSegTravPathROS::image_inference_srv_callback(semantic_segmentation_srvs::
   sensor_msgs::ImagePtr label_msg;
   sensor_msgs::ImagePtr color_label_msg;
   sensor_msgs::ImagePtr prob_msg;
+  sensor_msgs::ImagePtr uncertainty_msg;
   geometry_msgs::PointStampedPtr start_point_msg;
   geometry_msgs::PointStampedPtr end_point_msg;
-  std::tie(label_msg, color_label_msg, prob_msg, start_point_msg, end_point_msg) = inference(cv_ptr->image);
+  std::tie(label_msg, color_label_msg, prob_msg, uncertainty_msg, start_point_msg, end_point_msg) = inference(cv_ptr->image);
 
   res.label_img = *label_msg;
   res.colorlabel_img = *color_label_msg;
@@ -105,7 +111,8 @@ PyTorchSegTravPathROS::image_inference_srv_callback(semantic_segmentation_srvs::
  * @param[in] res  Response
  * @return    True if the service succeeded
  */
-std::tuple<sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, geometry_msgs::PointStampedPtr, geometry_msgs::PointStampedPtr>
+std::tuple<sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, 
+  geometry_msgs::PointStampedPtr, geometry_msgs::PointStampedPtr>
 PyTorchSegTravPathROS::inference(cv::Mat & input_img)
 {
 
@@ -118,16 +125,9 @@ PyTorchSegTravPathROS::inference(cv::Mat & input_img)
   cv::resize(input_img, input_img, s);
 
   at::Tensor input_tensor;
-  pt_wrapper_.img2tensor(input_img, input_tensor);
+  pt_wrapper_ptr_->img2tensor(input_img, input_tensor);
 
-  // Normalize from [0, 255] -> [0, 1]
-  input_tensor /= 255.0;
-  // z-normalization
-  std::vector<float> mean_vec{0.485, 0.456, 0.406};
-  std::vector<float> std_vec{0.229, 0.224, 0.225};
-  for(int i = 0; i < mean_vec.size(); i++) {
-    input_tensor[0][i] = (input_tensor[0][i] - mean_vec[i]) / std_vec[i];
-  }
+  normalize_tensor(input_tensor);
 
   // Execute the model and turn its output into a tensor.
   at::Tensor segmentation;
@@ -136,25 +136,32 @@ PyTorchSegTravPathROS::inference(cv::Mat & input_img)
   // segmentation: raw output for segmentation (before softmax)
   // prob: traversability
   // points: coordinates of the line points
-  std::tie(segmentation, prob, points) = pt_wrapper_.get_output(input_tensor);
+  std::tie(segmentation, prob, points) = pt_wrapper_ptr_->get_output(input_tensor);
 
   // Get class label map by taking argmax of 'segmentation'
-  at::Tensor output_args = pt_wrapper_.get_argmax(segmentation);
+  at::Tensor output_args = pt_wrapper_ptr_->get_argmax(segmentation);
 
   // Uncertainty of segmentation
-  at::Tensor uncertainty = pt_wrapper_.get_entropy(segmentation);
+  at::Tensor uncertainty = pt_wrapper_ptr_->get_entropy(segmentation, true);
 
   // 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);
+  cv::Mat uncertainty_cv;
+  // Segmentation label
+  pt_wrapper_ptr_->tensor2img(output_args[0], label);
+  // Traverability
+  pt_wrapper_ptr_->tensor2img((prob[0][0]*255).to(torch::kByte), prob_cv);
+  // Segmentation label
+  pt_wrapper_ptr_->tensor2img((uncertainty[0]*255).to(torch::kByte), uncertainty_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::resize(uncertainty_cv, uncertainty_cv, s_orig, cv::INTER_LINEAR);
+
   // Generate color label image
   cv::Mat color_label;
   label_to_color(label, color_label);
@@ -163,10 +170,11 @@ PyTorchSegTravPathROS::inference(cv::Mat & input_img)
   sensor_msgs::ImagePtr label_msg = cv_bridge::CvImage(std_msgs::Header(), "mono8", label).toImageMsg();
   sensor_msgs::ImagePtr color_label_msg = cv_bridge::CvImage(std_msgs::Header(), "rgb8", color_label).toImageMsg();
   sensor_msgs::ImagePtr prob_msg = cv_bridge::CvImage(std_msgs::Header(), "mono8", prob_cv).toImageMsg();
+  sensor_msgs::ImagePtr uncertainty_msg = cv_bridge::CvImage(std_msgs::Header(), "mono8", uncertainty_cv).toImageMsg();
   geometry_msgs::PointStampedPtr start_point_msg(new geometry_msgs::PointStamped), end_point_msg(new geometry_msgs::PointStamped);
   std::tie(start_point_msg, end_point_msg) = tensor_to_points(points, width_orig, height_orig);
 
-  return std::forward_as_tuple(label_msg, color_label_msg, prob_msg, start_point_msg, end_point_msg);
+  return std::forward_as_tuple(label_msg, color_label_msg, prob_msg, uncertainty_msg, start_point_msg, end_point_msg);
 }
 
 /** 
@@ -259,3 +267,20 @@ PyTorchSegTravPathROS::msg_to_cv_bridge(sensor_msgs::Image msg)
 
   return cv_ptr;
 }
+
+/** 
+ * @brief Normalize a tensor to feed in a model
+ * @param[in]  input        Tensor
+ */
+void 
+PyTorchSegTravPathROS::normalize_tensor(at::Tensor & input_tensor)
+{
+  // Normalize from [0, 255] -> [0, 1]
+  input_tensor /= 255.0;
+  // z-normalization
+  std::vector<float> mean_vec{0.485, 0.456, 0.406};
+  std::vector<float> std_vec{0.229, 0.224, 0.225};
+  for(int i = 0; i < mean_vec.size(); i++) {
+    input_tensor[0][i] = (input_tensor[0][i] - mean_vec[i]) / std_vec[i];
+  }
+}