From 4a6888ccf617271333a30e00f239209a26aeef3a Mon Sep 17 00:00:00 2001 From: Alexander Alekhin Date: Tue, 25 Jun 2019 13:42:04 +0300 Subject: [PATCH] imgproc: fix kmeans() call from grabCut() --- modules/core/src/kmeans.cpp | 2 +- modules/imgproc/src/grabcut.cpp | 24 ++++++++++++++++++------ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/modules/core/src/kmeans.cpp b/modules/core/src/kmeans.cpp index 3f10780cd3..bca8b9bd79 100644 --- a/modules/core/src/kmeans.cpp +++ b/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(data0.step)); diff --git a/modules/imgproc/src/grabcut.cpp b/modules/imgproc/src/grabcut.cpp index d68077de1e..7442066eea 100644 --- a/modules/imgproc/src/grabcut.cpp +++ b/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, - 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, - TermCriteria( CV_TERMCRIT_ITER, kMeansItCount, 0.0), 0, kMeansType ); + { + Mat _bgdSamples( (int)bgdSamples.size(), 3, CV_32FC1, &bgdSamples[0][0] ); + 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] ); + 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++ )