1 #include "opencv2/imgcodecs.hpp"
2 #include "opencv2/highgui/highgui.hpp"
3 #include "opencv2/imgproc/imgproc.hpp"
12 cout << "\nThis program demonstrates GrabCut segmentation -- select an object in a region\n"
13 "and then grabcut will attempt to segment it out.\n"
15 "./grabcut <image_name>\n"
16 "\nSelect a rectangular area around the object you want to segment\n" <<
18 "\tESC - quit the program\n"
19 "\tr - restore the original image\n"
20 "\tn - next iteration\n"
22 "\tleft mouse button - set rectangle\n"
24 "\tCTRL+left mouse button - set GC_BGD pixels\n"
25 "\tSHIFT+left mouse button - set CG_FGD pixels\n"
27 "\tCTRL+right mouse button - set GC_PR_BGD pixels\n"
28 "\tSHIFT+right mouse button - set CG_PR_FGD pixels\n" << endl;
31 const Scalar RED = Scalar(0,0,255);
32 const Scalar PINK = Scalar(230,130,255);
33 const Scalar BLUE = Scalar(255,0,0);
34 const Scalar LIGHTBLUE = Scalar(255,255,160);
35 const Scalar GREEN = Scalar(0,255,0);
37 const int BGD_KEY = EVENT_FLAG_CTRLKEY;
38 const int FGD_KEY = EVENT_FLAG_SHIFTKEY;
40 static void getBinMask( const Mat& comMask, Mat& binMask )
42 if( comMask.empty() || comMask.type()!=CV_8UC1 )
43 CV_Error( Error::StsBadArg, "comMask is empty or has incorrect type (not CV_8UC1)" );
44 if( binMask.empty() || binMask.rows!=comMask.rows || binMask.cols!=comMask.cols )
45 binMask.create( comMask.size(), CV_8UC1 );
46 binMask = comMask & 1;
52 enum{ NOT_SET = 0, IN_PROCESS = 1, SET = 2 };
53 static const int radius = 2;
54 static const int thickness = -1;
57 void setImageAndWinName( const Mat& _image, const string& _winName );
58 void showImage() const;
59 void mouseClick( int event, int x, int y, int flags, void* param );
61 int getIterCount() const { return iterCount; }
64 void setLblsInMask( int flags, Point p, bool isPr );
66 const string* winName;
69 Mat bgdModel, fgdModel;
71 uchar rectState, lblsState, prLblsState;
75 vector<Point> fgdPxls, bgdPxls, prFgdPxls, prBgdPxls;
79 void GCApplication::reset()
82 mask.setTo(Scalar::all(GC_BGD));
83 bgdPxls.clear(); fgdPxls.clear();
84 prBgdPxls.clear(); prFgdPxls.clear();
86 isInitialized = false;
89 prLblsState = NOT_SET;
93 void GCApplication::setImageAndWinName( const Mat& _image, const string& _winName )
95 if( _image.empty() || _winName.empty() )
99 mask.create( image->size(), CV_8UC1);
103 void GCApplication::showImage() const
105 if( image->empty() || winName->empty() )
111 image->copyTo( res );
114 getBinMask( mask, binMask );
115 image->copyTo( res, binMask );
118 vector<Point>::const_iterator it;
119 for( it = bgdPxls.begin(); it != bgdPxls.end(); ++it )
120 circle( res, *it, radius, BLUE, thickness );
121 for( it = fgdPxls.begin(); it != fgdPxls.end(); ++it )
122 circle( res, *it, radius, RED, thickness );
123 for( it = prBgdPxls.begin(); it != prBgdPxls.end(); ++it )
124 circle( res, *it, radius, LIGHTBLUE, thickness );
125 for( it = prFgdPxls.begin(); it != prFgdPxls.end(); ++it )
126 circle( res, *it, radius, PINK, thickness );
128 if( rectState == IN_PROCESS || rectState == SET )
129 rectangle( res, Point( rect.x, rect.y ), Point(rect.x + rect.width, rect.y + rect.height ), GREEN, 2);
131 imshow( *winName, res );
134 void GCApplication::setRectInMask()
136 CV_Assert( !mask.empty() );
137 mask.setTo( GC_BGD );
138 rect.x = max(0, rect.x);
139 rect.y = max(0, rect.y);
140 rect.width = min(rect.width, image->cols-rect.x);
141 rect.height = min(rect.height, image->rows-rect.y);
142 (mask(rect)).setTo( Scalar(GC_PR_FGD) );
145 void GCApplication::setLblsInMask( int flags, Point p, bool isPr )
147 vector<Point> *bpxls, *fpxls;
148 uchar bvalue, fvalue;
163 if( flags & BGD_KEY )
166 circle( mask, p, radius, bvalue, thickness );
168 if( flags & FGD_KEY )
171 circle( mask, p, radius, fvalue, thickness );
175 void GCApplication::mouseClick( int event, int x, int y, int flags, void* )
177 // TODO add bad args check
180 case EVENT_LBUTTONDOWN: // set rect or GC_BGD(GC_FGD) labels
182 bool isb = (flags & BGD_KEY) != 0,
183 isf = (flags & FGD_KEY) != 0;
184 if( rectState == NOT_SET && !isb && !isf )
186 rectState = IN_PROCESS;
187 rect = Rect( x, y, 1, 1 );
189 if ( (isb || isf) && rectState == SET )
190 lblsState = IN_PROCESS;
193 case EVENT_RBUTTONDOWN: // set GC_PR_BGD(GC_PR_FGD) labels
195 bool isb = (flags & BGD_KEY) != 0,
196 isf = (flags & FGD_KEY) != 0;
197 if ( (isb || isf) && rectState == SET )
198 prLblsState = IN_PROCESS;
201 case EVENT_LBUTTONUP:
202 if( rectState == IN_PROCESS )
204 rect = Rect( Point(rect.x, rect.y), Point(x,y) );
207 CV_Assert( bgdPxls.empty() && fgdPxls.empty() && prBgdPxls.empty() && prFgdPxls.empty() );
210 if( lblsState == IN_PROCESS )
212 setLblsInMask(flags, Point(x,y), false);
217 case EVENT_RBUTTONUP:
218 if( prLblsState == IN_PROCESS )
220 setLblsInMask(flags, Point(x,y), true);
225 case EVENT_MOUSEMOVE:
226 if( rectState == IN_PROCESS )
228 rect = Rect( Point(rect.x, rect.y), Point(x,y) );
229 CV_Assert( bgdPxls.empty() && fgdPxls.empty() && prBgdPxls.empty() && prFgdPxls.empty() );
232 else if( lblsState == IN_PROCESS )
234 setLblsInMask(flags, Point(x,y), false);
237 else if( prLblsState == IN_PROCESS )
239 setLblsInMask(flags, Point(x,y), true);
246 int GCApplication::nextIter()
249 grabCut( *image, mask, rect, bgdModel, fgdModel, 1 );
252 if( rectState != SET )
255 if( lblsState == SET || prLblsState == SET )
256 grabCut( *image, mask, rect, bgdModel, fgdModel, 1, GC_INIT_WITH_MASK );
258 grabCut( *image, mask, rect, bgdModel, fgdModel, 1, GC_INIT_WITH_RECT );
260 isInitialized = true;
264 bgdPxls.clear(); fgdPxls.clear();
265 prBgdPxls.clear(); prFgdPxls.clear();
272 static void on_mouse( int event, int x, int y, int flags, void* param )
274 gcapp.mouseClick( event, x, y, flags, param );
277 int main( int argc, char** argv )
284 string filename = argv[1];
285 if( filename.empty() )
287 cout << "\nDurn, couldn't read in " << argv[1] << endl;
290 Mat image = imread( filename, 1 );
293 cout << "\n Durn, couldn't read image filename " << filename << endl;
299 const string winName = "image";
300 namedWindow( winName, WINDOW_AUTOSIZE );
301 setMouseCallback( winName, on_mouse, 0 );
303 gcapp.setImageAndWinName( image, winName );
312 cout << "Exiting ..." << endl;
320 int iterCount = gcapp.getIterCount();
321 cout << "<" << iterCount << "... ";
322 int newIterCount = gcapp.nextIter();
323 if( newIterCount > iterCount )
326 cout << iterCount << ">" << endl;
329 cout << "rect must be determined>" << endl;
335 destroyWindow( winName );