imgproc: fix kmeans() call from grabCut()

modules/core/src/kmeans.cpp

@@ -238,7 +238,7 @@ double cv::kmeans( InputArray _data, int K,
 
     attempts = std::max(attempts, 1);
     CV_Assert( data0.dims <= 2 && type == CV_32F && K > 0 );
-    CV_Assert( N >= K );
+    CV_CheckGE(N, K, "Number of clusters should be more than number of elements");
 
     Mat data(N, dims, CV_32F, data0.ptr(), isrow ? dims * sizeof(float) : static_cast<size_t>(data0.step));
 

modules/imgproc/src/grabcut.cpp

@@ -45,6 +45,8 @@
 
 using namespace cv;
 
+namespace {
+
 /*
 This is implementation of image segmentation algorithm GrabCut described in
 "GrabCut - Interactive Foreground Extraction using Iterated Graph Cuts".
@@ -228,6 +230,8 @@ void GMM::calcInverseCovAndDeterm(int ci, const double singularFix)
     }
 }
 
+} // namespace
+
 /*
   Calculate beta - parameter of GrabCut algorithm.
   beta = 1/(2*avg(sqr(||color[i] - color[j]||)))
@@ -379,12 +383,20 @@ static void initGMMs( const Mat& img, const Mat& mask, GMM& bgdGMM, GMM& fgdGMM
         }
     }
     CV_Assert( !bgdSamples.empty() && !fgdSamples.empty() );
+    {
         Mat _bgdSamples( (int)bgdSamples.size(), 3, CV_32FC1, &bgdSamples[0][0] );
-    kmeans( _bgdSamples, GMM::componentsCount, bgdLabels,
+        int num_clusters = GMM::componentsCount;
+        num_clusters = std::min(num_clusters, (int)bgdSamples.size());
+        kmeans( _bgdSamples, num_clusters, bgdLabels,
                 TermCriteria( CV_TERMCRIT_ITER, kMeansItCount, 0.0), 0, kMeansType );
+    }
+    {
         Mat _fgdSamples( (int)fgdSamples.size(), 3, CV_32FC1, &fgdSamples[0][0] );
-    kmeans( _fgdSamples, GMM::componentsCount, fgdLabels,
+        int num_clusters = GMM::componentsCount;
+        num_clusters = std::min(num_clusters, (int)fgdSamples.size());
+        kmeans( _fgdSamples, num_clusters, fgdLabels,
                 TermCriteria( CV_TERMCRIT_ITER, kMeansItCount, 0.0), 0, kMeansType );
+    }
 
     bgdGMM.initLearning();
     for( int i = 0; i < (int)bgdSamples.size(); i++ )