augmented_object_selection: augmented_object_selection: GCApplication.cpp Source File

00001 #include "augmented_object_selection/GCApplication.hpp"
00002 
00003 using namespace std;
00004 using namespace cv;
00005 
00006 static const Scalar RD = Scalar(0,0,255);
00007 static const Scalar PK = Scalar(230,130,255);
00008 static const Scalar BL = Scalar(255,0,0);
00009 static const Scalar LB = Scalar(255,255,160);
00010 static const Scalar GR = Scalar(0,255,0);
00011 static const Scalar WH = Scalar(255,255,255);
00012 static const Scalar GY = Scalar(128,128,128);
00013 static const Scalar BK = Scalar(0,0,0);
00014 static const int DOT_RADIUS = 2;
00015 static const int LINE_THICKNESS = -1;
00016 
00017 // Initialize default rectangle marker (not actually used as a rectangle)
00018 const Rect GCApplication::DEFAULT_RECT(0,0,0,0);
00019 // This scale factor is what is actually used to calculate the default
00020 // rectangle.
00021 const int DEFAULT_RECT_SCALE_FACTOR = 2;
00022 
00023 const int BGD_KEY = CV_EVENT_FLAG_CTRLKEY;
00024 const int FGD_KEY = CV_EVENT_FLAG_SHIFTKEY;
00025 
00026 /* NB: GrabCut pixel classes from imgproc.hpp for reference:
00027  * GC_BGD    = 0
00028  * GC_FGD    = 1
00029  * GC_PR_BGD = 2
00030  * GC_PR_FGD = 3
00031  */
00032 
00033 /* Mouse callback trampoline for OpenCV window */
00034 void on_mouse( int event, int x, int y, int flags, void* obj )
00035 {
00036   static_cast<GCApplication*>(obj)->mouseClick(event, x, y, flags);
00037 }
00038 
00039 /* Helper function to convert mask in 4-value GrabCut format to a binary mask
00040  * with foreground = 1, background = 0. */
00041 static void 
00042 binaryMaskFromGCMask(const Mat &_gc_mask, Mat &_bin_mask)
00043 {
00044     if( _gc_mask.empty() || _gc_mask.type() != CV_8UC1 )
00045         CV_Error( CV_StsBadArg, "_gc_mask is empty or has incorrect type (not CV_8UC1)" );
00046 
00047     if(    _bin_mask.empty() 
00048         || _bin_mask.rows!=_gc_mask.rows 
00049         || _bin_mask.cols!=_gc_mask.cols )
00050         _bin_mask.create( _gc_mask.size(), CV_8UC1 );
00051 
00052     // Convert GC_BGD and GC_PR_BGD to 0, GC_FGD and GC_PR_FGD to 1
00053     _bin_mask = _gc_mask & 1;
00054 }
00055 
00056 GCApplication::GCApplication(const string& _name, Mat _img)
00057   : win_name_(_name)
00058 {
00059   // Create OpenCV window
00060   namedWindow( win_name_, CV_WINDOW_AUTOSIZE );
00061 
00062   setMouseCallback( win_name_, on_mouse, this );
00063   ROS_DEBUG("GCApplication: creating window");
00064 
00065   // Initialize window
00066   imageIs(_img);
00067   winImageStateIs(GCApplication::UPDATED);
00068 
00069   iter_count_ = 0;
00070 }
00071 
00072 GCApplication::~GCApplication()
00073 {
00074   destroyWindow( win_name_ );
00075 }
00076 
00077 /* Mutator for initialized state of GCApp object. Client can set
00078  * initialized=true with no side effects; setting initialized=false clears
00079  * background and foreground pixel arrays and returns application state
00080  * variables to NOT_SET. */
00081 // TODO: Should this be public?
00082 void 
00083 GCApplication::initializedIs(bool _init)
00084 {
00085   if (_init)
00086     initialized_=_init;
00087   else
00088   {
00089     winImageStateIs(STALE);
00090     if ( !mask_.empty() ) mask_.setTo(Scalar::all(GC_BGD));
00091     bgd_pxls_.clear();   fgd_pxls_.clear();
00092     pr_bgd_pxls_.clear();  pr_fgd_pxls_.clear();
00093     initialized_   = _init;
00094     rect_state_ = NOT_SET;
00095     lbls_state_ = NOT_SET;
00096     pr_lbls_state_ = NOT_SET;
00097     iterCountIs(0);
00098     winImageStateIs(UPDATED);
00099   }
00100 }
00101 
00102 /* Mutator for GCApp stored image. */
00103 void GCApplication::imageIs(const Mat& _image)
00104 {
00105     if( _image.empty() ) {
00106       ROS_WARN("GCApp: imageIs called with empty image");
00107       return;
00108     }
00109     _image.copyTo(image_);
00110     mask_.create( image_.size(), CV_8UC1);
00111     initializedIs(false);
00112 }
00113 
00114 /* Mutator for GC rectangle. */
00115 void GCApplication::rectIs(const Rect &_r)
00116 {
00117   if (_r == DEFAULT_RECT) 
00118     rect_ = Rect(image_.cols/(DEFAULT_RECT_SCALE_FACTOR*2), 
00119                  image_.rows/(DEFAULT_RECT_SCALE_FACTOR*2), 
00120                  image_.cols/DEFAULT_RECT_SCALE_FACTOR, 
00121                  image_.rows/DEFAULT_RECT_SCALE_FACTOR);
00122   else 
00123     rect_=_r;
00124 
00125   rect_state_ = SET;
00126   setRectInMask();
00127   assert( bgd_pxls_.empty() 
00128   && fgd_pxls_.empty() 
00129   && pr_bgd_pxls_.empty() 
00130   && pr_fgd_pxls_.empty() );
00131   winImageStateIs(UPDATED);
00132 
00133 }
00134 
00135 /* Accessor for binary mask. */
00136 Mat GCApplication::binaryMask()
00137 {
00138   Mat bin_mask;
00139   /*
00140   cv::imwrite("/home/stm1pal/ros/stacks/bosch-ros-pkg/mas/grabcut_app/full_mask.pbm", mask_);
00141   double fm_max;
00142   double fm_min;
00143   minMaxLoc(mask_, &fm_min, &fm_max);
00144   std::cout << "fm_max: " << fm_max << std::endl;
00145   std::cout << "fm_min: " << fm_min << std::endl;
00146   */
00147   binaryMaskFromGCMask(mask_, bin_mask);
00148   /*
00149   cv::imwrite("/home/stm1pal/ros/stacks/bosch-ros-pkg/mas/grabcut_app/bin_mask.pbm", bin_mask);
00150   double bm_max;
00151   double bm_min;
00152   minMaxLoc(bin_mask, &bm_min, &bm_max);
00153   std::cout << "bm_max: " << bm_max << std::endl;
00154   std::cout << "bm_min: " << bm_min << std::endl;
00155   */
00156   return bin_mask;
00157 }
00158 
00159 /* Mutator for window background color. */
00160 void GCApplication::winColorIs(WinColor _c) 
00161 {
00162   winImageStateIs(STALE);
00163   win_color_=_c; 
00164   winImageStateIs(UPDATED);
00165 }
00166 /* Mutator for state of image window. Client can set state=STALE with no side
00167  * effects. Setting state=UPDATED when state=STALE causes new image to be drawn
00168  * in window. */
00169 void GCApplication::winImageStateIs(WinImageState _wis)
00170 {
00171   // Set state=STALE with no side effects
00172   if (_wis == STALE) 
00173   {
00174     win_image_state_=_wis;
00175     return;
00176   }
00177 
00178   // Setting state=UPDATED when already UPDATED is a no-op.
00179   else if (win_image_state_ == UPDATED)
00180     return;
00181 
00182   // Otherwise setting state=UPDATED generates a new image and updates the
00183   // window.
00184   else
00185   {
00186     if( image_.empty() ) 
00187     {
00188       ROS_WARN("GCApp: window image state updated with empty image");
00189       return;
00190     }
00191 
00192     // Generate the display image
00193     Scalar color = BK;
00194     if (win_color_ == WHITE) color = WH;
00195     else if (win_color_ == GRAY) color = GY;
00196     else if (win_color_ == GREEN) color = GR;
00197     else if (win_color_ == BLUE) color = BL;
00198 
00199     Mat display_image(image_.size().height, image_.size().width, 
00200                       image_.type(), color);
00201     Mat bin_mask;
00202     if ( !initialized_ )
00203       // If we haven't created a mask, just copy the base image
00204       image_.copyTo( display_image );
00205     else
00206     {
00207       binaryMaskFromGCMask( mask_, bin_mask );
00208       image_.copyTo( display_image, bin_mask );
00209     }
00210 
00211     // Overlay the marked points, BG and FG, on the display image
00212     vector<Point>::const_iterator it;
00213     for( it = bgd_pxls_.begin(); it != bgd_pxls_.end(); ++it )
00214       circle( display_image, *it, DOT_RADIUS, BL, LINE_THICKNESS );
00215     for( it = fgd_pxls_.begin(); it != fgd_pxls_.end(); ++it )
00216       circle( display_image, *it, DOT_RADIUS, RD, LINE_THICKNESS );
00217     for( it = pr_bgd_pxls_.begin(); it != pr_bgd_pxls_.end(); ++it )
00218       circle( display_image, *it, DOT_RADIUS, LB, LINE_THICKNESS );
00219     for( it = pr_fgd_pxls_.begin(); it != pr_fgd_pxls_.end(); ++it )
00220       circle( display_image, *it, DOT_RADIUS, PK, LINE_THICKNESS );
00221 
00222     // Add the rectangle
00223     if( rect_state_ == IN_PROCESS || rect_state_ == SET )
00224       rectangle( display_image, 
00225           Point( rect_.x, rect_.y ), 
00226           Point(rect_.x + rect_.width, rect_.y + rect_.height ), 
00227           GR, 2);
00228 
00229     // Display the image
00230     imshow( win_name_, display_image );
00231   }
00232 }
00233 
00234 /* Private helper method to transfer pixels from designated rectangle to
00235  * Grabcut mask. */
00236 void GCApplication::setRectInMask()
00237 {
00238     assert( !mask_.empty() );
00239     mask_.setTo( GC_BGD );
00240     rect_.x = max(0, rect_.x);
00241     rect_.y = max(0, rect_.y);
00242     rect_.width = min(rect_.width, image_.cols-rect_.x);
00243     rect_.height = min(rect_.height, image_.rows-rect_.y);
00244     (mask_(rect_)).setTo( Scalar(GC_PR_FGD) );
00245 }
00246 
00247 /* Private helper method to transfer labels designated in UI to mask. */
00248 void GCApplication::setLblsInMask( int flags, Point p, bool isPr )
00249 {
00250     vector<Point> *bpxls, *fpxls;
00251     uchar bvalue, fvalue;
00252     if( !isPr )
00253     {
00254       bpxls = &bgd_pxls_;
00255       fpxls = &fgd_pxls_;
00256       bvalue = GC_BGD;
00257       fvalue = GC_FGD;
00258     }
00259     else
00260     {
00261       bpxls = &pr_bgd_pxls_;
00262       fpxls = &pr_fgd_pxls_;
00263       bvalue = GC_PR_BGD;
00264       fvalue = GC_PR_FGD;
00265     }
00266 
00267     if( flags & BGD_KEY )
00268     {
00269       bpxls->push_back(p);
00270       circle( mask_, p, DOT_RADIUS, bvalue, LINE_THICKNESS );
00271     }
00272     if( flags & FGD_KEY )
00273     {
00274       fpxls->push_back(p);
00275       circle( mask_, p, DOT_RADIUS, fvalue, LINE_THICKNESS );
00276     }
00277 }
00278 
00279 /* Mouse callback function passed to OpenCV window. Handles marking the ROI
00280  * rectangle and foreground and background pixel hints. */
00281 void GCApplication::mouseClick( int event, int x, int y, int flags)
00282 {
00283     // TODO add bad args check
00284   switch( event )
00285   {
00286     case CV_EVENT_LBUTTONDOWN: // set rect or GC_BGD(GC_FGD) labels
00287       {
00288         bool isb = (flags & BGD_KEY) != 0,
00289              isf = (flags & FGD_KEY) != 0;
00290         if( rect_state_ == NOT_SET && !isb && !isf )
00291         {
00292           rect_state_ = IN_PROCESS;
00293           rect_ = Rect( x, y, 1, 1 );
00294         }
00295         if ( (isb || isf) && rect_state_ == SET )
00296           lbls_state_ = IN_PROCESS;
00297       }
00298       break;
00299 
00300     case CV_EVENT_RBUTTONDOWN: // set GC_PR_BGD(GC_PR_FGD) labels
00301       {
00302         bool isb = (flags & BGD_KEY) != 0;
00303         bool isf = (flags & FGD_KEY) != 0;
00304         if ( (isb || isf) && rect_state_ == SET )
00305           pr_lbls_state_ = IN_PROCESS;
00306       }
00307       break;
00308 
00309     case CV_EVENT_LBUTTONUP:
00310       if( rect_state_ == IN_PROCESS )
00311       {
00312         rect_ = Rect( Point(rect_.x, rect_.y), Point(x,y) );
00313         rect_state_ = SET;
00314         setRectInMask();
00315         assert( bgd_pxls_.empty() 
00316                && fgd_pxls_.empty() 
00317                && pr_bgd_pxls_.empty() 
00318                && pr_fgd_pxls_.empty() );
00319         winImageStateIs(UPDATED);
00320       }
00321       if( lbls_state_ == IN_PROCESS )
00322       {
00323         setLblsInMask(flags, Point(x,y), false);
00324         lbls_state_ = SET;
00325         winImageStateIs(UPDATED);
00326       }
00327       break;
00328 
00329     case CV_EVENT_RBUTTONUP:
00330       if( pr_lbls_state_ == IN_PROCESS )
00331       {
00332         setLblsInMask(flags, Point(x,y), true);
00333         pr_lbls_state_ = SET;
00334         winImageStateIs(UPDATED);
00335       }
00336       break;
00337       
00338     case CV_EVENT_MOUSEMOVE:
00339       if( rect_state_ == IN_PROCESS )
00340       {
00341         rect_ = Rect( Point(rect_.x, rect_.y), Point(x,y) );
00342         assert( bgd_pxls_.empty() 
00343                && fgd_pxls_.empty() 
00344                && pr_bgd_pxls_.empty() 
00345                && pr_fgd_pxls_.empty() );
00346         winImageStateIs(UPDATED);
00347       }
00348       else if( lbls_state_ == IN_PROCESS )
00349       {
00350         setLblsInMask(flags, Point(x,y), false);
00351         winImageStateIs(UPDATED);
00352       }
00353       else if( pr_lbls_state_ == IN_PROCESS )
00354       {
00355         setLblsInMask(flags, Point(x,y), true);
00356         winImageStateIs(UPDATED);
00357       }
00358       break;
00359   }
00360 }
00361 
00362 void GCApplication::iterCountIs(int _icnt)
00363 {
00364   if (_icnt == 0) {
00365     iter_count_ = 0;
00366     return;
00367   }
00368 
00369   int iter_inc = _icnt - iter_count_;
00370   if (iter_inc <= 0) return;
00371   
00372   if( initialized() )
00373     grabCut( image_, mask_, rect_, bgd_model_, fgd_model_, iter_inc );
00374   else
00375   {
00376     if( rect_state_ != SET ) return;
00377 
00378     if( lbls_state_ == SET || pr_lbls_state_ == SET )
00379       grabCut( image_, mask_, rect_, bgd_model_, fgd_model_, 
00380                iter_inc, GC_INIT_WITH_MASK );
00381     else
00382       grabCut( image_, mask_, rect_, bgd_model_, fgd_model_, 
00383                iter_inc, GC_INIT_WITH_RECT );
00384 
00385     initializedIs(true);
00386   }
00387   iter_count_ = _icnt;
00388 
00389   bgd_pxls_.clear(); fgd_pxls_.clear();
00390   pr_bgd_pxls_.clear(); pr_fgd_pxls_.clear();
00391 
00392   winImageStateIs(UPDATED);
00393 
00394   return;
00395 }
00396 
00397