Skip to content

Commit

Permalink
Feature/enet image segmenter (#784)
Browse files Browse the repository at this point in the history
* Initial Release of ENEt Image Segmenter

* Node Publishes
-segmentation results in /image_segmented
-blended result with sensor /image_segmented_blended

Added README

* Removed unused params

* Added More Instructions to README

* Added extra instructions

* Package renamed to image_segmenter
Node renamed to image_segmenter_enet

* Modified package README to reflect updated naming

* Fixed dependencies

* Added curly braces to one line ifs

* Added Further compilation instructions

* Path in readme file updated
  • Loading branch information
amc-nu committed Aug 29, 2017
1 parent 4a052e2 commit 5c3adaf
Show file tree
Hide file tree
Showing 7 changed files with 486 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
cmake_minimum_required(VERSION 2.8.12)
project(image_segmenter)
execute_process(
COMMAND rosversion -d
OUTPUT_VARIABLE ROS_VERSION
OUTPUT_STRIP_TRAILING_WHITESPACE
)

include(FindPkgConfig)

if ("${ROS_VERSION}" MATCHES "(indigo|jade)")
FIND_PACKAGE(catkin REQUIRED COMPONENTS
cv_bridge
image_transport
roscpp
sensor_msgs
)
elseif ("${ROS_VERSION}" MATCHES "(kinetic)")
FIND_PACKAGE(catkin REQUIRED COMPONENTS
cv_bridge
image_transport
roscpp
sensor_msgs
)
endif()

FIND_PACKAGE(CUDA)
FIND_PACKAGE(OpenCV REQUIRED)

catkin_package(CATKIN_DEPENDS
cv_bridge
image_transport
roscpp
sensor_msgs
)
###########
## Build ##
###########

SET(CMAKE_CXX_FLAGS "-std=c++11 -O2 -g -Wall ${CMAKE_CXX_FLAGS}")

INCLUDE_DIRECTORIES(
${catkin_INCLUDE_DIRS}
)


#####ENET########
##############################ENet's CAFFE FORK NEEDS TO BE PREVIOUSLY COMPILED####################
set(ENET_CAFFE_PATH "$ENV{HOME}/ENet/caffe-enet/distribute")
####################################################################################################
if(EXISTS "${ENET_CAFFE_PATH}")

ADD_EXECUTABLE(image_segmenter_enet
src/image_segmenter_enet_node.cpp
src/image_segmenter_enet.cpp
)

TARGET_LINK_LIBRARIES(image_segmenter_enet
${catkin_LIBRARIES}
${OpenCV_LIBS}
${CUDA_LIBRARIES}
${CUDA_CUBLAS_LIBRARIES}
${CUDA_curand_LIBRARY}
${ENET_CAFFE_PATH}/lib/libcaffe.so
glog
)

TARGET_INCLUDE_DIRECTORIES(image_segmenter_enet PRIVATE
${CUDA_INCLUDE_DIRS}
${ENET_CAFFE_PATH}/include
include
)

ADD_DEPENDENCIES(image_segmenter_enet
${catkin_EXPORTED_TARGETS}
)
else()
message("'ENet/Caffe' is not installed. 'image_segmenter_enet' will not be built.")
endif()
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
# How to install Caffe for ENet

1. Complete the [pre-requisites](http://caffe.berkeleyvision.org/install_apt.html)

2. Clone ENet

```
cd ~
git clone --recursive https://github.com/TimoSaemann/ENet.git
cd ENet/caffe-enet
```

3. Compile ENet fork of Caffe using **Make** *(Don't use CMake to compile Caffe)*
Create Makefile.config from Makefile.config.example and setup your paths
as indicated in http://caffe.berkeleyvision.org/installation.html#compilation
```
make && make distribute
```

4. Download pretrained models as pointed in
https://github.com/TimoSaemann/ENet/tree/master/Tutorial#kick-start or use your own.

If you didn't install ENet Caffe in `ENet` inside your home for some reason,
modify the Autoware ENet's node CMakeFile and point the paths to match your system.

Once compiled, run from the terminal

```
% roslaunch image_segmenter image_segmenter_enet.launch
```
Remember to modify the launch file located inside
`computing/perception/detection/packages/image_segmenter/launch/image_segmenter_enet.launch`
and configure the network configuration file, the pretrained models and the LUT file.
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
/*
* image_segmenter_enet.hpp
*
* Created on: Aug 23, 2017
* Author: ne0
*/

#ifndef ENET_IMAGE_SEGMENTER_HPP_
#define ENET_IMAGE_SEGMENTER_HPP_

#include <algorithm>
#include <iosfwd>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <caffe/caffe.hpp>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

class ENetSegmenter
{
public:
ENetSegmenter(const std::string& in_model_file, const std::string& in_trained_file, const std::string& in_lookuptable_file);

void Predict(const cv::Mat& in_img, cv::Mat& out_segmented);

private:
void SetMean(const std::string& in_mean_file);

void WrapInputLayer(std::vector<cv::Mat>* in_input_channels);

void Preprocess(const cv::Mat& img, std::vector<cv::Mat>* in_input_channels);

cv::Mat Visualization(cv::Mat in_prediction_map, std::string in_lookuptable_file);

private:
std::shared_ptr<caffe::Net<float> > net_;
cv::Size input_geometry_;
int num_channels_;
std::string lookuptable_file_;
cv::Scalar pixel_mean_;

};

#endif /* ENET_IMAGE_SEGMENTER_HPP_ */
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
<launch>
<!-- arguments list -->
<arg name="network_definition_file" default="$(env HOME)/ENet/prototxts/enet_deploy_final.prototxt"/>
<arg name="pretrained_model_file" default="$(env HOME)/ENet/enet_weights_zoo/cityscapes_weights.caffemodel"/>
<arg name="lookuptable_file" default="$(env HOME)/ENet/scripts/cityscapes19.png"/>

<arg name="camera_id" default="/" />
<arg name="image_src" default="/image_raw"/>

<!-- ENET -->
<node pkg="image_segmenter" name="image_segmenter_enet" type="image_segmenter_enet" output="screen">
<param name="network_definition_file" type="str" value="$(arg network_definition_file)"/>
<param name="pretrained_model_file" type="str" value="$(arg pretrained_model_file)"/>
<param name="lookuptable_file" type="str" value="$(arg lookuptable_file)"/>
<param name="image_raw_node" type="str" value="$(arg camera_id)$(arg image_src)"/>
</node>

</launch>
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
<?xml version="1.0"?>
<package>
<name>image_segmenter</name>
<version>1.0.0</version>
<description>Image Segmentation</description>
<maintainer email="abrahammonrroy@yahoo.com">Abraham Monrroy</maintainer>
<license>BSD</license>
<buildtool_depend>catkin</buildtool_depend>
<build_depend>cv_bridge</build_depend>
<build_depend>image_transport</build_depend>
<build_depend>roscpp</build_depend>
<build_depend>sensor_msgs</build_depend>
<run_depend>cv_bridge</run_depend>
<run_depend>image_transport</run_depend>
<run_depend>roscpp</run_depend>
<run_depend>cv_bridge</run_depend>
<run_depend>sensor_msgs</run_depend>
</package>
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
/*
* image_segmenter_enet.cpp
*
* Created on: Aug 23, 2017
* Author: amc
*/

#include "image_segmenter_enet.hpp"

ENetSegmenter::ENetSegmenter(const std::string& in_model_file,
const std::string& in_trained_file,
const std::string& in_lookuptable_file)
{

caffe::Caffe::set_mode(caffe::Caffe::GPU);

/* Load the network. */
net_.reset(new caffe::Net<float>(in_model_file, caffe::TEST));
net_->CopyTrainedLayersFrom(in_trained_file);

CHECK_EQ(net_->num_inputs(), 1)<< "Network should have exactly one input.";
CHECK_EQ(net_->num_outputs(), 1)<< "Network should have exactly one output.";

caffe::Blob<float>* input_layer = net_->input_blobs()[0];
num_channels_ = input_layer->channels();
CHECK(num_channels_ == 3 || num_channels_ == 1) << "Input layer should have 1 or 3 channels.";
input_geometry_ = cv::Size(input_layer->width(), input_layer->height());

lookuptable_file_ = in_lookuptable_file;

pixel_mean_ = cv::Scalar(102.9801, 115.9465, 122.7717);
}

void ENetSegmenter::Predict(const cv::Mat& in_image_mat, cv::Mat& out_segmented)
{
caffe::Blob<float>* input_layer = net_->input_blobs()[0];
input_layer->Reshape(1, num_channels_, input_geometry_.height,
input_geometry_.width);
/* Forward dimension change to all layers. */
net_->Reshape();

std::vector<cv::Mat> input_channels;
WrapInputLayer(&input_channels);

Preprocess(in_image_mat, &input_channels);

net_->Forward();

/* Copy the output layer to a std::vector */
caffe::Blob<float>* output_layer = net_->output_blobs()[0];

int width = output_layer->width();
int height = output_layer->height();
int channels = output_layer->channels();

// compute argmax
cv::Mat class_each_row(channels, width * height, CV_32FC1,
const_cast<float *>(output_layer->cpu_data()));
class_each_row = class_each_row.t(); // transpose to make each row with all probabilities
cv::Point maxId; // point [x,y] values for index of max
double maxValue; // the holy max value itself
cv::Mat prediction_map(height, width, CV_8UC1);
for (int i = 0; i < class_each_row.rows; i++)
{
minMaxLoc(class_each_row.row(i), 0, &maxValue, 0, &maxId);
prediction_map.at<uchar>(i) = maxId.x;
}

out_segmented = Visualization(prediction_map, lookuptable_file_);

cv::resize(out_segmented, out_segmented, cv::Size(in_image_mat.cols, in_image_mat.rows));


}

cv::Mat ENetSegmenter::Visualization(cv::Mat in_prediction_map, std::string in_lookuptable_file)
{

cv::cvtColor(in_prediction_map.clone(), in_prediction_map, CV_GRAY2BGR);
cv::Mat label_colours = cv::imread(in_lookuptable_file, 1);
cv::cvtColor(label_colours, label_colours, CV_RGB2BGR);
cv::Mat output_image;
LUT(in_prediction_map, label_colours, output_image);

return output_image;
}

void ENetSegmenter::WrapInputLayer(std::vector<cv::Mat>* in_input_channels)
{
caffe::Blob<float>* input_layer = net_->input_blobs()[0];

int width = input_layer->width();
int height = input_layer->height();
float* input_data = input_layer->mutable_cpu_data();
for (int i = 0; i < input_layer->channels(); ++i)
{
cv::Mat channel(height, width, CV_32FC1, input_data);
in_input_channels->push_back(channel);
input_data += width * height;
}
}

void ENetSegmenter::Preprocess(const cv::Mat& in_image_mat,
std::vector<cv::Mat>* in_input_channels)
{
/* Convert the input image to the input image format of the network. */
cv::Mat sample;
if (in_image_mat.channels() == 3 && num_channels_ == 1){
cv::cvtColor(in_image_mat, sample, cv::COLOR_BGR2GRAY);}
else if (in_image_mat.channels() == 4 && num_channels_ == 1){
cv::cvtColor(in_image_mat, sample, cv::COLOR_BGRA2GRAY);}
else if (in_image_mat.channels() == 4 && num_channels_ == 3){
cv::cvtColor(in_image_mat, sample, cv::COLOR_BGRA2BGR);}
else if (in_image_mat.channels() == 1 && num_channels_ == 3){
cv::cvtColor(in_image_mat, sample, cv::COLOR_GRAY2BGR);}
else{
sample = in_image_mat;
}

cv::Mat sample_resized;
if (sample.size() != input_geometry_){
cv::resize(sample, sample_resized, input_geometry_);}
else{
sample_resized = sample;
}

cv::Mat sample_float;
if (num_channels_ == 3)
{
sample_resized.convertTo(sample_float, CV_32FC3);
}
else
{
sample_resized.convertTo(sample_float, CV_32FC1);
}

cv::split(sample_float, *in_input_channels);

CHECK(reinterpret_cast<float*>(in_input_channels->at(0).data)
== net_->input_blobs()[0]->cpu_data())
<< "Input channels are not wrapping the input layer of the network.";
}

0 comments on commit 5c3adaf

Please sign in to comment.