/*
 *
 * Created on: Jan 21, 2017
 * Author: Timo Sämann
 *
 * The basis for the creation of this script was the classification.cpp example (caffe/examples/cpp_classification/classification.cpp)
 *
 * This script visualize the semantic segmentation for your input image.
 *
 * To compile this script you can use a IDE like Eclipse. To include Caffe and OpenCV in Eclipse please refer to
 * http://tzutalin.blogspot.de/2015/05/caffe-on-ubuntu-eclipse-cc.html
 * and http://rodrigoberriel.com/2014/10/using-opencv-3-0-0-with-eclipse/ , respectively
 *
 *
 */



#define USE_OPENCV 1
#include <caffe/caffe.hpp>

#ifdef USE_OPENCV
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#endif // USE_OPENCV

#include <algorithm>
#include <iosfwd>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <chrono> //Just for time measurement

#ifdef USE_OPENCV
using namespace caffe; // NOLINT(build/namespaces)
using std::string;



class Classifier {
 public:
 Classifier(const string& model_file,
 const string& trained_file);


 void Predict(const cv::Mat& img, string LUT_file);

 private:
 void SetMean(const string& mean_file);

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

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

 void Visualization(Blob<float>* output_layer, string LUT_file);

 private:
 shared_ptr<Net<float> > net_;
 cv::Size input_geometry_;
 int num_channels_;

};

Classifier::Classifier(const string& model_file,
 const string& trained_file) {


 Caffe::set_mode(Caffe::GPU);

 /* Load the network. */
 net_.reset(new Net<float>(model_file, TEST));
 net_->CopyTrainedLayersFrom(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.";

 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());
}


void Classifier::Predict(const cv::Mat& img, string LUT_file) {
 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(img, &input_channels);


 std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now(); //Just for time measurement

 net_->Forward();

 std::chrono::steady_clock::time_point end= std::chrono::steady_clock::now();
 std::cout << "Processing time = " << (std::chrono::duration_cast<std::chrono::microseconds>(end - begin).count())/1000000.0 << " sec" <<std::endl; //Just for time measurement


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

 Visualization(output_layer, LUT_file);
}


void Classifier::Visualization(Blob<float>* output_layer, string LUT_file) {

 std::cout << "output_blob(n,c,h,w) = " << output_layer->num() << ", " << output_layer->channels() << ", "
			 << output_layer->height() << ", " << output_layer->width() << std::endl;

 cv::Mat merged_output_image = cv::Mat(output_layer->height(), output_layer->width(), CV_32F, const_cast<float *>(output_layer->cpu_data()));
 //merged_output_image = merged_output_image/255.0;

 merged_output_image.convertTo(merged_output_image, CV_8U);
 cv::cvtColor(merged_output_image.clone(), merged_output_image, CV_GRAY2BGR);
 cv::Mat label_colours = cv::imread(LUT_file,1);
 cv::Mat output_image;
 LUT(merged_output_image, label_colours, output_image);

 cv::imshow( "Display window", output_image);
 cv::waitKey(0);
}


/* Wrap the input layer of the network in separate cv::Mat objects
 * (one per channel). This way we save one memcpy operation and we
 * don't need to rely on cudaMemcpy2D. The last preprocessing
 * operation will write the separate channels directly to the input
 * layer. */
void Classifier::WrapInputLayer(std::vector<cv::Mat>* input_channels) {
 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);
 input_channels->push_back(channel);
 input_data += width * height;
 }
}

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

 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);

 /* This operation will write the separate BGR planes directly to the
 * input layer of the network because it is wrapped by the cv::Mat
 * objects in input_channels. */
 cv::split(sample_float, *input_channels);

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

int main(int argc, char** argv) {
 if (argc != 5) {
 std::cerr << "Usage: " << argv[0]
 << " \ndeploy.prototxt \nnetwork.caffemodel"
 << " \nimg.jpg" << " \ncamvid12.png (for example: /SegNet-Tutorial/Scripts/camvid12.png)" << std::endl;
 return 1;
 }

 ::google::InitGoogleLogging(argv[0]);

 string model_file = argv[1];
 string trained_file = argv[2]; //for visualization


 Classifier classifier(model_file, trained_file);

 string file = argv[3];
 string LUT_file = argv[4];

 std::cout << "---------- Semantic Segmentation for "
 << file << " ----------" << std::endl;

 cv::Mat img = cv::imread(file, 1);
 CHECK(!img.empty()) << "Unable to decode image " << file;
 cv::Mat prediction;

 classifier.Predict(img, LUT_file);
}
#else
int main(int argc, char** argv) {
 LOG(FATAL) << "This example requires OpenCV; compile with USE_OPENCV.";
}
#endif // USE_OPENCV