Merge pull request #2925 from flx42/fix_classification_example

Make classification.bin support models with less than 5 classes
BVLC · Aug 14, 2015 · c6b9f58 · c6b9f58
2 parents 65c7fa6 + 98593e3
commit c6b9f58
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/cpp_classification/classification.cpp b/examples/cpp_classification/classification.cpp
@@ -2,6 +2,7 @@
 #include <opencv2/core/core.hpp>
 #include <opencv2/highgui/highgui.hpp>
 #include <opencv2/imgproc/imgproc.hpp>
+#include <algorithm>
 #include <iosfwd>
 #include <memory>
 #include <string>
@@ -101,6 +102,7 @@ static std::vector<int> Argmax(const std::vector<float>& v, int N) {
 std::vector<Prediction> Classifier::Classify(const cv::Mat& img, int N) {
   std::vector<float> output = Predict(img);
 
+  N = std::min<int>(labels_.size(), N);
   std::vector<int> maxN = Argmax(output, N);
   std::vector<Prediction> predictions;
   for (int i = 0; i < N; ++i) {