@@ -25,14 +25,14 @@ class Classifier {
2525 const string& mean_file,
2626 const string& label_file);
2727
28- std::vector<Prediction> Classify (const cv::Mat& img , int N = 5 );
28+ std::vector<std::vector< Prediction> > Classify (const std::vector< cv::Mat>& imgs , int N = 5 );
2929
3030 private:
3131 void SetMean (const string& mean_file);
3232
33- std::vector<float > Predict (const cv::Mat& img );
33+ std::vector<std::vector< float > > Predict (const std::vector< cv::Mat>& imgs );
3434
35- void WrapInputLayer (std::vector<cv::Mat>* input_channels);
35+ void WrapInputLayer (std::vector<cv::Mat>* input_channels, int n );
3636
3737 void Preprocess (const cv::Mat& img,
3838 std::vector<cv::Mat>* input_channels);
@@ -102,18 +102,24 @@ static std::vector<int> Argmax(const std::vector<float>& v, int N) {
102102}
103103
104104/* Return the top N predictions. */
105- std::vector<Prediction> Classifier::Classify (const cv::Mat& img, int N) {
106- std::vector<float > output = Predict (img);
107-
108- N = std::min<int >(labels_.size (), N);
109- std::vector<int > maxN = Argmax (output, N);
110- std::vector<Prediction> predictions;
111- for (int i = 0 ; i < N; ++i) {
112- int idx = maxN[i];
113- predictions.push_back (std::make_pair (labels_[idx], output[idx]));
114- }
115-
116- return predictions;
105+ std::vector<std::vector<Prediction> > Classifier::Classify (const std::vector<cv::Mat>& imgs, int N) {
106+ std::vector<std::vector<float > > outputs = Predict (imgs);
107+
108+ std::vector<std::vector<Prediction> > all_predictions;
109+ for (int j = 0 ; j < outputs.size (); ++j) {
110+ std::vector<float > output = outputs[j];
111+
112+ N = std::min<int >(labels_.size (), N);
113+ std::vector<int > maxN = Argmax (output, N);
114+ std::vector<Prediction> predictions;
115+ for (int i = 0 ; i < N; ++i) {
116+ int idx = maxN[i];
117+ predictions.push_back (std::make_pair (labels_[idx], output[idx]));
118+ }
119+ all_predictions.push_back (predictions);
120+ }
121+
122+ return all_predictions;
117123}
118124
119125/* Load the mean file in binaryproto format. */
@@ -147,42 +153,48 @@ void Classifier::SetMean(const string& mean_file) {
147153 mean_ = cv::Mat (input_geometry_, mean.type (), channel_mean);
148154}
149155
150- std::vector<float > Classifier::Predict (const cv::Mat& img ) {
156+ std::vector<std::vector< float > > Classifier::Predict (const std::vector< cv::Mat>& imgs ) {
151157 Blob<float >* input_layer = net_->input_blobs ()[0 ];
152- input_layer->Reshape (1 , num_channels_,
158+ input_layer->Reshape (imgs. size () , num_channels_,
153159 input_geometry_.height , input_geometry_.width );
154160 /* Forward dimension change to all layers. */
155161 net_->Reshape ();
156162
157- std::vector<cv::Mat> input_channels;
158- WrapInputLayer (& input_channels) ;
159-
160- Preprocess (img , &input_channels);
161-
163+ for ( int i = 0 ; i < imgs. size (); ++i) {
164+ std::vector<cv::Mat> input_channels;
165+ WrapInputLayer (&input_channels, i);
166+ Preprocess (imgs[i] , &input_channels);
167+ }
162168 net_->ForwardPrefilled ();
163169
164- /* Copy the output layer to a std::vector */
170+ std::vector<std::vector<float > > outputs;
171+
165172 Blob<float >* output_layer = net_->output_blobs ()[0 ];
166- const float * begin = output_layer->cpu_data ();
167- const float * end = begin + output_layer->channels ();
168- return std::vector<float >(begin, end);
173+ for (int i = 0 ; i < output_layer->num (); ++i) {
174+ const float * begin = output_layer->cpu_data () + i * output_layer->channels ();
175+ const float * end = begin + output_layer->channels ();
176+ /* Copy the output layer to a std::vector */
177+ outputs.push_back (std::vector<float >(begin, end));
178+ }
179+ return outputs;
169180}
170181
171182/* Wrap the input layer of the network in separate cv::Mat objects
172183 * (one per channel). This way we save one memcpy operation and we
173184 * don't need to rely on cudaMemcpy2D. The last preprocessing
174185 * operation will write the separate channels directly to the input
175186 * layer. */
176- void Classifier::WrapInputLayer (std::vector<cv::Mat>* input_channels) {
187+ void Classifier::WrapInputLayer (std::vector<cv::Mat>* input_channels, int n ) {
177188 Blob<float >* input_layer = net_->input_blobs ()[0 ];
178189
179190 int width = input_layer->width ();
180191 int height = input_layer->height ();
181- float * input_data = input_layer->mutable_cpu_data ();
182- for (int i = 0 ; i < input_layer->channels (); ++i) {
183- cv::Mat channel (height, width, CV_32FC1, input_data);
184- input_channels->push_back (channel);
185- input_data += width * height;
192+ int channels = input_layer->channels ();
193+ float * input_data = input_layer->mutable_cpu_data () + n * width * height * channels;
194+ for (int i = 0 ; i < channels; ++i) {
195+ cv::Mat channel (height, width, CV_32FC1, input_data);
196+ input_channels->push_back (channel);
197+ input_data += width * height;
186198 }
187199}
188200
@@ -221,13 +233,15 @@ void Classifier::Preprocess(const cv::Mat& img,
221233 * objects in input_channels. */
222234 cv::split (sample_normalized, *input_channels);
223235
236+ /*
224237 CHECK(reinterpret_cast<float*>(input_channels->at(0).data)
225238 == net_->input_blobs()[0]->cpu_data())
226239 << "Input channels are not wrapping the input layer of the network.";
240+ */
227241}
228242
229243int main (int argc, char ** argv) {
230- if (argc != 6 ) {
244+ if (argc < 6 ) {
231245 std::cerr << " Usage: " << argv[0 ]
232246 << " deploy.prototxt network.caffemodel"
233247 << " mean.binaryproto labels.txt img.jpg" << std::endl;
@@ -242,20 +256,27 @@ int main(int argc, char** argv) {
242256 string label_file = argv[4 ];
243257 Classifier classifier (model_file, trained_file, mean_file, label_file);
244258
245- string file = argv[ 5 ] ;
246-
247- std::cout << " ---------- Prediction for "
248- << file << " ---------- " << std::endl ;
249-
250- cv::Mat img = cv::imread (file, - 1 );
251- CHECK (!img. empty ()) << " Unable to decode image " << file;
252- std::vector<Prediction> predictions = classifier.Classify (img );
259+ std::vector<cv::Mat> imgs ;
260+ for ( int i = 5 ; i < argc; ++i)
261+ {
262+ cv::Mat img = cv::imread (argv[i], - 1 ) ;
263+ CHECK (!img. empty ()) << " Unable to decode image " << argv[i];
264+ imgs. push_back (img );
265+ }
266+ std::vector<std::vector< Prediction> > all_predictions = classifier.Classify (imgs );
253267
254268 /* Print the top N predictions. */
255- for (size_t i = 0 ; i < predictions.size (); ++i) {
256- Prediction p = predictions[i];
257- std::cout << std::fixed << std::setprecision (4 ) << p.second << " - \" "
258- << p.first << " \" " << std::endl;
269+ for (size_t i = 0 ; i < all_predictions.size (); ++i) {
270+ std::cout << " ---------- Prediction for "
271+ << argv[5 + i] << " ----------" << std::endl;
272+
273+ std::vector<Prediction>& predictions = all_predictions[i];
274+ for (size_t j = 0 ; j < predictions.size (); ++j) {
275+ Prediction p = predictions[j];
276+ std::cout << std::fixed << std::setprecision (4 ) << p.second << " - \" "
277+ << p.first << " \" " << std::endl;
278+ }
279+ std::cout << std::endl;
259280 }
260281}
261282#else
0 commit comments