1 #include <stdio.h>  // for snprintf
  2 #include <string>
  3 #include <vector>
  4 
  5 #include "boost/algorithm/string.hpp"
  6 #include "google/protobuf/text_format.h"
  7 #include "leveldb/db.h"
  8 #include "leveldb/write_batch.h"
  9 
 10 #include "caffe/blob.hpp"
 11 #include "caffe/common.hpp"
 

 12 #include "caffe/net.hpp"
 13 #include "caffe/proto/caffe.pb.h"
 14 #include "caffe/util/io.hpp"
 15 #include "caffe/vision_layers.hpp"
 16 
 17 // liu 
 18 #include <opencv2/core/core.hpp>
 19 #include <opencv2/highgui/highgui.hpp>
 20 #include <opencv2/highgui/highgui_c.h>
 21 #include <opencv2/imgproc/imgproc.hpp>
 

 22 #include <opencv2/opencv.hpp>
 23 #include <sys/stat.h>  // for mkdir
 24 
 25 using namespace caffe;  // NOLINT(build/namespaces)
 26 
 27 template<typename Dtype>
 28 int feature_extraction_pipeline(int argc, char** argv);
 29 
 30 int main(int argc, char** argv) {
 31   return feature_extraction_pipeline
 
<float>(argc, argv);
 32 //  return feature_extraction_pipeline<double>(argc, argv);
 33 }
 34 
 35 template<typename Dtype>
 36 int feature_extraction_pipeline(int argc, char** argv) {
 37   ::google::InitGoogleLogging(argv[0]);
 38   const int num_required_args = 5;
 39   if (argc < num_required_args) {
 40     LOG(ERROR)<<
 41     "This program takes in a trained network and an input data layer, and then"
 42     " extract features of the input data produced by the net.\n"
 43     "Usage: test_det_net  pretrained_net_param"
 44     "  feature_extraction_proto_file num_mini_batches"
 45         "  output_dir"
 46     "  [CPU/GPU] [DEVICE_ID=0]\n"
 47     "Note: the feature blob names is fixed as 'fc_8_det' in code\n";
 48     return 1;
 49   }
 50   int arg_pos = num_required_args;
 51 
 52   arg_pos = num_required_args;
 53   if (argc > arg_pos && strcmp(argv[arg_pos], "GPU") == 0) {
 54     LOG(ERROR)<< "Using GPU";
 55     uint device_id = 0;
 56     if (argc > arg_pos + 1) {
 57       device_id = atoi(argv[arg_pos + 1]);
 58       CHECK_GE(device_id, 0);
 59     }
 60     LOG(ERROR) << "Using Device_id=" << device_id;
 61     Caffe::SetDevice(device_id);
 62     Caffe::set_mode(Caffe::GPU);
 63   } else {
 64     LOG(ERROR) << "Using CPU";
 65     Caffe::set_mode(Caffe::CPU);
 66   }
 67   Caffe::set_phase(Caffe::TEST);
 68 
 69   arg_pos = 0;  // the name of the executable
 70   string pretrained_binary_proto(argv[++arg_pos]);
 71 
 72   string feature_extraction_proto(argv[++arg_pos]);
 73   shared_ptr<Net<Dtype> > feature_extraction_net(
 74       new Net<Dtype>(feature_extraction_proto));
 75   feature_extraction_net->CopyTrainedLayersFrom(pretrained_binary_proto);
 76     // to get image_paths
 77     const vector<shared_ptr<Layer<float> > > layers = feature_extraction_net->layers();
 78     const caffe::ImageDataLayer<float> *image_layer = dynamic_cast<caffe::ImageDataLayer<float>* >(layers[0].get());
 79     CHECK(image_layer);
 80             
 81   const string blob_name = "fc_8_det";
 82   
 83   CHECK(feature_extraction_net->has_blob(blob_name))   \
 84         << "Unknown feature blob name " << blob_name      \
 85         << " in the network " << feature_extraction_proto;
 86 
 87 
 88   int num_mini_batches = atoi(argv[++arg_pos]);
 89     string output_dir = argv[++arg_pos];
 90     CHECK_EQ(mkdir(output_dir.c_str(),0744), 0) << "mkdir " << output_dir << " failed";
 91 
 92   LOG(ERROR)<< "Extracting Features";
 93 
 94   vector<Blob<float>*> input_vec;
 95   int image_index=0;
 96   for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index) {
 97     feature_extraction_net->Forward(input_vec);
 98         
 99         const shared_ptr<Blob<Dtype> > feature_blob = feature_extraction_net->blob_by_name(blob_name);
100     
101         int batch_size = feature_blob->num();
102         
103         int dim_features = feature_blob->count() / batch_size;
104         CHECK_EQ(dim_features, 4) << "the dim of feature is not equal to 4";
105         
106         Dtype* feature_blob_data;
107         int x1, y1, x2, y2;
108         for (int n = 0; n < batch_size; ++n) {
109             feature_blob_data = feature_blob->mutable_cpu_data() + feature_blob->offset(n);
110 
111             x1 = feature_blob_data[0];
112             y1 = feature_blob_data[1];
113             x2 = feature_blob_data[2];
114             y2 = feature_blob_data[3];
115             
116             string image_path = image_layer->lines_[image_index].first;
117             //LOG(ERROR) << "image_index " << image_index << " " <<  image_path   \
118                                 << " x1 " << feature_blob_data[0] << " y1 " << feature_blob_data[1] \
119                                 << " x2 " << feature_blob_data[2] << " y2 " << feature_blob_data[3];
120             
121             cv::Mat img_origin = cv::imread(image_path);
122             
123             std::vector<string> part_names;
124             boost::split(part_names, image_path, boost::is_any_of("/"));
125             string subname = part_names[part_names.size()-1];             // the last element is the image name.
126             string out_path(output_dir + "/" + subname);
127             
128             //LOG(ERROR) << subname;
129             line(img_origin, cv::Point(x1, y1), cv::Point(x2, y1), cv::Scalar(0, 0, 255), 3);
130             line(img_origin, cv::Point(x2, y1), cv::Point(x2, y2), cv::Scalar(0, 0, 255), 3);
131             line(img_origin, cv::Point(x2, y2), cv::Point(x1, y2), cv::Scalar(0, 0, 255), 3);
132             line(img_origin, cv::Point(x1, y2), cv::Point(x1, y1), cv::Scalar(0, 0, 255), 3);
133             CHECK_EQ(imwrite(output_dir + "/" + subname, img_origin), true) << "write image " + out_path + " failed";
134             
135             image_index ++ ;
136             if (image_index>=image_layer->lines_.size()){
137                 LOG(ERROR) << "Restarting data prefetching from start.";
138                 image_index = 0;
139             }
140             // (image_index>image_layer->lines_.size()-1)?(image_index=0):(image_index++);
141         }
142         
143   }  // for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index)
144   // write the last batch
145   
146   LOG(ERROR)<< "Successfully extracted the features!";
147   return 0;  
148 }