From: Felix Abecassis <fabecassis@nvidia.com>
Date: Fri, 14 Aug 2015 18:15:56 +0000 (-0700)
Subject: Make classification.bin support models with less than 5 classes
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=98593e3aa64d9a8f42723fb03fa46a1343e12c33;p=platform%2Fupstream%2Fcaffe.git

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