[update] Separate the base class and child classes

Author: ShigemichiMatsuzaki

CMakeLists.txt

@@ -45,7 +45,10 @@ include_directories(
 ## Declare a C++ library
 add_library(${PROJECT_NAME}
   src/pytorch_enet_ros.cpp
-  src/pytorch_cpp_wrapper.cpp
+  src/pytorch_seg_trav_path.cpp
+  src/pytorch_cpp_wrapper_seg_trav.cpp
+  src/pytorch_cpp_wrapper_seg_trav_path.cpp
+  src/pytorch_cpp_wrapper_base.cpp
 )
 
 target_link_libraries(${PROJECT_NAME}
@@ -60,10 +63,17 @@ set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 14)
 ## With catkin_make all packages are built within a single CMake context
 ## The recommended prefix ensures that target names across packages don't collide
 add_executable(${PROJECT_NAME}_node src/pytorch_enet_ros_node.cpp)
+add_executable(pytorch_seg_trav_path_node src/pytorch_seg_trav_path_node.cpp)
 
 ## Specify libraries to link a library or executable target against
 target_link_libraries(${PROJECT_NAME}_node
   ${catkin_LIBRARIES}
   ${PROJECT_NAME}
   ${TORCH_LIBRARIES}
 )
+
+target_link_libraries(pytorch_seg_trav_path_node
+  ${catkin_LIBRARIES}
+  ${PROJECT_NAME}
+  ${TORCH_LIBRARIES}
+)

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper.h

@@ -0,0 +1,63 @@
+#ifndef PYTORCH_CPP_WRAPPER_BASE
+#define PYTORCH_CPP_WRAPPER_BASE
+
+#include <torch/script.h> // One-stop header.
+#include <torch/data/transforms/tensor.h> // One-stop header.
+#include <c10/util/ArrayRef.h>
+#include <opencv2/opencv.hpp>
+
+#include <iostream>
+#include <memory>
+
+/**
+ * @brief this class is a base class of C++ wrapper of PyTorch
+ */
+class PyTorchCppWrapperBase {
+protected :
+  torch::jit::script::Module module_;
+
+public:
+  PyTorchCppWrapperBase();
+  PyTorchCppWrapperBase(const std::string & filename);
+  PyTorchCppWrapperBase(const char* filename);
+//  virtual ~PyTorchCppWrapperBase();
+
+  /**
+   * @brief import a network 
+   * @param filename
+   * @return true if import succeeded 
+   */
+  bool import_module(const std::string & filename);
+
+  /**
+   * @brief convert an image(cv::Mat) to a tensor (at::Tensor)
+   * @param[in] img
+   * @param[out] tensor
+   * @param[in] whether to use GPU
+   */
+  void img2tensor(cv::Mat & img, at::Tensor & tensor, const bool & use_gpu = true);
+
+  /**
+   * @brief convert a tensor (at::Tensor) to an image (cv::Mat)
+   * @param[in] tensor
+   * @param[out] img
+   */
+  void tensor2img(at::Tensor tensor, cv::Mat & img);
+
+  /**
+   * @brief convert a tensor (at::Tensor) to an image (cv::Mat)
+   * @param[in] tensor
+   * @param[out] img
+   */
+  at::Tensor get_argmax(at::Tensor input_tensor);
+
+  /**
+   * @brief convert a tensor (at::Tensor) to an image (cv::Mat)
+   * @param[in] input_tensor
+   * @param[out] Output from the network (depends on the implementation)
+   */
+//  virtual auto get_output(at::Tensor & input_tensor) = 0;
+  
+};
+//}
+#endif

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h

@@ -0,0 +1,22 @@
+#ifndef PYTORCH_CPP_WRAPPER
+#define PYTORCH_CPP_WRAPPER
+
+#include <torch/script.h> // One-stop header.
+#include <torch/data/transforms/tensor.h> // One-stop header.
+#include <c10/util/ArrayRef.h>
+#include <opencv2/opencv.hpp>
+#include "opencv2/highgui/highgui.hpp"
+#include "pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h"
+
+#include <iostream>
+#include <memory>
+
+class PyTorchCppWrapperSegTrav : public PyTorchCppWrapperBase {
+private :
+  // c = P(s|y=1) in PU learning, calculated during training
+  float c_{0.3};
+
+public:
+  std::tuple<at::Tensor, at::Tensor> get_output(at::Tensor input_tensor);
+};
+#endif

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper_seg_trav.h

@@ -0,0 +1,23 @@
+#ifndef PYTORCH_CPP_WRAPPER
+#define PYTORCH_CPP_WRAPPER
+
+#include <torch/script.h> // One-stop header.
+#include <torch/data/transforms/tensor.h> // One-stop header.
+#include <c10/util/ArrayRef.h>
+#include <opencv2/opencv.hpp>
+#include "opencv2/highgui/highgui.hpp"
+#include "pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h"
+
+#include <iostream>
+#include <memory>
+
+
+class PyTorchCppWrapperSegTravPath : public PyTorchCppWrapperBase {
+private :
+  // c = P(s|y=1) in PU learning, calculated during training
+  float c_{0.3};
+
+public:
+  std::tuple<at::Tensor, at::Tensor, at::Tensor> get_output(at::Tensor input_tensor);
+};
+#endif

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper_seg_trav_path.h

@@ -15,7 +15,7 @@
 //#include<semantic_segmentation_srvs/GetLabelImage.h>
 #include<semantic_segmentation_srvs/GetLabelAndProbability.h>
 
-#include"pytorch_cpp_wrapper/pytorch_cpp_wrapper.h"
+#include"pytorch_cpp_wrapper/pytorch_cpp_wrapper_seg_trav.h"
 
 #include <iostream>
 #include <memory>
@@ -34,7 +34,7 @@ class PyTorchENetROS {
   image_transport::Publisher  pub_color_image_;
   image_transport::Publisher  pub_prob_image_;
 
-  PyTorchCppWrapper pt_wrapper_;
+  PyTorchCppWrapperSegTrav pt_wrapper_;
 
   cv::Mat colormap_;
 

include/pytorch_enet_ros/pytorch_enet_ros.h

@@ -0,0 +1,55 @@
+/*
+ * A ROS node to do inference using PyTorch model
+ * Shigemichi Matsuzaki
+ *
+ */
+
+#ifndef PYTORCH_SEG_TRAV_PATH
+#define PYTORCH_SEG_TRAV_PATH
+
+#include <ros/ros.h>
+
+#include <opencv2/opencv.hpp>
+#include<image_transport/image_transport.h>
+#include<cv_bridge/cv_bridge.h>
+#include<geometry_msgs/Point.h>
+//#include<semantic_segmentation_srvs/GetLabelImage.h>
+#include<semantic_segmentation_srvs/GetLabelAndProbability.h>
+
+#include"pytorch_cpp_wrapper/pytorch_cpp_wrapper_seg_trav_path.h"
+
+#include <iostream>
+#include <memory>
+#include <tuple>
+
+class PyTorchSegTravPathROS {
+private:
+  ros::NodeHandle nh_;
+
+  ros::ServiceServer get_label_image_server_;
+
+  image_transport::ImageTransport it_;
+
+  image_transport::Subscriber sub_image_;
+  image_transport::Publisher  pub_label_image_;
+  image_transport::Publisher  pub_color_image_;
+  image_transport::Publisher  pub_prob_image_;
+
+//  PyTorchCppWrapper pt_wrapper_;
+  PyTorchCppWrapperSegTravPath pt_wrapper_;
+
+  cv::Mat colormap_;
+
+public:
+  PyTorchSegTravPathROS(ros::NodeHandle & nh); 
+
+  void image_callback(const sensor_msgs::ImageConstPtr& msg); 
+  std::tuple<sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, sensor_msgs::ImagePtr, geometry_msgs::PointPtr, geometry_msgs::PointPtr> inference(cv::Mat & input_image);
+  bool image_inference_srv_callback(semantic_segmentation_srvs::GetLabelAndProbability::Request  & req,
+                                    semantic_segmentation_srvs::GetLabelAndProbability::Response & res);
+  cv_bridge::CvImagePtr msg_to_cv_bridge(sensor_msgs::ImageConstPtr msg);
+  cv_bridge::CvImagePtr msg_to_cv_bridge(sensor_msgs::Image msg);
+  void label_to_color(cv::Mat& label, cv::Mat& color_label);
+};
+
+#endif

include/pytorch_enet_ros/pytorch_seg_trav_path.h

@@ -2,14 +2,14 @@
 <launch>
   <arg name="image" default="/camera/rgb/image_rect_color" />
 
-  <node pkg="pytorch_enet_ros" type="pytorch_enet_ros_node" name="pytorch_enet_ros_node" output="screen">
+  <node pkg="pytorch_enet_ros" type="pytorch_seg_trav_path_node" name="pytorch_seg_trav_path_node" output="screen">
     <!--
     <remap from="~image" to="/kinect2/qhd/image_color_rect" />
     <param name="model_file" value="$(find pytorch_enet_ros)/models/ENet_camvid.pt" />
     <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_trav_20210115-151110.pt" />
+    <param name="model_file" value="$(find pytorch_enet_ros)/models/espdnet_ue_trav_path_20210518-221714.pt" />
     <!--
     <param name="model_file" value="$(find pytorch_enet_ros)/models/ENet_greenhouse.pt" />
     -->

launch/pytorch_enet_ros.launch

@@ -6,41 +6,28 @@
 
 
 #include <torch/torch.h>
-#include "pytorch_cpp_wrapper/pytorch_cpp_wrapper.h"
+#include "pytorch_cpp_wrapper/pytorch_cpp_wrapper_base.h"
 #include <torch/script.h> // One-stop header.
 #include <torch/data/transforms/tensor.h> // One-stop header.
 #include <c10/util/ArrayRef.h>
 #include <opencv2/opencv.hpp>
 #include "opencv2/highgui/highgui.hpp"
 #include <typeinfo>
 
-//namespace mpl {
+PyTorchCppWrapperBase::PyTorchCppWrapperBase() {}
 
-PyTorchCppWrapper::PyTorchCppWrapper() {
-  std::vector<float> mean_vec{0.485, 0.456, 0.406};
-  std::vector<float> std_vec{0.229, 0.224, 0.225};
-  
-}
-
-PyTorchCppWrapper::PyTorchCppWrapper(const std::string filename) {
+PyTorchCppWrapperBase::PyTorchCppWrapperBase(const std::string & filename) {
   // Import
   import_module(filename);
-  std::vector<float> mean_vec{0.485, 0.456, 0.406};
-  std::vector<float> std_vec{0.229, 0.224, 0.225};
-  
 }
 
-PyTorchCppWrapper::PyTorchCppWrapper(const char* filename) {
+PyTorchCppWrapperBase::PyTorchCppWrapperBase(const char* filename) {
   // Import
   import_module(std::string(filename));
-
-  std::vector<float> mean_vec{0.485, 0.456, 0.406};
-  std::vector<float> std_vec{0.229, 0.224, 0.225};
-  
 }
 
 bool
-PyTorchCppWrapper::import_module(const std::string filename)
+PyTorchCppWrapperBase::import_module(const std::string & filename)
 {
   try {
     // Deserialize the ScriptModule from a file using torch::jit::load().
@@ -59,7 +46,7 @@ PyTorchCppWrapper::import_module(const std::string filename)
 }
 
 void
-PyTorchCppWrapper::img2tensor(cv::Mat & img, at::Tensor & tensor, const bool use_gpu)
+PyTorchCppWrapperBase::img2tensor(cv::Mat & img, at::Tensor & tensor, const bool & use_gpu)
 {
   // Get the size of the input image
   int height = img.size().height;
@@ -77,7 +64,7 @@ PyTorchCppWrapper::img2tensor(cv::Mat & img, at::Tensor & tensor, const bool use
 }
 
 void
-PyTorchCppWrapper::tensor2img(at::Tensor tensor, cv::Mat & img)
+PyTorchCppWrapperBase::tensor2img(at::Tensor tensor, cv::Mat & img)
 {
   // Get the size of the input image
   int height = tensor.sizes()[0];
@@ -89,36 +76,11 @@ PyTorchCppWrapper::tensor2img(at::Tensor tensor, cv::Mat & img)
   img = cv::Mat(height, width, CV_8U, tensor. template data<uint8_t>());
 }
 
-//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.
-  auto outputs_tmp = module_.forward({input_tensor}); //.toTuple();
-
-  auto outputs = outputs_tmp.toTuple();
-
-  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 std::forward_as_tuple(segmentation, prob);
-}
-
 at::Tensor 
-PyTorchCppWrapper::get_argmax(at::Tensor input_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);
 
   return output;
 }
-//} // namespace mpl