samples/cpp/meanshift_segmentation.cpp

   1 #include "opencv2/highgui/highgui.hpp"
   2 #include "opencv2/core/core.hpp"
   3 #include "opencv2/imgproc/imgproc.hpp"
   4
   5 #include <iostream>
   6
   7 using namespace cv;
   8 using namespace std;
   9
  10 static void help(char** argv)
  11 {
  12     cout << "\nDemonstrate mean-shift based color segmentation in spatial pyramid.\n"
  13     << "Call:\n   " << argv[0] << " image\n"
  14     << "This program allows you to set the spatial and color radius\n"
  15     << "of the mean shift window as well as the number of pyramid reduction levels explored\n"
  16     << endl;
  17 }
  18
  19 //This colors the segmentations
  20 static void floodFillPostprocess( Mat& img, const Scalar& colorDiff=Scalar::all(1) )
  21 {
  22     CV_Assert( !img.empty() );
  23     RNG rng = theRNG();
  24     Mat mask( img.rows+2, img.cols+2, CV_8UC1, Scalar::all(0) );
  25     for( int y = 0; y < img.rows; y++ )
  26     {
  27         for( int x = 0; x < img.cols; x++ )
  28         {
  29             if( mask.at<uchar>(y+1, x+1) == 0 )
  30             {
  31                 Scalar newVal( rng(256), rng(256), rng(256) );
  32                 floodFill( img, mask, Point(x,y), newVal, 0, colorDiff, colorDiff );
  33             }
  34         }
  35     }
  36 }
  37
  38 string winName = "meanshift";
  39 int spatialRad, colorRad, maxPyrLevel;
  40 Mat img, res;
  41
  42 static void meanShiftSegmentation( int, void* )
  43 {
  44     cout << "spatialRad=" << spatialRad << "; "
  45          << "colorRad=" << colorRad << "; "
  46          << "maxPyrLevel=" << maxPyrLevel << endl;
  47     pyrMeanShiftFiltering( img, res, spatialRad, colorRad, maxPyrLevel );
  48     floodFillPostprocess( res, Scalar::all(2) );
  49     imshow( winName, res );
  50 }
  51
  52 int main(int argc, char** argv)
  53 {
  54     if( argc !=2 )
  55     {
  56         help(argv);
  57         return -1;
  58     }
  59
  60     img = imread( argv[1] );
  61     if( img.empty() )
  62         return -1;
  63
  64     spatialRad = 10;
  65     colorRad = 10;
  66     maxPyrLevel = 1;
  67
  68     namedWindow( winName, WINDOW_AUTOSIZE );
  69
  70     createTrackbar( "spatialRad", winName, &spatialRad, 80, meanShiftSegmentation );
  71     createTrackbar( "colorRad", winName, &colorRad, 60, meanShiftSegmentation );
  72     createTrackbar( "maxPyrLevel", winName, &maxPyrLevel, 5, meanShiftSegmentation );
  73
  74     meanShiftSegmentation(0, 0);
  75     waitKey();
  76     return 0;
  77 }