Make classification.bin support models with less than 5 classes

The example program would crash if the number of classes was less than
5, since it was still attempting to get the top 5 predictions.

Close #2585

diff --git a/examples/cpp_classification/classification.cpp b/examples/cpp_classification/classification.cpp
index 1c6371e..dc8b863 100644 (file)
--- 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) {