WIP: Implement uncertainty estimation

Author: ShigemichiMatsuzaki

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h

@@ -49,6 +49,14 @@ protected :
    * @param[out] tensor that has index of max value in each element
    */
   at::Tensor get_argmax(at::Tensor input_tensor);
+
+  /**
+   * @brief Take element-wise entropy 
+   * @param[in]  tensor
+   * @param[out] tensor that has index of max value in each element
+   */
+  at::Tensor get_entropy(at::Tensor input_tensor);
+
 };
 //}
 #endif

launch/pytorch_enet_ros.launch

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <launch>
   <arg name="image" default="/camera/rgb/image_rect_color" />
-  <arg name="model_name" default="$(find pytorch_ros)/models/espdnet_ue_trav_path_20210518-221714.pt" />
+  <arg name="model_name" default="$(find pytorch_ros)/models/espdnet_ue_trav_path_20210712-134315.pt" />
 
   <node pkg="pytorch_ros" type="pytorch_seg_trav_path_node" name="pytorch_seg_trav_path_node" output="screen">
     <remap from="~image" to="$(arg image)" />

script/visualizer.py

@@ -83,4 +83,4 @@ def main():
 
 
 if __name__=='__main__':
-    main()
+    main()

src/pytorch_cpp_wrapper_base.cpp

@@ -105,3 +105,21 @@ PyTorchCppWrapperBase::get_argmax(at::Tensor input_tensor)
 
   return output;
 }
+
+  /**
+   * @brief Take element-wise entropy 
+   * @param[in]  tensor
+   * @param[out] tensor that has index of max value in each element
+   */
+at::Tensor
+PyTorchCppWrapperBase::get_entropy(at::Tensor input_tensor)
+{
+  input_tensor.to(torch::kCUDA);
+  // Calculate the entropy at each pixel
+  at::Tensor log_p = torch::log_softmax(input_tensor, /*dim=*/1);//at::argmax(input_tensor, 1).to(torch::kCPU).to(at::kByte);
+  at::Tensor p = torch::log_softmax(input_tensor, /*dim=*/1);
+
+  at::Tensor entropy = torch::sum(p * log_p, /*dim=*/1);
+
+  return entropy;
+}

src/pytorch_seg_trav_path_ros.cpp

@@ -133,10 +133,17 @@ PyTorchSegTravPathROS::inference(cv::Mat & input_img)
   at::Tensor segmentation;
   at::Tensor prob;
   at::Tensor points;
+  // 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);
 
+  // Get class label map by taking argmax of 'segmentation'
   at::Tensor output_args = pt_wrapper_.get_argmax(segmentation);
 
+  // Uncertainty of segmentation
+  at::Tensor uncertainty = pt_wrapper_.get_entropy(segmentation);
+
   // Convert to OpenCV
   cv::Mat label;
   cv::Mat prob_cv;