samples/cpp/grabcut.cpp

   1 #include "opencv2/highgui/highgui.hpp"
   2 #include "opencv2/imgproc/imgproc.hpp"
   3
   4 #include <iostream>
   5
   6 using namespace std;
   7 using namespace cv;
   8
   9 static void help()
  10 {
  11     cout << "\nThis program demonstrates GrabCut segmentation -- select an object in a region\n"
  12             "and then grabcut will attempt to segment it out.\n"
  13             "Call:\n"
  14             "./grabcut <image_name>\n"
  15         "\nSelect a rectangular area around the object you want to segment\n" <<
  16         "\nHot keys: \n"
  17         "\tESC - quit the program\n"
  18         "\tr - restore the original image\n"
  19         "\tn - next iteration\n"
  20         "\n"
  21         "\tleft mouse button - set rectangle\n"
  22         "\n"
  23         "\tCTRL+left mouse button - set GC_BGD pixels\n"
  24         "\tSHIFT+left mouse button - set CG_FGD pixels\n"
  25         "\n"
  26         "\tCTRL+right mouse button - set GC_PR_BGD pixels\n"
  27         "\tSHIFT+right mouse button - set CG_PR_FGD pixels\n" << endl;
  28 }
  29
  30 const Scalar RED = Scalar(0,0,255);
  31 const Scalar PINK = Scalar(230,130,255);
  32 const Scalar BLUE = Scalar(255,0,0);
  33 const Scalar LIGHTBLUE = Scalar(255,255,160);
  34 const Scalar GREEN = Scalar(0,255,0);
  35
  36 const int BGD_KEY = EVENT_FLAG_CTRLKEY;
  37 const int FGD_KEY = EVENT_FLAG_SHIFTKEY;
  38
  39 static void getBinMask( const Mat& comMask, Mat& binMask )
  40 {
  41     if( comMask.empty() || comMask.type()!=CV_8UC1 )
  42         CV_Error( Error::StsBadArg, "comMask is empty or has incorrect type (not CV_8UC1)" );
  43     if( binMask.empty() || binMask.rows!=comMask.rows || binMask.cols!=comMask.cols )
  44         binMask.create( comMask.size(), CV_8UC1 );
  45     binMask = comMask & 1;
  46 }
  47
  48 class GCApplication
  49 {
  50 public:
  51     enum{ NOT_SET = 0, IN_PROCESS = 1, SET = 2 };
  52     static const int radius = 2;
  53     static const int thickness = -1;
  54
  55     void reset();
  56     void setImageAndWinName( const Mat& _image, const string& _winName );
  57     void showImage() const;
  58     void mouseClick( int event, int x, int y, int flags, void* param );
  59     int nextIter();
  60     int getIterCount() const { return iterCount; }
  61 private:
  62     void setRectInMask();
  63     void setLblsInMask( int flags, Point p, bool isPr );
  64
  65     const string* winName;
  66     const Mat* image;
  67     Mat mask;
  68     Mat bgdModel, fgdModel;
  69
  70     uchar rectState, lblsState, prLblsState;
  71     bool isInitialized;
  72
  73     Rect rect;
  74     vector<Point> fgdPxls, bgdPxls, prFgdPxls, prBgdPxls;
  75     int iterCount;
  76 };
  77
  78 void GCApplication::reset()
  79 {
  80     if( !mask.empty() )
  81         mask.setTo(Scalar::all(GC_BGD));
  82     bgdPxls.clear(); fgdPxls.clear();
  83     prBgdPxls.clear();  prFgdPxls.clear();
  84
  85     isInitialized = false;
  86     rectState = NOT_SET;
  87     lblsState = NOT_SET;
  88     prLblsState = NOT_SET;
  89     iterCount = 0;
  90 }
  91
  92 void GCApplication::setImageAndWinName( const Mat& _image, const string& _winName  )
  93 {
  94     if( _image.empty() || _winName.empty() )
  95         return;
  96     image = &_image;
  97     winName = &_winName;
  98     mask.create( image->size(), CV_8UC1);
  99     reset();
 100 }
 101
 102 void GCApplication::showImage() const
 103 {
 104     if( image->empty() || winName->empty() )
 105         return;
 106
 107     Mat res;
 108     Mat binMask;
 109     if( !isInitialized )
 110         image->copyTo( res );
 111     else
 112     {
 113         getBinMask( mask, binMask );
 114         image->copyTo( res, binMask );
 115     }
 116
 117     vector<Point>::const_iterator it;
 118     for( it = bgdPxls.begin(); it != bgdPxls.end(); ++it )
 119         circle( res, *it, radius, BLUE, thickness );
 120     for( it = fgdPxls.begin(); it != fgdPxls.end(); ++it )
 121         circle( res, *it, radius, RED, thickness );
 122     for( it = prBgdPxls.begin(); it != prBgdPxls.end(); ++it )
 123         circle( res, *it, radius, LIGHTBLUE, thickness );
 124     for( it = prFgdPxls.begin(); it != prFgdPxls.end(); ++it )
 125         circle( res, *it, radius, PINK, thickness );
 126
 127     if( rectState == IN_PROCESS || rectState == SET )
 128         rectangle( res, Point( rect.x, rect.y ), Point(rect.x + rect.width, rect.y + rect.height ), GREEN, 2);
 129
 130     imshow( *winName, res );
 131 }
 132
 133 void GCApplication::setRectInMask()
 134 {
 135     CV_Assert( !mask.empty() );
 136     mask.setTo( GC_BGD );
 137     rect.x = max(0, rect.x);
 138     rect.y = max(0, rect.y);
 139     rect.width = min(rect.width, image->cols-rect.x);
 140     rect.height = min(rect.height, image->rows-rect.y);
 141     (mask(rect)).setTo( Scalar(GC_PR_FGD) );
 142 }
 143
 144 void GCApplication::setLblsInMask( int flags, Point p, bool isPr )
 145 {
 146     vector<Point> *bpxls, *fpxls;
 147     uchar bvalue, fvalue;
 148     if( !isPr )
 149     {
 150         bpxls = &bgdPxls;
 151         fpxls = &fgdPxls;
 152         bvalue = GC_BGD;
 153         fvalue = GC_FGD;
 154     }
 155     else
 156     {
 157         bpxls = &prBgdPxls;
 158         fpxls = &prFgdPxls;
 159         bvalue = GC_PR_BGD;
 160         fvalue = GC_PR_FGD;
 161     }
 162     if( flags & BGD_KEY )
 163     {
 164         bpxls->push_back(p);
 165         circle( mask, p, radius, bvalue, thickness );
 166     }
 167     if( flags & FGD_KEY )
 168     {
 169         fpxls->push_back(p);
 170         circle( mask, p, radius, fvalue, thickness );
 171     }
 172 }
 173
 174 void GCApplication::mouseClick( int event, int x, int y, int flags, void* )
 175 {
 176     // TODO add bad args check
 177     switch( event )
 178     {
 179     case EVENT_LBUTTONDOWN: // set rect or GC_BGD(GC_FGD) labels
 180         {
 181             bool isb = (flags & BGD_KEY) != 0,
 182                  isf = (flags & FGD_KEY) != 0;
 183             if( rectState == NOT_SET && !isb && !isf )
 184             {
 185                 rectState = IN_PROCESS;
 186                 rect = Rect( x, y, 1, 1 );
 187             }
 188             if ( (isb || isf) && rectState == SET )
 189                 lblsState = IN_PROCESS;
 190         }
 191         break;
 192     case EVENT_RBUTTONDOWN: // set GC_PR_BGD(GC_PR_FGD) labels
 193         {
 194             bool isb = (flags & BGD_KEY) != 0,
 195                  isf = (flags & FGD_KEY) != 0;
 196             if ( (isb || isf) && rectState == SET )
 197                 prLblsState = IN_PROCESS;
 198         }
 199         break;
 200     case EVENT_LBUTTONUP:
 201         if( rectState == IN_PROCESS )
 202         {
 203             rect = Rect( Point(rect.x, rect.y), Point(x,y) );
 204             rectState = SET;
 205             setRectInMask();
 206             CV_Assert( bgdPxls.empty() && fgdPxls.empty() && prBgdPxls.empty() && prFgdPxls.empty() );
 207             showImage();
 208         }
 209         if( lblsState == IN_PROCESS )
 210         {
 211             setLblsInMask(flags, Point(x,y), false);
 212             lblsState = SET;
 213             showImage();
 214         }
 215         break;
 216     case EVENT_RBUTTONUP:
 217         if( prLblsState == IN_PROCESS )
 218         {
 219             setLblsInMask(flags, Point(x,y), true);
 220             prLblsState = SET;
 221             showImage();
 222         }
 223         break;
 224     case EVENT_MOUSEMOVE:
 225         if( rectState == IN_PROCESS )
 226         {
 227             rect = Rect( Point(rect.x, rect.y), Point(x,y) );
 228             CV_Assert( bgdPxls.empty() && fgdPxls.empty() && prBgdPxls.empty() && prFgdPxls.empty() );
 229             showImage();
 230         }
 231         else if( lblsState == IN_PROCESS )
 232         {
 233             setLblsInMask(flags, Point(x,y), false);
 234             showImage();
 235         }
 236         else if( prLblsState == IN_PROCESS )
 237         {
 238             setLblsInMask(flags, Point(x,y), true);
 239             showImage();
 240         }
 241         break;
 242     }
 243 }
 244
 245 int GCApplication::nextIter()
 246 {
 247     if( isInitialized )
 248         grabCut( *image, mask, rect, bgdModel, fgdModel, 1 );
 249     else
 250     {
 251         if( rectState != SET )
 252             return iterCount;
 253
 254         if( lblsState == SET || prLblsState == SET )
 255             grabCut( *image, mask, rect, bgdModel, fgdModel, 1, GC_INIT_WITH_MASK );
 256         else
 257             grabCut( *image, mask, rect, bgdModel, fgdModel, 1, GC_INIT_WITH_RECT );
 258
 259         isInitialized = true;
 260     }
 261     iterCount++;
 262
 263     bgdPxls.clear(); fgdPxls.clear();
 264     prBgdPxls.clear(); prFgdPxls.clear();
 265
 266     return iterCount;
 267 }
 268
 269 GCApplication gcapp;
 270
 271 static void on_mouse( int event, int x, int y, int flags, void* param )
 272 {
 273     gcapp.mouseClick( event, x, y, flags, param );
 274 }
 275
 276 int main( int argc, char** argv )
 277 {
 278     if( argc!=2 )
 279     {
 280         help();
 281         return 1;
 282     }
 283     string filename = argv[1];
 284     if( filename.empty() )
 285     {
 286         cout << "\nDurn, couldn't read in " << argv[1] << endl;
 287         return 1;
 288     }
 289     Mat image = imread( filename, 1 );
 290     if( image.empty() )
 291     {
 292         cout << "\n Durn, couldn't read image filename " << filename << endl;
 293         return 1;
 294     }
 295
 296     help();
 297
 298     const string winName = "image";
 299     namedWindow( winName, WINDOW_AUTOSIZE );
 300     setMouseCallback( winName, on_mouse, 0 );
 301
 302     gcapp.setImageAndWinName( image, winName );
 303     gcapp.showImage();
 304
 305     for(;;)
 306     {
 307         int c = waitKey(0);
 308         switch( (char) c )
 309         {
 310         case '\x1b':
 311             cout << "Exiting ..." << endl;
 312             goto exit_main;
 313         case 'r':
 314             cout << endl;
 315             gcapp.reset();
 316             gcapp.showImage();
 317             break;
 318         case 'n':
 319             int iterCount = gcapp.getIterCount();
 320             cout << "<" << iterCount << "... ";
 321             int newIterCount = gcapp.nextIter();
 322             if( newIterCount > iterCount )
 323             {
 324                 gcapp.showImage();
 325                 cout << iterCount << ">" << endl;
 326             }
 327             else
 328                 cout << "rect must be determined>" << endl;
 329             break;
 330         }
 331     }
 332
 333 exit_main:
 334     destroyWindow( winName );
 335     return 0;
 336 }