Refactoring

Author: ShigemichiMatsuzaki

CMakeLists.txt

@@ -1,8 +1,7 @@
 cmake_minimum_required(VERSION 2.8.3)
 project(pytorch_enet_ros)
 
-## Compile as C++11, supported in ROS Kinetic and newer
-add_compile_options(-std=c++11)
+add_compile_options(-std=c++14)
 
 # Locate the cmake file of torchlib
 set(Torch_DIR "/opt/pytorch/pytorch/torch/share/cmake/Torch/")
@@ -21,6 +20,7 @@ find_package(catkin REQUIRED COMPONENTS
 )
 
 find_package(Torch REQUIRED)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
 find_package(OpenCV REQUIRED)
 
 catkin_package(
@@ -54,6 +54,7 @@ target_link_libraries(${PROJECT_NAME}
   ${Open_CV_LIBS}
   opencv_core opencv_highgui opencv_imgcodecs
 )
+set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 14)
 
 ## Declare a C++ executable
 ## With catkin_make all packages are built within a single CMake context
@@ -64,4 +65,5 @@ add_executable(${PROJECT_NAME}_node src/pytorch_enet_ros_node.cpp)
 target_link_libraries(${PROJECT_NAME}_node
   ${catkin_LIBRARIES}
   ${PROJECT_NAME}
+  ${TORCH_LIBRARIES}
 )

include/pytorch_cpp_wrapper/pytorch_cpp_wrapper.h

@@ -2,6 +2,8 @@
 #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"
 
@@ -11,7 +13,9 @@
 //namespace mpl {
 class PyTorchCppWrapper {
 private :
-  std::shared_ptr<torch::jit::script::Module> module_;
+//  std::shared_ptr<torch::jit::script::Module> module_;
+  torch::jit::script::Module module_;
+//  torch::data::transforms::Normalize<at::Tensor> normalizer_;
 
 public:
   PyTorchCppWrapper();

launch/pytorch_enet_ros.launch

@@ -9,7 +9,10 @@
     <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/ENet_greenhouse.pt" />
+    -->
     <param name="colormap" value="$(find pytorch_enet_ros)/images/greenhouse4.png" />
     <param name="model_name" value="greenhouse" />
   </node>

src/pytorch_cpp_wrapper.cpp

@@ -6,21 +6,36 @@
 
 
 #include "pytorch_cpp_wrapper/pytorch_cpp_wrapper.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 {
 
 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) {
   // 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) {
   // 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
@@ -30,8 +45,8 @@ PyTorchCppWrapper::import_module(const std::string filename)
     // Deserialize the ScriptModule from a file using torch::jit::load().
     module_ = torch::jit::load(filename);
     // Set evaluation mode
-    module_->eval();
-    std::cout << module_->is_training() << std::endl;
+    module_.eval();
+    std::cout << module_.is_training() << std::endl;
 
     std::cout << "Import succeeded" << std::endl;
     return true;
@@ -76,9 +91,14 @@ at::Tensor
 PyTorchCppWrapper::get_output(at::Tensor input_tensor)
 {
   // Execute the model and turn its output into a tensor.
-  at::Tensor output = module_->forward({input_tensor}).toTensor();
+  auto outputs_tmp = module_.forward({input_tensor}); //.toTuple();
 
-  return output;
+  auto outputs = outputs_tmp.toTuple();
+
+  at::Tensor output1 = outputs->elements()[0].toTensor();
+  at::Tensor output2 = outputs->elements()[1].toTensor();
+
+  return output1 + 0.5 * output2;
 }
 
 at::Tensor 

src/pytorch_enet_ros.cpp

@@ -9,10 +9,6 @@
 PyTorchENetROS::PyTorchENetROS(ros::NodeHandle & nh) 
   : it_(nh), nh_(nh)
 {
-//  nh_ = ros::NodeHandle(nh);
-//  it_ = image_transport::ImageTransport(nh_);
-//  ROS_INFO("[PyTorchENetROS] Constructor");
-
   sub_image_ = it_.subscribe("image", 1, &PyTorchENetROS::image_callback, this);
   pub_label_image_ = it_.advertise("label", 1);
   pub_color_image_ = it_.advertise("color_label", 1);
@@ -67,7 +63,6 @@ PyTorchENetROS::image_inference_srv_callback(semantic_segmentation_srvs::GetLabe
 {
   ROS_INFO("[PyTorchENetROS image_inference_srv_callback] Start");
 
-
   // Convert the image message to a cv_bridge object
   cv_bridge::CvImagePtr cv_ptr = msg_to_cv_bridge(req.img);
 
@@ -93,28 +88,30 @@ PyTorchENetROS::inference(cv::Mat & input_img)
   int height_orig = input_img.size().height;
   int width_orig  = input_img.size().width;
 
-  cv::Size s(480, 264);
+  cv::Size s(480, 256);
   // Resize the input image
   cv::resize(input_img, input_img, s);
 
-//  ROS_INFO("[PyTorchENetROS inference] Start");
   at::Tensor input_tensor;
   pt_wrapper_.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];
+  }
   std::cout << input_tensor.sizes() << std::endl;
+
   // Execute the model and turn its output into a tensor.
-//  at::Tensor output = module->forward({input_tensor}).toTensor();
-//  ROS_INFO("[PyTorchENetROS inference] get_output");
   at::Tensor output = pt_wrapper_.get_output(input_tensor);
   // Calculate argmax to get a label on each pixel
-  // at::Tensor output_args = at::argmax(output, 1).to(torch::kCPU).to(at::kByte);
-//  ROS_INFO("[PyTorchENetROS inference] get_argmax");
   at::Tensor output_args = pt_wrapper_.get_argmax(output);
 
   // Convert to OpenCV
-  // cv::Mat mat(height, width, CV_8U, output_args[0]. template data<uint8_t>());
   cv::Mat label;
-//  ROS_INFO("[PyTorchENetROS inference] tensor2img");
   pt_wrapper_.tensor2img(output_args[0], label);
 
   // Set the size
@@ -123,18 +120,12 @@ PyTorchENetROS::inference(cv::Mat & input_img)
   cv::resize(label, label, s_orig, cv::INTER_NEAREST);
 
   cv::Mat color_label;
-//  cv::applyColorMap(mat, color_label, cv::COLORMAP_JET);
   label_to_color(label, color_label);
 
   // Generate an image message
-//  ROS_INFO("[PyTorchENetROS inference] cv_bridge to image msg");
   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 color_msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", image).toImageMsg();
-  
-//  ROS_INFO("[PyTorchENetROS inference] Publish");
   return std::forward_as_tuple(label_msg, color_label_msg);
 }
 
@@ -155,11 +146,11 @@ cv_bridge::CvImagePtr
 PyTorchENetROS::msg_to_cv_bridge(sensor_msgs::ImageConstPtr msg)
 {
   cv_bridge::CvImagePtr cv_ptr;
+
   // Convert the image message to a cv_bridge object
   try
   {
     cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
-//    ROS_INFO("[PyTorchENetROS image_callback] Convert to cv_bridge object");
   }
   catch (cv_bridge::Exception& e)
   {
@@ -177,11 +168,11 @@ cv_bridge::CvImagePtr
 PyTorchENetROS::msg_to_cv_bridge(sensor_msgs::Image msg)
 {
   cv_bridge::CvImagePtr cv_ptr;
+
   // Convert the image message to a cv_bridge object
   try
   {
     cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
-//    ROS_INFO("[PyTorchENetROS image_callback] Convert to cv_bridge object");
   }
   catch (cv_bridge::Exception& e)
   {