[update] Update

Author: ShigemichiMatsuzaki

.gitignore

@@ -5,8 +5,8 @@ project(pytorch_enet_ros)
 add_compile_options(-std=c++11)
 
 # Locate the cmake file of torchlib
-#set(Torch_DIR "/opt/pytorch/pytorch/torch/share/cmake/Torch/")
-set(Torch_DIR "$ENV{TORCH_ROOT}/torch/share/cmake/Torch/")
+set(Torch_DIR "/opt/pytorch/pytorch/torch/share/cmake/Torch/")
+#set(Torch_DIR "$ENV{TORCH_ROOT}/torch/share/cmake/Torch/")
 
 ## Find catkin macros and libraries
 ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz)

CMakeLists.txt

@@ -1,14 +1,15 @@
 <?xml version="1.0"?>
-
 <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">
     <!--
     <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="/camera/rgb/image_rect_color" />
-    <param name="model_file" value="$(find pytorch_enet_ros)/models/ENet_scenes_train_c2-501_2_200.pt" />
+    <remap from="~image" to="$(arg image)" />
+    <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>

README.md

@@ -93,9 +93,9 @@ PyTorchENetROS::inference(cv::Mat & input_img)
   int height_orig = input_img.size().height;
   int width_orig  = input_img.size().width;
 
-//  cv::Size s(320, 240);
+  cv::Size s(480, 264);
   // Resize the input image
-//  cv::resize(input_img, input_img, s);
+  cv::resize(input_img, input_img, s);
 
 //  ROS_INFO("[PyTorchENetROS inference] Start");
   at::Tensor input_tensor;
@@ -120,7 +120,7 @@ PyTorchENetROS::inference(cv::Mat & input_img)
   // 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(label, label, s_orig, cv::INTER_NEAREST);
 
   cv::Mat color_label;
 //  cv::applyColorMap(mat, color_label, cv::COLORMAP_JET);